#!/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.table import Table, Column
from astropy import units as u

import matplotlib.pyplot as plt
import math, sys, os
import pickle

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 statsmodels.robust.scale import huber
from astropy.stats import sigma_clip
from astropy.stats import sigma_clipped_stats

from scipy.stats import norm
from scipy.stats import describe
from scipy.stats import sem
import subprocess

from numpy import absolute
from numpy import arange

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

inkey="ALL"



plotme=False

"""
ebands0={#'E02':[0.0,0.0],
         'E03':[0.0,0.0],
         'E04':[0.0,0.0],
         'E05':[0.0,0.0],
         'E06':[0.0,0.0],
         'E07':[0.0,0.0],
         'E08':[0.0,0.0],
         'E09':[0.0,0.0],
         'E10':[0.0,0.0],
         'E11':[0.0,0.0],
         'E12':[0.0,0.0],}
"""

ebands0={'B01':[0.0,0.0],
         'B02':[0.0,0.0],
         'B03':[0.0,0.0],
         'B04':[0.0,0.0],
         'B05':[0.0,0.0],
         'B06':[0.0,0.0],
         'B07':[0.0,0.0],
         'B08':[0.0,0.0],
         'B09':[0.0,0.0],
         'B10':[0.0,0.0],
         'B11':[0.0,0.0],
         'B12':[0.0,0.0],
         'B13':[0.0,0.0],
         'B14':[0.0,0.0],
         'B15':[0.0,0.0],
         'B16':[0.0,0.0],
         'B17':[0.0,0.0],
         'B18':[0.0,0.0],
         'B19':[0.0,0.0],
         'B20':[0.0,0.0],
         'B21':[0.0,0.0],
         }

skew0={'B01':[0.0,0.0],
         'B02':[0.0,0.0],
         'B03':[0.0,0.0],
         'B04':[0.0,0.0],
         'B05':[0.0,0.0],
         'B06':[0.0,0.0],
         'B07':[0.0,0.0],
         'B08':[0.0,0.0],
         'B09':[0.0,0.0],
         'B10':[0.0,0.0],
         'B11':[0.0,0.0],
         'B12':[0.0,0.0],
         'B13':[0.0,0.0],
         'B14':[0.0,0.0],
         'B15':[0.0,0.0],
         'B16':[0.0,0.0],
         'B17':[0.0,0.0],
         'B18':[0.0,0.0],
         'B19':[0.0,0.0],
         'B20':[0.0,0.0],
         'B21':[0.0,0.0],
         }

mcrab=u.def_unit('mCrab')
ctss=u.def_unit('cts/s')
u.add_enabled_units([mcrab,ctss])

#skey='Geminga'
if len(sys.argv) > 1:
    skeys = [sys.argv[1]]
else:
    skeys = list(skyreg.keys())
    # don't plot, if all regions are taken
    plotme=False


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

fitsdir = "{}fits/".format(specdir)
if not os.path.exists(fitsdir):
    os.makedirs(fitsdir)

with open(ignored_rev_file, 'rb') as fp:
    ignored_rev = pickle.load(fp)
    print(ignored_rev)
    print("{} orbits ignored".format(len(ignored_rev)))

ign=ignored_rev.tolist()



for skey in skeys:
    if not skey in skyreg.keys():
        print("{} not found in {}".format(skey,list(skyreg.keys())))
        sys.exit()

    for enkey in ebands0.keys():
        
        fn="detcnts.{}.{}.resid.fits".format(enkey,inkey)
        #d1 = fits.getdata(proddir+fn)
        #d2=np.array(d1)
        #df=pd.DataFrame(d2.view(d2.dtype.newbyteorder()))
        #df=pd.DataFrame(np.array(d).byteswap().newbyteorder())

        #with fits.open(proddir+fn) as data:
        #    df = pd.DataFrame(data[1].data)

        dat = Table.read(proddir+fn)
        df = dat.to_pandas()

        #print(df.columns)
        #sys.exit()
        #df = df.query("REV == @ign")

        
        query = "LON > {} & LON < {} & LAT > {} & LAT < {} & REV != @ign".format(
            skyreg[skey]['lon'] - skyreg[skey]['wlon']/2,
            skyreg[skey]['lon'] + skyreg[skey]['wlon']/2,
            skyreg[skey]['lat'] - skyreg[skey]['wlat']/2,
            skyreg[skey]['lat'] + skyreg[skey]['wlat']/2)

    
        df = df.query(query)
        print("{}, {}: {} N={}".format(skey, enkey, query, df.shape[0]))


        t = Table.from_pandas(df)
        t.write("{}fits/{}.{}.fits".format(specdir,skey,enkey),overwrite=True)
        
        texp = np.array(df['TEXP'])
        with open("{}fits/{}.{}.livetime".format(specdir,skey,enkey), 'w') as fp:
            fp.write("{} {} ScWs: {} Texp: {:.2f} Ms\n".format(skey,enkey,df.shape[0],np.sum(texp)/1e6))
        
        if not (df.shape[0]>0):
            continue
        
        sg_mean,sg_sem,skew_val,skew_err = get_spec(df, sigma=3,  grxe_err_cut=grxe_err_cut, skey=skey, enkey=enkey, plotme=False, bootstrap=False, gaussfit=True)
        ebands0[enkey]=[sg_mean,sg_sem]   
        skew0[enkey]=[skew_val,skew_err]   

        

    fspec="{}{}.spec".format(specdir,skey)
    with open(fspec, 'w') as fp:
        for enkey in ebands0.keys():
            flux=ebands0[enkey][0]
            err=ebands0[enkey][1]
            print("[DATA] {}: {} {:.6f} {:.6f}".format(skey,enkey,flux,err))
            fp.write("0 {} {:.6f} {:.6f} 0.0\n".format(bands[enkey],flux,err))    
    subprocess.run(["perl", "do_pha_igr_v3_mCrab.pl", fspec])

    fspec="{}{}.skew".format(specdir,skey)
    with open(fspec, 'w') as fp:
        fp.write("read serr 4\n")
        count=1
        for enkey in skew0.keys():
            fp.write("{} {} {:.6f} {:.6f}\n".format(count,bands[enkey],skew0[enkey][0],skew0[enkey][1]))    
            count+=1
    
    try:
        for remfile in ["cols","cols1","cols2","header",]:
            os.remove(remfile)
    except OSError:
        pass