#!/net/python/bin/python

import pylab as p
import numpy as num
from Data1D import *
from GetParams import *
from LightCurve import *
from Header import i_max

from mathfuncs import splrep,splev,smooth

def rawFit(filename, data=True):
    D = Data1D(filename)
    if data:
        D.pointgraph()
    D.resample().graph()

def fourierFit(filename,data=True):
    D = Data1D(filename)
    if data:
        D.pointgraph()
    D.fourier_fit().graph()

def smoothfit(filename,data=True,s=1.0):
    D = Data1D(filename)
    if data:
        D.pointgraph()
    D.smooth(s).graph()

def compareS(filename,data=True):
    D = Data1D(filename)

    for s in [0.1,0.5,2.0]:
        D.smooth(s).graph()

    if data:
        D.pointgraph()

    p.legend(['s=0.1','s=0.5','s=2.0'])
    p.xlabel('time (MJD)')
    p.ylabel('flux (normalized CCD counts)')
    p.title('Effect of s in fitting algorithm')

def Find_t0(cid,filters='ugriz',filterset = 'SDSS',cosmo='FLCDM',graph=False,fourier=False,raw=False):
    """
    fits the light curves found at the given file location, and determines
    an estimate of the B-band maximum
    """

    path = GetPath(filterset,filters) + cid
    fitfile = GetFitFile(filterset,filters,cosmo)
    
    print path
    print fitfile
    
    B_CENTRAL = 4400 #Johnson B central wavelength in angstroms

    Params = GetParams(path,fitfile)

    LC = ImportLightCurve(path,filterset)

    extincted_filters = LC.filterset.extincted(Params['MWEBV'])

    filter_means = []
    filter_names = extincted_filters.keys[:]
    
    for f in filter_names:
        blueshifted_mean = extincted_filters[f].mean()/(1.+Params['Redshift'])
        filter_means.append( blueshifted_mean )
    
    print filter_names
    print filter_means
    
    upper_index = len(filter_means)

    for i in range(len(filter_means)):
        #if filters are not in wavelength order, the algorithm will not work
        if i != 0:
            assert filter_means[i] > filter_means[i-1]
        if B_CENTRAL < filter_means[i]:
            upper_index = i
            break
    
    print 'MWEBV = %.3f' % Params['MWEBV']
    print 'z = %.3f' % Params['Redshift']
    print 'Previously computed t0 = %.3f +/- %.3f'\
          % (Params['DayMax'],Params['dDayMax'])

    if upper_index > 0: #this is good: we can interpolate
        lower_filter = filter_names[upper_index-1]
        upper_filter = filter_names[upper_index]

        mean_L = filter_means[upper_index-1]
        mean_U = filter_means[upper_index]

        D_L = LC[lower_filter]
        D_U = LC[upper_filter]

        #here is where the fit happens.
        # one of three fits can be specified
        if fourier:
            print "fourier"
            S_L = D_L.fourier_fit()
            S_U = D_U.fourier_fit()
        elif raw:
            print "raw"
            S_L = D_L
            S_U = D_U
        else:
            print "my method"
            S_L = D_L.smooth()
            S_U = D_U.smooth()

        if graph:
            D_L.pointgraph('b')
            D_U.pointgraph('r')
            S_L.graph('-b')
            S_U.graph('-r')
        
        tmax_L = S_L.maximum()
        tmax_U = S_U.maximum()

        t_max = tmax_L + (B_CENTRAL - mean_L)*(tmax_U-tmax_L)/(mean_U-mean_L)

        dt_max = 0.25 * num.abs(tmax_U-tmax_L)

        print "new guess: %.3f +/- %.3f" % (t_max,dt_max)

        return (cid,t_max,dt_max,Params['DayMax'],Params['dDayMax'])

    else:
        print "Warning: out of range!"

        return (cid,0,0,Params['DayMax'],Params['dDayMax'])

    

def compare_methods():
    outfile = open('out.txt','w')
    outfile.write('#cid\traw t0\t\tdt\tfourier t0\t\tdt0\tMy t0\t\tdt0\tSALT t0\t\tdt0\n')
    for cid in Type120:
        out = Find_t0(cid,raw=True)
        out2 = Find_t0(cid,fourier=True)
        out3 = Find_t0(cid)
        print out[:3]+out2[1:3]+out3[1:]
        outfile.write('%s\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f\n'\
                      % (out[:3]+out2[1:3]+out3[1:]) )

    outfile.close()


if __name__ == '__main__':

    #these are all the confirmed Ia supernovae in SDSS2 year 1
    Type120 = ['762', '1032', '1112', '1241', '1253', '1371', '1580', '1688', '2017', '2031', '2165', '2246', '2308', '2330', '2372', '2422', '2440', '2533', '2561', '2635', '2689', '2789', '2916', '2992', '3080', '3087', '3199', '3256', '3317', '3377', '3451', '3452', '3592', '3901', '4000', '4046', '4241', '4524', '4577', '4679', '5103', '5183', '5391', '5395', '5533', '5549', '5550', '5635', '5717', '5736', '5751', '5844', '5944', '5957', '5966', '5994', '6057', '6100', '6108', '6127', '6192', '6196', '6249', '6295', '6406', '6558', '6649', '6699', '6780', '6924', '6933', '7143', '7147', '7243', '7473', '7475', '7512', '7779', '7847', '7876', '8030', '8046', '8151', '8598', '9045', '9207', '2943', '6304', '6315', '6422', '6936']
    
    #  Crazy amounts of data
    #filename='/astro/users/vanderplas/research/SDSSdata/4064/lc2fit_i.dat'

    # used for plots:
    filename='/astro/users/vanderplas/research/SDSSdata/6936/lc2fit_g.dat'

    #  Used for plots:
    #filename = '/astro/users/vanderplas/research/SDSSdata/4064/lc2fit_g.dat'

    #extreme outlier
    #filename = '/astro/users/vanderplas/research/SDSSdata/6196/lc2fit_r.dat'

    #compareS(filename,data=True)
    #rawfit(filename)
    #fourierFit(filename)
    #smoothfit(filename,s=0.1)
    #compare_methods()
    
    #p.show()