#!/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"


sigma=3
plotme=True

ebands0={
    'E01':[0.0,0.0], # 25-60 keV
    'E14':[0.0,0.0], # 60-80 keV
    'E13':[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)

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()

    # generate array for bootstrap
    ebands_sim={}
    for enkey in ebands0.keys():
        ebands_sim[enkey] = []

    for enkey in ebands0.keys():
        #bkg_fn="detcnts.{}.BKG.resid.fits".format(enkey,inkey)
        #syserr, bkg_sem = get_syserror(proddir+bkg_fn)
        
        fn="detcnts.{}.{}.resid.fits".format(enkey,inkey)
        print("Reading {}".format(proddir+fn))
        dat = Table.read(proddir+fn, unit_parse_strict='silent')
        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/{}.{}.fits".format(fluxdir,skey,enkey),overwrite=True)
        
        texp = np.array(df['TEXP'])
        with open("{}fits/{}.{}.livetime".format(fluxdir,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 = get_spec(df, sigma=sigma,  grxe_err_cut=grxe_err_cut, skey=skey, enkey=enkey, plotme=plotme)
        ebands0[enkey]=[sg_mean,sg_sem]   

        nsel = int(df.shape[0]/simfrac)
        for n in range(nsim):
            df0=df.sample(nsel)
            sg_mean,sg_sem = get_spec(df0, grxe_err_cut=grxe_err_cut, skey=skey, enkey=enkey)
            ebands_sim[enkey].append(sg_mean)   
            

        
    ###
    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))    
    

    ###
    fspec="{}{}.sim.dat".format(fluxdir,skey)
    with open(fspec, 'w') as fp:
        for enkey in ebands_sim.keys():
            data=ebands_sim[enkey]

            (mu, sg) = norm.fit(data)
            fp.write("{} {} {} {:.6f} {:.6f}\n".format(skey,enkey,bands[enkey],mu,sg))    
            print("[BOOT] {}: {} {:.6f} {:.6f}".format(skey,enkey,mu,sg))

            if(plotme):
                n, bins, patches = plt.hist(data, 60, density=True, facecolor='green', alpha=0.75)
                # add a 'best fit' line
                y = norm.pdf(bins, mu, sg)
                l = plt.plot(bins, y, 'r--', linewidth=2)
                #plot
                plt.axvline(mu, color="black")
                plt.axvline(ebands0[enkey][0], color="black", linestyle="dashed")
                #plt.axvline(mu+sg_sem, color="black", linestyle="dashed")
                #plt.axvline(mu-sg_sem, color="black", linestyle="dashed")
                plt.axvline(mu+sg, color="blue", linestyle="dashed")
                plt.axvline(mu-sg, color="blue", linestyle="dashed")

                plt.xlabel('Flux, mCrab')
                plt.ylabel('Probability')
                plt.title("[BOOT] {}: {:.2f} {:.2f}".format(enkey, mu, sg))
                plt.grid(True)
                plt.show()