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

"""

Build stitched AKARI/IRC-Spitzer/IRS spectral cube

"""

# import logging, sys
# logging.basicConfig(stream=sys.stderr, level=logging.ERROR)
# print(logging.getLogger())
# logging.disable(sys.maxsize) # disable IDL print
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)

from tqdm import tqdm, trange

import os
import numpy as np
from scipy.optimize import curve_fit
from matplotlib.ticker import ScalarFormatter, NullFormatter

## rapyuta
from rapyuta.inout import fclean, fitsext, read_fits, write_fits
from rapyuta.imaging import ( iconvolve, iswarp, iuncert, concatenate, icrop,
                              imontage, improve, Jy_per_pix_to_MJy_per_sr )
from rapyuta.astrom import fixwcs
from rapyuta.arrays import closest
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_fp,
                        path_idl, path_phot, path_cal, path_ker, csv_ker,
                        out_irc, out_irs, path_tmp, path_out, path_fig
)

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

print('        MIRAGE - Mid-infrared cube builder - '+src)

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

# Nmc = 2

phot = ['IRAC2', 'IRAC3']
Nphot = len(phot)

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

##                 Concatenate IRC & IRS

##----------------------------------------------------------
concat_mir = input("Stitch MIR spectra (y/n)? ")
keep_frag = input(" - Keep fragmentary spectra (y/n)? ")

if keep_frag=='y':
    keepfrag = True
    cropedge = False
else:
    keepfrag = False
    crop_edge_frag = input(" - Crop edge if not keeping frag spectra (y/n)? ")
    if crop_edge_frag=='y':
        cropedge = True
    else:
        cropedge = False

if concat_mir=='y':
    concatenate([out_irc,out_irs], path_tmp+src+'_ma', wsort=False)
    data0 = read_fits(path_tmp+src+'_ma').data
    mask_any = np.isnan(data0).any(axis=0)
    mask_all = np.isnan(data0).all(axis=0)

    Nw, Ny, Nx = data0.shape
    for j in trange(Nmc+1,
                    desc='MIR concatenation [MC]'):
        if j==0:
            concatenate([out_irc,out_irs],
                        path_out+src+'_0', wsort=False,
                        keepfrag=keepfrag, cropedge=cropedge)
            concatenate([out_irc+'_unc',out_irs+'_unc'],
                        path_out+src+'_unc', wsort=False,
                        keepfrag=keepfrag, cropedge=cropedge)
        else:
            concatenate([out_irc+'_'+str(j),out_irs+'_'+str(j)],
                        path_out+src+'_'+str(j), wsort=False,
                        keepfrag=keepfrag, cropedge=cropedge)
        hd = read_fits(path_out+src+'_'+str(j))
        data = hd.data
        header = hd.header
        wvl = hd.wave
        for k in range(Nw):
            data[k][mask_all] = np.nan
        write_fits(path_out+src+'_'+str(j), header, data, wvl)
        if j==0:
            unc = read_fits(path_out+src+'_unc').data
            for k in range(Nw):
                unc[k][mask_all] = np.nan
            write_fits(path_out+src+'_unc', header, unc, wvl)

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

##                   Inter-Calibration

##----------------------------------------------------------
pre_calib1 = input("Run inter-calibration prepipeline (DustPedia) (y/n)? ")
pre_calib2 = input("Run inter-calibration prepipeline (SINGS) (y/n)? ")
int_calib = input("Run inter-calibration (y/n)? ")

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

for i in range(Nphot):

    ## Prepare photometry
    ##--------------------
    if phot[i]=='IRAC2':
        phot_ker = path_ker+'Kernel_HiRes_IRAC_4.5_to_Gauss_06.0'
    elif phot[i]=='IRAC3':
        phot_ker = path_ker+'Kernel_HiRes_IRAC_5.8_to_Gauss_06.0'

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

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

    fits_spec = path_cal+src+'_'+phot[i]+'_MIR'
    tmp_spec = path_tmp+'calib/'+src+'_'+phot[i]+'_MIR'

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

        ## Convert phot unit (Jy/pix -> MJy/sr)
        ## This step should be before iswarp reprojection (pixscale changed)
        Jy_per_pix_to_MJy_per_sr(raw_phot1, filOUT=path_tmp+src+'_'+phot[i]+'_DP')

        if phot[i]=='IRAC2':
            ## DustPedia IRAC2 of M82 has no uncertainty map
            ## Create homogeneous uncertainty map
            bg = read_fits(path_tmp+src+'_'+phot[i]+'_DP').data[1720:1760,3110:3150]
            iuncert(path_tmp+src+'_'+phot[i]+'_DP',
                    filOUT=path_tmp+src+'_'+phot[i]+'_DP_unc', BG_image=bg)
        elif phot[i]=='IRAC3':
            Jy_per_pix_to_MJy_per_sr(raw_phot1+'_unc',
                                     filOUT=path_tmp+src+'_'+phot[i]+'_DP_unc')

        ## Reproject phot
        mtg = imontage('exact', tmpdir=path_tmp)
        mtg.reproject_mc(path_tmp+src+'_'+phot[i]+'_DP', refheader=coadd_fp,
                         dist='norm', Nmc=Nmc, filOUT=tmp_phot1)

        ## Convolve phot
        for j in trange(Nmc+1, leave=False,
                        desc=phot[i]+': DP conv [MC]'):
            if j==0:
                conv = iconvolve(tmp_phot1, phot_ker, csv_ker, filOUT=fits_phot1)
            else:
                conv = iconvolve(tmp_phot1+'_'+str(j),
                                 kfile=phot_ker, klist=csv_ker,
                                 filOUT=tmp_phot1+'_'+str(j))
            conv.do_conv(idldir=path_idl)

        ## Calculate unc
        mcimage = []
        for j in range(Nmc+1):
            if j==0:
                hdr = read_fits(fits_phot1).header
            else:
                mcimage.append(read_fits(tmp_phot1+'_'+str(j)).data)
        if Nmc>1:
            mcimage = np.array(mcimage)
            unc = np.nanstd(mcimage, axis=0)
            write_fits(fits_phot1+'_unc', hdr, 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 IRAC2/IRAC3 uncertainty map (y/n)? ")
        # if gen_unc2=='y':
        #     ds_wgt = read_fits(wgt_phot2)
        #     if phot[i]== 'IRAC2':
        #         bg = read_fits(raw_phot2).data[542:582,1102:1142]
        #         bg_wgt = ds_wgt.data[542:582,1102:1142]
        #     elif phot[i]=='IRAC3':
        #         bg = read_fits(raw_phot2).data[402:442,1022:1062]
        #         bg_wgt = ds_wgt.data[402:442,1022:1062]
        #     wfac = .1
        #     iuncert(raw_phot2, filOUT=raw_phot2+'_unc',
        #             filWGT=wgt_phot2, wfac=wfac,
        #             BG_image=bg, BG_weight=bg_wgt)

        ## Reproject phot
        mtg = imontage('exact', tmpdir=path_tmp)
        mtg.reproject_mc(raw_phot2, refheader=coadd_fp,
                         dist='norm', Nmc=Nmc, filOUT=tmp_phot2)

        ## Convolve phot
        for j in trange(Nmc+1,# leave=False,
                        desc=phot[i]+': SINGS conv [MC]'):
            if j==0:
                conv = iconvolve(tmp_phot2, phot_ker, csv_ker, filOUT=fits_phot2)
            else:
                conv = iconvolve(tmp_phot2+'_'+str(j),
                                 kfile=phot_ker, klist=csv_ker,
                                 filOUT=tmp_phot2+'_'+str(j))
            conv.do_conv(idldir=path_idl)

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

        ## Synthetic photometry (spec)
        ##=============================
        mcimage = []
        for j in trange(Nmc+1, leave=False,
                        desc='MIR sythetic photometry [MC]'):
            ic = intercalib(path_out+src+'_'+str(j))
            sp = ic.synthetic_photometry(phot[i])
            if j==0:
                write_fits(fits_spec, ic.hdr, sp.Fnu_filt)
            else:
                ## Cal unc
                mcimage.append(sp.Fnu_filt)
        if Nmc>1:
            # print('Calculating uncertainty cube...')
            mcimage = np.array(mcimage)
            unc = np.nanstd(mcimage, axis=0)
            write_fits(fits_spec+'_unc', ic.hdr, unc)

    ## Fit correlation
    ##-----------------
    ## Data
    data_phot1 = read_fits(fits_phot1).data
    unc_phot1 = read_fits(fits_phot1+'_unc').data
    pix_phot1 = data_phot1[:,:].reshape((-1,))
    pix_phot1_unc = unc_phot1[:,:].reshape((-1,))

    data_phot2 = read_fits(fits_phot2).data
    unc_phot2 = read_fits(fits_phot2+'_unc').data
    data_spec = read_fits(fits_spec).data
    unc_spec = read_fits(fits_spec+'_unc').data
    if i==0:
        d1_phot2 = data_phot2
        u1_phot2 = unc_phot2
        d1_spec = data_spec
        u1_spec = unc_spec
    elif i==1:
        d2_phot2 = data_phot2
        u2_phot2 = unc_phot2
        d2_spec = data_spec
        u2_spec = unc_spec

    if int_calib=='y':

        ## Plot
        pix_phot2 = data_phot2[:,:].reshape((-1,))
        pix_spec = data_spec[:,:].reshape((-1,))
        pix_phot2_unc = unc_phot2[:,:].reshape((-1,))
        pix_spec_unc = unc_spec[:,:].reshape((-1,))

        # xgrid = np.arange(0,1e3,1)
        xgrid = np.logspace(-6,4,10000)

        ## NaNs mask
        mask = ~np.logical_or( np.isnan(pix_spec), np.isnan(pix_phot2) )

        # if i==0:
        #     mask_all = mask
        # else:
        #     mask_all = np.ma.array(mask,
        #                            mask=np.logical_and(mask,mask_all)).mask

        ## S/N ratio
        # print('S/N ('+phot[i]+' - MIR) = \n', pix_spec[mask]/pix_spec_unc[mask])
        # print('S/N ('+phot[i]+' - SINGS) = \n', pix_phot2[mask]/pix_phot2_unc[mask])
        # exit()

        ## SINGS - DP
        p0 = pplot(pix_phot1, pix_phot2,
                   yerr=pix_phot2_unc, xerr=pix_phot1_unc,
                   fmt='.', c='k', ec='r',elw=1,
                   xlog=1, ylog=1, nonposx='clip', nonposy='clip',
                   xlim=(1e-3,1e4), ylim=(1e-3,1e4),
                   xlabel='DustPedia (MJy/sr)', ylabel='SINGS (MJy/sr)',
                   figsize=(8,8), title=src+'_'+phot[i],
                   titlesize=20, labelsize=10, ticksize=10, legendsize=10)
        p0.save(path_cal+'SINGS-DP_'+phot[i])

        ## MIR spec - SINGS
        p = pplot(pix_phot2, pix_spec,
                  yerr=pix_spec_unc, xerr=pix_phot2_unc,
                  fmt='.', c='k', ec='r', legend='upper left',
                  xlog=1, ylog=1, nonposx='clip', nonposy='clip',
                  xlim=(1e-6,1e4), ylim=(1e-6,1e4),
                  xlabel='SINGS (MJy/sr)', ylabel='MIR spec (MJy/sr)',
                  figsize=(8,8), title=src+'_'+phot[i],
                  titlesize=20, labelsize=10, ticksize=10, legendsize=10)

        ## Linear fit
        if i==0:
            # mask_irac2 = pix_spec/pix_phot2>1.e-2
            # mask = np.logical_and( mask, mask_irac3 )

            popt, pcov = curve_fit(f_lin0, pix_spec[mask], pix_phot2[mask],
                                   sigma=pix_phot2_unc[mask])
            calib_gain = popt[0]
            # calib_off = popt[1]
            calib_off = 0.
            print('Inter-Calibration ('+phot[i]+') gain = {:.4}'.format(calib_gain))
            print('Inter-Calibration ('+phot[i]+') offset = {:.4}'.format(calib_off))
            label = phot[i]+': y={0:.4}x+{1:.4}'.format(calib_gain, calib_off)
            p.add_plot(f_lin0(xgrid, *popt), xgrid,
                       c='y', ls='-', label=label)

            # popt, pcov = curve_fit(f_lin0, pix_phot1[mask], pix_spec[mask],)
            #                        # sigma=pix_spec_unc[mask])
            # calib_gain = 1. / popt[0]
            # # calib_off = -popt[1] / popt[0]
            # calib_off = 0.
            # print('Inter-Calibration ('+phot[i]+') gain = {:.4}'.format(calib_gain))
            # print('Inter-Calibration ('+phot[i]+') offset = {:.4}'.format(calib_off))
            # label = phot[i]+': y={0:.4}x+{1:.4}'.format(calib_gain, calib_off)
            # p.add_plot(xgrid, f_lin0(xgrid, *popt),
            #            c='y', ls='-', label=label)
        else:
            # mask_irac3 = pix_spec/pix_phot2>1.e-2
            # mask = np.logical_and( mask, mask_irac3 )

            popt, pcov = curve_fit(f_lin0, pix_spec[mask], pix_phot2[mask],
                                   sigma=pix_phot2_unc[mask])
            calib_gain = popt[0]
            # calib_off = popt[1]
            calib_off = 0.
            print('Inter-Calibration ('+phot[i]+') gain = {:.4}'.format(calib_gain))
            print('Inter-Calibration ('+phot[i]+') offset = {:.4}'.format(calib_off))
            label = phot[i]+': y={0:.4}x+{1:.4}'.format(calib_gain, calib_off)
            p.add_plot(f_lin0(xgrid, *popt), xgrid,
                       c='y', ls='-', label=label)

            # popt, pcov = curve_fit(f_lin0, pix_phot1[mask], pix_spec[mask],)
            #                        # sigma=pix_spec_unc[mask])
            # calib_gain = 1. / popt[0]
            # # calib_off = -popt[1] / popt[0]
            # calib_off = 0.
            # print('Inter-Calibration ('+phot[i]+') gain = {:.4}'.format(calib_gain))
            # print('Inter-Calibration ('+phot[i]+') offset = {:.4}'.format(calib_off))
            # label = phot[i]+': y={0:.4}x+{1:.4}'.format(calib_gain, calib_off)
            # p.add_plot(xgrid, f_lin0(xgrid, *popt),
            #            c='y', ls='-', label=label)

        p.save(path_cal+'intercalib_'+phot[i])

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

##                      Plot spectra

##----------------------------------------------------------
do_plot = input("Plot MIR spectra (y/n)? ")
if do_plot=='y':
    ds1 = read_fits(out_irc+'_0')
    ds2 = read_fits(out_irs+'_0')
    ds = read_fits(path_out+src+'_0', path_out+src+'_unc')
    Nw, Ny, Nx = ds.data.shape
    pp = intercalib(path_out+src+'_0')
    pp.read_filter(phot)
    wcen = pp.wcen
    for x in range(Nx):
        for y in range(Ny):
            if ~np.isnan(ds.data[:,y,x]).any():
            # if ~np.isnan(ds1.data[:,y,x]).any():
            #     if ~np.isnan(ds2.data[:,y,x]).any():
                p = pplot(ds.wave, ds.data[:,y,x], yerr=ds.unc[:,y,x],
                          c='k', lw=.7, ec='r', xlog=1, ylog=1,
                          legend='upper left', figsize=(12,10),
                          title=src+'_('+str(x)+','+str(y)+')',
                          titlesize=20, labelsize=10, ticksize=10, legendsize=10)
                p.add_plot(ds1.wave, ds1.data[:,y,x],
                           c='y', lw=.7, ls='-', zorder=-1, label='IRC raw')
                p.add_plot(ds2.wave, ds2.data[:,y,x],
                           c='y', lw=.7, ls='-', zorder=-1, label='IRS raw')

                p.add_plot(wcen[0], d1_phot2[y,x], yerr=u1_phot2[y,x],
                           c='m', marker='o', ms=10, zorder=100, label=phot[0])
                p.add_plot(wcen[0], d1_spec[y,x], yerr=u1_spec[y,x],
                           c='g', marker='^', ms=10, zorder=101, label='MIR spec-'+phot[0])
                p.add_plot(wcen[1], d2_phot2[y,x], yerr=u2_phot2[y,x],
                           c='orange', marker='o', ms=10, zorder=100, label=phot[1])
                p.add_plot(wcen[1], d2_spec[y,x], yerr=u2_spec[y,x],
                           c='c', marker='^', ms=10, zorder=101, label='MIR spec-'+phot[1])

                ph = ['IRAC1','IRAC4','MIPS1']
                pp.read_filter(ph)
                wc = pp.wcen
                c1 = ['pink','tomato','maroon']
                c2 = ['lightgreen','lime','olive']
                for i,ip in enumerate(ph):
                    data_phot = read_fits(path_cal+src+'_'+ip+'_SINGS').data
                    unc_phot = read_fits(path_cal+src+'_'+ip+'_SINGS_unc').data
                    if i==0:
                        data_spec = read_fits(path_cal+src+'_'+ip+'_IRC').data
                        unc_spec = read_fits(path_cal+src+'_'+ip+'_IRC_unc').data
                    else:
                        data_spec = read_fits(path_cal+src+'_'+ip+'_IRS').data
                        unc_spec = read_fits(path_cal+src+'_'+ip+'_IRS_unc').data
                    p.add_plot(wc[i], data_phot[y,x], yerr=unc_phot[y,x],
                               c=c1[i], marker='o', ms=10, zorder=100, label=ph[i])
                    if i==0:
                        p.add_plot(wc[i], data_spec[y,x], yerr=unc_spec[y,x],
                                   c=c2[i], marker='^', ms=10, zorder=101, label='IRC-'+ph[i])
                    else:
                        p.add_plot(wc[i], data_spec[y,x], yerr=unc_spec[y,x],
                                   c=c2[i], marker='^', ms=10, zorder=101, label='IRS-'+ph[i])

                xtic = [2, 3, 4, 5, 6, 7, 8, 10, 12, 15, 20, 30, 40]
                # xtic = []
                xtic_min = np.arange(2., 41., 1.)
                p.ax.set_xticks(xtic, minor=False) # major
                p.ax.set_xticks(xtic_min, minor=True) # minor
                p.ax.xaxis.set_major_formatter(ScalarFormatter()) # major
                p.ax.xaxis.set_minor_formatter(NullFormatter()) # minor

                p.save(path_fig+src+'_('+str(x+1)+','+str(y+1)+')', transparent=1)

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

##                     Final cut spectra

##----------------------------------------------------------
final_cut = input("MIR spectral map final cut with NaN edge, manuel mask, etc. (y/n)? ")
if final_cut=='y':
    if not os.path.exists(path_out+'final/'):
        os.makedirs(path_out+'final/')

    imp = improve(path_out+src+'_0')
    im = imp.im
    imp_unc = improve(path_out+src+'_unc')
    unc = imp_unc.im
    Nw, Ny, Nx = imp.im.shape
    for x in range(Nx):
        for y in range(Ny):
            if np.isnan(im[:,y,x]).any():
                imp.im[:,y,x] = np.nan
                imp_unc.im[:,y,x] = np.nan
    xlist = []
    for x in range(Nx):
        if not np.isnan(imp.im[:,:,x]).all():
            xlist.append(x)
    ylist = []
    for y in range(Ny):
        if not np.isnan(imp.im[:,y,:]).all():
            ylist.append(y)
    xmin = min(xlist)
    xmax = max(xlist)+1
    ymin = min(ylist)
    ymax = max(ylist)+1
    dx = xmax-xmin
    dy = ymax-ymin
    x0 = xmin+dx/2
    y0 = ymin+dy/2

    imp.crop(sizpix=(dx,dy), cenpix=(x0,y0))
    imp_unc.crop(sizpix=(dx,dy), cenpix=(x0,y0))

    ## Manual mask
    delcol = [1,5,9,11]
    for x in delcol:
        imp.im[:,:,x] = np.nan
        imp_unc.im[:,:,x] = np.nan
    imp.im[:3,:,:] = np.nan # AKARI wvl edge
    imp_unc.im[:3,:,:] = np.nan
    imp.im[225:265,:,:] = np.nan # AKARI-SL2 wvl edge
    imp_unc.im[225:265,:,:] = np.nan
    imp.im[:,:6,2] = np.nan
    imp_unc.im[:,:6,2] = np.nan
    imp.im[:,55:,2] = np.nan
    imp_unc.im[:,55:,2] = np.nan

    ## Print spectral map info
    Nw, Ny, Nx = imp.im.shape
    Npix = 0
    for y in range(Ny):
        for x in range(Nx):
            if not np.isnan(imp.im[:,y,x]).all():
                Npix += 1
    print('Total number of spectra: '+str(Npix))
    print('Size of spectral sampling: '+str(Nw))

    write_fits(path_out+'final/'+src, imp.hdr, imp.im, imp.wvl, imp.wmod)
    write_fits(path_out+'final/'+src+'_unc', imp.hdr, imp_unc.im, imp.wvl, imp.wmod)

    ds = read_fits(path_out+'final/'+src, path_out+'final/'+src+'_unc')

    Nw, Ny, Nx = ds.data.shape
    for x in range(Nx):
        for y in range(Ny):
            if ~np.isnan(ds.data[:,y,x]).all():
                p = pplot(ds.wave, ds.data[:,y,x], yerr=ds.unc[:,y,x],
                          c='k', lw=.7, ec='r', xlog=0, ylog=0,
                          xlim=(3.,22.), ylim=(-.1,np.nanmax(ds.data[:503,y,x])),
                          legend='upper left', figsize=(12,10),
                          title=src+'_('+str(x)+','+str(y)+')',
                          titlesize=20, labelsize=20, ticksize=20, legendsize=20)
                # p.add_plot(ds.wave, ds.data[:,y,x]/ds.unc[:,y,x], c='c')

                # xtic = [2, 3, 4, 5, 6, 7, 8, 10, 12, 15, 20, 30, 40]
                # xtic_min = np.arange(2., 41., 1.)
                xtic = np.arange(3., 21., 1)
                xtic_min = np.arange(3., 21., .1)
                p.ax.set_xticks(xtic, minor=False) # major
                p.ax.set_xticks(xtic_min, minor=True) # minor
                p.ax.xaxis.set_major_formatter(ScalarFormatter()) # major
                p.ax.xaxis.set_minor_formatter(NullFormatter()) # minor

                ## Line check
                cline = [
                         # 4.052,     # Bra
                         # 5.128657,  # HI6-10
                         # 5.51116,   # H2S7
                         # 5.6098,    # MgV1
                         # 5.908213,  # Huc
                         # 5.981,     # KIV
                         # 6.10856,   # H2S6
                         # 6.709,     # ClV
                         # 6.90952,   # H2S5
                         # 6.947984,  # HeII1
                         # 6.985274,  # ArII
                         # 7.3178,    # NaIII
                         # 7.459858,  # Pfa
                         # 7.502493,  # Hub
                         # 7.6524,    # NeVI
                         # 7.8145,    # FeVII
                         # 7.90158,   # ArV
                         # 8.02505,   # H2S4
                         # 8.760064,  # HI7-10
                         # 8.99103,   # ArIII1
                         # 9.042,     # NiVI
                         # 9.5267,    # FeVII
                         # 9.66491,   # H2S3
                         # 9.713475,  # HeII2
                         # 10.3385,   # SiI
                         # 10.5105,   # SIV
                         # 12.27861,  # H2S2
                         # 12.368527, # Hua
                         # 12.81355,  # NeII
                         # 13.10219,  # ArV
                         # 13.521,    # MgV2
                         # 14.32168,  # NeV1
                         # 15.555,    # NeIII1
                         # 17.0398,   # H2S1
                         # 17.608246, # HI11-18
                         # 18.7129,   # SIII1
                         # 19.061898, # HI7-8
                         # 21.8291,   # ArIII2
                         # 24.3175,   # NeV2
                         # 25.8903,   # OIV
                         # 25.9882,   # FeII1
                         # 28.21883,  # H2S0
                         # 33.4810,   # SIII2
                         # 34.8152,   # SiII
                         # 35.3491,   # FeII2
                         # 36.0135    # NeIII2
                         ]
                for clin in cline:
                    p.ax.axvline(clin, color='g')
                ## Band check
                cband = [
                         # 3.291,     # Main 3.3
                         # 3.399,     # Main 3.4
                         # 3.499,     # Small 3.5
                         # 5.2394667, # Small 5.2
                         5.6437461, # Small 5.7 (1)
                         5.7490305, # Small 5.7 (2)
                         6.0106598, # Small 6.0
                         6.2034273, # Main 6.2 (1)
                         6.2672596, # Main 6.2 (2)
                         6.6273788, # Small 6.6
                         6.8548833, # Small 6.8
                         7.0791725, # Small 7.1
                         7.6000000, # Plateau 7.7
                         7.6171123, # Main 7.7 (1)
                         7.8704769, # Main 7.7 (2)
                         8.3623706, # Small 8.3
                         8.6204540, # Main 8.6
                         # 9.5244838, # Small 9.5
                         10.707132, # Small 10.7
                         11.038349, # Small 11.0
                         11.237893, # Main 11.2
                         11.400432, # Plateau 11.3
                         # 11.796389, # Small 11.8
                         11.949674, # Small 11.9
                         12.626842, # Main 12.7 (1)
                         12.760273, # Main 12.7 (2)
                         13.559342, # Small 13.6
                         14.257133, # Small 14.2
                         # 15.893117, # Small 15.6
                         16.482868, # Small 16.4
                         17.082868, # Plateau 17.0
                         # 17.428485, # Small 17.4
                         17.771096, # Small 17.8
                         # 18.925630, # Small 18.9
                         ]
                for cban in cband:
                    p.ax.axvline(cban, color='b')

                p.save(path_out+'final/'+src+'_('+str(x+1)+','+str(y+1)+')', transparent=1)


#!/usr/bin/python3
# -*- coding: utf-8 -*-
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~----->>>
#                    _          _
#                .__(.)<  ??  >(.)__.
#                 \___)        (___/
# @Time    : 2022/3/16 下午10:11
# @Author  : wds -->> hellowds2014@gmail.com
# @File    : main.py
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~----->>>
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import argparse
from dataloader import dataloader
from model import WeTe
from Trainer import Trainer

parser = argparse.ArgumentParser(description='Representing mixtures of word embeddings with mixtures of topic embeddings (WeTe)')
# WeTe options
parser.add_argument('--embedding_dim', type=int, default=100, metavar='N', help='embeddings dimension (default: 100)')
parser.add_argument('--K', type=int, default=100, help='topic size (default: 100)')
# Training options
parser.add_argument('--epoch', type=int, default=500, help='number of epochs to train WeTe (default: 500)')
parser.add_argument('--batchsize', type=int, default=500, help='batch size (default: 500)')
parser.add_argument('--lr', type=float, default=1e-2, help='learning rate for Adam to train WeTe (default: 1e-3)')
parser.add_argument('--test_every', type=int, default=10, help='do test every test_every epoches (default: 10)')
parser.add_argument('--beta', type=float, default=0.5, help='balance coefficient for bidirectional transport cost (default: 0.5)')
parser.add_argument('--epsilon', type=float, default=1.0, help='trade-off between transport cost and likelihood (default: 1.0)')
parser.add_argument('--device', type=str, default="0", help='which device for training: 0, 1, 2, 3 (GPU) or cpu')
parser.add_argument('--init_alpha', type=bool, default=True, help='Using K-means to initialise topic embeddings or not, setting to True will make faster and better performance.')
parser.add_argument('--train_word', type=bool, default=True, help='Finetuning word embedding or not, seting to True will make better performance.')
# Dataset options
parser.add_argument('--dataname', type=str, default='20ng_6', help='Dataset: 20ng_6|20ng_20|reuters|rcv2|web|tmn|dp')

# Other options
parser.add_argument('--glove', type=str, default="./glove.6B/glove.6B.100d.txt", help='load pretrained word embedding')
# parser.add_argument('--glove', type=str, default=None, help='load pretrained word embedding')
parser.add_argument('--every_test', type=int, default=5, help='test every test_num epoch')
parser.add_argument('--save_phi', type=int, default=10, help='save phi every save_phi epoch')
parser.add_argument('--save_path', type=str, default='WeTe_result', help='path for saving topics')

args = parser.parse_args()
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
args.device = device

if __name__ == '__main__':

    # Load dataset
    if args.dataname == '20ng_6':
        data_path = './dataset/20ng.pkl'
        args.clc_num = 20
    elif args.dataname == '20ng_20':
        data_path = './dataset/20ng.pkl'
        args.clc_num = 20
    elif args.dataname == 'reuters':
        data_path = './dataset/reuters.pkl'
        args.clc_num = 20
    elif args.dataname == 'rcv2':
        data_path ='dataset/rcv2.pkl'
        args.clc_num = 52
    elif args.dataname == 'web':
        data_path = './dataset/web.pkl'
        args.clc_num = 20
    elif args.dataname == 'tmn':
        data_path = './dataset/tmn.pkl'
        args.clc_num = 20
    elif args.dataname == 'dp':
        data_path = './dataset/dp.pkl'
        args.clc_num = 20
    else:
        print(f'Unknown dataset: {args.dataname}')
        exit()
    train_loader, voc = dataloader(data_path, dataname=args.dataname, mode='train', batch_size=args.batchsize)
    test_loader, _ = dataloader(data_path, dataname=args.dataname, mode='test', batch_size=args.batchsize)
    args.vocsize = len(voc)
    print(f'=============================   Setting   =============================\n {args}')
    print(f'len train: {len(train_loader)}, len test: {len(test_loader)}, voc_len: {len(voc)}')

    model = WeTe(args, voc=voc)
    trainer = Trainer(args, model, voc=voc)
    trainer.train(train_loader, test_loader)


# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
from astropy.utils import lazyproperty
import astropy.units as u
from ..utils.fitting import Fit, Parameters
from ..stats import cash
from ..maps import Map, MapAxis

__all__ = ["MapEvaluator", "MapDataset"]

class MapDataset:
    """Perform sky model likelihood fit on maps.

    Parameters
    ----------
    model : `~gammapy.cube.models.SkyModel`
        Fit model
    counts : `~gammapy.maps.WcsNDMap`
        Counts cube
    exposure : `~gammapy.maps.WcsNDMap`
        Exposure cube
    mask : `~gammapy.maps.WcsNDMap`
        Mask to apply to the likelihood.
    psf : `~gammapy.cube.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EnergyDispersion`
        Energy dispersion
    background_model: `~gammapy.cube.models.BackgroundModel`
        Background model to use for the fit.
    likelihood : {"cash"}
	    Likelihood function to use for the fit.
    """

    def __init__(
        self,
        model,
        counts=None,
        exposure=None,
        mask=None,
        psf=None,
        edisp=None,
        background_model=None,
    ):
        if mask is not None and mask.data.dtype != np.dtype("bool"):
            raise ValueError("mask data must have dtype bool")

        self.model = model
        self.counts = counts
        self.exposure = exposure
        self.mask = mask
        self.psf = psf
        self.edisp = edisp
        self.background_model = background_model
        if background_model:
            self.parameters = Parameters(
                self.model.parameters.parameters +
                self.background_model.parameters.parameters
            )
        else:
            self.parameters = Parameters(self.model.parameters.parameters)

        self.evaluator = MapEvaluator(
            model=self.model, exposure=exposure, psf=self.psf, edisp=self.edisp
        )

    @property
    def data_shape(self):
        """Shape of the counts data"""
        return self.counts.data.shape

    def npred(self):
        """Returns npred map (model + background)"""
        model_npred = self.evaluator.compute_npred()
        back_npred = self.background_model.evaluate()
        total_npred = model_npred.data + back_npred.data
        return back_npred.copy(data=total_npred)
        # TODO: return model_npred + back_npred
        # There is some bug: edisp.e_reco.unit is dimensionless
        # thus map arithmetic does not work.

    def likelihood_per_bin(self):
        """Likelihood per bin given the current model parameters"""
        return cash(n_on=self.counts.data, mu_on=self.npred().data)

    def likelihood(self, parameters, mask=None):
        """Total likelihood given the current model parameters.

        Parameters
        ----------
        mask : `~numpy.ndarray`
            Mask to be combined with the dataset mask.
        """
        if self.mask is None and mask is None:
            stat = self.likelihood_per_bin()
        elif self.mask is None:
            stat = self.likelihood_per_bin()[mask]
        elif mask is None:
            stat = self.likelihood_per_bin()[self.mask.data]
        else:
            stat = self.likelihood_per_bin()[mask & self.mask.data]
        return np.sum(stat, dtype=np.float64)

class MapEvaluator:
    """Sky model evaluation on maps.

    This evaluates a sky model on a 3D map and convolves with the IRFs,
    and returns a map of the predicted counts.
    Note that background counts are not added.

    For now, we only make it work for 3D WCS maps with an energy axis.
    No HPX, no other axes, those can be added later here or via new
    separate model evaluator classes.

    Parameters
    ----------
    model : `~gammapy.cube.models.SkyModel`
        Sky model
    exposure : `~gammapy.maps.Map`
        Exposure map
    background : `~gammapy.maps.Map`
        Background map
    psf : `~gammapy.cube.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EnergyDispersion`
        Energy dispersion
    """

    def __init__(self, model=None, exposure=None, psf=None, edisp=None):
        self.model = model
        self.exposure = exposure
        self.psf = psf
        self.edisp = edisp

        self.parameters = Parameters(self.model.parameters.parameters)

    @lazyproperty
    def geom(self):
        """This will give the energy axes in e_true"""
        return self.exposure.geom

    @lazyproperty
    def geom_image(self):
        return self.geom.to_image()

    @lazyproperty
    def energy_center(self):
        """True energy axis bin centers (`~astropy.units.Quantity`)"""
        energy_axis = self.geom.get_axis_by_name("energy")
        energy = energy_axis.center * energy_axis.unit
        return energy[:, np.newaxis, np.newaxis]

    @lazyproperty
    def energy_edges(self):
        """Energy axis bin edges (`~astropy.units.Quantity`)"""
        energy_axis = self.geom.get_axis_by_name("energy")
        energy = energy_axis.edges * energy_axis.unit
        return energy[:, np.newaxis, np.newaxis]

    @lazyproperty
    def energy_bin_width(self):
        """Energy axis bin widths (`astropy.units.Quantity`)"""
        return np.diff(self.energy_edges, axis=0)

    @lazyproperty
    def lon_lat(self):
        """Spatial coordinate pixel centers.

        Returns ``lon, lat`` tuple of `~astropy.units.Quantity`.
        """
        lon, lat = self.geom_image.get_coord()
        return (u.Quantity(lon, "deg", copy=False), u.Quantity(lat, "deg", copy=False))

    @lazyproperty
    def lon(self):
        return self.lon_lat[0]

    @lazyproperty
    def lat(self):
        return self.lon_lat[1]

    @lazyproperty
    def solid_angle(self):
        """Solid angle per pixel"""
        return self.geom.solid_angle()

    @lazyproperty
    def bin_volume(self):
        """Map pixel bin volume (solid angle times energy bin width)."""
        omega = self.solid_angle
        de = self.energy_bin_width
        return omega * de

    def compute_dnde(self):
        """Compute model differential flux at map pixel centers.

        Returns
        -------
        model_map : `~gammapy.maps.Map`
            Sky cube with data filled with evaluated model values.
            Units: ``cm-2 s-1 TeV-1 deg-2``
        """
        coord = (self.lon, self.lat, self.energy_center)
        dnde = self.model.evaluate(*coord)
        return dnde

    def compute_flux(self):
        """Compute model integral flux over map pixel volumes.

        For now, we simply multiply dnde with bin volume.
        """
        dnde = self.compute_dnde()
        volume = self.bin_volume
        flux = dnde * volume
        return flux

    def apply_exposure(self, flux):
        """Compute npred cube

        For now just divide flux cube by exposure
        """
        npred = (flux * self.exposure.quantity).to_value("")
        return self.exposure.copy(data=npred)

    def apply_psf(self, npred):
        """Convolve npred cube with PSF"""
        return npred.convolve(self.psf)

    def apply_edisp(self, npred):
        """Convolve map data with energy dispersion.

        Parameters
        ----------
        npred : `~gammapy.maps.Map`
            Predicted counts in true energy bins

        Returns
        ---------
        npred_reco : `~gammapy.maps.Map`
            Predicted counts in reco energy bins
        """
        loc = npred.geom.get_axis_index_by_name("energy")
        data = np.rollaxis(npred.data, loc, len(npred.data.shape))
        data = np.dot(data, self.edisp.pdf_matrix)
        data = np.rollaxis(data, -1, loc)
        e_reco_axis = MapAxis.from_edges(
            self.edisp.e_reco.bins, unit=self.edisp.e_reco.unit, name="energy"
        )
        geom_ereco = self.exposure.geom.to_image().to_cube(axes=[e_reco_axis])
        npred = Map.from_geom(geom_ereco, unit="")
        npred.data = data
        return npred

    def compute_npred(self):
        """
        Evaluate model predicted counts.

        Returns
        -------
        npred : `~gammapy.maps.Map`
            Predicted counts on the map (in reco energy bins)
        """
        flux = self.compute_flux()
        npred = self.apply_exposure(flux)
        if self.psf is not None:
            npred = self.apply_psf(npred)
        if self.edisp is not None:
            npred = self.apply_edisp(npred)
        return npred

import numpy,import astropy,from astropy


import numpy as np
import pandas as pd
from tensorflow import keras
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input as resnet50_preprocess

from tensorflow.keras.applications.efficientnet import EfficientNetB7
from tensorflow.keras.applications.efficientnet import preprocess_input as eff_preprocess

from tensorflow.keras.applications.inception_v3 import InceptionV3
from tensorflow.keras.applications.inception_v3 import preprocess_input as incept_preprocess

from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Conv2D, Flatten, Dense

class Predict:
    """
    This class uses the trained models to classify if an image can be classified as ec or nec
    ec means eczema while nec means not eczema
    Only the get_result method is public, it expects an image path and returns the result
    """
    def __init__(self):
        self.height = 300
        self.width = 300
        self.resnet_model = keras.models.load_model('models/res50_model.h5')
        self.eff_model = keras.models.load_model('models/eff11_model.h5')
        self.inception_model = keras.models.load_model('models/inception_model.h5')

    def __resnet(self, x):
        x = resnet50_preprocess(x)
        preds = self.resnet_model.predict(x)
        return self.format_res(preds)

    def __efficient(self, x):
        x = eff_preprocess(x)
        preds = self.eff_model.predict(x)
        return self.format_res(preds)

    def __inception(self, x):
        x = incept_preprocess(x)
        preds = self.inception_model.predict(x)
        return self.format_res(preds)

    @staticmethod
    def format_res(pred):  # formats output result to percentage
        pred = list((pred*100)[0])
        pred = [float(round(i, 2)) for i in pred]
        return pred

    def __get_x(self, image_path, case=None):
        if case == 'inception':
            img = load_img(image_path, target_size=(224, 224))
            x = img_to_array(img)
            x = np.expand_dims(x, axis=0)
            return x
        else:
            img = load_img(image_path, target_size=(self.height, self.width))
            x = img_to_array(img)
            x = np.expand_dims(x, axis=0)
            return x

    def get_result(self, image_path):
        x = self.__get_x(image_path=image_path)
        xx = self.__get_x(case='inception', image_path=image_path)
        return {'resnet': self.__resnet(x),
                'efficient': self.__efficient(x),
                'inception': self.__inception(xx)}


# -*- coding: UTF-8 -*-
import copy
import numpy as np
import os, re
import matplotlib
matplotlib.use("Agg")

import matplotlib.pyplot as plt
import shutil
import sys
import argparse
import xml.dom.minidom
import xmltodict
import cv2
#import matplotlib
#matplotlib.use('Agg')

prob_thres = [0.05, 0.1, 0.15, 0.20, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.905, 0.910, 0.915, 0.920, 0.925, 0.930, 0.935, 0.940, 0.945, 0.950, 0.955, 0.960, 0.965, 0.970, 0.975, 0.980, 0.985, 0.990, 0.995, 0.999]
prob_thres = [0.01, 0.03, 0.05, 0.08, 0.1, 0.15, 0.20, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.825, 0.85, 0.9, 0.922, 0.94, 0.96, 0.980, 0.985, 0.990, 0.995, 0.999]

def iou(b1, b2):#[xmin, ymin, xmax, ymax]
    iou_val = 0.0

    b1_w = b1[2] - b1[0]
    b1_h = b1[3] - b1[1]
    b2_w = b2[2] - b2[0]
    b2_h = b2[3] - b2[1]

    x1 = np.max([b1[0], b2[0]])
    y1 = np.max([b1[1], b2[1]])
    x2 = np.min([b1[2], b2[2]])
    y2 = np.min([b1[3], b2[3]])
    w = np.max([0, x2 - x1])
    h = np.max([0, y2 - y1])
    if w != 0 and h != 0:
        iou_val = float(w * h) / (b1_w * b1_h + b2_w * b2_h - w * h)

    return iou_val

def precision_recal(pred_origin, gt_origin, thres):

    p_num = len(pred_origin)  #pred total num
    r_num = len(gt_origin)    #gt total num
    p_cnt = 0
    r_cnt = 0
    if p_num == 0 or r_num == 0: return p_cnt, r_cnt, p_num, r_num
    gt = np.copy(gt_origin)
    pred = np.copy(pred_origin)

    flag = 0
    for b1 in pred:
        for idx, b2 in enumerate(gt):
            if iou(b1, b2) > thres:
                p_cnt += 1
                gt = np.delete(gt, [idx], axis=0)
                flag = 1
                break

    gt = np.copy(gt_origin)
    pred = np.copy(pred_origin)
    for b1 in gt:
        for idx, b2 in enumerate(pred):
            if iou(b1, b2) > thres:
                r_cnt += 1
                pred = np.delete(pred, [idx], axis=0)
                break

    return p_cnt, r_cnt, p_num, r_num

def get_gt(file_path, anno_key, resized_height = 0):
    gt = []
    #print file_path
    with open(file_path, 'r') as f:
        d = xmltodict.parse(f.read())
        anno = d['annotation']
        # folder = anno['folder']
        filename = anno['filename']
        width = anno['size']['width']
        height = anno['size']['height']
        depth = anno['size']['depth']

        if not 'object' in anno.keys():
            return np.array(gt)

        objs = anno['object']
        if not isinstance(objs, list):   #if len(objs) is one, objs will not be list
            objs = [objs]
        for obj in objs:
            name = obj['name'].lower().strip()
            if str(name) != anno_key: continue

            xmin = obj['bndbox']['xmin']
            ymin = obj['bndbox']['ymin']
            xmax = obj['bndbox']['xmax']
            ymax = obj['bndbox']['ymax']
            gt.append([int(float(xmin)), int(float(ymin)), int(float(xmax)), int(float(ymax))])

    return np.array(gt)

def get_pred(pred_dict, key_name):
    probs = []
    bboxs = []
    if not key_name in pred_dict.keys():
        return np.array(bboxs), np.array(probs)

    pred_lists = pred_dict[key_name]
    probs = []
    bboxs = []
    for item in pred_lists:
        score = item[0]
        pred = item[1]

        probs.append(score)
        bboxs.append(pred)

    return np.array(bboxs), np.array(probs)

def roc(test_gt_list, pred_dict, anno_key, iou_thres, save_dir):
    precisions = np.zeros((len(prob_thres), ), dtype=np.int32)
    recalls = np.zeros((len(prob_thres), ), dtype=np.int32)
    phits = np.zeros((len(prob_thres), ), dtype=np.int32)
    rhits = np.zeros((len(prob_thres), ), dtype=np.int32)

    error_list = []
    for idx, line in enumerate(test_gt_list):
        key_name, gt_file= line[0], line[1] #img_name, xml_file
        #print gt_file
        gt = get_gt(gt_file, anno_key)
        #print gt[0]
        #pred, probs = get_pred_txt(pred_file, key_name)
        pred_box, confs = get_pred(pred_dict, key_name)
        #print(pred_box)
        for k in range(len(prob_thres)):
            # print k
            if len(pred_box) > 0:
                I = confs > prob_thres[k]
                #prob_k = confs[I]
                pred_k = pred_box[I, :]
            else:
                pred_k = [] #pred_ori, gt_ori, gt_side_ori, thres

            pm, rm, pn, rn = precision_recal(pred_k, gt, iou_thres)
            #print(pm, rm, pn, rn)
            # if prob_thres[k] == 0.6 and pm < pn:
            #     error_list.append(key_name+'\n')

            phits[k] += pm
            rhits[k] += rm
            precisions[k] += pn
            recalls[k] += rn

        #if idx == 3: break
        if idx % 1000 == 0:
            print("processed {:d}/{:d}".format(idx, len(test_gt_list)))

    # fw = open('error_list', 'w')
    # fw.writelines(error_list)
    # fw.close()

    mean_p = np.zeros((len(prob_thres), ), dtype=np.float32)
    mean_r = np.zeros((len(prob_thres), ), dtype=np.float32)
    for k in range(len(prob_thres)):
        if precisions[k] == 0:  #pred_num = 0
            mean_r[k] = 0.0
            continue
        if recalls[k] == 0: #gt_num = 0
            mean_r[k] = 1.0
            continue

        mean_p[k] = phits[k] * 1.0 / precisions[k]
        mean_r[k] = rhits[k] * 1.0 / recalls[k]

    return mean_p, mean_r

#gen img_xml_list and read pred results to dict
def get_img_xml_pred(data_path, test_file, pred_file):
    img_xml_list = []
    pred_dict = {}

    lines = open(test_file, 'r').readlines()
    for _, line in enumerate(lines):
        image_path, xml_path = line.strip().split()
        abs_xml_path = os.path.join(data_path, xml_path)
        img_xml_list.append([image_path, abs_xml_path])

    lines = open(pred_file, 'r').readlines()
    for line in lines:
        split_item = line.strip().split()
        image_path = split_item[0]
        #image key first occur
        if not image_path in pred_dict.keys():
            pred_dict[image_path] = list()
        score = float(split_item[1])
        box = list(map(float, split_item[2:6]))
        pred_dict[image_path].append([score, list(map(int, box))])

    return img_xml_list, pred_dict

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='roc')
    parser.add_argument('--data_path', dest='data_path', help='gt path', type=str)
    parser.add_argument('--test_file', dest='test_file', help='gt file path', type=str)
    parser.add_argument('--pred_file', dest='pred_file', help='pred file path', type=str)

    parser.add_argument('--save_dir', dest='save_dir', help='save dir', type=str)
    parser.add_argument('--anno_key', dest='anno_key', help='dataset key, use it to get gt', type=str)
    parser.add_argument('--resized_height', dest='resized_height', default=None, help='img height when test',
                        type=float)
    parser.add_argument('--ious', dest='ious', default="(0.3,0.4)", help='ious to compute roc, should be tuple or list',
                        type=str)
    parser.add_argument('--show', dest='show', help='whether to show plot figure',
                        action='store_true')
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    args = parser.parse_args()

    test_gt_list, pred_dict = get_img_xml_pred(args.data_path, args.test_file, args.pred_file)  #name, gt, pred file

    iou_list = eval(args.ious)  #convert str into compute value(e.g. list)
    if not isinstance(iou_list, list) and not isinstance(iou_list, tuple):
        iou_list = (iou_list,)

    stats = []
    for iou_thres in iou_list:
        mean_p, mean_r = roc(test_gt_list, pred_dict, args.anno_key, iou_thres, args.save_dir)
        stats.append([mean_p, mean_r])

    #plt.figure(figsize=(8, 7))
    for mean_p, mean_r in stats:
        print("error1")
        plt.plot(mean_r,  mean_p, 'o-')
        print("error2")
        for idx, x, y in zip(range(len(mean_p)), mean_r, mean_p):
            print(idx, 'thres =', prob_thres[idx], 'R =',x, 'P =',y)
            if idx % 2 == 0:
                plt.text(x+0.02, y+0.02, '%d(%.3f)'%(idx, prob_thres[idx]), rotation=60, ha='center', va='bottom', fontsize=7)

    plt.grid(True)
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.xlabel("recall")
    plt.ylabel("precision")
    #plt.title(get_title(args.save_dir))
    plt.title(args.anno_key)
    #if args.show:
    #plt.show()
    plt.savefig("water.png")

# def get_test_gt_pred(data_path, test_file, pred_file):
#     test_gt_list = []
#     lines = open(test_file, 'r').readlines()
#     for _, line in enumerate(lines):
#         split_t = line.strip().split()
#         if len(split_t) == 2:
#             name, gt_xml = line.strip().split()
#             gt_xml = '{}/{}'.format(data_path, gt_xml)
#         test_gt_list.append([name, gt_xml])

#     pred_dict = {}
#     lines = open(pred_file, 'r').readlines()
#     for line in lines:
#         list_ = line.strip().split()
#         name = list_[0]     #image id
#         if not name in pred_dict.keys():
#             pred_dict[name] = list()
#         score = float(list_[1])
#         box = list(map(float, list_[2:6]))
#         pred_dict[name].append([score, list(map(int, box))])
#     return test_gt_list, pred_dict


import copy
import math
import random
import time
import uuid
import builtins
from copy import deepcopy
from itertools import count
from types import MethodType

import numpy as np
import torch
from gym.spaces import Box
from torch import nn
import torch.nn.functional as F
from torch.distributions import Bernoulli, Categorical
from trident.backend.common import *
from trident.backend.pytorch_backend import *
from trident.backend.pytorch_ops import *
from trident.backend.tensorspec import TensorSpec, ObjectType
from trident.data.image_common import image_backend_adaption
from trident.misc.visualization_utils import loss_metric_curve
from trident.layers.pytorch_initializers import kaiming_normal
from trident.optims.pytorch_trainer import Model
from trident.reinforcement.utils import ReplayBuffer, ActionStrategy, ObservationType

import_or_install('gym')
import gym

_session = get_session()
__all__ = ['PolicyBase', 'DqnPolicy', 'A2CPolicy']

def modify_env(env: gym.Env):
    def render(env: gym.Env, mode='human'):
        _env = None
        if isinstance(env, gym.core.Wrapper):
            _env = env.env
        elif isinstance(env, gym.core.Env):
            _env = env

        if mode == 'human':
            return _env.render('human')
        elif mode == 'observation':
            if env.obsetvation_type == ObservationType.Image and hasattr(env, 'observation') and callable(env.observation):
                return env.observation(_env.render('rgb_array').copy())
            else:
                return _env.render('rgb_array')
        elif mode == 'rgb_array':
            return _env.render('rgb_array')

    if isinstance(env.observation_space, Box) and len(env.observation_space.shape) in [2, 3] and (
            (env.observation_space.low.max() == 0 and env.observation_space.high.min() == 255)
            or env.observation_space.dtype == np.uint8):
        env.obsetvation_type = ObservationType.Image
    elif isinstance(env.observation_space, Box) and env.observation_space.dtype == np.float32:
        env.obsetvation_type = ObservationType.Box

    setattr(env, 'render', MethodType(render, env))

    return env

class PolicyBase(Model):
    """The base class for any RL policy.
    """

    def __init__(self, network: Layer, env: gym.Env, action_strategy=None, gamma=0.99, use_experience_replay=False, replay_unit='step', memory_length: int = 10000,
                 name=None) -> None:
        super().__init__()
        self.network = network
        if name is not None:
            self.network._name = name

        self.env = modify_env(env)
        self.env.reset()
        self.observation_space = env.observation_space
        self.action_space = env.action_space
        self.agent_id = uuid.uuid4().node
        self.gamma = gamma
        self.use_experience_replay = use_experience_replay
        self.memory = None
        if self.use_experience_replay:
            self.memory = ReplayBuffer(memory_length)
        else:
            self.memory = ReplayBuffer(1)

        if replay_unit not in ['step', 'episode']:
            raise ValueError('Only [step,episode] are valid unit options.')
        self.replay_unit = replay_unit
        self.name = name
        self.action_strategy = action_strategy
        self.state_pool = []
        self.reward_pool = []
        self.action_pool = []
        self.value_pool = []
        self.setting_network()

    def setting_network(self):
        super()._initial_graph(inputs=self.get_observation(), output=deepcopy(self.network))

    def get_observation(self):
        if hasattr(self.env, 'state'):
            return expand_dims(self.data_preprocess(self.env.state), 0)
        else:
            return expand_dims(self.data_preprocess(self.env.render('observation')), 0)

    def select_action(self, state, model_only=False, **kwargs):
        return self.env.action_space.samples()

    def get_rewards(self, action):
        return self.env.step(action)

    def experience_replay(self, batch_size):
        return NotImplemented

    def collect_samples(self, min_replay_samples, need_render=False) -> bool:

        if self.memory is None:
            self.memory = ReplayBuffer(10000)
        progress_inteval = int(min_replay_samples / 50) * 5
        self.state_pool = []
        self.reward_pool = []
        self.action_pool = []

        for i_episode in range(min_replay_samples):
            self.env.reset()
            state = self.get_observation()
            for t in count():
                action = self.select_action(state, model_only=True if self.action_strategy == ActionStrategy.OnPolicy else False)
                _observation, reward, done, info = self.get_rewards(action)
                if need_render:
                    self.env.render()
                next_state = None
                if not done:
                    next_state = self.get_observation()
                if self.replay_unit == 'step':
                    self.memory.push(state, action, next_state, reward)
                    if len(self.memory) < min_replay_samples and len(self.memory) % progress_inteval == 0:
                        print("Replay Samples:{0}".format(len(self.memory)))
                    if len(self.memory) == min_replay_samples:
                        # n1 = self.action_logs['model'][0]
                        # n2 = self.action_logs['model'][1]
                        # n3 = self.action_logs['random'][0]
                        # n4 = self.action_logs['random'][1]
                        # print('model: 0:{0} 1:{1}  random: 0:{2} 1:{3}  random: {4}'.format(float(n1) / (n1 + n2), float(n2) / (n1 + n2), float(n3) / (n3 + n4),
                        #                                                                                float(n4) / (n3 + n4), float(n3 + n4) / builtins.max(n1 + n2 + n3 + n4,1)))
                        #
                        # self.action_logs = OrderedDict()
                        # self.action_logs['model'] = OrderedDict()
                        # self.action_logs['random'] = OrderedDict()
                        # self.action_logs['model'][0] = 0
                        # self.action_logs['model'][1] = 0
                        # self.action_logs['random'][0] = 0
                        # self.action_logs['random'][1] = 0
                        return True
                elif self.replay_unit == 'episode':
                    self.state_pool.append(state)
                    self.action_pool.append(action)
                    self.reward_pool.append(reward)
                    if done:
                        self.memory.push(self.state_pool, self.action_pool, None, self.reward_pool)
                        if len(self.memory) < min_replay_samples and len(self.memory) % progress_inteval == 0:
                            print("Replay Samples:{0}".format(len(self.memory)))
                        self.state_pool = []
                        self.action_pool = []
                        self.reward_pool = []

                        if len(self.memory) == min_replay_samples:
                            return True
                        break
                state = next_state
                if done:
                    break

        return False

    def push_into_memory_criteria(self, *args, **kwargs) -> bool:
        return True

    def episode_complete_criteria(self, *args, **kwargs) -> bool:
        return False

    def task_complete_criteria(self, *args, **kwargs) -> bool:
        return False

    def estimate_future_return(self, *args, **kwargs):
        return NotImplemented

    def save_or_sync_weights(self):
        self.save_model(save_path=self.training_context['save_path'])

    def training_model(self, current_episode=0, current_step=0, num_episodes=100, train_timing=None, done=False, batch_size=1, repeat_train=1):

        is_collect_data = False
        for i in range(repeat_train):
            data = None
            if self.use_experience_replay:
                data = self.experience_replay(batch_size)
            else:
                data = self.memory.memory[0]
            self.estimate_future_return(*data)
            self.training_context['skip_generate_output'] = True
            if 'step' in train_timing:
                current_step = current_step * repeat_train + i
                if done:
                    total_batch = current_step * repeat_train + i + 1
                    is_collect_data = True
                else:
                    total_batch = current_step * repeat_train + i + 10
            elif 'episode' in train_timing:
                current_step = i
                total_batch = repeat_train

            super(PolicyBase, self).train_model(self.training_context['train_data'], self.training_context['test_data'],
                                                current_epoch=current_episode,
                                                current_batch=current_step,
                                                total_epoch=num_episodes,
                                                total_batch=total_batch,
                                                is_collect_data=True,
                                                is_print_batch_progress=False,
                                                is_print_epoch_progress=False,
                                                log_gradients=False, log_weights=False,
                                                accumulate_grads=(current_step * repeat_train + 1) % self.accumulation_steps != 0)
            self.save_or_sync_weights()

    def play(self, num_episodes, batch_size=1, min_replay_samples=1, print_progess_frequency=5, training=True, train_timing='on_episode_start', train_every_nstep=1, repeat_train=1,
             need_render=True):
        if train_timing not in ['on_episode_start', 'on_step_end', 'on_step_start']:
            raise ValueError('Only on_episode_start,on_step_end are valid  train_timing options')

        if training:
            self._model.train()
        else:
            self._model.eval()
        if self.use_experience_replay:
            self.collect_samples(min_replay_samples=min_replay_samples)
        else:
            self.collect_samples(min_replay_samples=1, need_render=True if self.replay_unit == 'episode' else False)
            print('start train....')
        self.state_pool = []
        self.reward_pool = []
        self.action_pool = []

        self.total_reward = 0
        self.t = 0
        self.i_episode = 0
        if hasattr(self.env, 'recording_enabled'):
            self.env.recording_enabled = True
        for i_episode in range(num_episodes):
            self.i_episode = i_episode

            if training and train_timing == 'on_episode_start' and i_episode % train_every_nstep == 0:
                self.training_model(i_episode, 0, num_episodes=num_episodes, repeat_train=repeat_train, train_timing=train_timing, batch_size=batch_size)
            self.env.reset()
            self.total_rewards = 0
            state = self.get_observation()
            for t in count():
                self.t = t
                # # Train on_step_start
                # if training and train_timing == 'on_step_start' and t % train_every_nstep == 0:
                #     self.training_model(i_episode, t,num_episodes=num_episodes, repeat_train=repeat_train, batch_size=batch_size)

                action = self.select_action(state, model_only=True)
                observation, reward, done, info = self.get_rewards(action)

                self.total_rewards += reward

                next_state = self.get_observation() if not done else None

                if need_render:
                    self.env.render()
                if self.replay_unit == 'step':
                    if self.push_into_memory_criteria(state, action, next_state, reward) or done:
                        self.memory.push(state, action, next_state, reward)
                elif self.replay_unit == 'episode':
                    self.state_pool.append(state)
                    self.action_pool.append(action)
                    self.reward_pool.append(reward)
                    if done:
                        if self.push_into_memory_criteria(self.state_pool, self.action_pool, None, self.reward_pool):
                            self.memory.push(self.state_pool, self.action_pool, None, self.reward_pool)
                        self.state_pool = []
                        self.action_pool = []
                        self.reward_pool = []

                complete = self.episode_complete_criteria()
                # Train on_step_end
                if training and train_timing == 'on_step_end' and t % train_every_nstep == 0:
                    self.training_model(i_episode, t, num_episodes=num_episodes, done=done or complete, repeat_train=repeat_train, train_timing=train_timing, batch_size=batch_size,
                                        accumulate_grads=accumulate_grads)

                state = next_state
                if done or complete:
                    self.epoch_metric_history.collect('rewards', i_episode, float(self.total_rewards))
                    self.epoch_metric_history.collect('t', i_episode, float(t + 1))
                    if self.use_experience_replay:
                        self.epoch_metric_history.collect('replay_buffer_utility', i_episode, float(len(self.memory)) / self.memory.capacity)

                    if print_progess_frequency == 1 or (i_episode > 0 and (i_episode + 1) % print_progess_frequency == 0):
                        self.print_epoch_progress(print_progess_frequency)
                        # n1 = self.action_logs['model'][0]
                        # n2 = self.action_logs['model'][1]
                        # n3 = self.action_logs['random'][0]
                        # n4 = self.action_logs['random'][1]
                        # print('model: 0:{0} 1:{1}  random: 0:{2} 1:{3}  random: {4}'.format(float(n1) / (n1 + n2), float(n2) / (n1 + n2), float(n3) / builtins.max(n3 + n4,1),
                        #                                                                                float(n4) / builtins.max(n3 + n4,1), float(n3 + n4) / builtins.max(n1 + n2 + n3 + n4,1)))
                        #
                        # self.action_logs = OrderedDict()
                        # self.action_logs['model'] = OrderedDict()
                        # self.action_logs['random'] = OrderedDict()
                        # self.action_logs['model'][0] = 0
                        # self.action_logs['model'][1] = 0
                        # self.action_logs['random'][0] = 0
                        # self.action_logs['random'][1] = 0
                    # 定期繪製損失函數以及評估函數對時間的趨勢圖
                    if i_episode > 0 and (i_episode + 1) % (5 * print_progess_frequency) == 0:
                        loss_metric_curve(self.epoch_loss_history, self.epoch_metric_history, metrics_names=list(self.epoch_metric_history.keys()), calculate_base='epoch',
                                          imshow=True)

                    if self.task_complete_criteria():
                        self.save_model(save_path=self.training_context['save_path'])
                        print('episode {0} meet task complete criteria, training finish! '.format(i_episode))
                        return True

                    break

        print('Complete')
        self.env.render()
        self.env.close()

    def learn(self, num_episodes, batch_size=1, min_replay_samples=1, print_progess_frequency=5, train_timing='on_episode_start', train_every_nstep=1, repeat_train=1,
              accumulate_grads=False):
        self.play(num_episodes=num_episodes, batch_size=batch_size, min_replay_samples=min_replay_samples, print_progess_frequency=print_progess_frequency, training=True,
                  train_timing=train_timing, train_every_nstep=train_every_nstep,
                  repeat_train=repeat_train, need_render=True)

    def resume(self, num_episodes=3000, **kwargs):
        pass

    @property
    def preprocess_flow(self):
        return self._preprocess_flow

    @preprocess_flow.setter
    def preprocess_flow(self, value):
        self._preprocess_flow = value
        objecttype = None
        if isinstance(self.model.input_spec, TensorSpec):
            objecttype = self.model.input_spec.object_type
        # super()._initial_graph(inputs=to_tensor(self.get_observation()).repeat_elements(2, 0), output=deepcopy(self.network))
        self.setting_network()
        if objecttype is not None:
            self.inputs.value_list[0].object_type = objecttype
            self.model.input_spec.object_type = objecttype

        self.env.reset()

    def data_preprocess(self, img_data):
        if self._model is not None:
            self._model.input_spec.object_type = ObjectType.rgb
        if not hasattr(self, '_preprocess_flow') or self._preprocess_flow is None:
            self._preprocess_flow = []
        if img_data.ndim == 4:
            return to_tensor(to_numpy([self.data_preprocess(im) for im in img_data]))
        if len(self._preprocess_flow) == 0:
            return image_backend_adaption(img_data)
        if isinstance(img_data, np.ndarray):
            for fc in self._preprocess_flow:
                if self._model is not None and self.signature is not None and len(self.signature) > 1 and self._model.input_spec is not None:
                    img_data = fc(img_data, spec=self._model.input_spec)
                else:
                    img_data = fc(img_data)
            img_data = image_backend_adaption(img_data)
            if self._model.input_spec is None:
                self._model.input_spec = TensorSpec(shape=tensor_to_shape(to_tensor(img_data), need_exclude_batch_axis=True, is_singleton=True), object_type=ObjectType.rgb,
                                                    name='input')

                self.input_shape = self._model.input_spec.shape[1:]

            return img_data
        else:
            return img_data

    def do_on_batch_end(self):
        self.training_context['time_batch_progress'] += (time.time() - self.training_context['time_batch_start'])
        self.training_context['time_epoch_progress'] += (time.time() - self.training_context['time_batch_start'])
        self.training_context['steps'] += 1
        if (self.training_context['steps'] + 1) % _session.epoch_equivalent == 0:
            if self.warmup > 0 and self.warmup == (self.training_context['steps'] + 1) // _session.epoch_equivalent:
                self.adjust_learning_rate(self.training_context['base_lr'])
                self.warmup = 0

class DqnPolicy(PolicyBase):
    """The base class for any RL policy.
    """

    def __init__(self, network: Layer, env: gym.Env, memory_length: int = 10000
                 , gamma=0.99, max_epsilon=0.9, min_epsilon=0.01, decay=100
                 , target_update=10, name='dqn') -> None:
        super(DqnPolicy, self).__init__(network=network, env=env, action_strategy=ActionStrategy.OffPolicy, gamma=gamma, use_experience_replay=True, replay_unit='step',
                                        memory_length=memory_length, name=name)

        self.gamma = gamma
        self.max_epsilon = max_epsilon
        self.min_epsilon = min_epsilon
        self.decay = decay
        self.target_update = target_update
        self.steps_done = 0

    def setting_network(self):
        super()._initial_graph(inputs=self.get_observation(), output=deepcopy(self.network))
        self.policy_net = self._model
        self.policy_net.to(get_device())
        kaiming_normal(self.policy_net)
        self.policy_net.train()

        self.target_net = copy.deepcopy(self.network)
        self.target_net.to(get_device())
        self.target_net.eval()
        self.summary()

    def get_observation(self):

        if hasattr(self.env, 'state'):
            return expand_dims(self.data_preprocess(to_numpy(self.env.state)), 0).astype(np.float32)
        else:
            return expand_dims(self.data_preprocess(self.env.render('observation')), 0).astype(np.float32)

    def select_action(self, state, model_only=False, **kwargs):
        sample = random.random()
        self.epsilon = self.min_epsilon + (self.max_epsilon - self.min_epsilon) * math.exp(-1.0 * self.steps_done / self.decay)
        self.steps_done += 1
        if model_only == True or sample > self.epsilon:
            with torch.no_grad():
                selected_action = argmax(self.policy_net(to_tensor(state)), 1)
                # self.action_logs['model'][selected_action.item()]+=1
                return selected_action.item()
        else:
            selected_action = np.random.randint(low=0, high=100) % self.env.action_space.n
            # self.action_logs['random'][selected_action] += 1
            return selected_action

    def get_rewards(self, action):
        return self.env.step(action)

    def experience_replay(self, batch_size):
        batch = self.memory.sample(minimum(batch_size, len(self.memory)))
        state_batch = to_tensor(batch.state, requires_grad=True).squeeze(1)
        action_batch = to_tensor(batch.action).long().detach()
        reward_batch = to_tensor(batch.reward).squeeze(1).detach()
        return state_batch, action_batch, batch.next_state, reward_batch

    def estimate_future_return(self, *args, **kwargs):
        state_batch, action_batch, next_state_batch, reward_batch = args

        # 基於目前狀態所產生的獎賞期望值預測 Q(s_t)。
        # 策略網路評估行動
        predict_rewards = self.policy_net(state_batch).gather(1, action_batch)

        target_rewards = zeros(len(next_state_batch))
        # 至於未來狀態的部分，未來狀態主要是計算未來價值使用，但是只要是done，等於不再有未來狀態，則要視成功(1)是失敗(-3)來決定獎賞
        for i in range(len(next_state_batch)):
            s = next_state_batch[i]
            # 要計算未來獎賞預估
            if s is not None:
                # 目標網路評估Q值
                self.target_net.eval()
                next_q = self.target_net(to_tensor(s))
                q_next = max(next_q.detach(), axis=1)

                # 計算目標獎賞值(目前實際獎賞+折扣因子]*未來獎賞預估值)
                target_rewards[i] = reward_batch[i] + (q_next * self.gamma)
            # 不再有未來，只有當期獎賞
            else:
                # 如果是最後一筆表示沒有未來價值
                target_rewards[i] = reward_batch[i] + 0

        # 將計算衍生結果以及data_feed暫存於training_context
        train_data = OrderedDict()
        train_data['state'] = state_batch
        train_data['predict_rewards'] = squeeze(predict_rewards, -1)
        train_data['target_rewards'] = target_rewards.detach()
        train_data['reward_batch'] = reward_batch

        data_feed = OrderedDict()
        data_feed['input'] = 'state'
        data_feed['output'] = 'predict_rewards'
        data_feed['target'] = 'target_rewards'

        self.training_context['data_feed'] = data_feed
        self.training_context['train_data'] = train_data

    def save_or_sync_weights(self):
        steps = self.training_context['steps']
        if steps <= 2 * self.accumulation_steps or (steps % self.target_update == 0):
            self.target_net.load_state_dict(self.policy_net.state_dict(), strict=True)
            self.save_model(save_path=self.training_context['save_path'])

    def learn(self, num_episodes, batch_size=1, min_replay_samples=1, print_progess_frequency=5, train_timing='on_episode_start', train_every_nstep=1, repeat_train=1,
              accumulate_grads=False):
        self.play(num_episodes=num_episodes, batch_size=batch_size, min_replay_samples=min_replay_samples, print_progess_frequency=print_progess_frequency, training=True,
                  train_timing=train_timing, train_every_nstep=train_every_nstep,
                  repeat_train=repeat_train, need_render=True)

class PGPolicy(PolicyBase):
    """The base class for any RL policy.
    """

    def __init__(self, network: Layer, env: gym.Env, use_experience_replay=False, memory_length: int = 10000
                 , gamma=0.999, name='pg') -> None:
        super(PGPolicy, self).__init__(network=network, env=env, action_strategy=ActionStrategy.OnPolicy, gamma=gamma, use_experience_replay=use_experience_replay,
                                       replay_unit='episode',
                                       memory_length=memory_length, name=name)

        def pg_loss(output, action, rewards):
            loss = 0.0
            for i in range(len(action)):
                probs_i = output[i]
                m = Bernoulli(probs_i)
                action_i = torch.FloatTensor([action[i]]).to(get_device())
                reward_i = rewards[i]
                loss = loss - m.log_prob(action_i) * reward_i
            return loss

        self.with_loss(pg_loss, name='pg_loss')

        self.gamma = gamma

    def setting_network(self):
        super()._initial_graph(inputs=self.get_observation(), output=deepcopy(self.network))
        self.policy_net = self._model

        kaiming_normal(self._model)
        self.policy_net.to(get_device())
        self.summary()
        self.policy_net.train()
        self.steps_done = 0

    def get_observation(self):
        if hasattr(self.env, 'state'):
            return expand_dims(self.data_preprocess(to_numpy(self.env.state)), 0).astype(np.float32)
        else:
            return expand_dims(self.data_preprocess(self.env.render('observation')), 0).astype(np.float32)

    def select_action(self, state, **kwargs):
        # 只根據模型來決定action   on-policy
        probs = self.model(state)[0]
        m = Bernoulli(probs)
        action = m.sample()
        return int(action.item())

    def get_rewards(self, action):
        return self.env.step(action)

    def experience_replay(self, batch_size):
        batch = self.memory.sample(builtins.min(batch_size, len(self.memory)))
        state_batch = batch.state
        action_batch = batch.action
        reward_batch = batch.reward
        return state_batch, action_batch, None, reward_batch

    def estimate_future_return(self, *args, **kwargs):
        state_pool = []
        action_pool = []
        reward_pool = []
        state_batch, action_batch, _, reward_batch = args
        for i in range(len(state_batch)):
            state_pool.extend(state_batch[i])
            action_pool.extend(action_batch[i])
            reward_pool.extend(reward_batch[i])
        running_add = 0
        # 逆向計算期望獎賞
        for i in reversed(range(len(reward_pool))):
            if reward_pool[i] == -3:
                running_add = -3
            else:
                running_add = running_add * self.gamma + reward_pool[i]
                reward_pool[i] = running_add

        # 將期望獎賞標準化
        reward_pool = to_numpy(reward_pool).astype(np.float32)
        reward_mean = np.mean(reward_pool.copy())
        reward_std = np.std(reward_pool.copy())
        for i in range(len(reward_pool)):
            reward_pool[i] = (reward_pool[i] - reward_mean) / reward_std

        train_data = OrderedDict()
        train_data['state'] = to_tensor(np.array(state_pool))
        train_data['output'] = self.policy_net(train_data['state'])
        train_data['action'] = to_tensor(np.array(action_pool)).float()
        train_data['rewards'] = to_tensor(np.array(reward_pool))
        data_feed = OrderedDict()
        data_feed['input'] = 'state'
        data_feed['output'] = 'output'
        data_feed['action'] = 'action'
        data_feed['rewards'] = 'rewards'

        self.training_context['data_feed'] = data_feed
        self.training_context['train_data'] = train_data

class A2CPolicy(PolicyBase):
    """The base class for any RL policy.
    """

    def __init__(self, network: Layer, env: gym.Env, use_experience_replay=False, memory_length: int = 10000
                 , gamma=0.999, name='a2c') -> None:
        super(A2CPolicy, self).__init__(network=network, env=env, action_strategy=ActionStrategy.OnPolicy, gamma=gamma, use_experience_replay=use_experience_replay,
                                        replay_unit='episode',
                                        memory_length=memory_length, name=name)

        def policy_loss(policy_losses):
            return policy_losses.sum()

        def value_loss(value_losses):
            return value_losses.sum()

        self.with_loss(policy_loss, name='policy_loss')
        self.with_loss(value_loss, name='value_loss')

        self.gamma = gamma

    def setting_network(self):
        super()._initial_graph(inputs=self.get_observation(), output=deepcopy(self.network))
        self.actor_critic = self._model

        kaiming_normal(self._model)
        self.actor_critic.to(get_device())
        self.summary()
        self.actor_critic.train()
        self.steps_done = 0

    def get_observation(self):
        return expand_dims(self.data_preprocess(self.env.render('observation')), 0).astype(np.float32)

    def select_action(self, state, **kwargs):
        # 只根據模型來決定action   on-policy
        if ndim(state) == 3:
            state = expand_dims(state, 0)
        state = to_tensor(state)
        out = self.model(state)
        if isinstance(out, OrderedDict):
            out = out.value_list
        probs, qvalue = out
        try:
            probs = where(is_abnormal_number(probs), zeros_like(probs), probs)
            m = Categorical(probs=probs)
            action = m.sample()
            return int(action.item())
        except Exception  as e:
            print(probs)
            print(e)
            action = argmax(probs, -1)
            return int(action.item())

    def get_rewards(self, action):
        return self.env.step(action)

    def experience_replay(self, batch_size):
        batch = self.memory.sample(builtins.min(batch_size, len(self.memory)))
        state_batch = batch.state
        action_batch = batch.action
        reward_batch = batch.reward
        return state_batch, action_batch, batch.next_state, reward_batch

    def estimate_future_return(self, *args, **kwargs):
        policy_losses = []
        value_losses = []
        state_batch, action_batch, _, reward_batch = args
        reward_batch = reward_batch.astype(np.float32)
        running_add = 0.0
        # 逆向計算期望獎賞
        for i in reversed(range(len(reward_batch))):
            if i == len(reward_batch) - 1:
                running_add = reward_batch[i]
            else:
                running_add = running_add * self.gamma + reward_batch[i]
                reward_batch[i] = running_add
        # 將期望獎賞標準化
        reward_batch = np.array(reward_batch, dtype=np.float32)
        reward_mean = np.mean(reward_batch.copy())
        reward_std = np.std(reward_batch.copy())
        reward_batch = np.expand_dims(((reward_batch - reward_mean) / reward_std), -1)
        policy_losses = to_tensor(0.0)
        value_losses = to_tensor(0.0)
        for action, state, expect_reward in zip(action_batch, state_batch, reward_batch):
            state = to_tensor(state)
            out = self.model(state)
            action_log_prob = None
            if isinstance(out, OrderedDict):
                out = out.value_list
            probs, value = out
            try:
                probs = where(is_abnormal_number(probs), zeros_like(probs), probs)
                m = Categorical(probs=probs)
                action_log_prob = m.log_prob(to_tensor(action))

            except Exception  as e:
                print(probs)
                print(e)

            expect_reward = to_tensor(expect_reward)
            advantage = (expect_reward - value).detach()

            # calculate actor (policy) loss
            policy_losses = policy_losses + (-action_log_prob * (advantage.detach()))

            # calculate critic (value) loss using L1 smooth loss
            value_losses = value_losses + ((value - expect_reward) ** 2)

        train_data = OrderedDict()
        train_data['policy_losses'] = true_divide(policy_losses, builtins.max(len(reward_batch), 1))
        train_data['value_losses'] = true_divide(value_losses, builtins.max(len(reward_batch), 1))
        data_feed = OrderedDict()
        data_feed['input'] = 'state'
        data_feed['output'] = 'output'
        data_feed['value_losses'] = 'value_losses'
        data_feed['policy_losses'] = 'policy_losses'

        self.training_context['data_feed'] = data_feed
        self.training_context['train_data'] = train_data


# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ pretrain/mockingjay/dataset.py ]
#   Synopsis     [ The bucketing datasets that apply the masked reconstruction task on-the-fly ]
#   Author       [ Andy T. Liu (https://github.com/andi611) ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
"""*********************************************************************************************"""

###############
# IMPORTATION #
###############
import os
import numpy as np
#-------------#
import torch
from torch.nn.utils.rnn import pad_sequence
import torchaudio
import scipy.io.wavfile
#-------------#
from pretrain.mockingjay.task import generate_masked_acoustic_model_data
from pretrain.bucket_dataset import FeatDataset

HALF_BATCHSIZE_TIME = 99999

class KaldiAcousticDataset(FeatDataset):

    def __init__(self, extracter, task_config, bucket_size, file_path, sets,
                 max_timestep=0, libri_root=None, **kwargs):
        super(KaldiAcousticDataset, self).__init__(extracter, task_config, bucket_size, file_path, sets,
                                                   max_timestep, libri_root, **kwargs)
        import pdb; pdb.set_trace()

    def _load_feat(self, feat_path):
        if self.libri_root is None:
            return torch.FloatTensor(np.load(os.path.join(self.root, feat_path)))
        else:
            wav, _ = torchaudio.load(os.path.join(self.libri_root, feat_path))
            feat = self.extracter(wav.squeeze())
            return feat

    def __getitem__(self, index):
        # Load acoustic feature and pad
        x_batch = [self._sample(self._load_feat(x_file)) for x_file in self.X[index]]
        import pdb; pdb.set_trace()
        x_pad_batch = pad_sequence(x_batch, batch_first=True)
        return generate_masked_acoustic_model_data(spec=(x_pad_batch,), config=self.task_config)

class OnlineAcousticDataset(FeatDataset):

    def __init__(self, extracter, task_config, bucket_size, file_path, sets,
                 max_timestep=0, libri_root=None, target_level=-25, **kwargs):
        max_timestep *= 160 #why 160?
        super(OnlineAcousticDataset, self).__init__(extracter, task_config, bucket_size, file_path, sets,
                                                    max_timestep, libri_root, **kwargs)
        import pdb; pdb.set_trace()
        self.target_level = target_level
        self.sample_length = self.sample_length * 160

    def _normalize_wav_decibel(self, wav):
        '''Normalize the signal to the target level'''
        if self.target_level == 'None':
            return wav
        rms = wav.pow(2).mean().pow(0.5)
        scalar = (10 ** (self.target_level / 20)) / (rms + 1e-10)
        wav = wav * scalar
        return wav

    def _load_feat(self, feat_path):
        #import pdb; pdb.set_trace()
        if self.libri_root is None:
            return torch.FloatTensor(np.load(os.path.join(self.root, feat_path)))
        else:
            #wav, _ = torchaudio.load(os.path.join(self.libri_root, feat_path))
            _, wav = scipy.io.wavfile.read(os.path.join(self.libri_root, feat_path))
            wav = np.expand_dims(wav, axis=0)
            wav = torch.Tensor(wav)
            wav = self._normalize_wav_decibel(wav.squeeze())
            return wav # (seq_len)

    def _process_x_pad_batch(self, x_pad_batch):
        import pdb; pdb.set_trace()
        if self.libri_root is not None:
            x_pad_batch = x_pad_batch.unsqueeze(1) # (batch_size, channel=1, seq_len)
            feat_list = self.extracter(x_pad_batch)
        return generate_masked_acoustic_model_data(feat_list, config=self.task_config)

    def __getitem__(self, index):
        # Load acoustic feature and pad
        x_batch = [self._sample(self._load_feat(x_file)) for x_file in self.X[index]]
        import pdb; pdb.set_trace()
        x_pad_batch = pad_sequence(x_batch, batch_first=True)
        import pdb; pdb.set_trace()
        return self._process_x_pad_batch(x_pad_batch)


import os
import argparse
import cv2
import numpy as np
import pandas as pd

import sys
sys.path.append('..')

nose_rect_index = {
    'aflw':[8,9,12,13,14],
    'wflw':[51,52,53,54,55,56,57,58,58,64,68]
}

def get_rect(index_array,landmarks,img_shape):
    """
    args:
        landmarks: (-1,2)
        index_array: [int,]

    return:
        left_top: (int,int)
        right_bottom : (int,int)
    """

    point_landmarks = np.array([landmarks[i] for i in index_array],dtype=np.float32)
    a0_max, a1_max = np.argmax(point_landmarks,axis=0)
    a0_min, a1_min = np.argmin(point_landmarks,axis=0)

    left_top_x = int(max(0,point_landmarks[a0_min][0]-np.random.randint(0,15)/1000.0 * img_shape[1]))
    left_top_y = int(max(0,point_landmarks[a1_min][1]-np.random.randint(0,15)/1000.0 * img_shape[0]))

    right_bottom_x = int(min(img_shape[1],point_landmarks[a0_max][0]+np.random.randint(0,15)/1000.0 * img_shape[1]))
    right_bottom_y = int(min(img_shape[0],point_landmarks[a1_max][1]+np.random.randint(0,15)/1000.0 * img_shape[0]))

    return (left_top_x,left_top_y),(right_bottom_x,right_bottom_y)
    # return (point_landmarks[a0_min][0],point_landmarks[a1_min][1]),(point_landmarks[a0_max][0],point_landmarks[a1_max][1])

def create_folder(str_path):
    paths = str_path.split('/')
    temp_folder = paths[0]
    for i in range (len(paths)-2):
        temp_folder = os.path.join(temp_folder,paths[i+1])
        if not os.path.exists(temp_folder):
            print("{} not exist , created.".format(temp_folder))
            os.mkdir(temp_folder)
    if os.path.exists(str_path):
        print("{} exist more than one face.".format(str_path))

def parse_args():
    parser = argparse.ArgumentParser(description='Gen Nose image dataset')

    parser.add_argument('--Dataset', help='experiment dataset name',
                        type=str, default="wflw")
    parser.add_argument('--Train_type', help='experiment dataset type',
                        type=str, default="test")
    parser.add_argument('--Gen_folder', help='generate dataset folder',
                        type=str, default="nose_images")

    args = parser.parse_args()
    return args

def main():

    args = parse_args()
    csv_file = os.path.join('data',args.Dataset,"face_landmarks_{}_{}.csv".format(args.Dataset,args.Train_type))
    image_folder = os.path.join('data',args.Dataset,"images")
    Gen_folder = os.path.join('data',args.Dataset,args.Gen_folder)
    if not os.path.exists(Gen_folder):
        os.mkdir(Gen_folder)
        print("Create folder : {}".format(Gen_folder))

    landmark_frame = pd.read_csv(csv_file)
    for i in range(landmark_frame.shape[0]):
        image_name = landmark_frame.iloc[i, 0]
        image_path = os.path.join(image_folder,image_name)
        gen_image_path = os.path.join(Gen_folder,image_name)
        # print(image_path)
        # print("Process image to {}".format(gen_image_path))
        img = cv2.imread(image_path)

        if os.path.exists(gen_image_path):
            print("{} have processed one face, add face inplace.".format(gen_image_path))
            nose_img = cv2.imread(gen_image_path)
        else:
            nose_img = np.zeros(img.shape,dtype=np.int)

        if args.Dataset == 'wflw':
            landmarks = landmark_frame.iloc[i,4:].values
        else :
            landmarks = landmark_frame.iloc[i,5:].values

        landmarks = landmarks.astype('int').reshape(-1, 2)

        left_top, right_bottom = get_rect(nose_rect_index[args.Dataset],landmarks,img.shape)

        nose_img[left_top[1]:right_bottom[1],left_top[0]:right_bottom[0]] = img[left_top[1]:right_bottom[1],left_top[0]:right_bottom[0]]

        # cv2.rectangle(nose_img,left_top,right_bottom,(0,255,255),1)
        # for l_i in range(landmarks.shape[0]) :
        #     cv2.circle(nose_img,(landmarks[l_i][0],landmarks[l_i][1]),1,(0,255,255),1,1)
        # #     cv2.putText(nose_img,str(l_i),(landmarks[l_i][0],landmarks[l_i][1]),cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,255),1)

        create_folder(gen_image_path)
        cv2.imwrite(gen_image_path,nose_img)
        # exit()

    # print(image_folder)

if __name__=="__main__":
    main()


import json
import os
import sys

import dateutil as du
import numpy as np
import pandas as pd
import pickle

from model.training_utils import load_raw_tracks, tracks_by_tag
from support.data_model import CLASSES, TAG_INDEXES
from test.test_utils import load_tracks, convert_remote_path, flatten_tag_tracks, load_model
from test.model_test import ModelTest

def frame_count(tracks):
    return int(np.sum([t.frame_count for t in tracks]))

def all_frame_counts(tag_tracks):
    return int(np.sum([frame_count(tag_tracks[t]) for t in CLASSES]))

def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    argv = sys.argv
    infos_path = argv[1]
    with open(argv[2]) as f:
        training_config = json.load(f)
    predicts_path = argv[3]
    arg_offset = 4
    save_predicts = False
    start_date = None
    while arg_offset < len(argv):
        if argv[arg_offset] == '--save_predicts':
            save_predicts = True
            arg_offset += 1
        elif argv[arg_offset] == '--start_date':
            start_date = du.parser.parse(argv[arg_offset+1])
            arg_offset += 2
        else:
            break
    test_name = os.path.split(infos_path)[1].split('_')[0]
    with open(predicts_path, 'rb') as f:
        full_data = pickle.load(f)
    track_data = {(t.clip_key, t.track_key): t.data for t in full_data}
    test_tracks = load_raw_tracks(infos_path)
    if start_date is not None:
        test_tracks = [t for t in test_tracks if t.start_time >= start_date]
        test_name += '_after_' + start_date.isoformat(timespec='minutes')
    input_values = training_config['input_values']
    input_scales = training_config['input_scales']
    for track in test_tracks:
        track.data = build_sample_data(track, track_data[(track.clip_key, track.track_key)], input_values, **input_scales)
    actuals = [TAG_INDEXES[t.tag] for t in test_tracks]

    text = f'\nTesting with {len(track)} tracks'
    if start_date is not None:
        text += ' from ' + start_date.isoformat()
    print(text + '\n')
    for key in CLASSES:
        print(f'{key:12}  {frame_count(tag_tracks[key]):>7}')
    print()

    for weights_path in argv[arg_offset:]:
        weights_path = convert_remote_path(weights_path)
        model = load_model(weights_path)

        for track in test_tracks:
            predicts = model.predict(track.data)
            class_index = np.argmax(predicts)

        track_predicts, frame_predicts = model_test.test(test_tracks, True)
        if save_predicts:
            frame_data = {
                'clip_id': [t.clip_key for t in test_tracks],
                'track_id': [t.track_key for t in test_tracks],
                'device_name': [t.device_name for t in test_tracks],
                'time': [t.start_time for t in test_tracks],
                'tag': [t.tag for t in test_tracks]
            }
            for i in range(len(track_predicts)):
                frame_data[f'predict {i}'] = [CLASSES[p] for p in track_predicts[i]]
            df = pd.DataFrame(frame_data)
            df.to_csv(f'{weights_path}-{test_name}-predicts.csv', index=False)
            assert len(test_tracks) == len(frame_predicts)
            for track, predicts in zip(test_tracks, frame_predicts):
                track.data = predicts
            with open(f'{weights_path}-{test_name}-predicts.pk', 'wb') as f:
                pickle.dump(test_tracks, f)

if __name__ == '__main__':
    sys.exit(main())