#!/usr/bin/env python

__author__      = "Roman Krivonos"
__copyright__   = "Space Research Institute (IKI)"

import numpy as np
import pandas as pd
from astropy.io import fits
from astropy import wcs
import matplotlib.pyplot as plt
import math, sys, os
import pickle
from numpy.polynomial import Polynomial
from astropy.table import Table, Column

from sklearn.linear_model import LinearRegression
from sklearn.linear_model import HuberRegressor
from sklearn.linear_model import RANSACRegressor
from sklearn.linear_model import TheilSenRegressor
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import RepeatedKFold


from astropy.coordinates import SkyCoord # High-level coordinates
from astropy.coordinates import ICRS, Galactic, FK4, FK5 # Low-level frames
from astropy.coordinates import Angle, Latitude, Longitude # Angles
import astropy.units as u

#from statsmodels.robust.scale import huber
from astropy.stats import sigma_clip
from numpy import absolute
from numpy import arange

from ridge.utils import *
from ridge.config import *

if not os.path.exists(proddir):
    os.makedirs(proddir)

#enkey = sys.argv[1]
enkey='A01'
inkey="ALL"

fn="detcnts.{}.{}.resid.fits".format(enkey,inkey)
dat = Table.read(proddir+fn)
df = dat.to_pandas()

with open(proddir+"detcnts.{}.ignored_scw.pkl".format(enkey), 'rb') as fp:
    ignored_scw = pickle.load(fp)

sco_crd = SkyCoord(sco_ra, sco_dec, frame=FK5(), unit="deg")

plotme=False

npix = 50
sw   = 30.0 # deg
pix  = sw/npix # pixel size in degrees


crabmodel={}
rota_arr=[]
a0=[]
b0=[]
a_full0=[]
b_full0=[]
b_est0=[]
err0=[]
rev0=[]

totx=[]
toty=[]

for i,rec in df.iterrows():
    obsid = rec['OBSID']#.decode("utf-8")

    if (obsid in ignored_scw):
        print("Skip ScW",obsid)
        continue

# accumulate full data set
for rev in range(revmin,revmax):


    df0 = df.query('SRC > 0.0 & REV == {} & PHASE > {} & PHASE < {} & SCO_SEP < {}'.format(rev,phmin,phmax,crab_sep_max))
    nobs=len(df0)
    if not (nobs> crab_nmax):
        continue

    print(rev,nobs)
    
    for n in df0['SCO_SEP'].values:
        totx.append(n)
        
    for n in df0['SRC'].values:
        toty.append(n)
        
        
x = np.array(totx)
y = np.array(toty)
x = x.reshape((-1, 1))
model = LinearRegression()
#model = TheilSenRegressor()
results = evaluate_model(x, y, model)
a_full,b_full,err_full = plot_best_fit(x, y, model)
if(plotme):
    plot_ab(x, y, a_full, b_full, err_full, title="REGRESSION")


    
# go over orbits
poly_x=[]
poly_y=[]

ntotal=0
nrev=0
for rev in range(revmin,revmax):
    df0 = df.query('SRC > 0.0 & REV == {} & PHASE > {} & PHASE < {} & SCO_SEP < {}'.format(rev,phmin,phmax,crab_sep_max))
    nobs=len(df0)
    if not (nobs> crab_nmax):
        continue
    #cen_ra = np.array(df0['RA'].values)
    #cen_dec = np.array(df0['DEC'].values)
    print("*** Orbit ",rev)
    x = np.array(df0['SCO_SEP'].values)
    y = np.array(df0['SRC'].values)

    #rota_deg = np.array(df0['ROTA'].values)
    #rota = np.array(df0['ROTA'].values) * np.pi / 180. # in radians
    #detx = np.cos(rota)*x
    #dety = np.sin(rota)*x

    x = x.reshape((-1, 1))
    model = LinearRegression()
    #model = TheilSenRegressor()
    results = evaluate_model(x, y, model)
    a,b,err = plot_best_fit(x, y, model)

    b_est = np.mean(y - a_full*x)
      
    a_full0.append(a_full)
    b_full0.append(b_full)
    b_est0.append(b_est)

    
    a0.append(a)
    b0.append(b)
    err0.append(err)
    rev0.append(rev)

    poly_x.append(rev)
    poly_y.append(b_est)


    crabmodel[rev]={'a':a_full, 'b':b_est, 'err':err}
    if(plotme):
        plot_ab(x, y, a_full, b_est, err, title="REGRESSION rev {}".format(rev))



    print("ax+b: a={:.2e}, b={:.2e}, std={:.2e}\n".format(a,b,err))
    ntotal=ntotal+nobs
    nrev=nrev+1
    
print("Orbits: {}, obs: {}".format(nrev,ntotal))

npoly=4
    
z = np.polyfit(poly_x, poly_y, npoly)

p = np.poly1d(z)
poly_z=[]
for t in poly_x:
    poly_z.append(p(t))
plt.scatter(poly_x, poly_y)
plt.plot(poly_x, poly_z, color='r')
plt.title("Crab detector count rate evolution")
plt.ylabel("Crab count rate cts/s/pix")
plt.xlabel("INTEGRAL orbit")
#plt.show()
plt.savefig(proddir+fn.replace(".fits",".sco_rate.png"))

indices = sorted(
    range(len(rev0)),
    key=lambda index: rev0[index]
)

coldefs = fits.ColDefs([
    fits.Column(name='REV', format='J', unit='', array=[rev0[index] for index in indices]),
    fits.Column(name='A', format='D', unit='', array=[a0[index] for index in indices]),
    fits.Column(name='B', format='D', unit='', array=[b0[index] for index in indices]),
    fits.Column(name='ERR', format='D', unit='', array=[err0[index] for index in indices]),
    fits.Column(name='A_FULL', format='D', unit='', array=[a_full0[index] for index in indices]),
    fits.Column(name='B_FULL', format='D', unit='', array=[b_full0[index] for index in indices]),
    fits.Column(name='B_EST', format='D', unit='', array=[b_est0[index] for index in indices]),
    fits.Column(name='B_POLY', format='D', unit='', array=[poly_z[index] for index in indices]),
])

fout = fn.replace(".fits",".scox1.fits")
hdu = fits.BinTableHDU.from_columns(coldefs, name='GRXE')
hdu.header['MISSION'] = ('INTEGRAL', '')
hdu.header['TELESCOP'] = ('IBIS', '')
hdu.header['INSTITUT'] = ('IKI', 'Affiliation')
hdu.header['AUTHOR'] = ('Roman Krivonos', 'Responsible person')
hdu.header['EMAIL'] = ('krivonos@cosmos.ru', 'E-mail')
#hdu.add_checksum()
print(hdu.columns)
hdu.writeto(proddir+fout, overwrite=True)