#!/usr/bin/python3

import sys
import matplotlib.pyplot as plt
import numpy as np
from astropy.timeseries import LombScargle
from datetime import datetime
import scipy as sp
from scipy.optimize import curve_fit 
import astropy.time
from scipy.interpolate import splrep, splev
import math

plt.rcParams.update({'font.size': 16})

def fit_function(x, a, b, c, d, omega):
    return a * x + b + c * np.sin(omega*x) +  d* np.cos(omega*x)

def find_object_from_noradid(noradid):
    with open("/home/astro/NOVAS-BERTHELOT/resources/ALL.TLE") as tlef:
        linestle = tlef.readlines()
        for l in linestle:
            if (l.split()[0] == '2' and int(l.split()[1])==int(noradid)):
                index = linestle.index(l)
                cosparid = linestle[index-1].split()[2]
                name     = linestle[index-2]
                return cosparid, name


print ("APSIS Photometry")
print ("opening ", sys.argv[1])
with open(sys.argv[1]) as f:
    lines = f.readlines()

count=0
time=[]
per=[]
err=[]
weight=[]

noradid = sys.argv[1].split(".")[0]
cosparid, name = find_object_from_noradid(noradid)
for l in lines:
        print (l)
        if (count==0):
            t0 = float(l.split()[0])
        time.append(float(l.split()[0])-t0)
        per.append(float(l.split()[1]))
        err.append(float(l.split()[2]))
        weight.append(1./float(l.split()[2]))
        count+=1
print ("Found", count, "period values for", cosparid, name)
# Data for plott
spl = splrep(time, per, weight, s=1)#CubicSpline(time, per)

fig, ax = plt.subplots(figsize=(14,7), layout='constrained')

tt= np.arange(time[0], time[-1], 1)

initial_guess = [-0.001, np.mean(per), 1., 1., 0.0172]  # Better initial guesses
#popt, pcov = curve_fit(fit_function, time, per, sigma=err, method='lm', absolute_sigma=False, p0=initial_guess)

#perr = np.sqrt(np.diag(pcov))
#print ("Period ", 2*math.pi/popt[4], "days", "+-", 2*math.pi/(popt[4]*popt[4])*perr[4], "Slope", popt[0], "err", perr[0])
#time_interval = np.arange(0, time[-1]+720)
#fit_evolution = fit_function(time_interval, *popt)
#for jd, fitted_value, per in  zip(julian_date, fitted_values, period):
#    rms+=(per-fitted_value)*(per-fitted_value)
#rms=np.sqrt(rms);
#print (rms)


ax.scatter(time, per, c='red')
plt.errorbar(time, per, yerr=err, fmt='o', c='green', ecolor='red', label='Periods (with 1-sigma errors)')
#plt.plot(time, per, color="green")
plt.plot(tt, splev(tt, spl), c='brown', label='Interpolated spline')
#plt.plot(time_interval, fit_evolution, c="black", label="Periodic with slope")
ax.set_xlabel('Time (days)')
ax.set_ylabel('Period (s)')
#plt.ylim(min(per)-min(per)/100., max(per)+max(per)/100.)
#plt.ylim(min( min(fit_evolution), min, max(fit_evolution))
ax.grid()
ax.legend()
plt.title(cosparid+" "+name)
output=sys.argv[1][:-4]+"_"+cosparid+"_EVO.png"
fig.savefig(output)
plt.show()

#print ("\n")
#print ("t = jd-",t0) 
#print ("Period(t)=", popt[0], "*t+", popt[1], "+", popt[2], "*sin(", popt[4], "*t)+",popt[3],"*sin(",popt[4], "*t)" )