import unittest

import sys
import os
import math
from itertools import islice
import numpy as np
import copy
import ctypes
import galsim
import yaml
from astropy.io import fits

from observation_sim.instruments.chip import chip_utils
import matplotlib.pyplot as plt


# test FUNCTION --- START #
def add_crosstalk(GSimg):
    crosstalk = np.zeros([16, 16])
    crosstalk[0, :] = np.array([1., 1e-4, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[1, :] = np.array([1e-4, 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[2, :] = np.array([0., 0., 1., 1e-4, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[3, :] = np.array([0., 0., 1e-4, 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[4, :] = np.array([0., 0., 0., 0., 1., 1e-4, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[5, :] = np.array([0., 0., 0., 0., 1e-4, 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[6, :] = np.array([0., 0., 0., 0., 0., 0., 1., 1e-4, 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[7, :] = np.array([0., 0., 0., 0., 0., 0., 1e-4, 1., 0., 0., 0., 0., 0., 0., 0., 0.])
    crosstalk[8, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 1., 1e-4, 0., 0., 0., 0., 0., 0.])
    crosstalk[9, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 1e-4, 1., 0., 0., 0., 0., 0., 0.])
    crosstalk[10, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1e-4, 0., 0., 0., 0.])
    crosstalk[11, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1e-4, 1., 0., 0., 0., 0.])
    crosstalk[12, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1e-4, 0., 0.])
    crosstalk[13, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1e-4, 1., 0., 0.])
    crosstalk[14, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1e-4])
    crosstalk[15, :] = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1e-4, 1.])

    # 2*8 -> 1*16
    img = chip_utils.formatOutput(GSimg)
    ny, nx = img.array.shape
    nsecy = 1
    nsecx = 16
    dy = int(ny/nsecy)
    dx = int(nx/nsecx)

    newimg = galsim.Image(nx, ny, init_value=0)
    for i in range(16):
        for j in range(16):
            newimg.array[:, int(i*dx):int(i*dx+dx)] += crosstalk[i, j]*img.array[:, int(j*dx):int(j*dx+dx)]

    # 1*16 -> 2*8
    newimg = chip_utils.formatRevert(newimg)

    return newimg 
# test FUNCTION --- END #


class detModule_coverage(unittest.TestCase):
    def __init__(self, methodName='runTest'):
        super(detModule_coverage, self).__init__(methodName)
        self.dataPath = os.path.join(
            os.getenv('UNIT_TEST_DATA_ROOT'), 'csst_msc_sim/csst_fz_msc')

    def test_add_crosstalk(self):
        nsecy = 2
        nsecx = 8
        ny, nx = 1024,1024
        dy = int(ny/nsecy)
        dx = int(nx/nsecx)
        mapclip = np.zeros([dy, int(nsecx*nsecy*dx)])
        for i in range(int(nsecx*nsecy)):
            mapclip[:, i*dx:dx+i*dx] = np.random.randint(10)+np.random.rand(int(dy*dx)).reshape([dy, dx])
        mapclip = galsim.ImageF(mapclip)

        nsecy = 1
        nsecx = 16
        ny,nx = mapclip.array.shape
        dy = int(ny/nsecy)
        dx = int(nx/nsecx)
        for i in range(int(nsecy*nsecx)):
            gal = galsim.Gaussian(sigma=0.2, flux=500000).drawImage(nx=32, ny=32).array
            py = np.random.randint(450)+10
            mapclip.array[py:py+32, int(i*dx)+10:int(i*dx)+42] += gal
        
        tmap = chip_utils.formatRevert(mapclip, nsecy=1, nsecx=16)  # 1*16 -> 2*8
        temp = add_crosstalk(tmap)

        fig = plt.figure(figsize=(20,60))
        ax = plt.subplot(311)
        plt.imshow(np.log10(mapclip.array+1), origin='lower', cmap='gray')
        ax = plt.subplot(312)
        plt.imshow(np.log10(temp.array+1), origin='lower', cmap='gray')
        ax = plt.subplot(313)
        plt.imshow(np.log10(temp.array-tmap.array+1), origin='lower', cmap='gray')
        plt.savefig(os.path.join(self.dataPath, "./test_crosstalk.png"), dpi=300, bbox_inches='tight')
        self.assertTrue(True)


if __name__ == '__main__':
    unittest.main()