"""
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

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.pri_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.pri_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()
