"""
Unit test suite for the calibrate_kgain module.

Frames for unit tests are simulated SCI-size frames made with Gaussian and Poisson 
noises included. The assumed flux map is a realistic pupil image made from 
TVAC frames.
"""

import os
import pytest
import warnings
import numpy as np
import pandas as pd
from pathlib import Path
from astropy.io import fits

import test_check
from corgidrp import check
from corgidrp.data import Image, Dataset
from corgidrp.mocks import (create_default_headers, make_fluxmap_image, nonlin_coefs)
from corgidrp.calibrate_kgain import (calibrate_kgain, CalKgainException, kgain_params)

######################## function definitions ###############################

def count_contiguous_repeats(arr):
    """
    This function returns the count of the contiguous repeated exposure times.
    If input array is empty, returns empty array.

    Args:
        arr (np.array): 1d array of exposure times.

    Returns:
        counts (list): count of the contiguous repeated exposure times
    """

    if isinstance(arr, (np.ndarray, list)) and len(arr) == 0:
        return []
    counts = []
    current_count = 1
    for i in range(1, len(arr)):
        if arr[i] == arr[i - 1]:
            current_count += 1
        else:
            counts.append(current_count)
            current_count = 1
    counts.append(current_count)  # append the count of the last element
    return counts

############### make stacks with simulated frames ##################

# path to nonlin table made from running calibrate_nonlin.py on TVAC frames
# table used only to choose parameters to make analytic nonlin functions
here = os.path.abspath(os.path.dirname(__file__))
nonlin_table_path = Path(here,'test_data','nonlin_table_TVAC.txt')
nonlin_flag = False # True adds nonlinearity to simulated frames

# Load the arrays needed for calibrate_nonlin function from the .npz file
loaded = np.load(Path(here,'test_data','nonlin_arrays_ut.npz'))
# Access the arrays needed for calibrate_nonlin function
exp_time_stack_arr0 = loaded['array1']
time_stack_arr0 = loaded['array2']
len_list0 = loaded['array3']
# Reducing the number of frames used in unit tests (each has 5 substacks).
# Set to False to run the same test as in the IIT pipeline
rn_in = 130 # read noise in e-
if True:
    n_cal = 3
    exp_time_stack_arr0 = np.delete(exp_time_stack_arr0, np.s_[n_cal*5:])
    time_stack_arr0 = np.delete(time_stack_arr0, np.s_[n_cal*5:])
    # Update len_list0
    len_list0[0] = n_cal*5
    # Usual number of frames to deal with real rn values is ~200
    rn_in = 130/np.sqrt(200/n_cal) # read noise in e-

# Load the flux map
fluxmap_init =  np.load(Path(here,'test_data','FluxMap1024.npy'))
fluxmap_init[fluxmap_init < 50] = 0 # cleanup flux map a bit
fluxMap = 0.8*fluxmap_init # e/s/px, for G = 1

# assumed detector parameters
kgain_in = 8.7 # e-/DN
bias = 2000 # e-
actual_gain = 1.0

# cubic function nonlinearity for emgain of 1
if nonlin_flag:
    coeffs_1, DNs, _ = nonlin_coefs(nonlin_table_path,1.0,3)
else:
    coeffs_1 = [0.0, 0.0, 0.0, 1.0]
    _, DNs, _ = nonlin_coefs(nonlin_table_path,1.0,3)

frame_list = []
# make some uniform frames with emgain = 1 (must be unity) P.S. IIT would use ~30
n_mean = 3
for j in range(n_mean):
    image_sim = make_fluxmap_image(fluxMap,bias,kgain_in,rn_in, 1, 7.0,coeffs_1,
        nonlin_flag=nonlin_flag)
    # Datetime cannot be duplicated
    image_sim.ext_hdr['DATETIME'] = time_stack_arr0[j]
    # Temporary keyword value. Mean frame is TBD
    image_sim.ext_hdr['OBSTYPE'] = 'MNFRAME'
    frame_list.append(image_sim)

index = 0
iG = 0 # doing only the em gain = 1 case
g = actual_gain # Note: Same value for all frames used to calibrate K-gain
exp_time_loop = exp_time_stack_arr0[index:index+len_list0[iG]]
index = index + len_list0[iG]
if nonlin_flag:
    coeffs, _, vals = nonlin_coefs(nonlin_table_path,g,3)
else:
    coeffs = [0.0, 0.0, 0.0, 1.0]
    vals = np.ones(len(DNs))

exp_repeat_counts = count_contiguous_repeats(exp_time_loop)
for j in range(len(exp_repeat_counts)):
    for t in range(exp_repeat_counts[j]):
        # Simulate full frame
        exp_time = exp_time_loop[t+j*exp_repeat_counts[j]]
        image_sim = make_fluxmap_image(fluxMap,bias,kgain_in,rn_in,g,
                               exp_time,coeffs,nonlin_flag=nonlin_flag,
                               divide_em=True)
        image_sim.ext_hdr['DATETIME'] = time_stack_arr0[t+j*exp_repeat_counts[j]]
        # OBSTYPE has no KGAIN value, but NONLIN
        image_sim.ext_hdr['OBSTYPE'] = 'NONLIN'
        frame_list.append(image_sim)
dataset_kg = Dataset(frame_list)

# set input parameters for calibrate_kgain function
min_val = 800
max_val = 3000
binwidth = 68

def test_expected_results_sub():
    """Outputs are as expected, for imported frames."""
    kgain = calibrate_kgain(dataset_kg, n_cal, n_mean, min_val, max_val, binwidth)
        
    signal_bins_N = kgain_params['signal_bins_N']
    # kgain - should be close to the assumed value
    assert np.isclose(round(kgain.value,1), kgain_in, atol=0.5)
    assert np.all(np.equal(kgain.ptc.shape, (signal_bins_N,2)))

def test_psi():
    """These three below must be positive scalar integers."""
    check_list = test_check.psilist
    # min_val
    for perr in check_list:
        with pytest.raises(TypeError):
            calibrate_kgain(dataset_kg, n_cal, n_mean, perr, max_val, binwidth)
    # max_val
    for perr in check_list:
        with pytest.raises(TypeError):
            calibrate_kgain(dataset_kg, n_cal, n_mean, min_val, perr, binwidth)

    # binwidth
    for perr in check_list:
        with pytest.raises(TypeError):
            calibrate_kgain(dataset_kg, n_cal, n_mean, min_val, max_val, perr)
      
def test_binwidth():
    """binwidth must be >= 10."""
    with pytest.raises(CalKgainException):
        calibrate_kgain(dataset_kg, n_cal, n_mean, min_val, max_val, 9)
 
if __name__ == '__main__':
    print('Running test_expected_results_sub')
    test_expected_results_sub()
    print('Running test_psi')
    test_psi()
    print('Running test_binwidth')
    test_binwidth()
