#!/usr/bin/env python2.7

# ...plots a .cat file. see also fitsplot.
# usage: python catplot.py catfile 'mag' 'fwhm'

from __future__ import division
import sys
import warnings
import numpy as np
from matplotlib import *
from scipy.optimize import curve_fit
from pylab import *

warnings.filterwarnings('ignore', '', FutureWarning)
numbins = 14

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

def p(e, sig, es):
    """
    probability distribution from Viola et al 2013.
    es is the list of ellipticities.
    """

    return (len(es)/numbins)*(np.abs(e)/sig**2)*\
        np.exp(-np.abs(e)**2/(2*sig**2))

if __name__ == '__main__':
    if len(sys.argv) == 3: 
        cols = get_column_names(sys.argv[1])
        try:
            e_loc = cols.index('ellipticity')
        except (IndexError, ValueError) as e:
            e_loc = cols.index('ellip')
        es = []
        for line in open(sys.argv[1]):
            if not line.startswith('#'):
                es.append(float(line.split()[e_loc]))
        binned = np.histogram(es, numbins, range=(0,1.0))

        cols1 = get_column_names(sys.argv[2])
        try:
            e_loc1 = cols1.index('ellipticity')
        except (IndexError, ValueError) as e:
            e_loc1 = cols1.index('ellip')

        es1 = []
        for line in open(sys.argv[2]):
            if not line.startswith('#'):
                es1.append(float(line.split()[e_loc1]))
        binned1 = np.histogram(es1, numbins, range=(0,1.0))

        bar(binned[1][:-1], binned[0], width=1.1/numbins, color='purple',
            edgecolor='none', label='Galaxies')
        #bar(binned1[1][:-1], binned1[0], width=1.1/50, color='red',
        #    edgecolor='none', label='Stars')

        popt, pcov = curve_fit(lambda e, sig: p(e, sig, es=es),
            binned[1][:-1], binned[0], p0=0.3)
        #print popt, pcov
        print popt[0], np.sqrt(np.diag(pcov))

        x = np.linspace(0, 1.0, 200)
        #plot(x, [a*(len(es)/50) for a in p(x, *popt)], 'k-',
        #     label='Rayleigh distrubution')
        plot(x, p(x, *popt, es=es), 'k-',
             label='Rayleigh distrubution')
        xticks(np.linspace(0, binned[1][-1], 6))
        xlabel('Ellipticity')
        ylabel('Number of Galaxies')
        legend()
        title('Ellipticity Distribution for Galaxies in A1348')
        draw()
        pause(1)
        raw_input()
        close()

    else:
        print "usage: distribution catfile starfile"
