#!/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={
    'A01':[0.0,0.0], # 25-60 keV
    'A02':[0.0,0.0], # 60-80 keV
    'A03':[0.0,0.0], # 80-200 keV
}


if len(sys.argv) > 1:
    skeys = [sys.argv[1]]
else:
    skeys = list(skyreg.keys())

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

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

with open(ignored_rev_file, 'rb') as fp:
    ignored_rev = pickle.load(fp)
    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)
        print("Reading {}".format(proddir+fn))
        dat = Table.read(proddir+fn) 
        df = dat.to_pandas()

        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".format(fitsdir,skey,enkey),overwrite=True)
        
        texp = np.array(df['TEXP'])
        with open("{}{}.{}.livetime".format(fitsdir,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=4,  grxe_err_cut=grxe_err_cut, skey=skey, enkey=enkey, plotme=True, bootstrap=False, gaussfit=True)
        ebands0[enkey]=[sg_mean,sg_sem]


    fspec="{}{}.dat".format(fluxdir,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("{} {} {} {:.6f} {:.6f}\n".format(skey,enkey,bands[enkey],flux,err))