import glob
import os
import sys
import numpy as np
from astropy.io import fits
import pickle
from scipy.stats import sem
from astropy.wcs import WCS

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 matplotlib.pyplot as plt
from astropy.table import QTable, Table, Column

from mpdaf.obj import WCS as mWCS


#from sherpa.astro.ui import *

from coma.config import *

from scipy.optimize import minimize
from random import uniform

def sky2pix(wcs,dsrc):
    # converts ra,dec --> x,y   

    # checks
    #sk = wcs.pix2sky([0.0,0.0])
    #print(sk)
    #px = wcs.sky2pix([-6.86525105,36.54071803],nearest=False)
    #print(px)
    #sk = wcs.pix2sky([px[0][1],px[0][0]])
    #print(sk)
    
    for idx, s in enumerate(dsrc):
        ra0=s['ra']
        dec0=s['dec']
        
        px = wcs.sky2pix([dec0,ra0],nearest=False)
        s['x']=px[0][1]
        s['y']=px[0][0]
        
        #sk = wcs.pix2sky([y0,x0])
        #ra=sk[0][1]
        #dec=sk[0][0]
        
        #px = wcs.sky2pix([dec,ra],nearest=False)
        #x=px[0][1]
        #y=px[0][0]
        #print("{:.4f} {:.4f} --> {} {} --> {:.4f} {:.4f} --> {} {}".format(ra0,dec0,x0,y0,ra,dec,x,y))

        
def pix2sky(wcs,dsrc):
    # converts x,y --> ra1,dec1  
    for idx, s in enumerate(dsrc):
        sk = wcs.pix2sky([s['y'],s['x']])
        s['ra1']=sk[0][1]
        s['dec1']=sk[0][0]

def plot_resid(dsrc,dref,crval1,crval2):
    # calculates RMS of original ra,dec
    center_crd = SkyCoord(crval1, crval2, frame=FK5(), unit="deg")
    offset=[]
    resid=[]
    error=[]
    for idx, s in enumerate(dsrc):
        src_crd = SkyCoord(dsrc[idx]['ra'], dsrc[idx]['dec'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        offset.append(src_crd.separation(center_crd).arcmin)
        resid.append(src_crd.separation(ref_crd).arcsec)
        error.append(s['radec_err'])

    indices = sorted(
        range(len(offset)),
        key=lambda index: offset[index]
    )

    return ([offset[index] for index in indices],
            [resid[index] for index in indices],
            [error[index] for index in indices])

def plot_resid1(dsrc,dref,crval1,crval2):
    # calculates RMS of original ra,dec
    center_crd = SkyCoord(crval1, crval2, frame=FK5(), unit="deg")
    offset=[]
    resid=[]
    error=[]
    for idx, s in enumerate(dsrc):
        src_crd = SkyCoord(dsrc[idx]['ra1'], dsrc[idx]['dec1'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        offset.append(src_crd.separation(center_crd).arcmin)
        resid.append(src_crd.separation(ref_crd).arcsec)
        error.append(s['radec_err'])

    indices = sorted(
        range(len(offset)),
        key=lambda index: offset[index]
    )

    return ([offset[index] for index in indices],
            [resid[index] for index in indices],
            [error[index] for index in indices])

      
def get_rms(dsrc,dref):
    # calculates RMS of original ra,dec
    resid=0.0
    n=0
    for idx, s in enumerate(dsrc):
        radec_err=s['radec_err']
        src_crd = SkyCoord(dsrc[idx]['ra'], dsrc[idx]['dec'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        sep=src_crd.separation(ref_crd).arcsec
        resid = resid+sep**2
        n=n+1
    return np.sqrt(resid/n)

def get_rms_w(dsrc,dref):
    # calculates RMS of original ra,dec
    resid=0.0
    n=0
    A=0.0
    B=0.0
    for idx, s in enumerate(dsrc):
        radec_err=s['radec_err']
        src_crd = SkyCoord(dsrc[idx]['ra'], dsrc[idx]['dec'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        sep=src_crd.separation(ref_crd).arcsec
        A=A+sep/radec_err**2
        B=B+1.0/radec_err**2
        n=n+1
    return A/B

#    A=Total(table_lc_sav.rate/(table_lc_sav.err)^2)
#       B=Total(1.0/(table_lc_sav.err)^2)
def get_rms1_w(dsrc,dref):
    # calculates RMS of original ra,dec
    resid=0.0
    n=0
    A=0.0
    B=0.0
    for idx, s in enumerate(dsrc):
        radec_err=s['radec_err']
        src_crd = SkyCoord(dsrc[idx]['ra1'], dsrc[idx]['dec1'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        sep=src_crd.separation(ref_crd).arcsec
        A=A+sep/radec_err**2
        B=B+1.0/radec_err**2
        n=n+1
    return A/B


def get_chi2_radec(dsrc,dref):
    # calculates RMS of original ra,dec
    resid=0.0
    n=0
    for idx, s in enumerate(dsrc):
        radec_err=s['radec_err']
        src_crd = SkyCoord(dsrc[idx]['ra'], dsrc[idx]['dec'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        sep=src_crd.separation(ref_crd).arcsec
        resid = resid+sep**2/radec_err**2
        n=n+1
    return resid

def get_rms1(dsrc,dref):
    # calculates RMS of modified ra1,dec1
    resid=0.0
    n=0
    for idx, s in enumerate(dsrc):
        radec_err=s['radec_err']
        src_crd = SkyCoord(dsrc[idx]['ra1'], dsrc[idx]['dec1'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        sep=src_crd.separation(ref_crd).arcsec
        resid = resid+sep*sep
        n=n+1
    return np.sqrt(resid/n)
    
def get_chi2_radec1(dsrc,dref):
    # calculates RMS of modified ra1,dec1
    resid=0.0
    n=0
    for idx, s in enumerate(dsrc):
        radec_err=s['radec_err']
        src_crd = SkyCoord(dsrc[idx]['ra1'], dsrc[idx]['dec1'], frame=FK5(), unit="deg")
        ref_crd = SkyCoord(dref[idx]['ra'], dref[idx]['dec'], frame=FK5(), unit="deg")
        sep=src_crd.separation(ref_crd).arcsec
        resid = resid+sep**2/radec_err**2
        n=n+1
    return resid


    
def do_transform_resid(params, wcs, dsrc, dref, crval1, crval2):
    
    angle_in, crval1_in, crval2_in = params

    # fills x,y
    sky2pix(wcs,dsrc)
    
    wcs1=wcs.copy()

    # do transform here
    #wcs1.rotate(angle_in/60)

    wcs1.set_crval1(crval1_in)
    wcs1.set_crval2(crval2_in)
    
    # fills ra1,dec1
    pix2sky(wcs1,dsrc)

    # caclulate new statistics
    resid1 = get_chi2_radec1(dsrc,dref)
    #print("--> {:+.4f} {:+.4f} {:.4f}".format((crval1-wcs.get_crval1())*3600,(crval2-wcs.get_crval2())*3600,resid1))
    return resid1


def do_astro_corr(src=None, ref=None, catalog=None):
    """ calculates astronomy correction """

    log=catalog.replace(".fits", ".shift.log")
    log0=catalog.replace(".fits", ".orig.qdp")
    log1=catalog.replace(".fits", ".shift.qdp")
    fout=catalog.replace(".fits", ".shift-map.fits")

    title=os.path.split(catalog)[1]
    key=title.replace("_SourceCatalog_en00.fits","")

    #png=catalog.replace(".fits", ".opti.png")
    #png=png.replace(key,"{}_{}".format(obslist[key],key))
    t = Table(names=('rota', 'shift_ra','shift_dec','chi2',), dtype=('f8','f8','f8','f8'))

    print(src)
    hdul = fits.open(src)
    src_map_data = hdul[0].data
    src_map_hdr = hdul[0].header
    src_tab_data = hdul[1].data
    src_tab_hdr = hdul[1].header
    hdul.close()

    dsrc=[]
    for rec in src_tab_data:
        d = dict(ra = rec['RA'], dec = rec['DEC'], ra_err = rec['RA_ERR'], dec_err = rec['DEC_ERR'], radec_err = rec['RADEC_ERR'])
        dsrc.append(d)
    
    print(ref)
    hdul = fits.open(ref)
    ref_map_data = hdul[0].data
    ref_map_hdr = hdul[0].header
    ref_tab_data = hdul[1].data
    ref_tab_hdr = hdul[1].header
    hdul.close()

    dref=[]
    for rec in ref_tab_data:
        d = dict(ra = rec['RA'], dec = rec['DEC'], ra_err = rec['RA_ERR'], dec_err = rec['DEC_ERR'])
        dref.append(d)

    rms = get_rms(dsrc,dref)
    chi2_radec = get_chi2_radec(dsrc,dref)
    wcs_src=mWCS(src_map_hdr)
    wcs_src.info()
    crval1=wcs_src.get_crval1()
    crval2=wcs_src.get_crval2()
    #print(crval1,crval2)

    offset0, sep0, err0 = plot_resid(dsrc,dref,crval1,crval2)
    with open(log0, 'w') as writer:
        
        writer.write("label Y Separation, arcsec\nlabel X Offset, arcmin\nr y -1 15\n")
        writer.write("ma 9 on\n")
        writer.write("read serr 2\n")
        for idx, s in enumerate(offset0):
            writer.write("{:.2f} {:.2f} {:.2f}\n".format(offset0[idx],sep0[idx],err0[idx]))


    
    wcs_ref=mWCS(ref_map_hdr)
    wcs_ref.info()

    chi2_radec1_opt=1000
    rota_opt=0.0
    crval1_opt=crval1
    crval2_opt=crval2
    
    rota_min=-30
    rota_max=30
    Rsim=10.0
    Nsim=20000
    crd = SkyCoord(crval1, crval2, frame=FK5(), unit="deg")
    i=0
    while i < Nsim:
        #rota_rand = uniform(rota_min, rota_max)
        rota_rand=0.0
        crval1_rand = uniform(crval1-Rsim/3600, crval1+Rsim/3600)
        crval2_rand = uniform(crval2-Rsim/3600, crval2+Rsim/3600)
        #crval1_rand = crval1
        #crval2_rand = crval2
        
        sim_crd = SkyCoord(crval1_rand, crval2_rand, frame=FK5(), unit="deg")
        if(sim_crd.separation(crd).arcsec > Rsim):
            continue
        chi2_radec1 = do_transform_resid([rota_rand, crval1_rand, crval2_rand], wcs_src, dsrc, dref, crval1, crval2)
        t.add_row((rota_rand,(crval1-crval1_rand)*3600,(crval2-crval2_rand)*3600,chi2_radec1))
        if(chi2_radec1 < chi2_radec1_opt):
            chi2_radec1_opt = chi2_radec1
            rota_opt   = rota_rand
            crval1_opt = crval1_rand
            crval2_opt = crval2_rand
        i += 1

    # set optimal parameters
    r = do_transform_resid([rota_opt, crval1_opt, crval2_opt], wcs_src, dsrc, dref, crval1, crval2)
    print(chi2_radec1_opt,'==',r)

    shift_ra=(crval1-crval1_opt)*3600
    shift_dec=(crval2-crval2_opt)*3600

    t.write(fout, format='fits', overwrite=True)

    rms1 = get_rms1(dsrc,dref)

    offset1, sep1, err1 = plot_resid1(dsrc,dref,crval1,crval2)
    with open(log1, 'w') as writer:
        writer.write("label Y Separation, arcsec\nlabel X Offset, arcmin\nr y -1 15\n")
        writer.write("ma 9 on\n")
        writer.write("read serr 2\n")
        for idx, s in enumerate(offset1):
            writer.write("{:.2f} {:.2f} {:.2f}\n".format(offset1[idx],sep1[idx],err1[idx]))

    with open(log, 'w') as writer:
        writer.write("{} {:+.2f} {:+.2f} {:+.2f} ({:.2f} - {:.2f}) = {:.2f} {:.2f} ({:.2f}) {:.2f} ({:.2f}) N={} -- {}\n".format(key,
                                                                                                                            rota_opt,
                                                                                                                            shift_ra,
                                                                                                                            shift_dec,
                                                                                                                            get_chi2_radec(dsrc,dref),
                                                                                                                            get_chi2_radec1(dsrc,dref),
                                                                                                                            chi2_radec-chi2_radec1_opt,
                                                                                                                            rms,get_rms_w(dsrc,dref),
                                                                                                                            rms1,get_rms1_w(dsrc,dref),
                                                                                                                            len(dsrc),key))
        


    """
    Nsim=20000 Shift only
    N01 +0.00 +4.76 +0.07 (171.22 -  28.06) = 143.15  6.35 (5.11) 3.92 (1.40) N=22 -- tm1_obs_1
    N02 +0.00 +2.77 -2.74 (354.15 - 225.27) = 128.88 10.10 (5.34) 9.38 (2.61) N=31 -- tm5_obs_1
    N03 +0.00 +3.88 -3.65 (542.25 - 303.25) = 239.00  8.74 (6.87) 6.73 (3.07) N=36 -- tm5_obs_2
    N04 +0.00 +0.28 +0.03  (43.16 -  43.04) =   0.12  6.04 (3.39) 6.07 (3.36) N=17 -- tm6_park_1
    N05 +0.00 +0.65 -0.78 (175.10 - 169.55) =   5.56  6.77 (4.47) 6.54 (4.51) N=38 -- tm6_scan_1
    N06 +0.00 -0.38 +0.27  (26.60 -  26.37) =   0.23  6.54 (4.27) 6.52 (4.23) N=11 -- tm6_park_2
    N07 +0.00 +0.04 -0.35 (128.82 - 127.98) =   0.84  5.76 (4.21) 5.75 (4.18) N=35 -- tm6_scan_2
    N08 +0.00 +0.58 -0.22  (10.07 -   9.79) =   0.29  3.94 (2.83) 3.89 (2.82) N= 9 -- tm6_park_3
    N09 +0.00 +1.35 -0.56  (95.06 -  87.04) =   8.02  5.73 (4.45) 5.60 (4.28) N=38 -- tm6_scan_3
    N10 +0.00 +0.56 +0.36  (15.71 -  15.41) =   0.30  4.22 (3.49) 4.22 (3.43) N=10 -- tm6_park_4
    N11 +0.00 +1.08 -0.62 (116.34 - 108.60) =   7.73  5.38 (4.41) 5.11 (4.26) N=37 -- tm6_scan_4
    N12 +0.00 -0.56 -1.56  (84.93 -  59.86) =  25.07  4.20 (2.45) 4.09 (1.88) N=30 -- tm6_obs_1
    N13 +0.00 +0.19 -1.03  (17.58 -  13.28) =   4.30  3.38 (1.90) 3.23 (1.31) N=13 -- tm6_obs_2_badpix
    N14 +0.00 +1.03 -1.99  (63.23 -  35.27) =  27.97  4.36 (2.73) 3.27 (2.05) N=23 -- tm6_obs_3_badpix
    N15 +0.00 +0.88 -1.66 (187.51 - 164.26) =  23.25  6.93 (3.29) 6.62 (2.68) N=33 -- tm7_obs_1
    N16 +0.00 +2.74 -0.52 (113.70 -  75.09) =  38.61  5.34 (4.23) 4.73 (3.42) N=26 -- tm7_obs_2

    Nsim=20000 Shift+Rotation
    N01 -1.50 +4.93 +0.00 (171.22 -  27.67) = 143.55  6.35 (5.11) 3.83 (1.38) N=22 -- tm1_obs_1
    N02 -2.15 +2.47 -3.05 (354.15 - 223.27) = 130.88 10.10 (5.34) 9.29 (2.71) N=31 -- tm5_obs_1
    N03 -5.08 +3.87 -4.04 (542.25 - 296.00) = 246.25  8.74 (6.87) 6.93 (2.82) N=36 -- tm5_obs_2
    N04 +5.90 -0.19 +0.50  (43.16 -  39.56) =   3.60  6.04 (3.39) 5.57 (3.57) N=17 -- tm6_park_1
    N05 -5.84 +1.37 -0.85 (175.10 - 111.02) =  64.08  6.77 (4.47) 5.53 (3.58) N=38 -- tm6_scan_1
    N06 -0.30 -0.27 +0.31  (26.60 -  26.36) =   0.24  6.54 (4.27) 6.50 (4.27) N=11 -- tm6_park_2
    N07 -5.49 +0.73 -0.98 (128.82 -  86.16) =  42.66  5.76 (4.21) 4.94 (3.62) N=35 -- tm6_scan_2
    N08 -0.28 +0.51 -0.29  (10.07 -   9.78) =   0.29  3.94 (2.83) 3.90 (2.82) N= 9 -- tm6_park_3
    N09 -4.56 +0.91 -0.40  (95.06 -  64.56) =  30.49  5.73 (4.45) 4.93 (3.56) N=38 -- tm6_scan_3
    N10 +2.01 +0.54 +0.63  (15.71 -  15.17) =   0.54  4.22 (3.49) 4.03 (3.39) N=10 -- tm6_park_4
    N11 -3.83 +1.01 -1.29 (116.34 -  82.25) =  34.08  5.38 (4.41) 4.75 (3.72) N=37 -- tm6_scan_4
    N12 -1.35 -0.23 -1.27  (84.93 -  60.30) =  24.63  4.20 (2.45) 4.20 (1.82) N=30 -- tm6_obs_1
    N13 -1.52 +0.17 -1.38  (17.58 -  13.30) =   4.29  3.38 (1.90) 3.21 (1.28) N=13 -- tm6_obs_2_badpix
    N14 -1.24 +0.63 -2.18  (63.23 -  34.42) =  28.81  4.36 (2.73) 3.28 (1.92) N=23 -- tm6_obs_3_badpix
    N15 -5.23 +0.82 -1.47 (187.51 - 161.20) =  26.31  6.93 (3.29) 6.31 (2.66) N=33 -- tm7_obs_1
    N16 -5.41 +2.47 -0.97 (113.70 -  60.40) =  53.30  5.34 (4.23) 4.59 (2.89) N=26 -- tm7_obs_2
    """
def do_astro_corr_rota(src=None, ref=None, catalog=None):
    """ calculates astronomy correction """

    title=os.path.split(catalog)[1]
    png=catalog.replace(".fits", ".rota.png")
    log=catalog.replace(".fits", ".rota.log")
    key=title.replace("_SourceCatalog_en0.fits","")
    #png=png.replace(key,"{}_{}".format(obslist[key],key))
    
    print(src)
    hdul = fits.open(src)
    src_map_data = hdul[0].data
    src_map_hdr = hdul[0].header
    src_tab_data = hdul[1].data
    src_tab_hdr = hdul[1].header
    hdul.close()

    dsrc=[]
    for rec in src_tab_data:
        d = dict(ra = rec['RA'], dec = rec['DEC'], ra_err = rec['RA_ERR'], dec_err = rec['DEC_ERR'], radec_err = rec['RADEC_ERR'])
        dsrc.append(d)
    
    

    print(ref)
    hdul = fits.open(ref)
    ref_map_data = hdul[0].data
    ref_map_hdr = hdul[0].header
    ref_tab_data = hdul[1].data
    ref_tab_hdr = hdul[1].header
    hdul.close()

    dref=[]
    for rec in ref_tab_data:
        d = dict(ra = rec['RA'], dec = rec['DEC'], ra_err = rec['RA_ERR'], dec_err = rec['DEC_ERR'])
        dref.append(d)

        
    resid = get_rms(dsrc,dref)
    print(resid)

    wcs_src=mWCS(src_map_hdr)
    wcs_src.info()
    
    wcs_ref=mWCS(ref_map_hdr)
    wcs_ref.info()

    x=[]
    y=[]
    rms_min=1000.0
    rota_min=1000.0
    for angle in np.arange(-0.5, 0.5, 0.01):
        resid1=do_transform(angle,wcs_src,dsrc,dref)
        x.append(angle*60)
        y.append(resid1/resid)
        if(resid1<rms_min):
            rms_min=resid1
            rota_min=angle*60

    with open(log, 'w') as writer:
        writer.write("{} {} {:.2f} {:.2f} {:.2f} {:.2f}%\n".format(key,title,rota_min, rms_min, resid, (1.0-rms_min/resid)*100.0))


    fig, ax = plt.subplots( nrows=1, ncols=1 )  # create figure & 1 axis
    ax.plot(x,y)
    plt.ylabel('RMS/RMS0, RMS0={:.2f} arcsec'.format(resid))
    plt.xlabel('Rotation, arcmin')
    plt.title("{} {} {} pairs, {:.1f}%".format(key, title,len(dsrc),(1.0-rms_min/resid)*100.0))
    plt.grid(True)
    plt.axvline(x = 0.0, color = 'b', label = 'zero')
    fig.savefig(png)
    plt.close(fig)    # close the figure window

def do_transform(angle,wcs,dsrc,dref):
    # fills x,y
    sky2pix(wcs,dsrc)

    # do transform here
    wcs.rotate(angle)
    
    # fills ra1,dec1
    pix2sky(wcs,dsrc)

    resid1 = get_rms1(dsrc,dref)
    return resid1
