#!/usr/bin/env python

# Exoplanet transit plotter
# v1.0 15 Jan 2016 Jacob Isbell 

import math
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from scipy import optimize
from sys import argv
from astropy.time import Time

def dflux(p,z):
    """Takes radius ratio and distance z as arguments.
    Returns the stellar intensity when planet is at distance z from center. Relative to 1."""
    lam = 0
    if p+1 < z:
        lam = 0
    elif z >abs(1-p) and z <= 1+p:
        k1 = math.acos((1-p**2 + z**2)/(2*z))
        k0 = math.acos((p**2 + z**2 -1)/(2*p*z))
        lam = (1/math.pi)*(k0*(p**2) + k1 - math.sqrt((4*z**2 - math.pow(1+(-p**2)+z**2,2))/4))
    elif z <= 1-p:
        lam = p**2
    elif z <= p-1:
        lam =1
    fraction = 1-lam
    return limb_darkening(p,z)*fraction

def limb_darkening(p,z):
    """Calculates and returns the intensity of the star at distance z from center
    which is affected by linear limb darkening"""
    i = 1
    if not (z > 1+p):
        q = math.atan(z/D_to_star)
        i = (1-w*(1-math.cos(q)))
    return i

def dmag(t,params):
    """Function to call from other scripts, houses the process of computing
    intensity over time"""
    p,b,v = params[:]
    z = np.sqrt((v*t)**2 + b**2) #distance from center of star to center of planet
    #print (p, z)
    f = dflux(p,z)
    mag=(2.5*math.log(f)) #converts fractional intensity to magnitude change
    
    return mag

def transit_center_index(smooth):
    """Returns duration of transit (time from "limb" to "limb" of lightcurve)
    as well as the time it takes to do from "limb" to bottom (called tau)"""
    obs_data = np.array(smooth)
    k = 5
    beg_limb = np.median(obs_data[:k])
    blimb_end = 0
    elimb_beg = 0
    for i in range(0,len(smooth)-(k-1),k):
        rng_med = np.median(obs_data[i:i+k])
        if abs(rng_med-beg_limb) >= 0.013: #mmag. Assumed minimum of Gemini's visibility
            blimb_end = i
            break
    l=len(smooth)
    end_limb = np.median([obs_data[l-1],obs_data[l-2],obs_data[l-3],obs_data[l-4],obs_data[l-5]])
    for i in range(l, 0, -k):
        rng_med = np.median(obs_data[i:(i-k-1):-1])
        if abs(rng_med-end_limb) >= 0.013: #mmag. Assumed minimum of Gemini's visibility
            elimb_beg = i
            break
    center = int((blimb_end + elimb_beg)/2)
    limb_avg = 0.5*(np.median(obs_data[:blimb_end]) + np.median(obs_data[l-1:elimb_beg:-1]))
    return center, limb_avg

def prepare_curve_data(filename, reduce=False):
    """ Reads data from file and processes it to use in other functions
    Takes a filename as argument. Reduce is optional, and takes the median of every 5
    data points to smooth out the data. It is usually unnecessary."""
    file = open(filename, 'r')
    line1 = file.readline() #skip first line
    all_data = file.readlines()
    file.close()
    time, obs_data, obs_err = [],[],[]
    for line in all_data:
        d = line.split()
        time.append(float(d[6]))
        obs_data.append(-1*float(d[9]))
        obs_err.append(float(d[10]))
    smooth = []; smooth_time = []; sigma = []
    """Takes a median of every five data points to smooth the data. Optional"""
    if reduce:
        for i in range(0,len(v0)-4,5):
            med = np.median([v0[i],v0[i+1],v0[i+2],v0[i+3],v0[i+4]])
            medTime = np.median([time[i],time[i+1],time[i+2],time[i+3],time[i+4]])
            medErr = np.median([err0[i],err0[i+1],err0[i+2],err0[i+3],err0[i+4]])
            smooth.append(med)
            smooth_time.append(medTime)
            sigma.append(medErr)
    else:
        smooth = obs_data; smooth_time = time; sigma = obs_err
    """Subtracts any overall slope the data might have by comparing
    the median of the first few smooth points to the last few"""
    med_beg = np.median([smooth[0],smooth[1], smooth[2]])
    mx = len(smooth)
    med_end = np.median([smooth[mx-1], smooth[mx-2], smooth[mx-3]])
    slope = (med_end - med_beg)/mx
    no_slope_smooth = [smooth[x] - slope*x for x in range(len(smooth))]
    """Estimate transit center using minimum value of smooth
    AND KEEP MHJD FOR LATER CONVERSION BACK TO UT"""
    MHJD_list = [x for x in smooth_time]
    c, limb_avg = transit_center_index(smooth)
    '''Move the light curve to y=0 because absolute height doesn't matter, only relative change.'''
    processed_obs_data = [no_slope_smooth[i] - limb_avg for i in range(len(no_slope_smooth))]
    '''Center of transit time'''
    midTime = (smooth_time[c]-smooth_time[0])*24*3600
    """Convert smooth_time to a difference from center"""
    dTime = [((smooth_time[x]-smooth_time[0])*24*3600)-midTime for x in range(len(smooth_time))]
    
    return processed_obs_data, dTime, sigma, MHJD_list
def chi_squared(params, *args):
    """Calculates the distance from the actual data using weighted least squares.
    This is the function that is minimized."""
    time, data, sigma = args
    error = 0
    for n in range(len(time)):
        error += ((dmag(time[n], params)-data[n])/sigma[n])**2
    error = error/(len(time)-len(params))
    return error
def uncertainty(residuals):
    variance = 0
    for x in residuals:
        variance = variance + x**2
    unc = math.sqrt(variance/(len(residuals)-3))
    return unc
def residuals(model_data, data):
    """Returns difference from each model data point to the corresponding observation data
    as an array"""
    a_md = np.asarray(model_data)
    a_d = np.asarray(data)
    resid = np.subtract(a_md, a_d)
    return resid
def MHJD_to_UT_list(time):
    """Takes a list of MHJD time points and converts them to UT time for
    representation on the plot. Returns list of UT times."""
    ut = [x for x in time]
    for i in range(len(ut)):
        ut[i] = ut[i] + 0.5 
        ut[i] = ut[i] - int(ut[i])
        ut[i] = ut[i] * 24.
    return ut
def calculate_date(JD):
    """Takes Julian Date (JD) of observation and converts it to calendar date.
    Returns calendar date and start time of observation"""
    t = Time(JD, format='jd')
    date = t.iso
    return date
def plot_values(title, subtitle, time, transit_values, transit_uncertainties, model_values, residual_values):
    """Function to plot observation data, model data, and residuals on the same plot with titles and legend.
    Takes main title, subtitle as strings. transit_values and transit_uncertainties are lists of the data points
    given by the photout file. model_values is a list of simulated data points using the calculated parameters.
    residual values is the list of differences between model and observational data."""
    #create plot
    fig = plt.figure(figsize=(15,10))
    #Obs data has known uncertainties, so plotting with errorbars. Alpha tweaked so model is visible.
    plt.errorbar(time, transit_values,yerr=transit_uncertainties, fmt='.', label='Observed', alpha=0.7)
    #Model and residuals plotted as distinct points that correspond to obs data
    plt.plot(time, model_values, '.', label = 'Model')
    plt.plot(time, residual_values, '.', label='Residuals')
    #Labeling axes
    plt.xlabel('UT Time (hours)')
    plt.ylabel('Differential Magnitude (mag)')
    #Adding titles, main title first
    plt.suptitle(title)
    plt.title(subtitle)
    #Add grid for clarity
    plt.grid()
    #Add plot legend to bottom right corner of graph
    plt.legend(bbox_to_anchor=(1,0.2))
    name = title.split()[0]+'model.png'
    plt.savefig(name, bbinches='tight')
    print(('Completed. Graph is called %s' % name))

# ============= MAIN ==============

#Define constants
R_SUN = 695500.0
D_to_star = 293*(9.46e15)/R_SUN #in units of solar radius
R_JUPITER = 71492.0
R_WASP10 = .7 * R_SUN
w = .5 #limb darkening coefficient

#Get file to be processed, either from initial call or from user input
fname = ''
if len(argv) == 1:
    fname = input("Please enter a photometric file (.photout) to analyze: ")
else:
    fname = argv[1]
    
#Read photout file, and return observation data for use in minimization
transit_values, time, transit_uncertainties, MHJD_list = prepare_curve_data(fname)

#Initial Parameter Values -- GUESSES
p = 0.15
b = 0.5
v = 0.001
x0 = (p, b, v)

# Minimizes parameter values and returns them along with least squares error. Uses a scipy minimzation function.
method = 'Nelder-Mead'
res = optimize.minimize(chi_squared, x0, args=(time, transit_values, transit_uncertainties), method=method)
params = [res.x[0], res.x[1], res.x[2]]
error = res.fun
print(('Weighted Least Squares Value: %f' % error))
print(('Radius ratio: %f of parent star \t Impact Parameter: %f of parent star\t Fractional Velocity: %f of parent star/second' % (params[0], params[1], params[2])))

# Uses calculated parameter values to calculate model values at each data point.
model_values = [dmag(x, params) for x in time]
# Subtracts model values from transit values to calculate the model's residuals
model_residuals = residuals(model_values, transit_values)
model_uncertainty = uncertainty(model_residuals)
model_residuals = model_residuals+0.04 #for clarity in plotting -- shifts residuals away from 0 so they're not in the way

# Converts MHJD to a list of UT times for plotting
ut_time_list = MHJD_to_UT_list(MHJD_list)
# From MHJD given in data file, calculate first the JD of the observation then the calendar date
obs_JD = MHJD_list[0] + 2449000 #adding 2449000 to change MJD to JD
obs_date = calculate_date(obs_JD)

#Plotting the model, observational data, and residuals.
TARGET = fname.split('.')[0]
title = TARGET + ' (Observation Start: ' + obs_date + ')'
subtitle = r'$R_p/R_*$ = %.2f, Impact Parameter = %.2f, P = %.2f days, $\sigma$ = %.2f mmag' % (params[0],abs(params[1]), params[2], model_uncertainty*1.e3)
plot_values(title, subtitle, ut_time_list, transit_values, transit_uncertainties, model_values, model_residuals)