#!/usr/bin/env python2.7

# meant to implement the most basic Jarvis correction factor. Makes a map of
# star ellipticities moments; uses this to calculate R = 1 - r_o/r_* (the 
# resolution parameter).

from __future__ import division
import sys
import numpy as np
from astropy.io import fits
from matplotlib import *
from pylab import *

def get_column_names(fiatfile):
    """
    returns list of column names from a .fiat file.
    """
    col_names = []
    with open(fiatfile) as catalog:
        for line in catalog:
            if line.startswith('#'):
                if 'ttype' in line.lower():
                    col_names.append(line.split('=')[1].split()[0].lower())
            else:
                return col_names

def make_map(fiatfile):
    cols = get_column_names(fiatfile)
    x_loc = cols.index('x')
    y_loc = cols.index('y')
    try:
        size_loc = cols.index('size')
    except ValueError:
        size_loc = cols.index('fwhm')
    xx_loc = cols.index('ixx')
    yy_loc = cols.index('iyy')
    xy_loc = cols.index('ixy') # moments
    stars = []
    with open(fiatfile) as catalog:
        for l in catalog:
            if not l.startswith('#'):
                line = l.split()
                stars.append([line[x_loc], line[y_loc], line[size_loc],
                              line[xx_loc], line[yy_loc], line[xy_loc]])

    fin = [[float(s) for s in a] for a in stars]
    return fin

def correct_galaxies(galaxyfile, starfile):
    # going to be pretty slow for now. should split into quarters
    stars = make_map(starfile)
    root = '.'.join(galaxyfile.split('.')[:-1])
    targetfile = root + '.galaxy.corrected' # .fiat format
    galaxy_cols = get_column_names(galaxyfile)
    x_loc = galaxy_cols.index('x')
    y_loc = galaxy_cols.index('y')
    try:
        size_loc = galaxy_cols.index('fwhm')
    except ValueError:
        size_loc = galaxy_cols.index('size')
    e_loc = galaxy_cols.index('ellipticity')
    xx_loc = galaxy_cols.index('ixx')
    yy_loc = galaxy_cols.index('iyy')
    xy_loc = galaxy_cols.index('ixy') # moments
    ct = 0
    with open(targetfile, 'w+') as target:
        print "writing to {0}...".format(targetfile)
        with open(galaxyfile) as galaxies:
            for l in galaxies:
                if l.startswith('#'):
                    target.write(l)
                else: # find the nearest star
                    line = l.split()
                    x = float(line[x_loc])
                    y = float(line[y_loc])
                    #size = float(line[size_loc])
                    xx = float(line[xx_loc])
                    yy = float(line[yy_loc])
                    xy = float(line[xy_loc])
                    size = 0.5*(xx + yy)
                    min_distance = 1000.0
                    for s in stars:
                        sep = np.sqrt((s[0] - x)**2 + (s[1] - y)**2)
                        if sep < min_distance:
                            min_distance = sep
                            #min_loc = [stars[0], stars[1]]
                            star_size = 0.5*(s[3]+s[4])
                    if star_size < size:
                        R = 1.0 - star_size/size
                    else:
                        R = 1.0
                    newline = line
                    newline[e_loc] = str(float(newline[e_loc])/R)
                    newline[xx_loc] = str(float(newline[xx_loc])/R)
                    newline[yy_loc] = str(float(newline[yy_loc])/R)
                    newline[xy_loc] = str(float(newline[xy_loc])/R)
                    if size > 2.0*star_size and float(newline[e_loc]) < 1:
                        print R
                        target.write(' '.join(newline) + '\n')
                        ct += 1

            print "{0} objects survived.".format(ct)

if __name__ == '__main__':
    if len(sys.argv) == 3:
        correct_galaxies(sys.argv[1], sys.argv[2])
    else:
        print "selection galaxyfile starfile"
