from corgidrp.mocks import create_psfsub_dataset,create_ct_cal, gaussian_array
from corgidrp.klip_fm import meas_klip_thrupt, get_closest_psf, inject_psf, measure_noise
from corgidrp.l3_to_l4 import do_psf_subtraction
from corgidrp.astrom import centroid, seppa2xy, create_circular_mask
from corgidrp.data import Image
from astropy.io import fits
from scipy.ndimage import shift, rotate
from pyklip.fakes import gaussfit2d


import pytest
import numpy as np
import os
import itertools
## Helper functions/quantities

lam = 573.8e-9 #m
d = 2.36 #m  
pixscale_mas = 0.0218 * 1000 # mas
fwhm_mas = 1.22 * lam / d * 206265. * 1000.
fwhm_pix = fwhm_mas / pixscale_mas  
res_elem = 3 * fwhm_pix # pix

owa_mas = 450. 
owa_pix = owa_mas / pixscale_mas   

iwa_mas = 140. 
iwa_pix = iwa_mas / pixscale_mas  

# Overriding separations for testing purposes
seps = np.array([10.,15.,20.,25.,30.])

# Mock CT calibration settings

nx,ny = (21,21)
cenx, ceny = (25.,30.)

# Mock PSF subtraction data settings
st_amp = 100.
noise_amp = 1e-3
pl_contrast = 0.0
pa_aper_degs = [0,10.,0,0]

# Injection test settings
inj_flux = 10.

# KLIP throughput calculation settings
inject_snr = 10.
outdir = 'klipcal_output'
fileprefix = 'FAKE'
annuli = 1
subsections = 1
movement = 1
numbasis = [1]
calibrate_flux = False

# max_thrupt_tolerance = 1 + (noise_amp * (2*fwhm_pix)**2)
max_thrupt_tolerance = 1.05

if not os.path.exists(outdir):
    os.mkdir(outdir)
   
## pyKLIP data class tests

def test_create_ct_cal():
    """Test that mocks.create_ct_cal() generates the correct number of PSFs, 
    each with the correct shape, and that each PSF is a predictable amplitude 
    for debugging purposes.
    """

    nx,ny = (3,5)
    
    cenx,ceny = (10.5,20.5)

    n_psfs = nx * ny

    ctcal = create_ct_cal(fwhm_mas, cfam_name='1F',
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny)
    
    # Check that the correct number of psfs are created
    assert ctcal.ct_excam.shape == (3,n_psfs)

    # Check that each psf is the correct amplitude 
    # (scales with i for debugging purposes)
    for i in range(1,n_psfs+1):
        psf = ctcal.data[i-1]
        assert np.max(psf) == pytest.approx(ctcal.ct_excam[2][i-1],rel=0.01)

        # Check that psf is odd shape and is centered
        assert np.all(np.array(psf.shape) % 2 == 1)
        assert np.array(centroid(psf)) == pytest.approx(np.array(psf.shape)/2.-0.5)


def test_get_closest_psf():
    """Test that the correct PSF is grabbed from the CT Calibration 
    object for each position.
    """
    # Should get the first PSF
    nx,ny = (3,5)
    ct_cenx,ct_ceny = (10.5,20.5)
    cenx,ceny = (10.0,20.0)
    ctcal = create_ct_cal(fwhm_mas, cfam_name='1F',
                  cenx=ct_cenx,ceny=ct_ceny,
                  nx=nx,ny=ny)
    
    goal_dxdy = (-100,-100)
    psf = get_closest_psf(ctcal,cenx,ceny,*goal_dxdy)
    assert psf == pytest.approx(ctcal.data[0])

    # Should get the last PSF
    goal_dxdy = (100,100)
    psf = get_closest_psf(ctcal,cenx,ceny,*goal_dxdy)
    assert psf == pytest.approx(ctcal.data[-1])

    # Should get the center PSF (at index 7)
    goal_dxdy = (0.,0.)
    psf = get_closest_psf(ctcal,cenx,ceny,*goal_dxdy)
    assert psf == pytest.approx(ctcal.data[7])

    # Should get the PSF at (9.5,21.5) (at index 3)
    goal_dxdy = (-1.2,1.1)
    psf = get_closest_psf(ctcal,cenx,ceny,*goal_dxdy)
    assert psf == pytest.approx(ctcal.data[3])


def test_inject_psf():
    """Test that PSFs are injected with the correct amplitude and peak pixel location
    """
    # Mock CT cal object
    nx,ny = (21,21)
    cenx, ceny = (25.,30.)
    ctcal = create_ct_cal(fwhm_mas, cfam_name='1F',
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny,
                  psfsize=5)

    # Test 0 separation to make sure we can scale and center
    sep = 0.0
    pa = 0.0
    pa_aper_deg = 0.0
    frame_shape_yx = (50,60)
    expected_peak = (ceny,cenx)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)
    
    assert np.max(frame_out.data) == pytest.approx(inj_flux)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak
    
    # Test separation of exactly 1 pixel
    sep = 1.0
    pa = 0.0
    pa_aper_deg = 0.0
    frame_shape_yx = (50,60)
    expected_peak = (ceny+1,cenx)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)
    
    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak

    # Test PA
    sep = 1.0
    pa = 90.0
    pa_aper_deg = 0.0
    frame_shape_yx = (50,60)
    expected_peak = (ceny,cenx-1)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)
    
    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak

    # Test Rotation
    sep = 5.0
    pa = 0.0
    pa_aper_deg = 90.0
    frame_shape_yx = (50,60)
    expected_peak = (ceny,cenx+5)

    pri_hdr = fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)
    
    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak

    # Test injecting a psf over the left edge of the array
    sep = 9.0
    pa = 90.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,21)
    ceny,cenx = (5,10)
    expected_peak_yx = (5,1)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx

    # Test injecting a psf over the right edge of the array
    sep = 9.0
    pa = -90.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,21)
    ceny,cenx = (5,10)
    expected_peak_yx = (5,19)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx

    
    # Test injecting a psf over the top edge of the array
    sep = 4.0
    pa = 0.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,21)
    ceny,cenx = (5,10)
    expected_peak_yx = (9,10)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx


    # Test injecting a psf over the bottom edge of the array
    sep = 4.0
    pa = 180.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,21)
    ceny,cenx = (5,10)
    expected_peak_yx = (1,10)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx


    # Test injecting a psf over the bottom left corner of the array
    sep = 4.0 * np.sqrt(2)
    pa = 135.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,11)
    ceny,cenx = (5,5)
    expected_peak_yx = (1,1)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx


    # Test injecting a psf over the bottom right corner of the array
    sep = 4.0 * np.sqrt(2)
    pa = -135.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,11)
    ceny,cenx = (5,5)
    expected_peak_yx = (1,9)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx

    # Test injecting a psf over the top left corner of the array
    sep = 4.0 * np.sqrt(2)
    pa = 45.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,11)
    ceny,cenx = (5,5)
    expected_peak_yx = (9,1)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx
    
    # Test injecting a psf over the top right corner of the array
    sep = 4.0 * np.sqrt(2)
    pa = -45.0
    pa_aper_deg = 0.0
    frame_shape_yx = (11,11)
    ceny,cenx = (5,5)
    expected_peak_yx = (9,9)

    pri_hdr=fits.Header()
    ext_hdr = fits.Header()
    pri_hdr['PA_APER'] = pa_aper_deg
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    frame = Image(np.zeros(frame_shape_yx),
                  pri_hdr=pri_hdr,
                  ext_hdr=ext_hdr)

    frame_out, psf_model, psf_cenxy = inject_psf(frame, ctcal,inj_flux, sep,pa)

    assert np.unravel_index(np.argmax(frame_out.data),frame_out.data.shape) == expected_peak_yx


def test_measure_noise():
    """Check that annular noise profile measurement produces arrays with 
    the correct shape and values for a case with zero noise, and for a case 
    with different noise levels within and without a given radius of pixels.
    Also checks that klip_fm.measure_noise() can return results for only a 
    specific KL mode truncation if desired.
    """

    cenx,ceny = (120.,130.)
    frame_shape_yx = (200,200)
    seps = np.arange(50.,71.,5.)
    border_sep = 62.5
    
    fwhm = 2. # pix

    ext_hdr = fits.Header()
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny
    
    image = np.zeros((2,*frame_shape_yx))

    frame = Image(image,
                  pri_hdr=fits.Header(),
                  ext_hdr=ext_hdr)

    # No KL mode specified
    noise_profile = measure_noise(frame, seps, fwhm)

    # Check data shape
    assert noise_profile.shape == (5,2)
    assert noise_profile == pytest.approx(0.)

    # Different noise values for each KL mode 
    # and inner/outer region

    std1 = 1.
    std2 = 2.
    std3 = 3.
    std4 = 4.

    image = np.zeros((2,*frame_shape_yx))
    frame = Image(image,
                  pri_hdr=fits.Header(),
                  ext_hdr=ext_hdr)

    y, x = np.indices(frame.data.shape[1:])
    sep_map = np.sqrt((y-ceny)**2 + (x-cenx)**2)
    

    rng = np.random.default_rng(0)
    frame.data[0,:,:] = np.where(sep_map>border_sep,rng.normal(0.,std1,frame_shape_yx),rng.normal(0.,std2,frame_shape_yx))
    frame.data[1,:,:] = np.where(sep_map>border_sep,rng.normal(0.,std3,frame_shape_yx),rng.normal(0.,std4,frame_shape_yx))

    # No KL mode specified
    noise_profile = measure_noise(frame, seps, fwhm)

    # Check data shape
    assert noise_profile.shape == (5,2)

    # Check that std measurement is correct
    assert noise_profile[:3,0] == pytest.approx(std2,rel=0.05)
    assert noise_profile[3:,0] == pytest.approx(std1,rel=0.05)
    assert noise_profile[:3,1] == pytest.approx(std4,rel=0.05)
    assert noise_profile[3:,1] == pytest.approx(std3,rel=0.05)
    
    # Specify a particular KL mode
    noise_profile = measure_noise(frame, seps, fwhm,klmode_index=0)

    # Check data shape
    assert noise_profile.shape == (5,)
    assert noise_profile[:3] == pytest.approx(std2,rel=0.05)
    assert noise_profile[3:] == pytest.approx(std1,rel=0.05)

    # Specify a different particular KL mode
    noise_profile = measure_noise(frame, seps, fwhm,klmode_index=1)

    # Check data shape
    assert noise_profile.shape == (5,)
    assert noise_profile[:3] == pytest.approx(std4,rel=0.05)
    assert noise_profile[3:] == pytest.approx(std3,rel=0.05)


def test_meas_klip_ADI():
    """Checks that KLIP throughput measurement for ADI is always between 0 and 1 
    (within 5%) and that throughput increases with separation from the mask.
    """
    global kt_adi 

    mode = 'ADI'
    nsci, nref = (2,0)

    nx,ny = (21,21)
    cenx, ceny = (25.,30.)
    ctcal = create_ct_cal(fwhm_mas, cfam_name='1F',
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny)

    klip_params = {
                'outdir':outdir,'fileprefix':fileprefix,
                'annuli':annuli, 'subsections':subsections, 
                'movement':movement, 'numbasis':numbasis,
                'mode':mode,'calibrate_flux':calibrate_flux}
    
    mock_sci,mock_ref = create_psfsub_dataset(nsci,nref,pa_aper_degs,
                                            fwhm_pix=fwhm_pix,
                                            st_amp=st_amp,
                                            noise_amp=noise_amp,
                                            pl_contrast=pl_contrast)

    # # Plot input dataset
    # import matplotlib.pyplot as plt
    # dataset_frames = mock_sci if mock_ref is None else (*mock_sci,*mock_ref)
    # fig,axes = plt.subplots(1,len(dataset_frames),sharey=True,layout='constrained',figsize=(3*len(dataset_frames),3))
    # for f,frame in enumerate(dataset_frames):
    #     im0 = axes[f].imshow(frame.data,origin='lower')
    #     plt.colorbar(im0,ax=axes[f],shrink=0.8)
    #     axes[f].scatter(frame.ext_hdr['STARLOCX'],frame.ext_hdr['STARLOCY'],s=1)
    #     name = 'Sci' if f < len(mock_sci) else 'Ref'
    #     n_frame = f if f < len(mock_sci) else f-len(mock_sci)
    #     axes[f].set_title(f'{name} Input Frame {n_frame}')
    # plt.show()
    klip_kwargs={
                'numbasis' : numbasis,
                'mode' : mode}
    psfsub_dataset = do_psf_subtraction(mock_sci,
                                reference_star_dataset=mock_ref,
                                fileprefix='test_KL_THRU',
                                measure_klip_thrupt=False,
                                measure_1d_core_thrupt=False,
                                **klip_kwargs)
    
    # # Plot Psf subtraction result
    # if psfsub_dataset[0].pri_hdr['KLIP_ALG'] == 'RDI':
    #     analytical_result = rotate(mock_sci[0].data - mock_ref[0].data,-pa_aper_degs[0],reshape=False,cval=np.nan)
    # elif psfsub_dataset[0].pri_hdr['KLIP_ALG'] == 'ADI':
    #     analytical_result = shift((rotate(mock_sci[0].data - mock_sci[1].data,-pa_aper_degs[0],reshape=False,cval=0) + rotate(mock_sci[1].data - mock_sci[0].data,-pa_aper_degs[1],reshape=False,cval=0)) / 2,
    #                     [0.5,0.5],
    #                     cval=np.nan)
    # elif psfsub_dataset[0].pri_hdr['KLIP_ALG'] == 'ADI+RDI':
    #     analytical_result = (rotate(mock_sci[0].data - (mock_sci[1].data/2+mock_ref[0].data/2),-pa_aper_degs[0],reshape=False,cval=0) + rotate(mock_sci[1].data - (mock_sci[0].data/2+mock_ref[0].data/2),-pa_aper_degs[1],reshape=False,cval=0)) / 2
    # import matplotlib.pyplot as plt
    # fig,axes = plt.subplots(1,3,sharey=True,layout='constrained',figsize=(12,3))
    # im0 = axes[0].imshow(psfsub_dataset[0].data[0],origin='lower')
    # plt.colorbar(im0,ax=axes[0],shrink=0.8)
    # axes[0].set_title(f'Output data')
    # im1 = axes[1].imshow(analytical_result,origin='lower')
    # plt.colorbar(im1,ax=axes[1],shrink=0.8)
    # axes[1].set_title('Analytical result')
    # diff = psfsub_dataset[0].data[0] - analytical_result
    # im2 = axes[2].imshow(diff,origin='lower')
    # plt.colorbar(im2,ax=axes[2],shrink=0.8)
    # axes[2].set_title('Difference')
    # plt.suptitle(f'PSF Subtraction {psfsub_dataset[0].pri_hdr["KLIP_ALG"]} ({psfsub_dataset[0].ext_hdr["KLMODE0"]} KL Modes)')
    # plt.show()   

    klip_params['mode'] = mode
    kt_adi = meas_klip_thrupt(mock_sci, mock_ref, # pre-psf-subtracted dataset
                     psfsub_dataset, # post-subtraction dataset
                     ctcal,
                     klip_params,
                     inject_snr,
                     seps = seps, # in pixels from mask center
                     cand_locs=[])

    # import matplotlib.pyplot as plt
    # fig,ax = plt.subplots(figsize=(6,4))
    # plt.plot(kt_adi[0],kt_adi[1],label='ADI')
    # plt.title('KLIP throughput')
    # plt.legend()
    # plt.xlabel('separation (pixels)')
    # plt.show()

    # Check KL thrupt is <= 1 within noise tolerance
    assert np.all(kt_adi[1:,:,0] < max_thrupt_tolerance)

    # Check KL thrupt is > 0
    assert np.all(kt_adi[1:,:,0] > 0.)

    # Check KL thrupt increases with separation
    for i in range(1,len(kt_adi[0])):
        assert np.all(kt_adi[1:,i,0] > kt_adi[1:,i-1,0] - 0.05) # add a fudge factor of 5%

    # Check recovered FWHM increases with separation
    for i in range(1,len(kt_adi[0])):
        assert np.all(kt_adi[1:,i,1] > kt_adi[1:,i-1,1] - 0.1) # add fudge factor of 0.1 pixels


def test_meas_klip_RDI():
    """Checks that KLIP throughput measurement for RDI is always between 0.8 and 1 
    (within 5%).
    """
    global kt_rdi
    mode = 'RDI'
    nsci, nref = (1,1)

    nx,ny = (21,21)
    cenx, ceny = (25.,30.)
    ctcal = create_ct_cal(fwhm_mas, cfam_name='1F',
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny)

    klip_params = {
                'outdir':outdir,'fileprefix':fileprefix,
                'annuli':annuli, 'subsections':subsections, 
                'movement':movement, 'numbasis':numbasis,
                'mode':mode,'calibrate_flux':calibrate_flux}
    
    mock_sci,mock_ref = create_psfsub_dataset(nsci,nref,pa_aper_degs,
                                            fwhm_pix=fwhm_pix,
                                            st_amp=st_amp,
                                            noise_amp=noise_amp,
                                            pl_contrast=pl_contrast)

    # # Plot input dataset
    # import matplotlib.pyplot as plt
    # dataset_frames = mock_sci if mock_ref is None else (*mock_sci,*mock_ref)
    # fig,axes = plt.subplots(1,len(dataset_frames),sharey=True,layout='constrained',figsize=(3*len(dataset_frames),3))
    # for f,frame in enumerate(dataset_frames):
    #     im0 = axes[f].imshow(frame.data,origin='lower')
    #     plt.colorbar(im0,ax=axes[f],shrink=0.8)
    #     axes[f].scatter(frame.ext_hdr['STARLOCX'],frame.ext_hdr['STARLOCY'],s=1)
    #     name = 'Sci' if f < len(mock_sci) else 'Ref'
    #     n_frame = f if f < len(mock_sci) else f-len(mock_sci)
    #     axes[f].set_title(f'{name} Input Frame {n_frame}')
    # plt.show()
    klip_kwargs={
                'numbasis' : numbasis,
                'mode' : mode}
    psfsub_dataset = do_psf_subtraction(mock_sci,
                                reference_star_dataset=mock_ref,
                                fileprefix='test_KL_THRU',
                                measure_klip_thrupt=False,
                                measure_1d_core_thrupt=False,
                                **klip_kwargs)

    klip_params['mode'] = mode
    kt_rdi = meas_klip_thrupt(mock_sci, mock_ref, # pre-psf-subtracted dataset
                     psfsub_dataset, # post-subtraction dataset
                     ctcal,
                     klip_params,
                     inject_snr,
                     seps=seps)

    # import matplotlib.pyplot as plt
    # fig,ax = plt.subplots(figsize=(6,4))
    # plt.plot(kt_rdi[0],kt_rdi[1],label='RDI')
    # plt.title('KLIP throughput')
    # plt.legend()
    # plt.xlabel('separation (pixels)')
    # plt.show()

    # Check KL thrupt is <= 1 within noise tolerance
    assert np.all(kt_rdi[1:,:,0] < max_thrupt_tolerance)

    # Check KL thrupt > 0.8
    assert np.all(kt_rdi[1:,:,0] > 0.8)


def test_meas_klip_ADIRDI():
    """Checks that KLIP throughput measurement for ADI+RDI is always between 0 and 1.
    """
    global kt_adirdi
    
    mode = 'ADI+RDI'
    nsci, nref = (2,1)

    ctcal = create_ct_cal(fwhm_mas, cfam_name='1F',
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny)
    
    klip_params = {
                'outdir':outdir,'fileprefix':fileprefix,
                'annuli':annuli, 'subsections':subsections, 
                'movement':movement, 'numbasis':numbasis,
                'mode':mode,'calibrate_flux':calibrate_flux}
    
    mock_sci,mock_ref = create_psfsub_dataset(nsci,nref,pa_aper_degs,
                                            fwhm_pix=fwhm_pix,
                                            st_amp=st_amp,
                                            noise_amp=noise_amp,
                                            pl_contrast=pl_contrast)

    klip_kwargs={
        'numbasis' : numbasis,
        'mode' : mode}
    psfsub_dataset = do_psf_subtraction(mock_sci,
                                reference_star_dataset=mock_ref,
                                fileprefix='test_KL_THRU',
                                measure_klip_thrupt=False,
                                measure_1d_core_thrupt=False,
                                **klip_kwargs)

    kt_adirdi = meas_klip_thrupt(mock_sci, mock_ref, # pre-psf-subtracted dataset
                     psfsub_dataset, # post-subtraction dataset
                     ctcal,
                     klip_params,
                     inject_snr,
                     seps=seps)

    # Check KL thrupt is <= 1 within noise tolerance
    assert np.all(kt_adirdi[1:,:,0] < max_thrupt_tolerance)

    # Check KL thrupt is > 0
    assert np.all(kt_adirdi[1:,:,0] > 0.)


def test_compare_RDI_ADI():
    """Check that mean ADI throughput is lower than mean RDI throughput.
    """

    # import matplotlib.pyplot as plt
    # fig,ax = plt.subplots()
    # ax.plot(kt_adi[0],kt_adi[1],label='ADI')
    # ax.plot(kt_rdi[0],kt_adirdi[1],label='ADI+RDI')
    # ax.plot(kt_rdi[0],kt_rdi[1],label='RDI')
    # plt.legend()
    # ax.set_ylim(-0.1,1.1)
    # plt.title('KLIP Throughput')
    # plt.xlabel('Separation (pixels)')
    # plt.show()

    # Check that ADI thrupt < RDI thrupt
    mean_adi = np.mean(kt_adi[1:,:,0])
    mean_rdi = np.mean(kt_rdi[1:,:,0])

    assert mean_adi < mean_rdi

@pytest.mark.parametrize("coro_type,FOV,band", [("HLC", "NFOV", "1F"), ("SPC", "WFOV", "4F"), ("SPC", "WFOV", "1F")])
def test_psfsub_withklipandctmeas_adi(coro_type,FOV,band):
    """Check that KLIP throughput and CT calibration can be run as part 
    of the PSF subtraction step function. Check that KLIP throughput and 
    core throughput sample the same separations by default. Check that an 
    input planet flux (after CT) can be recovered via the KLIP throughput 
    within 1% error after ADI.
    """
    nsci, nref = (2,0) # Mode ADI
 
    st_amp = 100.
    noise_amp = 1e-3
    pl_contrast = 3e-4 
    pl_amp = st_amp * pl_contrast
    pl_loc = (20.,0.)
    est_pl_snr = pl_amp / noise_amp
    data_shape = [101,101]
    mock_sci,mock_ref = create_psfsub_dataset(nsci,nref,pa_aper_degs,
                                            fwhm_pix=fwhm_pix,
                                            st_amp=st_amp,
                                            noise_amp=noise_amp,
                                            pl_contrast=pl_contrast,
                                            pl_sep=pl_loc[0],
                                            data_shape=data_shape,
                                )
    if coro_type == "SPC":
        for im in mock_sci:
            im.ext_hdr['SPAMNAME'] = FOV
            im.ext_hdr['LSAMNAME'] = FOV
            im.ext_hdr['CFAMNAME'] = band
            if band == "1F":
                im.ext_hdr['FPAMNAME'] = 'SPC12_R1C1'
            if band == "4F":
                im.ext_hdr['FPAMNAME'] = 'SPC34_R5C1'
    nx,ny = (21,21)
    cenx, ceny = (25.,30.)
    ctcal = create_ct_cal(fwhm_mas, cfam_name=band,
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny)
    
    klip_kwargs={
                'numbasis' : numbasis}
                                
    psfsub_dataset = do_psf_subtraction(mock_sci,ctcal,
                                reference_star_dataset=mock_ref,
                                fileprefix='test_KL_THRU',
                                measure_klip_thrupt=True,
                                measure_1d_core_thrupt=True,
                                cand_locs= [pl_loc],
                                kt_seps=[pl_loc[0]],
                                **klip_kwargs)
    

    # # Plot Psf subtraction result
    # if psfsub_dataset[0].pri_hdr['KLIP_ALG'] == 'RDI':
    #     analytical_result = rotate(mock_sci[0].data - mock_ref[0].data,-pa_aper_degs[0],reshape=False,cval=np.nan)
    # elif psfsub_dataset[0].pri_hdr['KLIP_ALG'] == 'ADI':
    #     analytical_result = (rotate(mock_sci[0].data - mock_sci[1].data,-pa_aper_degs[0],reshape=False,cval=0) + rotate(mock_sci[1].data - mock_sci[0].data,-pa_aper_degs[1],reshape=False,cval=0)) / 2
    # elif psfsub_dataset[0].pri_hdr['KLIP_ALG'] == 'ADI+RDI':
    #     analytical_result = (rotate(mock_sci[0].data - (mock_sci[1].data/2+mock_ref[0].data/2),-pa_aper_degs[0],reshape=False,cval=0) + rotate(mock_sci[1].data - (mock_sci[0].data/2+mock_ref[0].data/2),-pa_aper_degs[1],reshape=False,cval=0)) / 2
    # import matplotlib.pyplot as plt
    # fig,axes = plt.subplots(1,3,sharey=True,layout='constrained',figsize=(12,3))
    # im0 = axes[0].imshow(psfsub_dataset[0].data[0],origin='lower')
    # plt.colorbar(im0,ax=axes[0],shrink=0.8)
    # axes[0].set_title(f'Output data')
    # im1 = axes[1].imshow(analytical_result,origin='lower')
    # plt.colorbar(im1,ax=axes[1],shrink=0.8)
    # axes[1].set_title('Analytical result')
    # diff = psfsub_dataset[0].data[0] - analytical_result
    # im2 = axes[2].imshow(diff,origin='lower')
    # plt.colorbar(im2,ax=axes[2],shrink=0.8)
    # axes[2].set_title('Difference')
    # plt.suptitle(f'PSF Subtraction {psfsub_dataset[0].pri_hdr["KLIP_ALG"]} ({psfsub_dataset[0].ext_hdr["KLMODE0"]} KL Modes)')
    # plt.show()    

    # Check that klip and ct separations are the same
    kt = psfsub_dataset[0].hdu_list['KL_THRU'].data
    kt_seps = kt[0,:,0]
    kt_thru = kt[1,:,0]

    # # Plot KL throughput measurement
    # import matplotlib.pyplot as plt
    # fig,ax = plt.subplots(figsize=(6,4))
    # plt.scatter(kt_seps,kt_thru,label=psfsub_dataset[0].pri_hdr["KLIP_ALG"])
    # plt.title('KLIP throughput')
    # plt.legend()
    # plt.xlabel('separation (pixels)')
    # plt.show()

    ct = psfsub_dataset[0].hdu_list['CT_THRU'].data
    ct_seps = ct[0]

    assert np.all(kt_seps == ct_seps)

    # Fit 2d gaussian to planet location:

    # Crop data around location to be same as psf_model cutout
    locxy = seppa2xy(*pl_loc,psfsub_dataset[0].ext_hdr['STARLOCX'],psfsub_dataset[0].ext_hdr['STARLOCY'])

    # Measure background via sigma clip 
    n_loops = 5
    masked_data = psfsub_dataset[0].data.copy()[0]
    for n in range(n_loops):
        std = np.nanstd(masked_data)
        med = np.nanmedian(masked_data)
        clip_thresh = 3 * std
        masked_data = np.where(np.abs(masked_data-med)>clip_thresh,np.nan,masked_data)
    
    # Subtract median
    bg_level = np.nanmedian(masked_data)
    medsubtracted_data = psfsub_dataset[0].data[0] - bg_level

    # Crop the data, pad with nans if we're cropping over the edge
    cutout_shape = np.array([39,39])
    cutout = np.zeros(cutout_shape)
    cutoutcenyx = cutout_shape/2. - 0.5
    cutout[:] = np.nan
    cutout_starty, cutout_startx = (0,0)
    cutout_endy, cutout_endx = cutout.shape

    data_shape = medsubtracted_data.shape
    data_center_indyx = np.array([locxy[1],locxy[0]]).astype(int)
    data_start_indyx = (data_center_indyx - cutout_shape//2)
    data_end_indyx = (data_start_indyx + cutout_shape)
    data_starty,data_startx = data_start_indyx
    data_endy,data_endx = data_end_indyx
    
    if data_starty < 0:
        cutout_starty = -data_starty
        data_starty = 0
    
    if data_startx < 0:
        cutout_startx = -data_startx
        data_startx = 0
    
    if data_endy >= data_shape[0]:
        y_overhang = data_endy - medsubtracted_data.shape[0]
        cutout_endy = cutout_shape[0] - y_overhang
        data_endy = data_shape[0]

    if data_endx >= data_shape[1]:
        x_overhang = data_endx - medsubtracted_data.shape[1]
        cutout_endx = cutout_shape[1] - x_overhang
        data_endx = data_shape[1]

    cutout[cutout_starty:cutout_endy,
                cutout_startx:cutout_endx] = medsubtracted_data[data_starty:data_endy,
                                                    data_startx:data_endx]
    
    postklip_peak, post_fwhm, post_xfit, post_yfit = gaussfit2d(
                cutout, 
                cutoutcenyx[1], 
                cutoutcenyx[0], 
                searchrad=5, 
                guessfwhm=fwhm_pix, 
                guesspeak=pl_amp, 
                refinefit=True) 

    # # Plot final PSF fit
    # import matplotlib.pyplot as plt
    # post_sigma = post_fwhm / (2 * np.sqrt(2. * np.log(2.)))
    # final_model = gaussian_array(array_shape=cutout_shape,
    #                              sigma=post_sigma,
    #                              amp=postklip_peak,
    #                              xoffset=post_xfit-cutout_shape[1]//2,
    #                              yoffset=post_yfit-cutout_shape[0]//2)
    # fig,axes = plt.subplots(1,3,sharey=True,layout='constrained',figsize=(12,3))
    # im0 = axes[0].imshow(cutout,origin='lower')
    # plt.colorbar(im0,ax=axes[0],shrink=0.8)
    # axes[0].set_title(f'Data')
    # im1 = axes[1].imshow(final_model,origin='lower')
    # plt.colorbar(im1,ax=axes[1],shrink=0.8)
    # axes[1].set_title('Model')
    # diff = cutout-final_model
    # im2 = axes[2].imshow(diff,origin='lower')
    # plt.colorbar(im2,ax=axes[2],shrink=0.8)
    # axes[2].set_title('Residuals')
    # plt.suptitle(f'Final PSF Fit')
    # plt.show()    

    
    pl_kt = kt[1,np.argmin(np.abs(kt_seps-pl_loc[0])),0]
    
    pl_counts = np.pi * pl_amp * fwhm_pix**2 / 4. / np.log(2.)
    recovered_pl_counts = np.pi * postklip_peak * post_fwhm**2 / 4. / np.log(2.)
    recovered_pl_counts_ktcorrected = recovered_pl_counts / pl_kt


    assert pl_counts == pytest.approx(recovered_pl_counts_ktcorrected,rel = 0.10) 

@pytest.mark.parametrize("coro_type,FOV,band", [("HLC", "NFOV", "1F"), ("SPC", "WFOV", "4F"),("SPC", "WFOV", "1F")])
def test_psfsub_withklipandctmeas_rdi(coro_type,FOV,band):
    """Check that KLIP throughput and CT calibration can be run as part 
    of the PSF subtraction step function. Check that KLIP throughput and 
    core throughput sample the same separations by default. Check that an 
    input planet flux (after CT) can be recovered via the KLIP throughput 
    within 5% error after RDI.
    """
    nsci, nref = (1,1) # Mode RDI
 
    st_amp = 100.
    noise_amp = 1e-3
    pl_contrast = 3e-3 
    pl_amp = st_amp * pl_contrast
    pl_loc = (20.,0.)
    est_pl_snr = pl_amp / noise_amp
    data_shape = (101,101)
    mock_sci,mock_ref = create_psfsub_dataset(nsci,nref,pa_aper_degs,
                                            fwhm_pix=fwhm_pix,
                                            st_amp=st_amp,
                                            noise_amp=noise_amp,
                                            pl_contrast=pl_contrast,
                                            pl_sep=pl_loc[0],
                                            data_shape=data_shape
                                )
    complete_dataset = [mock_sci,mock_ref]
    if coro_type == "SPC":
        for im in itertools.chain(*complete_dataset):                                
            im.ext_hdr['SPAMNAME'] = FOV
            im.ext_hdr['LSAMNAME'] = FOV
            im.ext_hdr['CFAMNAME'] = band
            if band == "1F":
                im.ext_hdr['FPAMNAME'] = 'SPC12_R1C1'
            if band == "4F":
                im.ext_hdr['FPAMNAME'] = 'SPC34_R5C1'
    nx,ny = (21,21)
    cenx, ceny = (25.,30.)
    ctcal = create_ct_cal(fwhm_mas, cfam_name=band,
                  cenx=cenx,ceny=ceny,
                  nx=nx,ny=ny)
    
    klip_kwargs={
                'numbasis' : numbasis}
                                
    psfsub_dataset = do_psf_subtraction(mock_sci,ctcal,
                                reference_star_dataset=mock_ref,
                                fileprefix='test_KL_THRU',
                                measure_klip_thrupt=True,
                                measure_1d_core_thrupt=True,
                                cand_locs= [pl_loc],
                                kt_seps=[pl_loc[0]],
                                **klip_kwargs)
    



    # Check that klip and ct separations are the same
    kt = psfsub_dataset[0].hdu_list['KL_THRU'].data
    kt_seps = kt[0,:,0]

    # # Plot KLIP throughput
    # import matplotlib.pyplot as plt
    # fig,ax = plt.subplots(figsize=(6,4))
    # plt.scatter(kt[0],kt[1],label=psfsub_dataset[0].pri_hdr["KLIP_ALG"])
    # plt.title('KLIP throughput')
    # plt.legend()
    # plt.xlabel('separation (pixels)')
    # plt.show()

    ct = psfsub_dataset[0].hdu_list['CT_THRU'].data
    ct_seps = ct[0]

    assert np.all(kt_seps == ct_seps)

    # Fit 2d gaussian to planet location:

    # Crop data around location to be same as psf_model cutout
    locxy = seppa2xy(*pl_loc,psfsub_dataset[0].ext_hdr['STARLOCX'],psfsub_dataset[0].ext_hdr['STARLOCY'])

    # Measure background via sigma clip 
    n_loops = 5
    masked_data = psfsub_dataset[0].data.copy()[0]
    for n in range(n_loops):
        std = np.nanstd(masked_data)
        med = np.nanmedian(masked_data)
        clip_thresh = 3 * std
        masked_data = np.where(np.abs(masked_data-med)>clip_thresh,np.nan,masked_data)
    
    # Subtract median
    bg_level = np.nanmedian(masked_data)
    medsubtracted_data = psfsub_dataset[0].data[0] #- bg_level

    # # Plot Psf subtraction result
    # analytical_result = rotate(mock_sci[0].data - mock_ref[0].data,-pa_aper_degs[0],reshape=False,cval=np.nan)
    # mask = create_circular_mask(analytical_result.shape[-2:],
    #                             r=3*fwhm_pix,
    #                             center=(psfsub_dataset[0].ext_hdr['STARLOCX'],
    #                                     psfsub_dataset[0].ext_hdr['STARLOCY']))
    # masked_analytical_result = np.where(mask,np.nan,analytical_result)
    # import matplotlib.pyplot as plt
    # fig,axes = plt.subplots(1,3,sharey=True,layout='constrained',figsize=(12,3))
    # im0 = axes[0].imshow(medsubtracted_data,origin='lower')
    # plt.colorbar(im0,ax=axes[0],shrink=0.8)
    # axes[0].set_title(f'BG-subtracted Output data')
    # im1 = axes[1].imshow(analytical_result,origin='lower')
    # plt.colorbar(im1,ax=axes[1],shrink=0.8)
    # axes[1].set_title('Analytical result')
    # diff = medsubtracted_data - masked_analytical_result
    # im2 = axes[2].imshow(diff,origin='lower')
    # plt.colorbar(im2,ax=axes[2],shrink=0.8)
    # axes[2].set_title('Difference')
    # plt.suptitle(f'PSF Subtraction {psfsub_dataset[0].pri_hdr["KLIP_ALG"]} ({psfsub_dataset[0].ext_hdr["KLMODE0"]} KL Modes)')
    # plt.show()    


    # Crop the data, pad with nans if we're cropping over the edge
    cutout_shape = np.array([21,21])
    cutout = np.zeros(cutout_shape)
    cutoutcenyx = cutout_shape/2. - 0.5
    cutout[:] = np.nan
    cutout_starty, cutout_startx = (0,0)
    cutout_endy, cutout_endx = cutout.shape

    data_shape = medsubtracted_data.shape
    data_center_indyx = np.array([locxy[1],locxy[0]]).astype(int)
    data_start_indyx = (data_center_indyx - cutout_shape//2)
    data_end_indyx = (data_start_indyx + cutout_shape)
    data_starty,data_startx = data_start_indyx
    data_endy,data_endx = data_end_indyx
    
    if data_starty < 0:
        cutout_starty = -data_starty
        data_starty = 0
    
    if data_startx < 0:
        cutout_startx = -data_startx
        data_startx = 0
    
    if data_endy >= data_shape[0]:
        y_overhang = data_endy - medsubtracted_data.shape[0]
        cutout_endy = cutout_shape[0] - y_overhang
        data_endy = data_shape[0]

    if data_endx >= data_shape[1]:
        x_overhang = data_endx - medsubtracted_data.shape[1]
        cutout_endx = cutout_shape[1] - x_overhang
        data_endx = data_shape[1]

    cutout[cutout_starty:cutout_endy,
                cutout_startx:cutout_endx] = medsubtracted_data[data_starty:data_endy,
                                                    data_startx:data_endx]
    

    postklip_peak, post_fwhm, post_xfit, post_yfit = gaussfit2d(
                cutout, 
                cutoutcenyx[1], 
                cutoutcenyx[0], 
                searchrad=5, 
                guessfwhm=fwhm_pix, 
                guesspeak=pl_amp, 
                refinefit=True) 
    post_sigma = post_fwhm / (2 * np.sqrt(2. * np.log(2.)))
    final_model = gaussian_array(array_shape=cutout_shape,
                                 sigma=post_sigma,
                                 amp=postklip_peak,
                                 xoffset=post_xfit-10.,
                                 yoffset=post_yfit-10.)
    

    diff = cutout-final_model

    # # Plot Final PSF Fit
    # import matplotlib.pyplot as plt
    # fig,axes = plt.subplots(1,3,sharey=True,layout='constrained',figsize=(12,3))
    # im0 = axes[0].imshow(cutout,origin='lower')
    # plt.colorbar(im0,ax=axes[0],shrink=0.8)
    # axes[0].set_title(f'Data')
    # im1 = axes[1].imshow(final_model,origin='lower')
    # plt.colorbar(im1,ax=axes[1],shrink=0.8)
    # axes[1].set_title('Model')
    # im2 = axes[2].imshow(diff,origin='lower')
    # plt.colorbar(im2,ax=axes[2],shrink=0.8)
    # axes[2].set_title('Residuals')
    # plt.suptitle(f'Final PSF Fit')
    # plt.show()    
    
    pl_kt = kt[1,np.argmin(np.abs(kt_seps-pl_loc[0])),0]
    
    pl_counts = np.pi * pl_amp * fwhm_pix**2 / 4. / np.log(2.)
    recovered_pl_counts = np.pi * postklip_peak * post_fwhm**2 / 4. / np.log(2.)
    recovered_pl_counts_ktcorrected = recovered_pl_counts / pl_kt

    assert pl_counts == pytest.approx(recovered_pl_counts_ktcorrected,rel = 0.05) 


def test_measure_noise_nsigma_and_correction():
    """Tests for nsigma scaling and small sample statistics correction
    in measure_noise().

    This is a unit test that directly tests measure_noise() with synthetic data.
    It verifies:
      - Simple nsigma scaling (noise * nsigma)
      - Backward compatibility (default behavior unchanged)
      - Mawet et al. (2014) small sample correction behavior
      - Convergence of the correction at large separations
      - Input validation (bad inputs raise errors)
      - FWHM input flexibility (scalar vs array)
    """
    from scipy.stats import t as t_dist, norm

    # --- Setup: create a synthetic frame with known Gaussian noise ---
    # Two KL-mode slices of pure Gaussian noise (mean=0, std=1) so that
    # the expected annular standard deviation is ~1.0 everywhere.
    cenx, ceny = (120., 130.)
    frame_shape_yx = (200, 200)
    seps = np.arange(50., 71., 5.)  # separations where we'll measure noise
    fwhm = 2.  # PSF FWHM in pixels (used as annulus halfwidth and for small sample correction)

    ext_hdr = fits.Header()
    ext_hdr['STARLOCX'] = cenx
    ext_hdr['STARLOCY'] = ceny

    rng = np.random.default_rng(42)
    image = rng.normal(0., 1., (2, *frame_shape_yx))  # 2 KL mode slices
    frame = Image(image, pri_hdr=fits.Header(), ext_hdr=ext_hdr)

    # --- Test 1: nsigma scaling is a pure multiplier ---
    # Without small sample correction, nsigma=5 should produce exactly 5x the nsigma=1 result,
    # since it's just multiplying the standard deviation by nsigma.
    noise_1sig = measure_noise(frame, seps, fwhm, nsigma=1)
    noise_5sig = measure_noise(frame, seps, fwhm, nsigma=5)
    assert noise_5sig == pytest.approx(5 * noise_1sig)

    # --- Test 2: backward compatibility ---
    # Calling without nsigma (default=1) should give the same result as nsigma=1,
    # ensuring old code that doesn't pass nsigma still works.
    noise_default = measure_noise(frame, seps, fwhm)
    assert noise_1sig == pytest.approx(noise_default)

    # --- Test 3: small sample correction makes noise MORE conservative ---
    # The Mawet et al. (2014) correction replaces the naive nsigma multiplier with a
    # t-distribution threshold that's always >= nsigma for finite n (number of resolution
    # elements). So corrected noise should be >= naive nsigma * std at every separation.
    noise_corr = measure_noise(frame, seps, fwhm, nsigma=5, fwhm=fwhm,
                               small_sample_correction=True)
    for i, sep in enumerate(seps):
        n = 2 * np.pi * sep / fwhm  # number of independent resolution elements in annulus
        if n > 2:
            # t-distribution threshold > Gaussian threshold for finite degrees of freedom
            assert np.all(noise_corr[i] >= noise_5sig[i] - 1e-10)

    # --- Test 4: convergence at large separations ---
    # With many resolution elements (large n), the t-distribution converges to the
    # Gaussian, so the corrected noise should be approximately equal to the naive
    # nsigma * std (within 5% tolerance).
    large_seps = np.array([500., 600., 700.])
    large_shape = (1500, 1500)  # need a bigger frame to fit these separations
    large_cenx, large_ceny = (750., 750.)
    ext_hdr_large = fits.Header()
    ext_hdr_large['STARLOCX'] = large_cenx
    ext_hdr_large['STARLOCY'] = large_ceny
    large_image = rng.normal(0., 1., (1, *large_shape))
    large_frame = Image(large_image, pri_hdr=fits.Header(), ext_hdr=ext_hdr_large)
    noise_large_naive = measure_noise(large_frame, large_seps, fwhm, nsigma=5)
    noise_large_corr = measure_noise(large_frame, large_seps, fwhm, nsigma=5, fwhm=fwhm,
                                     small_sample_correction=True)
    assert noise_large_corr == pytest.approx(noise_large_naive, rel=0.05)

    # --- Test 5: input validation ---
    # Requesting small_sample_correction without providing fwhm should fail
    with pytest.raises(ValueError):
        measure_noise(frame, seps, fwhm, nsigma=5, small_sample_correction=True)

    # nsigma must be positive
    with pytest.raises(ValueError):
        measure_noise(frame, seps, fwhm, nsigma=0)
    with pytest.raises(ValueError):
        measure_noise(frame, seps, fwhm, nsigma=-1)

    # fwhm array length must match number of separations
    with pytest.raises(ValueError):
        measure_noise(frame, seps, fwhm, nsigma=5, fwhm=np.array([1., 2.]),
                      small_sample_correction=True)

    # --- Test 6: fwhm input flexibility ---
    # Scalar fwhm should be broadcast to all separations
    noise_scalar_fwhm = measure_noise(frame, seps, fwhm, nsigma=5, fwhm=2.0,
                                      small_sample_correction=True)
    assert noise_scalar_fwhm.shape == noise_5sig.shape

    # An array of identical values should give the same result as the scalar
    fwhm_arr = np.full(len(seps), 2.0)
    noise_arr_fwhm = measure_noise(frame, seps, fwhm, nsigma=5, fwhm=fwhm_arr,
                                   small_sample_correction=True)
    assert noise_arr_fwhm == pytest.approx(noise_scalar_fwhm)


if __name__ == '__main__':
    # test_create_ct_cal()
    # test_get_closest_psf()
    # test_inject_psf()
    # test_measure_noise()
    # test_measure_noise_nsigma_and_correction()

    # test_meas_klip_ADI()
    # test_meas_klip_RDI()
    # test_meas_klip_ADIRDI()
    # test_compare_RDI_ADI()

    test_psfsub_withklipandctmeas_adi()
    test_psfsub_withklipandctmeas_rdi()
    
