import unittest

import sys,os,math
from itertools import islice

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.use('Agg')

import yaml
from ObservationSim.Config import Config
from ObservationSim.Config.Config import config_dir
from ObservationSim.Instrument import Chip

from ObservationSim.PSF.PSFInterp import PSFInterp


def defineCCD(iccd, config_file):
    with open(config_file, "r") as stream:
        try:
            config = yaml.safe_load(stream)
            #for key, value in config.items():
            #    print (key + " : " + str(value))
        except yaml.YAMLError as exc:
            print(exc)
    path_dict = config_dir(config=config, work_dir=config['work_dir'], data_dir=config['data_dir'])
    chip = Chip(chipID=iccd, config=config)
    #chip = Chip(chipID=iccd, ccdEffCurve_dir=path_dict["ccd_dir"], CRdata_dir=path_dict["CRdata_dir"], normalize_dir=path_dict["normalize_dir"], sls_dir=path_dict['sls_dir'], config=config)
    return chip

def psfSecondMoments(psfMat, cenX, cenY, pixSize=1):
    apr = 0.5 #arcsec, 0.5角秒内测量
    fl  = 28. #meters
    pxs = 2.5 #microns
    apr = np.deg2rad(apr/3600.)*fl*1e6
    apr = apr/pxs
    apr = np.int(np.ceil(apr))

    I = psfMat
    ncol = I.shape[1]
    nrow = I.shape[0]
    w   = 0.0
    w11 = 0.0
    w12 = 0.0
    w22 = 0.0
    for icol in range(ncol):
        for jrow in range(nrow):
            x = icol*pixSize - cenX
            y = jrow*pixSize - cenY
            rr = np.sqrt(x*x + y*y)
            wgt= 0.0
            if rr <= apr:
                wgt = 1.0
            w   += I[jrow, icol]*wgt
            w11 += x*x*I[jrow, icol]*wgt
            w12 += x*y*I[jrow, icol]*wgt
            w22 += y*y*I[jrow, icol]*wgt
    w11 /= w
    w12 /= w
    w22 /= w
    sz = w11 + w22
    e1 = (w11 - w22)/sz
    e2 = 2.0*w12/sz

    return sz, e1, e2


def test_psfEll(iccd, iwave, psfMat):
    psfMat_iwave = psfMat.psfMat[iwave-1, :,:,:]
    npsf = np.shape(psfMat_iwave)[0]

    imx = np.zeros(npsf)
    imy = np.zeros(npsf)
    psf_e1 = np.zeros(npsf)
    psf_e2 = np.zeros(npsf)
    psf_sz = np.zeros(npsf)

    for ipsf in range(1, npsf+1):
        print('ipsf-{:}'.format(ipsf), end='\r')
        imx[ipsf-1] = psfMat.cen_col[iwave-1, ipsf-1]
        imy[ipsf-1] = psfMat.cen_row[iwave-1, ipsf-1]
        psfMat_iwave_ipsf = psfMat_iwave[ipsf-1, :, :]
        cenX = 256
        cenY = 256
        sz, e1, e2  = psfSecondMoments(psfMat_iwave_ipsf, cenX, cenY, pixSize=1)
        psf_e1[ipsf-1] = e1
        psf_e2[ipsf-1] = e2
        psf_sz[ipsf-1] = sz
        #print('ell======', ipsf, np.sqrt(e1**2 + e2**2))
    #######
    arr = [imx, imy, psf_e1, psf_e2, psf_sz]
    np.save('data/psfEll{:}_{:}_{:}'.format(int(np.sqrt(npsf)),iccd, iwave), arr)


def test_psfEllPlot(OVERPLOT=False):
    #if ThisTask == 0:
    if True:
        prefix = 'psfEll30'
        iccd = 1
        iwave= 1
        data = np.load('data/'+prefix+'_1_1.npy')
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        #######
        plt.cla()
        plt.close("all")

        fig = plt.figure(figsize=(12, 12))
        plt.subplots_adjust(wspace=0.1, hspace=0.1)
        ax = plt.subplot(1, 1, 1)
        for ipsf in range(npsf):
            plt.plot(imx[ipsf], imy[ipsf], 'r.')

            ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
            ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
            ell *= 15
            lcos = ell*np.cos(ang)
            lsin = ell*np.sin(ang)
            plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'r', lw=2)
            ###########
        ang = 0.
        ell = 0.05
        ell*= 15
        lcos = ell*np.cos(ang)
        lsin = ell*np.sin(ang)
        plt.plot([imx[898]-lcos, imx[898]+lcos],[imy[898]+5.-lsin, imy[898]+5.+lsin],'k', lw=2)
        plt.annotate('{:}'.format(ell/15), (imx[898]-2., imy[898]+6.), xycoords='data', fontsize=10)
        plt.xlabel('CCD X (mm)')
        plt.ylabel('CCD Y (mm)')

        if OVERPLOT == True:
            prefix = 'psfEll20'
            data = np.load('data/'+prefix+'_1_1.npy')
            imx= data[0]
            imy= data[1]
            psf_e1 = data[2]
            psf_e2 = data[3]

            npsf = np.shape(imx)[0]
            for ipsf in range(npsf):
                plt.plot(imx[ipsf], imy[ipsf], 'b.')

                ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
                ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
                ell *= 15
                lcos = ell*np.cos(ang)
                lsin = ell*np.sin(ang)
                plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'b', lw=2)

        plt.gca().set_aspect(1)
        if OVERPLOT == True:
            prefix = 'psfEllOP'
        plt.savefig('figs/'+prefix+'_iccd{:}.pdf'.format(iccd))


def test_psfIDW(iccd, iwave, psfMatA, chip, psfMatB):
    bandpass = iwave-1
    class pos_img():
        def __init__(self,x, y):
            self.x = x*1e3/10.  #in unit of pixels
            self.y = y*1e3/10.

    psfMat_iwave = psfMatA.psfMat[iwave-1, :,:,:]
    npsf = np.shape(psfMat_iwave)[0]
    psf_e1 = np.zeros(npsf)
    psf_e2 = np.zeros(npsf)
    psf_sz = np.zeros(npsf)

    for ipsf in range(1, npsf+1):
        print('ipsf:', ipsf, end='\r', flush=True)
        tpos_img = pos_img(psfMatA.cen_col[iwave-1, ipsf-1], psfMatA.cen_row[iwave-1, ipsf-1])
        psfIDW = psfMatB.get_PSF(chip, tpos_img, bandpass, galsimGSObject=False, findNeighMode='treeFind')
        np.save('figs/psfIDW/psfIDW_{:}_{:}_{:}'.format(iccd, iwave, ipsf), psfIDW)

        cenX = 256
        cenY = 256
        sz, e1, e2  = psfSecondMoments(psfIDW, cenX, cenY, pixSize=1)
        psf_e1[ipsf-1] = e1
        psf_e2[ipsf-1] = e2
        psf_sz[ipsf-1] = sz
    arr = [psf_e1, psf_e2, psf_sz]
    np.save('data/psfEll20IDW_{:}_{:}'.format(iccd, iwave), arr)


def test_psfResidualPlot(iccd, iwave, ipsf, psfMatA):
    psfMat_iwave = psfMatA.psfMat[iwave-1, :,:,:]

    psfMatORG = psfMat_iwave[ipsf-1, :, :]
    psfMatIDW = np.load('figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf))

    npix = psfMatORG.shape[0]
    pixCutEdge= int(npix/2-15)

    img0 = psfMatORG[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge]
    img1 = psfMatIDW[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge]
    imgX = (img1 - img0)/img0

    img0 = np.log10(img0)
    img1 = np.log10(img1)
    imgX = np.log10(np.abs(imgX))

    fig = plt.figure(figsize=(18,4))
    ax = plt.subplot(1,3,1)
    plt.imshow(img0, origin='lower', vmin=-7, vmax=-1.3)
    plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--')
    plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--')
    plt.annotate('ORG', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    cticks=[-7, -6, -5, -4, -3, -2, -1]
    cbar = plt.colorbar(ticks=cticks)
    cbar.ax.set_yticklabels(['$10^{-7}$', '$10^{-6}$', '$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$'])
    print(img0.min(), img0.max())


    ax = plt.subplot(1,3,2)
    plt.imshow(img1, origin='lower', vmin=-7, vmax=-1.3)
    plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--')
    plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--')
    plt.annotate('IDW', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    cticks=[-7, -6, -5, -4, -3, -2, -1]
    cbar = plt.colorbar(ticks=cticks)
    cbar.ax.set_yticklabels(['$10^{-7}$', '$10^{-6}$', '$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$'])
    print(img1.min(), img1.max())


    ax = plt.subplot(1,3,3)
    plt.imshow(imgX, origin='lower', vmin =-3, vmax =np.log10(3e-1))
    plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--')
    plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--')
    #plt.annotate('(IDW-ORG)/ORG', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15)
    cticks=[-5, -4, -3, -2, -1]
    cbar = plt.colorbar(ticks=cticks)
    cbar.ax.set_yticklabels(['$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$'])

    print(np.max((psfMatORG-psfMatIDW)))
    plt.savefig('figs/psfResidual_iccd{:}.pdf'.format(iccd))


def test_psfEllIDWPlot(OVERPLOT=False):
    #if ThisTask == 0:
    if True:
        prefix = 'psfEll20'
        iccd = 1
        iwave= 1
        data = np.load('data/'+prefix+'_1_1.npy')
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        #######
        plt.cla()
        plt.close("all")

        fig = plt.figure(figsize=(12, 12))
        plt.subplots_adjust(wspace=0.1, hspace=0.1)
        ax = plt.subplot(1, 1, 1)
        for ipsf in range(npsf):
            plt.plot(imx[ipsf], imy[ipsf], 'b.')

            ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
            ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
            ell *= 15
            lcos = ell*np.cos(ang)
            lsin = ell*np.sin(ang)
            plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'b', lw=2)
            ###########
        ang = 0.
        ell = 0.05
        ell*= 15
        lcos = ell*np.cos(ang)
        lsin = ell*np.sin(ang)
        #plt.plot([imx[898]-lcos, imx[898]+lcos],[imy[898]+5.-lsin, imy[898]+5.+lsin],'k', lw=2)
        #plt.annotate('{:}'.format(ell/15), (imx[898]-2., imy[898]+6.), xycoords='data', fontsize=10)
        plt.xlabel('CCD X (mm)')
        plt.ylabel('CCD Y (mm)')

        if OVERPLOT == True:
            prefix = 'psfEll20IDW'
            data = np.load('data/'+prefix+'_1_1.npy')
            #imx= data[0]
            #imy= data[1]
            psf_e1 = data[0]
            psf_e2 = data[1]

            npsf = np.shape(imx)[0]
            for ipsf in range(npsf):
                #plt.plot(imx[ipsf], imy[ipsf], 'r.')

                ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2
                ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2)
                ell *= 15
                lcos = ell*np.cos(ang)
                lsin = ell*np.sin(ang)
                plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'r', lw=1)

        plt.gca().set_aspect(1)
        if OVERPLOT == True:
            prefix = 'psfEllOPIDW'
        plt.savefig('figs/'+prefix+'_iccd{:}.pdf'.format(iccd))


def test_psfdEllabsPlot(iccd):
    #if ThisTask == 0:
    if True:
        prefix = 'psfEll20'
        #iccd = 1
        #iwave= 1
        data = np.load('data/'+prefix+'_{:}_1.npy'.format(iccd))
        imx= data[0]
        imy= data[1]
        psf_e1 = data[2]
        psf_e2 = data[3]
        psf_sz = data[4]
        print(np.shape(imx))
        npsf = np.shape(imx)[0]

        ellX = np.sqrt(psf_e1**2 + psf_e2**2)
        angX = np.arctan2(psf_e2, psf_e1)/2
        angX = np.rad2deg(angX)
        szX  = psf_sz

        ##############################
        prefix = 'psfEll20IDW'
        data = np.load('data/'+prefix+'_{:}_1.npy'.format(iccd))
        #imx= data[0]
        #imy= data[1]
        psf_e1 = data[0]
        psf_e2 = data[1]
        psf_sz = data[2]

        ellY = np.sqrt(psf_e1**2 + psf_e2**2)
        angY = np.arctan2(psf_e2, psf_e1)/2
        angY = np.rad2deg(angY)
        szY  = psf_sz

        ##############################
        fig=plt.figure(figsize=(6, 5))
        grid = plt.GridSpec(3,1,left=0.15, right=0.95, wspace=None, hspace=0.02)
        #plt.subplots_adjust(left=None,bottom=None,right=None,top=None,wspace=None,hspace=0.02)
        ax = plt.subplot(grid[0:2,0])
        plt.plot([0.01,0.1],[0.01,0.1], 'k--', lw=1. )
        plt.scatter(ellX, ellY, color='b', alpha=1., s=3., edgecolors='None')
        #plt.xlim([0.015, 0.085])
        #plt.ylim([0.015, 0.085])
        plt.ylabel('$\epsilon_{\\rm IDW}$')
        plt.gca().axes.get_xaxis().set_visible(False)

        ax = plt.subplot(grid[2,0])
        plt.plot([0.015,0.085],[0.,0.], 'k--', lw=1. )
        plt.scatter(ellX, (ellY-ellX), color='b', s=3., edgecolors='None')
        #plt.xlim([0.015, 0.085])
        #plt.ylim([-0.0018, 0.0018])
        plt.xlabel('$\epsilon_{\\rm ORG}$')
        plt.ylabel('$\Delta$')
        plt.savefig('figs/psfEllOPIDWPDF_{:}.pdf'.format(iccd))

        fig=plt.figure(figsize=(6, 6))
        plt.hist((szY-szX)/szX, bins=20, color='r', alpha=0.5)
        plt.xlabel('$(R_{\\rm IDW}-R_{\\rm ORG})/R_{\\rm ORG}$')
        plt.ylabel('PDF')
        plt.savefig('figs/psfEllOPIDWPDF_dsz_{:}.pdf'.format(iccd))


class PSFInterpModule_coverage(unittest.TestCase):
    def test_psfEll_(self):
        iccd = 1
        iwave= 1

        config_file = "/public/home/weichengliang/CSST_git/newVersion/CSST/config/config_C3.yaml"
        chip = defineCCD(iccd, config_file)
        print(chip.chipID)
        print(chip.cen_pix_x, chip.cen_pix_y)

        ipath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/psfCubeTest'
        psfMatA = PSFInterp(chip, npsf=400, PSF_data_file=ipath, PSF_data_prefix="S20x20_")
        psfMatB = PSFInterp(chip, npsf=900, PSF_data_file=ipath, PSF_data_prefix="S30x30_")

        test_psfEll(iccd, iwave, psfMatA)
        test_psfEll(iccd, iwave, psfMatB)
        test_psfEllPlot(OVERPLOT=True)

        test_psfIDW(iccd, iwave, psfMatA, chip, psfMatB)
        ipsf = 1
        test_psfResidualPlot(iccd, iwave, ipsf, psfMatA)
        test_psfEllIDWPlot(OVERPLOT=True)
        test_psfdEllabsPlot(iccd)


if __name__ == '__main__':
    unittest.main()
    print('#####haha#####')