#!/usr/bin/env python

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

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

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 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.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


from numpy import absolute
from numpy import arange

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

enkey = sys.argv[1]
lon_nbin = int(sys.argv[2])
key="ALL"
fn='detcnts.{}.{}.resid.fits'.format(enkey,key)

print("Reading {}".format(proddir+fn))
dat = Table.read(proddir+fn)
df = dat.to_pandas()

#df = df.query('abs(LAT) < {} & abs(LON) < {}'.format(5,5))

n_bins = 80
minmax=[-300,800]
sigma=3
maxiters=10

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


glon, step = np.linspace(-lon_max, lon_max, num=lon_nbin, endpoint=False,retstep=True)

print("STEP",step)
#sys.exit()

#
# initiate 2d arrays
#
mean_map = np.array([0.0 for j in range(lon_nbin)])
sign_map = np.array([0.0 for j in range(lon_nbin)])
sem_map = np.array([0.0 for j in range(lon_nbin)])
cnt_map = np.array([0 for j in range(lon_nbin)])

mean_sim_map = np.array([0.0 for j in range(lon_nbin)])
error_sim_map = np.array([0.0 for j in range(lon_nbin)])

mean_sim={}
for i in range(lon_nbin):
    dkey="{:04d}".format(i)
    mean_sim[dkey] = []


obsid_map = {}
grxe_map = {}
grxe_err_map = {}

# redefine simfrac for low number of ScWs in pixel
#simfrac=2
#nsim=10

lon_arr=[]
lat_arr=[]

for i in range(lon_nbin):
    dkey="{:04d}".format(i)

    #if(i != 5):
    #    continue
    
    df0 = df.query('LON > {} & LON <= {} & LAT > {} & LAT < {} & REV != @ign'.format(glon[i],glon[i]+step,-lat_max,lat_max))

    if  (df0.shape[0] < nscw_min):
        continue

    sg_mean,sg_sem,skew_val,skew_err  = get_spec(df0, sigma=sigma,  grxe_err_cut=grxe_err_cut, enkey=enkey, plotme=True, nscw_min=nscw_min, gaussfit=True)
        
    #nsel = int(df0.shape[0]*simfrac/100)
    #print("lon {:.2f} ".format(glon[i]),"nsel=",nsel,df0.shape[0])

    #print('sg_sem',sg_sem)
    mean_map[i] = sg_mean
    sem_map[i]  = sg_sem
    cnt_map[i]  = df0.shape[0]

    for index, row in df0.iterrows():
        lon_arr.append(row['LON'])
        lat_arr.append(row['LAT'])



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

print("saving simulations")
for i in range(lon_nbin):
    dkey="{:04d}".format(i)
    data=mean_sim[dkey]
    #print("{} size {}".format(dkey,len(data)))
    #if(len(data)>10):
    (mu, sg) = norm.fit(data)
    mean_sim_map[i] = mu
    error_sim_map[i] = sg
    

coldefs = fits.ColDefs([
    fits.Column(name='LON1', format='D', array=glon), 
    fits.Column(name='LON2', format='D', array=glon+step), 
    fits.Column(name='GRXE_FLUX', format='D', array=mean_map), 
    fits.Column(name='GRXE_ERROR', format='D', array=sem_map), 
    fits.Column(name='GRXE_SIM_FLUX', format='D', array=mean_sim_map), 
    fits.Column(name='GRXE_SIM_ERROR', format='D', array=error_sim_map), 
    fits.Column(name='nScW', format='J', array=cnt_map), 
])

fout = fn.replace(".fits",".galprof.fits")
hdu = fits.BinTableHDU.from_columns(coldefs, name='GRXE')
hdu.header['MISSION'] = ('INTEGRAL', '')
#hdu.header['TELESCOP'] = (outkey, '')
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(profdir+fout, overwrite=True)

with fits.open(profdir+fout, mode='update') as hdus:
    hdus[1].add_checksum() 

    
# to check, which ScWs were used
coldefs = fits.ColDefs([
    fits.Column(name='LON', format='D', array=lon_arr), 
    fits.Column(name='LAT', format='D', array=lat_arr), 
])


fout = fn.replace(".fits",".gal.fits")
hdu = fits.BinTableHDU.from_columns(coldefs, name='GRXE')
hdu.header['MISSION'] = ('INTEGRAL', '')
#hdu.header['TELESCOP'] = (outkey, '')
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(profdir+fout, overwrite=True)

with fits.open(profdir+fout, mode='update') as hdus:
    hdus[1].add_checksum()