#!/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]
#key = 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()


hdulist = fits.open(modelsdir+modelrxte)
w = wcs.WCS(hdulist[0].header)
smap =hdulist[0].data

sx=int(hdulist[0].header['NAXIS1'])
sy=int(hdulist[0].header['NAXIS2'])

# fill AITOF map indexes
# -- Already done in 02_grxe_resid.py
"""
ds9x=[]
ds9y=[]
for i,row in df.iterrows():
    lon=row['LON']
    lat=row['LAT']
    world = SkyCoord(lon,lat, frame=Galactic, unit="deg")
    ra=world.fk5.ra.deg
    dec=world.fk5.dec.deg
    pixcrd = w.wcs_world2pix([(lon,lat)], 1)
    y=int(pixcrd[0][0])
    x=int(pixcrd[0][1])
    ds9x.append(x)
    ds9y.append(y)
    #print(x,y,smap[y-1,x-1])
df['DS9Y']=ds9x
df['DS9X']=ds9y
"""

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

# simulations
mean_sim_map = np.array([[0.0 for i in range(sx)] for j in range(sy)])
error_sim_map = np.array([[0.0 for i in range(sx)] for j in range(sy)])
sign_sim_map = np.array([[0.0 for i in range(sx)] for j in range(sy)])

mean_sim={}
for i in range(sx):
    for j in range(sy):
        dkey="{:04d}{:04d}".format(j,i)
        mean_sim[dkey] = []


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

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

for i in range(sx):
    for j in range(sy):
        dkey="{:04d}{:04d}".format(j,i)
        world = w.wcs_pix2world([(i+1,j+1)], 1)
        lon = world[0][0]
        lat = world[0][1]
        if(np.isnan(lon) or np.isnan(lat)):
            continue
        
        ds9i=i+1
        ds9j=j+1
        #df0 = df.query('DS9X == {} & DS9Y == {}'.format(ds9i,ds9j))
        df0 = df.query('DS9X == {} & DS9Y == {} & REV != @ign'.format(ds9i,ds9j))
        if  (df0.shape[0] < nscw_min):
            continue


        sg_mean,sg_sem,skew_val,skew_err = get_spec(df0, sigma=4,  grxe_err_cut=grxe_err_cut, enkey=enkey, plotme=False, bootstrap=False, gaussfit=True)

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

""" Filter by error map """
# Calculate the percentiles across the x and y dimension

perc = np.percentile(sem_map, sem_cut, axis=(0, 1), keepdims=False)
print("{} {}: {}% cut of SEM map: {:.2f} mCrab".format(key,enkey,sem_cut,perc))

idx=np.where(sem_map > perc)
print("index size {}".format(len(idx)))
mean_map[idx]=0.0
sem_map[idx]=0.0
cnt_map[idx]=0
sign_map[idx]=0.0

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

hdulist[0].data=mean_map
hdulist.writeto(mapsdir+fn.replace(".fits",".mean.fits"),overwrite=True)

hdulist[0].data=sem_map
hdulist.writeto(mapsdir+fn.replace(".fits",".error.fits"),overwrite=True)

hdulist[0].data=cnt_map
hdulist.writeto(mapsdir+fn.replace(".fits",".cnt.fits"),overwrite=True)

hdulist[0].data=sign_map
hdulist.writeto(mapsdir+fn.replace(".fits",".sign.fits"),overwrite=True)