#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""

Build AKARI/IRC spectral cube
Inter-calibrate AKARI/IRC spectra with Spitzer/IRAC1

"""

# import logging, sys
# logging.disable(sys.maxsize)
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", message="Skipping SYSTEM_VARIABLE record")

from tqdm import tqdm, trange

import os
import math
import numpy as np
from scipy.optimize import curve_fit
## rapyuta
from rapyuta.inout import fclean, fitsext, read_fits, write_fits, read_hdf5
from rapyuta.imaging import ( iconvolve, iswarp, cupid, iuncert,
                              icrop, igroupixel, ismooth,
                              imontage, improve, Jy_per_pix_to_MJy_per_sr )
from rapyuta.astrom import fixwcs
from rapyuta.arrays import closest, pix2sup, sup2pix
from rapyuta.calib import intercalib
from rapyuta.maths import f_lin, f_lin0, f_lin1
from rapyuta.plots import pplot

##----------------------------------------------------------

##                     Preliminaries

##----------------------------------------------------------

## Local
from buildinfo import ( src, Nmc, verbose, coadd_tool, colors, markers,
                        path_irc, path_build, parobs, # build
                        filog, fits_irc, out_irc, path_out,
                        path_phot, path_ker, path_cal, # calib
                        path_idl, csv_ker, path_tmp, path_fig
)

## Banner
print('\n============================================================\n')

print('        MIRAGE - AKARI/IRC cube builder - '+src)

print('\n============================================================\n')

Nobs = len(parobs)

# Nmc = 2

# exit()

##----------------------------------------------------------

##                      Build slits

##----------------------------------------------------------
do_build = input("(Re)build IRC slits (y/n)? ")
if do_build=='y':

    for i in trange(Nobs, #leave=False,
                    desc='<cupid> IRC slit building'):
        for j in trange(Nmc+1, leave=False,
                        desc=' - IRC slit unc [MC]'):
            cup = cupid(path_irc, obsid=parobs[i][1], slit=parobs[i][4],
                        spec=parobs[i][2], imref=parobs[i][3])
            if j==0:
                cup.spec_build(fits_irc[i]+'_'+str(j), tmpdir=path_build,
                               filRAW=fits_irc[i], fiLOG=filog+parobs[i][0],
                               Nx=parobs[i][7], Ny=parobs[i][5], Nsub=parobs[i][6],
                               pixscale=1, wmin=2.55, wmax=4.25, supix=True)
                if i==0:
                    wave0 = cup.wave()
            else:
                cup.spec_build(fits_irc[i]+'_'+str(j), tmpdir=path_build,
                               dist='splitnorm', sig_pt=3, fill_pt='med',
                               Nx=parobs[i][7], Ny=parobs[i][5], Nsub=parobs[i][6],
                               pixscale=1, wmin=2.55, wmax=4.25, supix=True)
            ## Interpolate wavelength grid of IRC slits
            ##------------------------------------------
            ## (A,B,C,D,G,H,I,J,L,N) (E,F,K) (M) have slightly different wgrid...
            ismooth(fits_irc[i]+'_'+str(j), wgrid=wave0, filOUT=fits_irc[i]+'_'+str(j))

        ## MC unc
        if Nmc>1:
            mcimage = []
            for j in range(Nmc):
                ds = read_fits(fits_irc[i]+'_'+str(j+1))
                mcimage.append(ds.data)
            mcimage = np.array(mcimage)
            unc = np.nanstd(mcimage, axis=0)
            write_fits(out_irc[i]+'_unc', ds.header, unc, ds.wave)

##----------------------------------------------------------

##                 Coadd slits (for footprint)

##----------------------------------------------------------
irc_coadd = input("Coadd IRC slits (y/n)? ")
if irc_coadd=='y':

    ## Coadding
    ##----------

    ## IRC grid
    refheader = fixwcs(fits_irc[0]+fitsext).header
    swp = iswarp(fits_irc, refheader=refheader,
                 tmpdir=path_build, verbose=verbose)
    footprint_irc = swp.refheader

    slice_irc = []
    for f in fits_irc:
        ds = read_fits(f)
        ds.data[0] = 1
        write_fits(f+'_slice', ds.header, ds.data[0])
        slice_irc.append(f+'_slice')

    if coadd_tool=='swarp':

        ## <iswarp> coadding
        ##===================
        swp = iswarp(refheader=footprint_irc,
                     tmpdir=path_build, verbose=verbose)
        swp.combine(slice_irc,# combtype='wgt_avg',
                    keepedge=True, #cropedge=True,
                    filOUT=path_out+src+'_footprint')

    elif coadd_tool=='reproject':

        ## <imontage> coadding
        ##=====================
        mtg = imontage('exact', tmpdir=path_tmp, verbose=verbose)
        mtg.coadd(slice_irc, refheader=footprint_irc,
                  filOUT=path_out+src+'_footprint')

    ## Crop NaN edge
    edge = 2 # leave an edge of some pixels
    data = read_fits(path_out+src+'_footprint').data
    Ny, Nx = data.shape
    xlist = []
    for x in range(Nx):
        if not np.isnan(data[:,x]).all():
            xlist.append(x)
    ylist = []
    for y in range(Ny):
        if not np.isnan(data[y,:]).all():
            ylist.append(y)
    xmin = min(xlist) - edge
    if xmin<0:
        xmin = 0
    xmax = max(xlist)+1 + edge
    if xmax>Nx:
        xmax = Nx
    ymin = min(ylist) - edge
    if ymin<0:
        ymin = 0
    ymax = max(ylist)+1 + edge
    if ymax>Ny:
        ymax = Ny
    dx = xmax-xmin
    dy = ymax-ymin
    x0 = xmin+dx/2
    y0 = ymin+dy/2

    icrop(path_out+src+'_footprint', filOUT=path_out+src+'_footprint',
          sizpix=(dx,dy), cenpix=(x0,y0), verbose=verbose)

## Header of th atlas IRC map
header_atlas = fixwcs(path_out+src+'_footprint'+fitsext).header

##----------------------------------------------------------

##                   Inter-calibration

##----------------------------------------------------------
pre_calib1 = input("Run inter-calibration prepipeline (DustPedia IRAC1) (y/n)? ")
pre_calib2 = input("Run inter-calibration prepipeline (SINGS IRAC1) (y/n)? ")
synt_phot = input("Run inter-calibration prepipeline (IRC synthetic photometry) (y/n)? ")

if not os.path.exists(path_tmp+'calib/'):
    os.makedirs(path_tmp+'calib/')

## Prepare photometry
##--------------------
phot = 'IRAC1'
phot_ker = path_ker+'Kernel_HiRes_IRAC_3.6_to_Gauss_06.0'

raw_phot1 = path_phot+src+'_'+phot+'_DP'
fits_phot1 = path_cal+src+'_'+phot+'_DP'
tmp_phot1 = path_tmp+'calib/'+src+'_'+phot+'_DP'

raw_phot2 = path_phot+src+'_'+phot+'_SINGS'
wgt_phot2 = path_phot+src+'_'+phot+'_SINGS_wt'
fits_phot2 = path_cal+src+'_'+phot+'_SINGS'
tmp_phot2 = path_tmp+'calib/'+src+'_'+phot+'_SINGS'

fits_spec = path_cal+src+'_'+phot+'_IRC'
tmp_spec = path_tmp+'calib/'+src+'_'+phot+'_IRC'

## DustPedia (phot1)
##===================
if pre_calib1=='y':

    ## Convert phot unit (Jy/pix -> MJy/sr)
    ## This step should be before reprojection (pixscale may change)
    Jy_per_pix_to_MJy_per_sr(raw_phot1, filOUT=tmp_phot1)

    ## DustPedia IRAC1 of M82 has no uncertainty map
    ## Create homogeneous uncertainty map
    bg = read_fits(tmp_phot1).data[1713:1753,3110:3150]*10
    iuncert(tmp_phot1, filOUT=tmp_phot1+'_unc', BG_image=bg)

    ## Reproject phot (atlas)
    if coadd_tool=='swarp':
        swp = iswarp(refheader=header_atlas, tmpdir=path_tmp)
        swp.combine_mc(tmp_phot1, keepedge=True, #cropedge=True,
                       dist='norm', Nmc=Nmc, filOUT=tmp_phot1+'_MC')
    elif coadd_tool=='reproject':
        mtg = imontage('exact', tmpdir=path_tmp)
        mtg.reproject_mc(tmp_phot1, refheader=header_atlas,
                         dist='norm', Nmc=Nmc, filOUT=tmp_phot1+'_MC')

    for i in trange(Nobs, #leave=False,
                    desc=phot+' (DustPedia) processing'):
        xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
        refheader = fixwcs(fits_irc[i]+'_0'+fitsext).header

        for j in trange(Nmc+1, leave=False,
                        desc=' - '+parobs[i][0]+' [MC]'):
            if i==0:
                ## Convolve phot
                if j==0:
                    conv = iconvolve(tmp_phot1+'_MC',
                                     kfile=phot_ker, klist=csv_ker,
                                     filOUT=tmp_phot1+'_MC')
                else:
                    conv = iconvolve(tmp_phot1+'_MC_'+str(j),
                                     kfile=phot_ker, klist=csv_ker,
                                     # filUNC=tmp_phot1+'_unc', dist='norm',
                                     filOUT=tmp_phot1+'_MC_'+str(j))
                conv.do_conv(idldir=path_idl)

            ## Reproject phot (slit)
            if coadd_tool=='swarp':
                swp = iswarp(refheader=refheader, tmpdir=path_tmp)
                if j==0:
                    swp.combine(tmp_phot1+'_MC',
                                filOUT=fits_phot1+'_'+parobs[i][0])
                else:
                    swp.combine(tmp_phot1+'_MC_'+str(j),
                                filOUT=tmp_phot1+'_'+parobs[i][0]+'_'+str(j))
            elif coadd_tool=='reproject':
                mtg = imontage('exact', tmpdir=path_tmp)
                if j==0:
                    mtg.reproject(tmp_phot1+'_MC', refheader=refheader,
                                  filOUT=fits_phot1+'_'+parobs[i][0])
                else:
                    mtg.reproject(tmp_phot1+'_MC_'+str(j), refheader=refheader,
                                  filOUT=tmp_phot1+'_'+parobs[i][0]+'_'+str(j))

            ## Calculate unc
            if j==0:
                igroupixel(fits_phot1+'_'+parobs[i][0],
                           xscale=xscale, yscale=yscale,
                           filOUT=fits_phot1+'_'+parobs[i][0])
            else:
                igroupixel(tmp_phot1+'_'+parobs[i][0]+'_'+str(j),
                           xscale=xscale, yscale=yscale,
                           filOUT=tmp_phot1+'_'+parobs[i][0]+'_'+str(j))
        if Nmc>1:
            mcimage = []
            for j in range(Nmc):
                ds = read_fits(tmp_phot1+'_'+parobs[i][0]+'_'+str(j+1))
                mcimage.append(ds.data)
            mcimage = np.array(mcimage)
            unc = np.nanstd(mcimage, axis=0)
            write_fits(fits_phot1+'_'+parobs[i][0]+'_unc', ds.header, unc)

## SINGS (phot2)
##===============
if pre_calib2=='y':

    ## Create uncertainty via weight map (suppose uniform contribution)

    ## (Sect. 3.2) The weight maps contain the information on the number of frames
    ## that were used to create the science mosaics at each pixel (value= # frames x10);
    ## the pixel size of the weight maps is the same as the science mosaics.
    ## -- SINGS v5 release doc
    ## https://irsa.ipac.caltech.edu/data/SPITZER/SINGS/doc/sings_fifth_delivery_v2.pdf

    ## Only need to run ONCE for each photometry (comment after using)
    # gen_unc2 = input("Create SINGS IRAC1 uncertainty map (y/n)? ")
    # if gen_unc2=='y':
    #     ds_wgt = read_fits(wgt_phot2)
    #     bg = read_fits(raw_phot2).data[694:734,56:96]
    #     bg_wgt = ds_wgt.data[694:734,56:96]
    #     wfac = .1
    #     iuncert(raw_phot2, filOUT=raw_phot2+'_unc',
    #             filWGT=wgt_phot2, wfac=wfac,
    #             BG_image=bg, BG_weight=bg_wgt)

    ## Reproject phot (atlas)
    if coadd_tool=='swarp':
        swp = iswarp(refheader=header_atlas, tmpdir=path_tmp)
        swp.combine_mc(raw_phot2, keepedge=True, #cropedge=True,
                       dist='norm', Nmc=Nmc, filOUT=tmp_phot2+'_MC')
    elif coadd_tool=='reproject':
        mtg = imontage('exact', tmpdir=path_tmp)
        mtg.reproject_mc(raw_phot2, refheader=header_atlas,
                         dist='norm', Nmc=Nmc, filOUT=tmp_phot2+'_MC')

    for i in trange(Nobs, #leave=False,
                    desc=phot+' (SINGS) processing'):
        xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
        refheader = fixwcs(fits_irc[i]+'_0'+fitsext).header

        for j in trange(Nmc+1, leave=False,
                        desc=' - '+parobs[i][0]+' [MC]'):
            if i==0:
                ## Convolve phot
                if j==0:
                    conv = iconvolve(tmp_phot2+'_MC',
                                     kfile=phot_ker, klist=csv_ker,
                                     filOUT=tmp_phot2+'_MC')
                else:
                    conv = iconvolve(tmp_phot2+'_MC_'+str(j),
                                     kfile=phot_ker, klist=csv_ker,
                                     # filUNC=tmp_phot2+'_unc', dist='norm',
                                     filOUT=tmp_phot2+'_MC_'+str(j))
                conv.do_conv(idldir=path_idl)

            ## Reproject phot (slit)
            if coadd_tool=='swarp':
                swp = iswarp(refheader=refheader, tmpdir=path_tmp)
                if j==0:
                    swp.combine(tmp_phot2+'_MC',
                                filOUT=fits_phot2+'_'+parobs[i][0])
                else:
                    swp.combine(tmp_phot2+'_MC_'+str(j),
                                filOUT=tmp_phot2+'_'+parobs[i][0]+'_'+str(j))
            elif coadd_tool=='reproject':
                mtg = imontage('exact', tmpdir=path_tmp)
                if j==0:
                    mtg.reproject(tmp_phot2+'_MC', refheader=refheader,
                                  filOUT=fits_phot2+'_'+parobs[i][0])
                else:
                    mtg.reproject(tmp_phot2+'_MC_'+str(j), refheader=refheader,
                                  filOUT=tmp_phot2+'_'+parobs[i][0]+'_'+str(j))

            ## Calculate unc
            if j==0:
                igroupixel(fits_phot2+'_'+parobs[i][0],
                           xscale=xscale, yscale=yscale,
                           filOUT=fits_phot2+'_'+parobs[i][0])
            else:
                igroupixel(tmp_phot2+'_'+parobs[i][0]+'_'+str(j),
                           xscale=xscale, yscale=yscale,
                           filOUT=tmp_phot2+'_'+parobs[i][0]+'_'+str(j))
        if Nmc>1:
            mcimage = []
            for j in range(Nmc):
                ds = read_fits(tmp_phot2+'_'+parobs[i][0]+'_'+str(j+1))
                mcimage.append(ds.data)
            mcimage = np.array(mcimage)
            unc = np.nanstd(mcimage, axis=0)
            write_fits(fits_phot2+'_'+parobs[i][0]+'_unc', ds.header, unc)

## Synthetic photometry (spec)
##=============================
if (synt_phot=='y' and do_build=='y'):
    ## fits_irc[i] can be changed by inter-calib,
    ## One need to rebuild before synthetic photometry
    for i in trange(Nobs, #leave=False,
                    desc=phot+' (IRC synt phot) slits'):
        xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
        for j in trange(Nmc+1, leave=False,
                        desc=' - '+parobs[i][0]+': IRC sythetic photometry [MC]'):
            ic = intercalib(fits_irc[i]+'_'+str(j))
            sp = ic.synthetic_photometry(phot, xscale=xscale, yscale=yscale)
            if j==0:
                write_fits(fits_spec+'_'+parobs[i][0], ic.hdr, sp.Fnu_filt)
            else:
                write_fits(tmp_spec+'_'+parobs[i][0]+'_'+str(j), ic.hdr, sp.Fnu_filt)
        if Nmc>1:
            mcimage = []
            for j in range(Nmc):
                ds = read_fits(tmp_spec+'_'+parobs[i][0]+'_'+str(j+1))
                mcimage.append(ds.data)
            mcimage = np.array(mcimage)
            unc = np.nanstd(mcimage, axis=0)
            write_fits(fits_spec+'_'+parobs[i][0]+'_unc', ds.header, unc)

##-----------------
## Fit correlation
##-----------------
## Data dictionary
## e.g. [{'A0': a0, 'A1': a1, ...}, {'B0': b0, 'B1': b1, ...}, ...]
dict_phot1 = []
dict_phot1_unc = []
dict_phot2 = []
dict_phot2_unc = []
dict_spec = []
dict_spec_unc = []
## Data 1D array
## e.g. [array([a1, a2, ...]), array([b1, b2, ...]), ...]
pix_phot1 = []
pix_phot1_unc = []
pix_phot2 = []
pix_phot2_unc = []
pix_spec = []
pix_spec_unc = []

# xgrid = np.arange(0,1e3,1)
xgrid = np.logspace(-2,3,1000)

for i in range(Nobs):
    # Nx = read_hdf5(filog+parobs[i][0], 'Slit width')
    Ny = read_hdf5(filog+parobs[i][0], 'Slit length')[0]
    xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
    # Nxs = math.ceil(Nx/xscale)
    Nys = math.ceil(Ny/yscale)

    ## Init dict
    data1 = {}
    unc1 = {}
    data2 = {}
    unc2 = {}
    data3 = {}
    unc3 = {}
    for ys in range(Nys):
        subname = parobs[i][0][0]+str(ys+1)
        yarr = sup2pix(ys, yscale, Npix=Ny, origin=0)
        ## Phot1
        ds = read_fits(fits_phot1+'_'+parobs[i][0], fits_phot1+'_'+parobs[i][0]+'_unc')
        data1[subname] = ds.data[yarr[0],0]
        unc1[subname] = ds.unc[yarr[0],0]
        ## Phot2
        ds = read_fits(fits_phot2+'_'+parobs[i][0], fits_phot2+'_'+parobs[i][0]+'_unc')
        data2[subname] = ds.data[yarr[0],0]
        unc2[subname] = ds.unc[yarr[0],0]
        ## spec
        ds = read_fits(fits_spec+'_'+parobs[i][0], fits_spec+'_'+parobs[i][0]+'_unc')
        data3[subname] = ds.data[yarr[0],0]
        unc3[subname] = ds.unc[yarr[0],0]
    dict_phot1.append(data1)
    dict_phot1_unc.append(unc1)
    dict_phot2.append(data2)
    dict_phot2_unc.append(unc2)
    dict_spec.append(data3)
    dict_spec_unc.append(unc3)

    ## Dictionary to 1D array
    pix_phot1.append(np.array(list(dict_phot1[i].values())))
    pix_phot1_unc.append(np.array(list(dict_phot1_unc[i].values())))
    pix_phot2.append(np.array(list(dict_phot2[i].values())))
    pix_phot2_unc.append(np.array(list(dict_phot2_unc[i].values())))
    pix_spec.append(np.array(list(dict_spec[i].values())))
    pix_spec_unc.append(np.array(list(dict_spec_unc[i].values())))

    ## Mask NaNs and negtive values
    mask_nan = ~np.logical_or( np.isnan(pix_spec[i]), np.isnan(pix_phot2[i]) )
    mask_neg = np.logical_and( pix_spec[i]>0, pix_phot2[i]>0 )
    mask = np.logical_and( mask_nan, mask_neg )

    ## Add calibration error
    pix_phot1_unc[i] += pix_phot1[i] * .03
    pix_phot2_unc[i] += pix_phot2[i] * .03
    pix_spec_unc[i] += pix_spec[i] * .05

    ## S/N ratio
    # print(parobs[i][0]+' S/N (IRC) = \n', pix_spec[i][mask]/pix_spec_unc[i][mask])
    # print(parobs[i][0]+' S/N (SINGS) = \n', pix_phot2[i][mask]/pix_phot2_unc[i][mask])

    if mask.any():
        ## DP - SINGS plot
        if i==0:
            p0 = pplot(fmt='s', ec='grey', elw=1,# clib=colors,
                       xlog=1, ylog=1, nonposx='clip', nonposy='clip',
                       # xlim=(1e-2,1e3), ylim=(1e-2,1e3),
                       xlabel='DustPedia (MJy/sr)', ylabel='SINGS (MJy/sr)',
                       title=src+' '+phot+' calibration',
                       figsize=(10,8), right=.8, left=.15, bottom=.15,
                       legend='upper left', anchor=(1,1),
                       titlesize=20, labelsize=20, ticksize=20, legendsize=20)
        p0.add_plot(pix_phot1[i][mask], pix_phot2[i][mask],
                    yerr=pix_phot2_unc[i][mask], xerr=pix_phot1_unc[i][mask],
                    fmt='s', ec='grey', c=colors[i+1],
                    marker=markers[i], markersize=10, capsize=2,
                    label=parobs[i][0])
        p0.add_plot(pix_phot1[i][mask], pix_phot2[i][mask],
                    yerr=pix_phot2_unc[i][mask], xerr=pix_phot1_unc[i][mask],
                    fmt='o', ec='grey', c=colors[i+1], zorder=100,
                    marker=markers[i], markersize=.1, capsize=2)
        p0.save(path_cal+'DP-SINGS_'+phot)

        if (synt_phot=='y' or pre_calib2=='y'):
            ## IRC - SINGS plot
            if i==0:
                p = pplot(fmt='s', ec='grey', elw=1,# clib=colors,
                          xlog=1, ylog=1, nonposx='clip', nonposy='clip',
                          # xlim=(1e-2,1e3), ylim=(1e-2,1e3),
                          # xlabel='IRC (MJy/sr)', ylabel='SINGS (MJy/sr)',
                          # title=src+' IRC-'+phot+' inter-calibration',
                          xlabel='IRC (MJy/sr)', ylabel=r'$\rm IRAC_{3.6\mu m}\ (MJy/sr)$',
                          title=None,
                          figsize=(11,8), right=.78, left=.12, bottom=.1, top=.95,
                          legend='upper left', anchor=(1,1),
                          titlesize=20, labelsize=20, ticksize=20, legendsize=20)

            p.add_plot(pix_spec[i][mask], pix_phot2[i][mask],
                       yerr=pix_phot2_unc[i][mask], xerr=pix_spec_unc[i][mask],
                       fmt='s', ec='grey', c=colors[i+1],
                       marker=markers[i], markersize=10, capsize=2,
                       label=parobs[i][0])
            p.add_plot(pix_spec[i][mask], pix_phot2[i][mask],
                       yerr=pix_phot2_unc[i][mask], xerr=pix_spec_unc[i][mask],
                       fmt='s', ec='grey', c=colors[i+1], zorder=100,
                       marker=markers[i], markersize=.1, capsize=2)
            p.save(path_cal+'IC_'+phot+'.png')

        ## Linear fit (IRC - SINGS, SLITS)
        ##=================================
        popt, pcov = curve_fit(f_lin0, pix_spec[i][mask], pix_phot2[i][mask],)
                               # sigma=pix_phot2_unc[i][mask])
        slit_gain = popt[0]
        # slit_gain = 1.
        # slit_off = popt[1]
        slit_off = 0.
        # slit_off = popt[0]

        if (synt_phot=='y' or pre_calib2=='y'):
            print(parobs[i][0]+' inter-calibration ('+phot+') gain = {:.4}'.format(slit_gain))
            # print(parobs[i][0]+' inter-Calibration ('+phot+') offset = {:.4}'.format(slit_off))
            # label = parobs[i][0]+': y={0:.4}x'.format(slit_gain)
            # label = parobs[i][0]+': y={0:.4}x+{1:.4}'.format(slit_gain, slit_off)
            # p.add_plot(xgrid, f_lin0(xgrid, *popt),
            #            c='k', ls='-', label=label)
            # p.save(path_cal+'IC_'+phot+'.png')

## Linear fit (DP - SINGS)
##=========================
fitx = np.concatenate(pix_phot1, axis=0)
uncx = np.concatenate(pix_phot1_unc, axis=0)
fity = np.concatenate(pix_phot2, axis=0)
uncy = np.concatenate(pix_phot2_unc, axis=0)
maskfit = ~np.logical_or( np.isnan(fitx), np.isnan(fity) )

popt, pcov = curve_fit(f_lin0, fitx[maskfit], fity[maskfit],)
                       # sigma=uncy[maskfit])
atlas_gain = popt[0]
# atlas_gain = 1.
# atlas_off = popt[1]
atlas_off = 0.
# atlas_off = popt[0]

print('DustPedia - SINGS calibration ('+phot+') gain = {:.4}'.format(atlas_gain))
# print('DustPedia - SINGS calibration ('+phot+') offset = {:.4}'.format(atlas_off))
label = 'y={0:.4}x'.format(atlas_gain)
# label = 'y={0:.4}x+{1:.4}'.format(atlas_gain, atlas_off)
p0.add_plot(xgrid, f_lin0(xgrid, *popt),
            c='k', ls='-', label=label)
p0.save(path_cal+'DP-SINGS_'+phot+'.png')

## Linear fit (IRC - SINGS, ATLAS)
##=================================
fitx = np.concatenate(pix_spec, axis=0)
uncx = np.concatenate(pix_spec_unc, axis=0)
fity = np.concatenate(pix_phot2, axis=0)
uncy = np.concatenate(pix_phot2_unc, axis=0)
maskfit = ~np.logical_or( np.isnan(fitx), np.isnan(fity) )

popt, pcov = curve_fit(f_lin0, fitx[maskfit], fity[maskfit],)
                       # sigma=uncy[maskfit])
atlas_gain = popt[0]
# atlas_gain = 1.
# atlas_off = popt[1]
atlas_off = 0.
# atlas_off = popt[0]

if (synt_phot=='y' or pre_calib2=='y'):
    print('Atlas inter-calibration ('+phot+') gain = {:.4}'.format(atlas_gain))
    # print('Atlas inter-Calibration ('+phot+') offset = {:.4}'.format(atlas_off))
    label = 'y={0:.4}x'.format(atlas_gain)
    # label = 'y={0:.4}x+{1:.4}'.format(atlas_gain, atlas_off)
    p.add_plot(xgrid, f_lin0(xgrid, *popt),
               c='k', ls='-', label=label)
    p.ax.text(.9,.05,'(a)',size=20,c='grey',transform=p.ax.transAxes)
    p.ax.legend(loc='upper left', bbox_to_anchor=(1,1),
                fontsize=20, framealpha=0,)
    p.save(path_cal+'IC_'+phot+'.png', transparent=True)

## Spectral correction
##---------------------
write_irc = input("Write IRC spectra (y/n)? ")
if write_irc=='y':
    if do_build=='y':
        correct_atlas = input(" - ATLAS level correction (y/n)? ")
        if correct_atlas=='y':
            correct_pixel = None
        else:
            correct_pixel = input(" - PIXEL level correction (y/n)? ")
    else:
        warnings.warn('Rebuild slits before spectral correction.')
        correct_atlas = None
        correct_pixel = None
else:
    correct_atlas = None
    correct_pixel = None

for i in range(Nobs):
    xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
    mask_nan = ~np.logical_or( np.isnan(pix_spec[i]), np.isnan(pix_phot2[i]) )
    mask_neg = np.logical_and( pix_spec[i]>0, pix_phot2[i]>0 )
    mask = np.logical_and( mask_nan, mask_neg )
    sc = intercalib(fits_irc[i]+'_0') # Attention: NOT fits_irc[i]
    Nw, Ny, Nx = sc.im.shape
    pixel_gain = np.ones((Ny,Nx))
    if mask.any():
        for y in range(Ny):
            ys = pix2sup(y, yscale, origin=0)
            if mask[ys]:
                pixel_gain[y,:] *= pix_phot2[i][ys]/pix_spec[i][ys]
            # if y%yscale==0:
            #     print(parobs[i][0]+' inter-calibration gain: {}'.format(pixel_gain[y,0]))
    if write_irc=='y':
        if correct_pixel=='y':
            calib_gain = pixel_gain
        elif correct_atlas=='y':
            calib_gain = atlas_gain
        else:
            calib_gain = 1.0
        sc.correct_spec(calib_gain, filOUT=out_irc[i])

##----------------------
## Fit correlation (MC)
##----------------------
for j in trange(Nmc, #leave=False,
                desc='IRC spectral correction [MC]'):
    ## Data dictionary
    ## e.g. [{'A0': a0, 'A1': a1, ...}, {'B0': b0, 'B1': b1, ...}, ...]
    dict_spec = []
    dict_spec_unc = []
    ## Data 1D array
    ## e.g. [array([a1, a2, ...]), array([b1, b2, ...]), ...]
    pix_spec = []
    pix_spec_unc = []

    for i in range(Nobs):
        # Nx = read_hdf5(filog+parobs[i][0], 'Slit width')
        Ny = read_hdf5(filog+parobs[i][0], 'Slit length')[0]
        xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
        # Nxs = math.ceil(Nx/xscale)
        Nys = math.ceil(Ny/yscale)

        ## Init dict
        data3 = {}
        unc3 = {}
        for ys in range(Nys):
            subname = parobs[i][0][0]+str(ys+1)
            yarr = sup2pix(ys, yscale, Npix=Ny, origin=0)
            ## spec
            ds = read_fits(tmp_spec+'_'+parobs[i][0]+'_'+str(j+1),
                           fits_spec+'_'+parobs[i][0]+'_unc') # intermediate unc
            data3[subname] = ds.data[yarr[0],0]
            unc3[subname] = ds.unc[yarr[0],0]
        dict_spec.append(data3)
        dict_spec_unc.append(unc3)

        ## Dictionary to 1D array
        pix_spec.append(np.array(list(dict_spec[i].values())))
        pix_spec_unc.append(np.array(list(dict_spec_unc[i].values())))

        ## Mask NaNs and negtive values
        mask_nan = ~np.logical_or( np.isnan(pix_spec[i]), np.isnan(pix_phot2[i]) )
        mask_neg = np.logical_and( pix_spec[i]>0, pix_phot2[i]>0 )
        mask = np.logical_and( mask_nan, mask_neg )

    ## Linear fit (IRC - SINGS, ATLAS)
    ##=================================
    fitx = np.concatenate(pix_spec, axis=0)
    uncx = np.concatenate(pix_spec_unc, axis=0)
    fity = np.concatenate(pix_phot2, axis=0)
    uncy = np.concatenate(pix_phot2_unc, axis=0)
    maskfit = ~np.logical_or( np.isnan(fitx), np.isnan(fity) )

    popt, pcov = curve_fit(f_lin0, fitx[maskfit], fity[maskfit],)
                           # sigma=uncy[maskfit])
    atlas_gain = popt[0]
    # atlas_gain = 1.
    # atlas_off = popt[1]
    atlas_off = 0.
    # atlas_off = popt[0]

    ## Spectral correction
    ##---------------------
    for i in range(Nobs):
        xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
        mask_nan = ~np.logical_or( np.isnan(pix_spec[i]), np.isnan(pix_phot2[i]) )
        mask_neg = np.logical_and( pix_spec[i]>0, pix_phot2[i]>0 )
        mask = np.logical_and( mask_nan, mask_neg )
        sc = intercalib(fits_irc[i]+'_'+str(j+1))
        Nw, Ny, Nx = sc.im.shape
        pixel_gain = np.ones((Ny,Nx))
        if mask.any():
            for y in range(Ny):
                ys = pix2sup(y, yscale, origin=0)
                if mask[ys]:
                    pixel_gain[y,:] *= pix_phot2[i][ys]/pix_spec[i][ys]
                # if y%yscale==0:
                #     print(parobs[i][0]+' inter-calibration gain: {}'.format(pixel_gain[y,0]))
        if (write_irc=='y'):
            ## fits_irc[i] can be changed by inter-calib,
            ## One need to rebuild before synthetic photometry
            if correct_pixel=='y':
                calib_gain = pixel_gain
            elif correct_atlas=='y':
                calib_gain = atlas_gain
            else:
                calib_gain = 1.0
            sc.correct_spec(calib_gain, filOUT=fits_irc[i]+'_'+str(j+1))

if (write_irc=='y' and Nmc>1):
    for i in trange(Nobs, leave=False,
                    desc='Calculating uncertainties for IRC slits'):
        mcimage = []
        for j in range(Nmc):
            ds = read_fits(fits_irc[i]+'_'+str(j+1))
            mcimage.append(ds.data)
        mcimage = np.array(mcimage)
        unc = np.nanstd(mcimage, axis=0)
        write_fits(out_irc[i]+'_unc', ds.header, unc, ds.wave)

##----------------------------------------------------------

##                      Plot spectra

##----------------------------------------------------------
plot_spec = input("Plot IRC spectra (y/n)? ")
if plot_spec=='y':

    for i in range(Nobs):
        ds = read_fits(out_irc[i], out_irc[i]+'_unc')
        Nw, Ny, Nx = ds.data.shape
        ds0 = read_fits(fits_irc[i]+'_0')
        pp = intercalib(out_irc[i])
        pp.read_filter(phot)
        wcen = pp.wcen

        xscale, yscale = read_hdf5(filog+parobs[i][0], 'Super pixel size')
        for y in range(Ny):
            ys = pix2sup(y, yscale, origin=0)
            subname = parobs[i][0][0]+str(ys+1)
            maskspec = ~np.isnan(ds.data[:,y,0]).any()
            if (maskspec and y%yscale==0):
                p = pplot(ds.wave, ds.data[:,y,0], yerr=ds.unc[:,y,0],
                          # xlog=1, ylog=1,
                          c='k', lw=.7, ec='r', label=subname,
                          xlabel=r'${\rm Wavelengths}\ \lambda\ (\mu m)$',
                          ylabel=r'${\rm Surface\ brightness}\ F_{\nu}\ (MJy/sr)$',
                          figsize=(8,8), legend='upper left',
                          title='IRC_'+subname,
                          titlesize=20, labelsize=10, ticksize=10, legendsize=10)
                ## Non-intercalib
                p.add_plot(ds.wave, ds0.data[:,y,0],
                           c='y', lw=.7, ls='--', zorder=-1)
                ## Photometry
                p.add_plot(wcen[0], dict_phot2[i][subname], yerr=dict_phot2_unc[i][subname],
                           c='m', marker='o', ms=10, zorder=100, label=phot)
                ## Synthetic photometry
                p.add_plot(wcen[0], dict_spec[i][subname], yerr=dict_spec_unc[i][subname],
                           c='g', marker='^', ms=10, zorder=101, label='IRC-'+phot)

                p.save(path_fig+'IRC_'+subname)


import numpy as np

def calculate_map_main(gt_results, pred_results, iou_gt_thr=0.5, class_num=5):
    '''
    说明: 此函数用于计算目标检测中的mAP
    输入:
          gt_results:  list, 每一个元素对应一个样本中的所有目标的真实值
          如gt_gt_result[0] = [
                                [xmin, ymin, xmax, ymax, label],
                                [xmin, ymin, xamx, yamx, label],
                                ...
                              ]
          pred_results: list, 每一个元素对应一个样本中的所有目标的预测值
          如pred_results[0] = [
                                [xmin, ymin, xmax, ymax, label, score],
                                [xmin, ymin, xamx, yamx, label, score],
                                ...
                              ]
          iou_gt_thr:   float, 用于判定正负样本, 如iou_gt_thr=0.5, 计算出的就是mAP 0.5
          class_num:    int, 类别数
    输出：
          calss_ap:     array, 各类别的AP
          mean_ap :     float, 各类别平均得到mAP
    '''

    all_tp = [ [] for i in range(class_num) ]    # 用于存放各类所有的tp
    all_fp = [ [] for i in range(class_num) ]    # 用于存放各类所有的fp
    all_score = [ [] for i in range(class_num) ] # 用于存放bbox对应的scores
    all_gt_num = np.zeros([class_num])           # 用于存放各类的真实目标数, 之后计算recall会用到
    data_num = len(gt_results)                        # 样本总数

    # 对于每一个样本, 计算tp, fp, 并且统计预测bbox的score以及真实目标的个数
    for i in range(data_num):
        gt_result = gt_results[i]
        pred_result = pred_results[i]
        # print(f'computing {i} tp and fp')
        tp, fp, score, gt_num = calculate_tpfp_single(gt_result, pred_result, iou_gt_thr, class_num)
        # 按类别更新到总数中
        for n in range(class_num):
            all_tp[n].extend(tp[n])
            all_fp[n].extend(fp[n])
            all_score[n].extend(score[n])
            all_gt_num[n] += gt_num[n]

    # 计算出各类的AP，进而得到mAP
    all_map = calculate_map(all_tp, all_fp, all_score, all_gt_num, class_num)
    mean_ap = np.mean(all_map[1:])
    print('mAP', mean_ap)
    return all_map, mean_ap

def calculate_tpfp_single(gt_result, pred_result, iou_gt_thr, class_num):
    '''
    说明: 此函数用于计算单个样本的tp和fp, 并且存放对应的score, 统计真实目标的个数
    输入:
          gt_result:  list, 一个样本中所有的目标的真实值
          如gt_result=[
                        [xmin, ymin, xmax, ymax, label],
                        [xmin, ymin, xmax, ymax, label],
                        ...
                      ]
          pred_result: list, 一个样本中所有的目标的预测值
          如pred_result=[
                          [xmin, ymin, xmax, ymax, label, score],
                          [xmin, ymin, xmax, ymax, label, score],
                          ...
                        ]
          iou_gt_thr:  float, 用于判断正负样本的iou阈值
          class_num:   类别数
    输出:
          all_tp:      list, 每一个元素对应该样本计算出的对应类别的tp
          all_fp:      list, 每一个元素对应该样本计算出的对应类别的fp
          all_score:   list, 每一个元素对应该样本bbox对应类别的score
          gt_num:      list, 每一个元素对应该样本对应类别的真实目标数
    '''

    all_tp = [[] for i in range(class_num)]
    all_fp = [[] for i in range(class_num)]
    all_score = [[] for i in range(class_num)]
    gt_num = np.zeros([class_num])

    # 逐个类别提取真实bbox和预测bbox
    for i in range(class_num):
        tp = []
        fp = []
        score = []

        match_gt_bbox = [obj[0:4] for obj in gt_result if int(obj[4]) == i]
        match_pred_bbox = [obj[0:4] for obj in pred_result if int(obj[4]) == i]
        match_pred_score = [obj[5] for obj in pred_result if int(obj[4]) == i]

        len_gt = len(match_gt_bbox)
        len_pred = len(match_pred_bbox)

        if len_gt == 0 and len_pred != 0:
            # 说明不存在该类目标，但是预测出来了，属于误检
            score.extend(match_pred_score)
            for k in range(len_pred):
                tp.extend([0])
                fp.extend([1])

        if len_gt != 0 and len_pred != 0:
            # 说明存在该目标，并且检测出来了,那么计算若干gt与若干pred的iou
            score.extend(match_pred_score)
            ious = calculate_iou(match_gt_bbox, match_pred_bbox)
            max_iou = np.max(ious, axis=0)  # [x,x,x...] 每一个预测框与某个gt最大的iou
            # if any(s > 0.65 for s in max_iou):
            #     print(f"The image has good IOU: {max_iou}")
            # if any(s < 0.1 for s in max_iou):
                # print(f"The image has bad IOU: {max_iou}")
            # 使用iou_gt_thr来进行正负样本的判定，若满足条件，则为tp，否则为fp
            for k in range(len_pred):
                if max_iou[k] >= iou_gt_thr:
                    tp.extend([1])
                    fp.extend([0])
                if max_iou[k] < iou_gt_thr:
                    tp.extend([0])
                    fp.extend([1])

        all_tp[i].extend(tp)
        all_fp[i].extend(fp)
        all_score[i].extend(score)
        gt_num[i] += len_gt

    return all_tp, all_fp, all_score, gt_num

def calculate_area(bbox):
    # 计算一个bbox的面积
    w = max(bbox[2] - bbox[0], 0)
    h = max(bbox[3] - bbox[1], 0)
    w = max(0, w)
    h = max(0, h)
    return w * h

def calculate_inter(bbox1, bbox2):
    # 计算两个bbox的交集面积
    xmin = max(bbox1[0], bbox2[0])
    ymin = max(bbox1[1], bbox2[1])
    xmax = min(bbox1[2], bbox2[2])
    ymax = min(bbox1[3], bbox2[3])
    return calculate_area([xmin, ymin, xmax, ymax])

def calculate_union(bbox1, bbox2):
    # 计算两个bbox的并集面积
    area1 = calculate_area(bbox1)
    area2 = calculate_area(bbox2)
    inter = calculate_inter(bbox1, bbox2)
    union = area1 + area2 - inter
    return union

def IOU(bbox1, bbox2):
    # 计算两个bbox的iou
    inter = calculate_inter(bbox1, bbox2)
    union = calculate_union(bbox1, bbox2)
    iou = inter / union
    return iou

def calculate_iou(bbox1, bbox2):
    '''
    说明: 此函数用于计算M个bbox与N个bbox的iou
    输入:
          bbox1: list, 每一个元素是一个bbox, 如bbox1=[
                                                     [xmin, ymin, xamx, ymax],
                                                     [xmin, ymin, xmax, ymax],
                                                     ...
                                                    ]
          bbox2: list, 每一个元素是一个bbox, 如bbox2=[
                                                     [xmin, ymin, xamx, ymax],
                                                     [xmin, ymin, xmax, ymax],
                                                     ...
                                                    ]
    输出:
          ans:   array, size=[M, N], 计算出的iou矩阵
    '''

    len_1 = len(bbox1)
    len_2 = len(bbox2)
    ans = np.zeros([len_1, len_2])
    for i in range(len_1):
        for j in range(len_2):
            # 计算bbox1[i]和bbox2[j]的iou
            ans[i, j] = IOU(bbox1[i], bbox2[j])
    return ans

def calculate_map(all_tp, all_fp, all_score, all_gt_num, class_num):
    '''
    说明: 此函数的输入为所有类别的tp, fp, score和真实目标数, 计算每一个类别的AP
    输入:
          all_tp:     list,  每个元素是该类别下的tp
          all_fp:     list,  每个元素是该类别下的fp
          all_score:  list,  每个元素是该类别下预测bbox对应的score
          all_gt_num: list,  每个元素是该类下真实母目标的个数
          class_num:  int,   类别数
    输出:
          all_map:    array, 每个元素是该类的AP
    '''

    all_map = np.zeros([class_num])
    for i in range(class_num):
    # 首先提取出每一类的信息
        class_tp = all_tp[i]
        class_fp = all_fp[i]
        class_score = all_score[i]
        class_gt_num = all_gt_num[i]
        # 计算每一类的PR曲线
        class_P, class_R = calculate_PR(class_tp, class_fp, class_score, class_gt_num)
        # 计算PR曲线的面积，即AP
        class_map = calculate_map_single(class_P, class_R)
        # 写入该类别下
        all_map[i] = class_map
    return all_map

def calculate_PR(class_tp, class_fp, class_score, class_gt_num):
    '''
    说明: 此函数用于计算某一类的PR曲线
    输入:
          class_tp:     list, 该类下的tp, 每个元素为0或1, 代表当前样本是否为正样本
          class_fp:     list, 该类下的fp, 每个元素为0或1, 代表当前样本是否为负样本
          class_score:  list, 该类下预测bbox对应的score
          class_gt_num: int,  类别数
    输出:
          P: list, 该类下的查准率曲线
          R: list, 该类下的查全率曲线
    '''

    # 按照score排序
    sort_inds = np.argsort(class_score)[::-1].tolist()
    tp = [class_tp[i] for i in sort_inds]
    fp = [class_fp[i] for i in sort_inds]
    # 累加
    tp = np.cumsum(tp).tolist()
    fp = np.cumsum(fp).tolist()
    # 计算PR
    P = [tp[i] / (tp[i] + fp[i]) for i in range(len(tp))]
    R = [tp[i] / class_gt_num for i in range(len(tp))]
    return P, R

def calculate_map_single(P, R):
    '''
    说明: 此函数用于计算PR曲线的面积, 即AP
    输入:
          P: list, 查准率曲线
          R: list, 查全率曲线
    输出:
          single_map: float, 曲线面积, 即AP
    '''
    mpre = np.concatenate(([0.], P, [0.]))
    mrec = np.concatenate(([0.], R, [1.]))
    for i in range(np.size(mpre) - 1, 0, -1):
        # mpre的平整化
        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
    # 寻找mrec变化的坐标
    i = np.where(mrec[1:] != mrec[:-1])[0]
    # 计算面积
    single_map = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return single_map

def NMS(bounding_boxes, S=7, img_size=224, confidence_threshold=0.5, iou_threshold=0.5):
    """Compute non max suppressing to reduce overlapping bounding box.

    Args:
        bounding_boxes (list): a list of bounding box.
        S (int): the number of grid cells.
        img_size (int): image size
        confidence_threshold (float): the threshold to select a box on has_obj_prob (has_object_probability).
        iou_threshold (float): the threshold of IOU to remove bounding boxes.

    Returns:
        a list of bounding boxes.
        bounding box format = [center_x (int), center_y (int), width (int), height (int), has_obj_prob (float),
                               class_probs (float), confident_score (float), class_id (int)].
            class_probs is a list of class probability. confident_score = has_obj_prob * max(class_prob).
            class_id is argmax(class_probs).

    """
    bounding_boxes = bounding_boxes.cpu().detach().numpy().tolist()
    nms_boxes_buf = []
    grid_size = img_size / S
    for batch in range(len(bounding_boxes)):
        predict_boxes = []
        nms_boxes = []
        for i in range(S):
            for j in range(S):
                gridX = grid_size * j
                gridY = grid_size * i
                if bounding_boxes[batch][i][j][4] < bounding_boxes[batch][i][j][9]:
                    bounding_box = bounding_boxes[batch][i][j][5:10]
                else:
                    bounding_box = bounding_boxes[batch][i][j][0:5]
                bounding_box.extend(bounding_boxes[batch][i][j][10:])
                if bounding_box[4] >= confidence_threshold:
                    predict_boxes.append(bounding_box)

                centerX = (int)(gridX + bounding_box[0] * grid_size)
                centerY = (int)(gridY + bounding_box[1] * grid_size)
                width = (int)(bounding_box[2] * img_size)
                height = (int)(bounding_box[3] * img_size)
                bounding_box[0] = max(0, (int)(centerX - width / 2))
                bounding_box[1] = max(0, (int)(centerY - height / 2))
                bounding_box[2] = min(img_size - 1, (int)(centerX + width / 2))
                bounding_box[3] = min(img_size - 1, (int)(centerY + height / 2))
                class_idx = np.argmax(bounding_box[5:])
                confident_score = bounding_box[4] * bounding_box[5 + class_idx]  # has_obj_prob * class_prob
                bounding_box.append(confident_score)
                bounding_box.append(class_idx)

        while len(predict_boxes) != 0:
            predict_boxes.sort(key=lambda box: box[4])
            assured_box = predict_boxes[0]
            curr_class = assured_box[-1]
            temp = []
            nms_boxes.append(assured_box)
            i = 1
            while i < len(predict_boxes):
                compared_box = predict_boxes[i]
                if compared_box[-1] != curr_class or IOU(assured_box, predict_boxes[i]) <= iou_threshold:
                    temp.append(predict_boxes[i])
                i = i + 1
            predict_boxes = temp

        nms_boxes_buf.append(nms_boxes)

    return nms_boxes_buf

def gt_std(gt_results, S=7, B=2, img_size=224):

    gt_results_all = []
    grid_size = img_size / S
    for instance_index in range(gt_results.shape[0]): # N
        gt_results_instance = []
        for index_i in range(gt_results.shape[1]): # 7
            for index_j in range(gt_results.shape[2]): # 7
                gridX = grid_size * index_j
                gridY = grid_size * index_i
                area = gt_results[instance_index, index_i, index_j, 9]
                if area > 0:
                    gt_results_patch = gt_results[instance_index, index_i, index_j].tolist()
                    centerX = (int)(gridX + gt_results_patch[0] * grid_size)
                    centerY = (int)(gridY + gt_results_patch[1] * grid_size)
                    width = (int)(gt_results_patch[2] * img_size)
                    height = (int)(gt_results_patch[3] * img_size)
                    class_idx = int(gt_results[instance_index, index_i, index_j, 10:].argmax())
                    gt_results_patch[0] = max(0, (int)(centerX - width / 2))
                    gt_results_patch[1] = max(0, (int)(centerY - height / 2))
                    gt_results_patch[2] = min(img_size - 1, (int)(centerX + width / 2))
                    gt_results_patch[3] = min(img_size - 1, (int)(centerY + height / 2))
                    gt_results_patch[4] = class_idx
                    gt_results_instance.append(gt_results_patch[0:5])

        gt_results_all.append(gt_results_instance)

    return gt_results_all


# Copyright 2021 United States Government as represented by the Administrator of the National Aeronautics and Space
# Administration.  No copyright is claimed in the United States under Title 17, U.S. Code. All Other Rights Reserved.

"""
This module defines the interface to the default GIANT star catalogue.

Catalogue Description
=====================

The default GIANT star catalogue is a blending of the UCAC4 and Tycho2 catalogues into a stripped down sqlite3 database.
This makes retrieval of stars very fast instead of using the UCAC4 or Tycho2 catalogue files themselves and also
includes blended stars (stars that are so close together they will appear as a single star in most cameras).  This
catalogue includes stars down to about 16th magnitude and should be sufficient for nearly all stellar OpNav needs.

The current implementation of the database storing this catalogue is a single table with a primary index column of rnm,
which contains the unique ID of the provided star from the UCAC4 catalogue.  For blended stars, the index number becomes
the negative of the unique ID of the brightest star in the group.  In addition to the index, the following columns are
provided

================= ====== ======== ======================================================================================
Column            Units  Type     Description
================= ====== ======== ======================================================================================
source            N/A    string   The original catalogue source of the star (UCAC4 or Tycho2)
zone              N/A    integer  The UCAC4 zone number for the star, if applicable
rnz               N/A    integer  The star number in the UCAC4 zone if applicable
ra                deg    double   The right ascension of the star in degrees
dec               deg    double   The declination of the star in degrees
distance          km     double   The distance to the star in km.  If not known then this is replaced with the average
                                  distance to a star
ra_proper_motion  deg/yr double   The proper motion of the right ascension of the star in degrees per SI year
dec_proper_motion deg/yr double   The proper motion of the declination of the star in degrees per SI year
mag               mag    double   The visual magnitude of the star.  This is the APASM_V magnitude if available,
                                  otherwise the MAGM (model magnitude) from the UCAC4 catalogue.
ra_sigma          deg    double   The right ascension uncertainty in units of degrees
dec_sigma         deg    double   The declination uncertainty in units of degrees
distance_sigma    km     double   The distance uncertainty in units of kilometers.  For stars for which this is not
                                  known it is set to a large number
ra_pm_sigma       deg/yr double   The right ascension proper motion uncertainty in degrees per SI year.
dec_pm_sigma      deg/yr double   The declination proper motion uncertainty in degrees per SI year.
================= ====== ======== ======================================================================================

While this implementation mirrors the :attr:`.GIANT_COLUMNS` currently, and probably will in the future, it isn't
guaranteed to stay that way.  In addition, while this is currently a blend of the UCAC4 and Tycho2 catalogues, in the
future it will likely be built from the GAIA DR2 catalogue since this provides much more accurate stars
positions/magnitudes for many more stars.

Use
===

The GIANT catalogue can be used anywhere that a star catalogue is required in GIANT and is generally the default
catalogue that is used if you do not override it.  It is stored in a sqlite file in a directory called data in the same
directory hosting this file, though it is possible to override this if desired (which you may want to do if you need
different versions of the catalogue for different cameras). To access the default catalogue simply initialize this class
with no arguments and then call :meth:`~.GIANTCatalogue.query_catalogue` to retrieve the star records that you want.

This implementation also provides 2 helper methods that can retrieve the original star record from either the Tycho 2 or
UCAC4 catalogue (if the star exists in them and is not a blended star).  These are
:meth:`~.GIANTCatalogue.get_ucac4_record` and :meth:`~.GIANTCatalogue.get_tycho2_record`.  They take in a pandas
DataFrame of stars retrieved from this catalogue and return a dataframe with the original records for those stars.
This can be useful if you need more information about a star than what GIANT typically considers.  Just note that these
methods will require that the entire UCAC4/Tycho2 star catalogues be downloaded if they aren't already.

This module also provides a few functions that can be used to build a new version of this catalogue,
:func:`build_catalogue`, :func:`find_star_pairs`, and :func:`blend_stars`.  Typically you won't interact with these
directly and instead will use the script :mod:`~.scripts.build_catalogue` which provides a command line interface,
however they're provided for those who are interested in what they do or in doing some more advanced things.
"""

from pathlib import Path
import time

from itertools import repeat, starmap

from warnings import filterwarnings, catch_warnings

import sqlite3

from datetime import datetime

from typing import Optional, Union

import numpy as np
import pandas as pd
from scipy.spatial import cKDTree

from giant.catalogues.utilities import radec_to_unit, apply_proper_motion, radec_distance, DEG2RAD
from giant.catalogues.meta_catalogue import Catalogue
from giant._typing import PATH, Real, ARRAY_LIKE

DEFAULT_CAT_FILE: Path = (Path(__file__).parent / "data") / 'giant_cat.db'
"""
This specifies the default location of the sqlite3 database file that contains the data for this catalogue.

This is stored as Path object.  It defaults to a file called "giant_cat.db" in a "data" directory in the same directory
containing this file. If you wish to use a different catalogue, typically you should simply provide a key
word argument to the class constructor instead of modifying this module attribute.
"""

_STARS_TABLE_SQL: str = '''CREATE TABLE IF NOT EXISTS "stars" (
  "rnm" INTEGER UNIQUE PRIMARY KEY NOT NULL ON CONFLICT REPLACE,
  "source" TEXT,
  "zone" INTEGER,
  "rnz" INTEGER,
  "ra" REAL,
  "dec" REAL,
  "distance" REAL,
  "ra_proper_motion" REAL,
  "dec_proper_motion" REAL,
  "mag" REAL,
  "ra_sigma" REAL,
  "dec_sigma" REAL,
  "distance_sigma" REAL,
  "ra_pm_sigma" REAL,
  "dec_pm_sigma" REAL,
  "epoch" REAL
)'''
"""
This SQL command creates a new table in the sqlite3 database file called stars for storing the GIANT catalogue

This shouldn't be used by the user and may change without notice.
"""

def build_catalogue(database_file: Optional[PATH] = None, limiting_magnitude: Real = 12, number_of_stars: int = 0,
                    use_tycho_mag: bool = False, limiting_separation: float = 0.04, blending_magnitude: Real = 8,
                    ucac_dir: Optional[PATH] = None):
    """
    Build a sqlite3 catalogue from the UCAC catalogue for faster query times.

    This function can be used to build a new GIANT catalogue (or overwrite an old one) from the UCAC4/Tycho2 star
    catalogues.  Typically a user will not use this directly and instead will use the command line utility
    :mod:`~.scripts.build_catalogue`.

    If you want to use this function, you can adjust how (and where) the catalogue is built by adjusting the key word
    arguments.  Note that this will require downloading the UCAC4 catalogue (if it isn't already available) and possibly
    the Tycho2 Catalogue.  In addition, building the catalogue can take a long time, so you will want to have a period
    where you can leave this run for a while without interruption to ensure that the catalogue is built successfully and
    not corrupted.

    :param database_file: The file to save the catalogue database to
    :param limiting_magnitude: The maximum magnitude to include in the catalogue
    :param number_of_stars: The maximum number of stars that can be blended together in any group.  To turn off star
                            blending, set this to 0.
    :param use_tycho_mag: A flag specifying whether to replace the UCAC4/APASM_V magnitudes with the Tycho VT magnitude.
                          The Tycho VTMag is more accurate, but this can make things take even longer to compile so by
                          default it is not used.
    :param limiting_separation: The maximum separation between stars for them to be considered for blending in degrees.
                                Typically this should be set to around the IFOV on the detector you are considering.
    :param blending_magnitude: The magnitude of the blended star for it to be included as a blended star in the
                               catalogue.
    :param ucac_dir: The directory containing the UCAC4 data files.  This is passed to the :class:`.UCAC4` class.
    """

    with catch_warnings():
        # ignore dumb warnings
        filterwarnings('ignore', message='The requested UCAC4')

        # import the UCAC4 interface and default directory
        from giant.catalogues.ucac import UCAC4, UCAC_DIR

        # use the default if the user didn't specify
        if database_file is None:
            database_file = DEFAULT_CAT_FILE
        else:
             database_file = Path(database_file)

        # make sure the directory for the database file exists
        database_file.parent.mkdir(exist_ok=True, parents=True)

        if database_file.exists():
            database_file.unlink()

        # connect to the database
        db_con = sqlite3.connect(str(database_file))

        # get the UCAC instance
        if ucac_dir is None:
            ucac = UCAC4(UCAC_DIR)
        else:
            ucac = UCAC4(ucac_dir)

        # determine the magnitude to query from the catalogue
        if number_of_stars == 0:
            blend_mag = -4.0
            query_mag = limiting_magnitude
        else:
            blend_mag = blending_magnitude + (2.5 * np.log10(number_of_stars))
            query_mag = max(limiting_magnitude, blend_mag)

        print('Query mag {}'.format(query_mag), flush=True)

        # Create the table/indices in the database
        db_con.execute("DROP TABLE IF EXISTS stars")
        db_con.execute(_STARS_TABLE_SQL)
        db_con.execute("CREATE UNIQUE INDEX idx_rnm on stars(rnm)")
        db_con.execute("CREATE INDEX idx_ra on stars(ra)")
        db_con.execute("CREATE INDEX idx_dec on stars(dec)")
        db_con.execute("CREATE INDEX idx_mag on stars(mag)")
        db_con.commit()

        print('creating database', flush=True)
        if number_of_stars == 0:
            # dump the UCAC4 stars to the database
            ucac.dump_to_sqlite(db_con, limiting_mag=query_mag, use_tycho_mag=use_tycho_mag,
                                return_locations=False)
        else:
            # dump the UCAC4 stars to the database
            records = ucac.dump_to_sqlite(db_con, limiting_mag=query_mag, use_tycho_mag=use_tycho_mag,
                                          return_locations=True, return_mag=blend_mag)

            # pair the stars based on distance
            print('finding close star pairs', flush=True)
            pairs = find_star_pairs(records, limiting_separation)

            # get rid of the old dataframe for memory reasons
            del records

            # blend the stars together
            print('blending stars', flush=True)
            combined_stars = blend_stars(pairs, db_con, limiting_magnitude)

            # rename the index column
            combined_stars.index.name = 'rnm'

            # dump to the stars table in the database
            print('adding blended stars to db', flush=True)
            combined_stars.to_sql('stars', db_con, if_exists='append')

def _repair(star_records: pd.DataFrame, pairs: pd.DataFrame) -> pd.DataFrame:
    """
    This helper function combines multiple pairs and gets rid of duplicates

    Don't use this yourself.

    :param star_records: The dataframe of star records
    :param pairs: The dataframe specifying groups of stars
    :return: The dataframe specifying groups of stars after correcting the groupings
    """
    # get the unique right hand sides
    unique_others = pairs.b.unique()

    paired_dict: pd.DataFrame = pairs.groupby('a')['b'].apply(set)
    removes = []
    for initial, others in paired_dict.iteritems():
        sets = []
        starts = []
        # look for where initial is also paired to another star
        if (initial in unique_others) and (initial not in removes):
            for local_initial, local_others in paired_dict.iteritems():
                if local_initial == initial:
                    continue
                elif initial in local_others:
                    # if the others is a subset of the first group we don't need to do anything
                    if others.issubset(local_others):
                        # should we also remove local_initial here?
                        continue
                    # otherwise store them for use
                    sets.append(local_others)
                    starts.append(local_initial)
            # if anything needs modified
            if starts:
                for local_initial, local_others in zip(starts, sets):
                    if len(starts) == 1:
                        # keep which ever has the higher magnitude on the left hand side and discard the other
                        if star_records.loc[local_initial, "mag"] < star_records.loc[initial, "mag"]:
                            local_others.update(others)
                            removes.append(initial)

                        else:
                            others.update(local_others)
                            removes.append(local_initial)

                    else:
                        # keep the brightest magnitude
                        best = initial
                        best_mag = star_records.loc[initial, "mag"]
                        best_set = others
                        for o, ls in zip(starts, sets):
                            if best_mag > star_records.loc[o, "mag"]:
                                best_mag = star_records.loc[o, "mag"]
                                best = o
                                best_set = ls

                        # get rid of whichever aren't the brightest and feed it
                        if best != initial:
                            removes.append(initial)
                            best_set.update(others)
                        for local_local_initial, local_local_others in zip(starts, sets):
                            if local_local_initial == best:
                                continue
                            best_set.update(local_local_others)
                            removes.append(local_local_initial)

    # get rid of the bad ones
    return paired_dict.drop(removes)

def find_star_pairs(star_records: pd.DataFrame, max_separation: float) -> pd.DataFrame:
    """
    This identifies possible star pairs based on separation.

    Stars are paired if their max separation is less that the input ``max_separation`` in degrees. This is done by
    creating unit vectors for all of the stars and then doing a pair query using a KDTree.  The pairs are sorted based
    on magnitude so that the first star in each pair is brighter.

    The result of this function will be a dataframe where the first column "a" is the primary star and the second column
    "b" is a set of stars that should be combined with "a".

    Generally this is not used directly by the user.  Instead see :func:`build_catalogue` or script
    :mod:`~.scripts.build_catalogue`.

    :param star_records: The dataframe containing the stars that are to be paired
    :param max_separation: The maximum separation in degrees between stars for them to be paired
    :return: A dataframe specifying stars to pair together.
    """

    # get the unit vectors
    units = radec_to_unit(star_records.ra.values * DEG2RAD, star_records.dec.values * DEG2RAD).T

    # build the kdtree
    # noinspection PyArgumentList
    tree = cKDTree(units, compact_nodes=False, balanced_tree=False)

    # find the pairs.  Tell pycharm to stop complaining because numpy/scipy don't document right
    # noinspection PyArgumentList,PyUnresolvedReferences
    pairs = tree.query_pairs(np.sin(max_separation * np.pi / 360) * 2, output_type='ndarray')

    # get the pairs
    pairs = pd.DataFrame(star_records.index.values[pairs], columns=['a', 'b'])

    # sort the pairs on magnitude
    for pair in pairs.itertuples():

        if float(star_records.loc[pair.a, "mag"]) > float(star_records.loc[pair.b, "mag"]):
            pairs.loc[pair.Index, "a"] = pair.b
            pairs.loc[pair.Index, "b"] = pair.a

    # condense pairs so that stars aren't in multiple pairs
    return _repair(star_records, pairs)

def _blend_stars(input_group: tuple, database_connection: sqlite3.Connection,
                 index: int, limiting_mag: Real, reference_mag: Real) -> Optional[pd.Series]:
    """
    This helper function computes a blended star from the input star indices.

    This function queries the star records from the database for memory reasons (so that we can use multiprocessing).

    The stars are blended into a single record with a combined magnitude, right ascension, declination, and proper
    motion (distance is not considered).  This is based off of an internal note on blending stars that Sean Semper sent.

    :param input_group: The group of stars to be blended
    :param database_connection: The database connection to retrieve the star records from
    :param index: The index of the group of stars that we're working on.  Purely for printing purposed
    :param limiting_mag: The magnitude that the blended stars must reach for them to be included
    :param reference_mag: The reference magnitude to use when blending the stars.
    :return: A series with the blended star, or ``None`` if the limiting magnitude wasn't met
    """

    # interpret the group
    initial = input_group[0]
    group = list(input_group[1])

    start = time.time()

    # all the ids we need to query
    star_ids = [initial] + group
    # the security risk is minimized here by calling int on all of the elements in the star_ids list
    # its already minimal because a user should never be directly interacting with this function anyway
    star_records = pd.read_sql('select * from stars where rnm in {}'.format(tuple(map(int, star_ids))),  # nosec
                               database_connection, index_col='rnm')

    # get the initial star from the dataframe
    initial_star = star_records.loc[initial]
    # get the rest of the stars from the dataframe
    other_stars = star_records.loc[group]

    # compute the weights for each star
    initial_weight = 1 / (10 ** (0.4 * (initial_star.mag - reference_mag)))
    other_weights = 1 / (10 ** (0.4 * (other_stars.mag.values - reference_mag)))

    # determine the reference declination from the brightest star
    if (initial_star.mag <= other_stars.mag).all():
        ref_dec = np.cos(initial_star.dec*DEG2RAD)
    else:
        ref_dec = np.cos(other_stars.loc[other_stars.mag == other_stars.mag.min(), "dec"].values[0]*DEG2RAD)

    # compute the combined magnitude
    combined_mag = -2.5 * np.log10((10 ** (-0.4 * initial_star.mag) + 10 ** (-0.4 * other_stars.mag.values)).sum())

    # if the blended star is too dim stop here
    if combined_mag > limiting_mag:
        return None

    # make the Series to return
    combined_star = initial_star.copy()
    # set the blended magnitude
    combined_star.mag = combined_mag
    # update the RNM to be negative
    combined_star.name *= -1

    # compute the combined position
    denominator = initial_weight + other_weights.sum()
    combined_star.ra = (initial_weight * initial_star.ra * ref_dec +
                        (other_weights * other_stars.ra.values * ref_dec).sum()) / denominator / ref_dec
    combined_star.dec = (initial_weight * initial_star.dec +
                         (other_weights * other_stars.dec.values).sum()) / denominator
    combined_star.ra_proper_motion = ((initial_weight * initial_star.ra_proper_motion * ref_dec +
                                      (other_weights * other_stars.ra_proper_motion.values * ref_dec).sum()) /
                                      denominator / ref_dec)
    combined_star.dec_proper_motion = ((initial_weight * initial_star.dec_proper_motion +
                                        (other_weights * other_stars.dec_proper_motion.values).sum()) / denominator)

    # give a status
    print('Pair {} blended in {:.3f}'.format(index, time.time() - start))

    return combined_star

def blend_stars(groups: pd.DataFrame, database_connection: sqlite3.Connection, limiting_mag: Real,
                ref_mag: Real = 4) -> pd.DataFrame:
    """
    Blends groups of stars together into a single "apparent" star as viewed by a camera.

    Star magnitude, right ascension, declination, and proper motion are all blended in the final product.  The blending
    is based off of an internal memo by Sean Semper.

    The groups input should provide 2 columns, the first column "a" should provide the primary (brightest) star in each
    group.  The second column "b" should provide a set of all of the stars that are to be blended to each other an "a".
    This is what is returned by :func:`find_star_pairs`.  This function uses the database to retrieve the individual
    star records for memory purposes.

    The blended star is given an id that is the negative of the brightest star in the group.  The blended stars are
    returned as a pandas dataframe.

    Typically this is not used directly by the user.  Instead se :func:`build_catalogue` or script
    :mod:`.scripts.build_catalogue`.

    :param groups: The dataframe specifying the groups to blend
    :param database_connection: The connection to the sqlite3 database to retrieve the stars from
    :param limiting_mag: The limiting magnitude that blended stars must achieve for them to be included
    :param ref_mag: The reference magnitude to use when blending the stars
    :return: The dataframe of the blended apparent stars
    """

    # get the number of groups we need to blend for reporting purposes
    number_of_groups = len(groups)

    # notify the user
    print('{} stars to blend'.format(number_of_groups), flush=True)

    # combine the stars.  Perhaps can use multiprocessing for this
    combined_stars = list(starmap(_blend_stars,
                                  zip(groups.iteritems(),
                                      repeat(database_connection),
                                      range(1, number_of_groups + 1),
                                      repeat(limiting_mag),
                                      repeat(ref_mag))))

    # return the dataframe by concatenating the individual blended series.  This will ignore Nones
    return pd.concat(combined_stars, axis=1).T

class GIANTCatalogue(Catalogue):
    """
    This class provides access to the default GIANT star catalogue built from the UCAC4 and Tycho2 Catalogues.

    This class is a fully functional catalogue for GIANT and can be used anywhere that GIANT expects a star catalogue.
    As such, it implements the :attr:`include_proper_motion` to turn proper motion on or off as well as the method
    :meth:`query_catalogue` which is how stars are queried into the GIANT format.  In addition, this catalogue provides
    2 additional methods :meth:`get_ucac4_record` and :meth:`get_tycho2_record` to get the original record that was
    used to create the GIANT catalogue record.  These methods aren't used anywhere by GIANT itself, but may be useful if
    you are doing some advanced analysis.

    To use this class simply initialize it, pointing to the file where the database is stored (if you are using one that
    is different from the default).  If the catalogue file does not exist it will ask you if you want to build it, and
    if you answer yes, it will dispatch to :func:`build_catalogue` (which takes a long time in most instances).  Once
    the class is initialized, you can query stars from it using :meth:`query_catalogue` which will return a dataframe of
    the star records with :attr:`.GIANT_COLUMNS` columns.
    """

    def __init__(self, db_file: PATH = DEFAULT_CAT_FILE, include_proper_motion: bool = True):
        """
        :param db_file: The file containing the sqlite3 database that the stars are stored in
        :param include_proper_motion: A boolean flag specifying whether to apply proper motion when retrieving the stars
        """

        super().__init__(include_proper_motion=include_proper_motion)

        self.catalogue_path: Path = Path(db_file)
        """
        The path to the catalogue file containing the database
        """

        self._catalogue: Optional[sqlite3.Connection] = None
        """
        The sqlite3 catalogue connection
        """

        if db_file.exists():
            try:
                self._catalogue = sqlite3.connect(str(self.catalogue_path))
                self._catalogue.execute("SELECT * FROM stars LIMIT 1")
            except (sqlite3.OperationalError, sqlite3.DatabaseError):
                print("GIANT catalogue corrupted at {}".format(db_file), flush=True)
                user_response = input("Would you like to build the GIANT catalogue from the UCAC catalogue (y/n)?\n"
                                      "    WARNING: THIS MAY TAKE A LONG TIME AND WILL USE UP 600 MB OF SPACE!\n    ")

                if user_response[:1].lower() == 'y':
                    build_catalogue(db_file)

                    self._catalogue = sqlite3.connect(str(self.catalogue_path))

                else:
                    raise sqlite3.DatabaseError('The GIANT catalogue database file is corrupted')

        else:
            print("GIANT catalogue not found at {}".format(db_file), flush=True)
            user_response = input("Would you like to build the GIANT catalogue from the UCAC catalogue (y/n)?\n"
                                  "    WARNING: THIS MAY TAKE A LONG TIME AND WILL USE UP 600 MB OF SPACE!\n    ")

            if user_response[:1].lower() == 'y':
                build_catalogue(db_file)

                self._catalogue = sqlite3.connect(str(self.catalogue_path))

            else:
                raise FileNotFoundError('The GIANT catalogue database file cannot be found')

    def __reduce__(self):
        return self.__class__, (self.catalogue_path, self.include_proper_motion)

    @property
    def catalogue(self) -> sqlite3.Connection:
        """
        This is a sqlite3 connection object which is used to read from the catalogue.

        It should not be used externally unless you really know what you're doing...
        """

        return self._catalogue

    @catalogue.setter
    def catalogue(self, err):
        raise AttributeError('You cannot set the catalogue directly.  It is purely for internal use.\n' +
                             'If you really know what you are doing, you can set the catalogue file by accessing\n' +
                             'the _catalogue attribute, however this is highly warned against.')

    @catalogue.deleter
    def catalogue(self):
        self._catalogue.close()

    def query_catalogue(self, ids: Optional[ARRAY_LIKE] = None, min_ra: Real = 0, max_ra: Real = 360,
                        min_dec: Real = -90, max_dec: Real = 90, min_mag: Real = -4, max_mag: Real = 20,
                        search_center: Optional[ARRAY_LIKE] = None, search_radius: Optional[Real] = None,
                        new_epoch: Optional[Union[datetime, Real]] = None) -> pd.DataFrame:
        """
        This method queries stars from the catalogue that meet specified constraints and returns them as a DataFrame
        with columns of :attr:`.GIANT_COLUMNS`.

        Stars can either be queried by ID directly or by right ascension/declination/magnitude. You cannot filter using
        both with this method.  If :attr:`apply_proper_motion` is ``True`` then this will shift the stars to the new
        epoch input by the user (``new_epoch``) using proper motion.

        :param ids: A sequence of star ids to retrieve from the catalogue.  What these ids are vary from catalogue to
                    catalogue so see the catalogue documentation for details.
        :param min_ra: The minimum ra bound to query stars from in degrees
        :param max_ra: The maximum ra bound to query stars from in degrees
        :param min_dec: The minimum declination to query stars from in degrees
        :param max_dec: The maximum declination to query stars from in degrees
        :param min_mag: The minimum magnitude to query stars from.  Recall that magnitude is inverse (so lower
                        magnitude is a dimmer star)
        :param max_mag: The maximum magnitude to query stars from.  Recall that magnitude is inverse (so higher
                        magnitude is a dimmer star)
        :param search_center: The center of a search cone as a ra/dec pair.
        :param search_radius: The radius about the center of the search cone
        :param new_epoch: The epoch to translate the stars to using proper motion if :attr:`apply_proper_motion` is
                          turned on
        :return: A Pandas dataframe with columns :attr:`GIANT_COLUMNS`.
        """

        if ids is not None:
            records = self.get_from_ids(ids)
        else:

            records = self.get_all_with_criteria(min_ra=min_ra, max_ra=max_ra, min_dec=min_dec, max_dec=max_dec,
                                                 min_mag=min_mag, max_mag=max_mag,
                                                 search_center=search_center, search_radius=search_radius)

        if self.include_proper_motion and (new_epoch is not None):
            apply_proper_motion(records, new_epoch, copy=False)

        return records

    def get_from_ids(self, ids: ARRAY_LIKE) -> pd.DataFrame:
        """
        This method queries star records from the database based off of ID (``rnm`` in the database).

        This can be used if you are interested in a particular set of stars.  The stars are returned in the GIANT
        DataFrame format according the :attr:`.GIANT_COLUMNS`.

        Note that this does not apply proper motion.  If you need to apply proper motion see :meth:`query_catalogue`

        :param ids: The ids of the stars to retrieve as an iterable
        :return: The dataframe of stars according to :attr:`.GIANT_COLUMNS`
        """

        # map(int, ids) protects against sql injection
        return pd.read_sql(f'select * from stars where rnm in {tuple(map(int, ids))}', self._catalogue,  # nosec
                           index_col='rnm')

    def get_all_with_criteria(self, min_ra: Real = 0, max_ra: Real = 360,
                              min_dec: Real = -90, max_dec: Real = 90, min_mag: Real = -4, max_mag: Real = 20,
                              search_center: Optional[ARRAY_LIKE] = None,
                              search_radius: Optional[Real] = None) -> pd.DataFrame:
        """
        This method queries star records from the database based off of location and magnitude requirements.

        Note that this does not apply proper motion.  If you need to apply proper motion see :meth:`query_catalogue`.

        :param min_ra: The minimum ra bound to query stars from in degrees
        :param max_ra: The maximum ra bound to query stars from in degrees
        :param min_dec: The minimum declination to query stars from in degrees
        :param max_dec: The maximum declination to query stars from in degrees
        :param min_mag: The minimum magnitude to query stars from.  Recall that magnitude is inverse (so lower
                        magnitude is a dimmer star)
        :param max_mag: The maximum magnitude to query stars from.  Recall that magnitude is inverse (so higher
                        magnitude is a dimmer star)
        :param search_center: The center of a search cone as a ra/dec pair.
        :param search_radius: The radius about the center of the search cone
        :return: A Pandas dataframe with columns :attr:`GIANT_COLUMNS`.
        """

        # make sure everything is a float (a) to validate input and (b) to protect against sql injection attacks
        min_ra = float(min_ra)
        max_ra = float(max_ra)
        min_dec = float(min_dec)
        max_dec = float(max_dec)
        min_mag = float(min_mag)
        max_mag = float(max_mag)

        # determine what the rectangular bounds should look like for the search center/radius
        if search_center is not None:
            min_ra = search_center[0] - search_radius
            max_ra = search_center[0] + search_radius

            min_dec = search_center[1] - search_radius
            max_dec = search_center[1] + search_radius

        # adjust for if we are at a corner case
        if min_dec < -90:
            min_dec = -90
            min_ra = 0
            max_ra = 360

        elif max_dec > 90:
            max_dec = 90
            min_ra = 0
            max_ra = 360

        # determine what the query should look like based on the rectangular bounds
        # the security risk is non-existent here due to calling float on all of the values before they are used as
        # parameters
        if min_ra < 0:
            query = ('SELECT * FROM stars WHERE ((ra >= {} AND ra <= {}) OR (ra >= {} AND ra <= {})) '   # nosec
                     'AND dec >= {} AND dec <= {} AND mag <={} '
                     'AND mag >= {}'.format(min_ra + 360, 360, 0, max_ra, min_dec, max_dec, max_mag, min_mag))

        elif max_ra > 360:
            query = ('SELECT * FROM stars WHERE ((ra >= {} AND ra <= {}) OR (ra >= {} AND ra <= {})) '   # nosec
                     'AND dec >= {} AND dec <= {} AND mag <={} '
                     'AND mag >= {}'.format(min_ra, 360, 0, max_ra - 360, min_dec, max_dec, max_mag, min_mag))

        else:
            query = ('SELECT * FROM stars WHERE ra >= {} AND ra <= {} '  # nosec
                     'AND dec >= {} AND dec <= {} AND mag <={} '
                     'AND mag >= {}'.format(min_ra, max_ra, min_dec, max_dec, max_mag, min_mag))

        # query the results from the database
        records = pd.read_sql(query, self.catalogue, index_col='rnm')

        # now do the real radial search if it is needed
        if search_center is not None:
            records = records.loc[radec_distance(records.ra * DEG2RAD, records.dec * DEG2RAD,
                                                 search_center[0] * DEG2RAD, search_center[1] * DEG2RAD) <=
                                  (search_radius * DEG2RAD)]
        if "epoch" not in records.columns:
            records = records.assign(epoch=2000.0)

        return records

    @staticmethod
    def get_tycho2_record(stars: pd.DataFrame, ucac_directory: Optional[PATH] = None,
                          tycho_directory: Optional[PATH] = None) -> pd.DataFrame:
        """
        This method can be used to retrieve the corresponding full (not GIANT) Tycho 2 records for a set of GIANT
        catalogue stars.

        This method requires that the UCAC4 and Tycho 2 catalogues be available, and will request to download them if
        they are not available (which takes a long time).  For a description of the columns refer to the Tycho 2
        documentation.

        Note that these records are not directly usable for GIANT, therefore only use this if you need the raw records
        yourself.

        Any stars that are not available in the Tycho 2 catalogue (blended stars or stars that are too dim) will be
        included in the output dataframe but with NANs for all columns

        :param stars: The stars to retrieve the Tycho records for
        :param ucac_directory: The directory containing the UCAC4 star catalogue files (or None to use the default)
        :param tycho_directory: The directory containing the Tycho 2 star catalogue files (or None to use the default)
        :return: The raw Tycho 2 records
        """

        from giant.catalogues.ucac import UCAC4
        from giant.catalogues.tycho import Tycho2

        ucac = UCAC4(directory=ucac_directory)

        tycho = Tycho2(directory=tycho_directory)

        return ucac.cross_ref_tycho(stars.loc[:, ['zone', 'rnz']], tycho_cat=tycho)

    @staticmethod
    def get_ucac4_record(stars: pd.DataFrame, ucac_directory: Optional[PATH] = None) -> pd.DataFrame:
        """
        This method can be used to retrieve the corresponding full (not GIANT) UCAC4 records for a set of GIANT
        catalogue stars.

        This method requires that the UCAC4 catalogue be available, and will request to download it if it is not (which
        takes a long time).  For a description of the columns refer to the UCAC4 documentation.

        Note that these records are not directly usable for GIANT, therefore only use this if you need the raw records
        yourself.

        Any stars that are not available in the UCAC4 catalogue (blended stars) will be included in the output dataframe
        but with NANs for all columns

        :param stars: The UCAC4 records for the stars
        :param ucac_directory: The directory containing the UCAC4 star catalogue files (or None to use the default)
        :return: The raw UCAC4 records
        """

        from giant.catalogues.ucac import UCAC4

        ucac = UCAC4(directory=ucac_directory)

        return ucac.query_catalogue_raw(ids=stars.loc[:, ['zone', 'rnz']].itertuples(False), generator=False)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# pack and unpack with mask function from: https://github.com/kevinkwl/AoAReader/blob/master/aoareader/AoAReader.py

# @Time    : 2017/12/6 11:13
# @From    : PyCharm
# @File    : hiso
# @Author  : Liujie Zhang
# @Email   : liujiezhangbupt@gmail.com
import json
import pickle
import torch
import torch.nn as nn
import numpy as np
from torch.autograd import Variable
import torch.nn.functional as F
from gensim.models import Word2Vec
from torch.nn.utils.rnn import pad_packed_sequence as unpack
from torch.nn.utils.rnn import pack_padded_sequence as pack

from utils.attention import Attention, cudaWrapper

def sort_batch(data, seq_len):
    sorted_seq_len, sorted_idx = torch.sort(seq_len, dim=0, descending=True)
    sorted_data = data[sorted_idx.data]
    _, reverse_idx = torch.sort(sorted_idx, dim=0, descending=False)

    return sorted_data, cudaWrapper(sorted_seq_len), cudaWrapper(reverse_idx)

def create_mask(seq_lens):
    mask = torch.zeros(seq_lens.data.size(0), torch.max(seq_lens.data))
    for i, seq_len in enumerate(seq_lens.data):
        mask[i][:seq_len] = 1

    return cudaWrapper(mask.float())

def getSeqLength(input, return_variable=True):
    seq_len = torch.LongTensor([torch.nonzero(input.data[i]).size(0) for i in range(input.data.size(0))])
    if return_variable:
        seq_len = Variable(seq_len)
    return cudaWrapper(seq_len)

class HISO(nn.Module):
    def __init__(self, opt):
        super(HISO, self).__init__()

        self.model_name = 'HISO'
        self.opt = opt
        # Embedding Layer
        self.wd_embed = nn.Embedding(opt.voc_size, opt.embed_dim)
        self.pos_embed = nn.Embedding(opt.pos_size, opt.embed_dim)
        self.initEmbedWeight()
        # Bi-GRU Layer
        self.wd_bi_gru = nn.GRU(input_size = opt.embed_dim,
                hidden_size = opt.ghid_size,
                num_layers = opt.glayer,
                bias = True,
                batch_first = True,
                dropout = 0.5,
                bidirectional = True
                )
        self.attention = Attention(opt)

        # Bi-GRU Layer
        self.pos_bi_gru = nn.GRU(input_size = opt.embed_dim,
                hidden_size = opt.ghid_size,
                num_layers = opt.glayer,
                bias = True,
                batch_first = True,
                dropout = 0.5,
                bidirectional = True
                )

        # output from pos hidden layer to predict middle labels
        pos_hidden_size = opt.ghid_size * opt.glayer
        self.pos_fc = nn.Sequential(
                nn.BatchNorm1d(pos_hidden_size),
                nn.ReLU(inplace=True),
                nn.Linear(pos_hidden_size, opt.auxiliary_labels),
                nn.Sigmoid()
                )
        # predict final labels
        combine_size = opt.ghid_size * opt.glayer + opt.auxiliary_labels
        self.fc = nn.Sequential(
                nn.Linear(combine_size, 128),
                nn.BatchNorm1d(128),
                nn.ReLU(inplace = True),
                nn.Linear(128, opt.label_dim),
                nn.Sigmoid()
                )
        self.softmax = nn.Softmax(dim=1)

    def initEmbedWeight(self):
        '''
        init embedding layer from random|word2vec|sswe
        '''
        if 'w2v' in self.opt.init_embed:
            weights = Word2Vec.load('../docs/data/w2v_word_100d_5win_5min')
            voc = json.load(open('../docs/data/voc.json','r'))['voc']
            print(weights[list(voc.keys())[3]])

            word_weight = np.zeros((len(voc),self.opt.embed_dim))
            for wd,idx in voc.items():
                vec = weights[wd] if wd in weights else np.random.randn(self.opt.embed_dim)
                word_weight[idx] = vec
            # print(word_weight[3])
            self.wd_embed.weight.data.copy_(torch.from_numpy(word_weight))

            weights = Word2Vec.load('../docs/data/w2v_pos_100d_5win_5min')
            pos = json.load(open('../docs/data/pos.json','r'))['voc']
            pos_weight = np.zeros((len(pos),self.opt.embed_dim))
            for ps,idx in pos.items():
                vec = weights[ps] if ps in weights else np.random.randn(self.opt.embed_dim)
                pos_weight[idx] = vec
            self.pos_embed.weight.data.copy_(torch.from_numpy(pos_weight))

        elif 'sswe' in self.opt.init_embed:
            word_weight = pickle.load(open('../docs/model/%s'% self.opt.embed_path,'rb'))
            self.wd_embed.weight.data.copy_(torch.from_numpy(word_weight))
        # random default

    def forward(self, wd, pos):
        wd_len, pos_len = getSeqLength(wd), getSeqLength(pos)
        wd_mask, pos_mask = create_mask(wd_len), create_mask(pos_len)

        s_wd, s_wd_len, reverse_wd_idx = sort_batch(wd, wd_len)
        s_pos, s_pos_len, reverse_pos_idx = sort_batch(pos, pos_len)

        wd_embedding = pack(self.wd_embed(s_wd), list(s_wd_len.data), batch_first=True)
        pos_embedding = pack(self.pos_embed(s_pos), list(s_pos_len.data), batch_first=True)

        # Bi-GRU
        wd_out, _ = self.wd_bi_gru(wd_embedding)
        pos_out, _ = self.pos_bi_gru(pos_embedding)

        wd_out, _ = unpack(wd_out, batch_first=True)
        pos_out,_ = unpack(pos_out, batch_first=True)

        wd_out = wd_out[reverse_wd_idx.data]
        pos_out = pos_out[reverse_pos_idx.data]

        # attention
        if 'word' in self.opt.attention:
            wd_atten = self.attention(wd_out, weight=wd_out, mask=wd_mask)
        elif 'pos' in self.opt.attention:
            wd_atten = self.attention(wd_out, weight=pos_out, mask=pos_mask)
        else:
            wd_atten = wd_out[:,-1,:]

        # pos_out to predict auxiliary label
        auxi_probs = self.pos_fc(pos_out[:, -1, :].contiguous())

        # combine wd_out with auxi_probs as feature
        combine_feature = torch.cat((wd_atten, auxi_probs), dim=1)
        logits = self.fc(combine_feature)

        return logits, auxi_probs

    def init_hidden(self, batch_size):
        h0 = torch.zeros(self.opt.glayer, batch_size, self.opt.ghid_size)
        return Variable(h0)

class HisoLoss(nn.Module):
    def __init__(self, opt):
        super(HisoLoss, self).__init__()
        self.opt = opt
        # self.reconstruction_loss = 0 // todo

    def forward(self,auxi_probs, auxi_labels, final_probs, final_labels):
        # calcu auxi_labels margin loss
        self.auxi_loss = self.marginLoss(auxi_probs, auxi_labels)

        # calcu final_labels margin loss
        self.final_loss = self.marginLoss(final_probs, final_labels)

        self.loss = self.opt.loss_alpha * self.auxi_loss + self.final_loss
        return self.loss

    def marginLoss(self, probs, labels):

        left = F.relu(self.opt.max_margin - probs, inplace=True)**2
        right = F.relu(probs - self.opt.min_margin, inplace=True)**2

        margin_loss = labels * left + (1. - labels) * right
        return margin_loss.sum() / labels.size(0)

class opt(object):
    voc_size = 100
    pos_size = 57
    embed_dim = 20
    ghid_size = 3
    seq_len = 4
    glayer = 2
    auxiliary_labels = 3
    label_dim = 6
    max_margin = 0.9
    min_margin = 0.1
    embed_path='lookup_01-22-19:10'
    init_embed='randn'
    loss_alpha=1e-2
    attention='word'

if __name__ == '__main__':
    import torch.optim as optim

    wd = Variable(torch.LongTensor([[2,45,75,0], [5,54,76,23]]))
    pos = Variable(torch.LongTensor([[3,45,8,0], [13,56,7,43]]))
    labels = Variable(torch.FloatTensor([[1,0,0,1,0,0],[0,0,1,0,1,0]]))
    auxi = Variable(torch.FloatTensor([[1,0,0],[0,1,0]]))

    model = HISO(opt)
    Loss = HisoLoss(opt)
    op = optim.SGD(model.parameters(),lr=0.1)
    model.train()
    for i in range(100):
        final_probs,auxi_probs = model(wd, pos)
        loss = Loss(auxi_probs, auxi, final_probs, labels)
        op.zero_grad()
        loss.backward()
        op.step()
        print(loss.data[0])