#%%
import inspect
from typing import List,Union,Optional,Tuple
import numpy as np
from multiprocessing import Pool, cpu_count
import bottleneck as bn
from tqdm import tqdm
import re
from astropy.io import fits
from astropy.time import Time
import os
import matplotlib.pyplot as plt
from tqdm.contrib.concurrent import process_map
import gc
import warnings
from astropy.wcs.wcs import FITSFixedWarning
from multiprocessing.shared_memory import SharedMemory


warnings.filterwarnings("ignore", category=FITSFixedWarning, message=".*SIP.*")

from ezphot.helper import Helper
from ezphot.imageobjects import ScienceImage, CalibrationImage, ReferenceImage, Errormap, Background 
from ezphot.methods import Reproject
from ezphot.methods import BackgroundGenerator
from ezphot.methods import PSFPhotometry
#%%
helper = Helper()
def combine_patch(patch_tuple, combine_method='mean', clip_method='sigma', sigma=3.0, nlow=1, nhigh=1):
    (i_start, i_end, j_start, j_end, tile_stack, bkgrms_stack) = patch_tuple

    if combine_method.lower() in ['weight', 'weighted'] and bkgrms_stack is None:
        raise ValueError("combine_method='weight' requires bkgrms_stack to be provided.")

    # --- Clipping ---
    if clip_method is None or clip_method.lower() == 'none':
        clipped = tile_stack
        clipped_rms = bkgrms_stack
    elif clip_method.lower() == 'sigma':
        mean = np.nanmean(tile_stack, axis=0)
        std = np.nanstd(tile_stack, axis=0)
        mask = np.abs(tile_stack - mean) < sigma * std
        clipped = np.where(mask, tile_stack, np.nan)
        clipped_rms = np.where(mask, bkgrms_stack, np.nan)
    elif clip_method.lower() == 'extrema':
        sorted_idx = np.argsort(tile_stack, axis=0)
        valid_idx = sorted_idx[nlow:len(tile_stack) - nhigh]

        # Apply indices to both image and RMS
        clipped = np.take_along_axis(tile_stack, valid_idx, axis=0)
        if bkgrms_stack is not None:
            clipped_rms = np.take_along_axis(bkgrms_stack, valid_idx, axis=0)
        else:
            clipped_rms = None
    else:
        raise ValueError(f"Unknown clip_method: {clip_method}")

    # -- Combine image ---
    if combine_method.lower() in ['weight', 'weighted']:
        if clipped_rms is None:
            raise ValueError("Weighted combination requires background RMS (clipped_rms).")

        # Compute weights as 1 / (rms^2), safely avoiding divide-by-zero
        weights = 1.0 / np.where(clipped_rms > 0, clipped_rms**2, np.nan)
        
        # Weighted average
        weighted_sum = np.nansum(clipped * weights, axis=0)
        weight_total = np.nansum(weights, axis=0)
        combined = np.divide(weighted_sum, weight_total, out=np.zeros_like(weight_total), where=weight_total > 0)

    elif combine_method.lower() == 'mean':
        combined = bn.nanmean(clipped, axis=0)

    elif combine_method.lower() == 'median':
        combined = bn.median(clipped, axis=0)

    elif combine_method.lower() == 'sum':
        combined = np.nansum(clipped, axis=0)

    else:
        raise ValueError(f"Unknown combine_method: {combine_method}")
    
    # --- Combine RMS ---
    if clipped_rms is not None:
        N = clipped_rms.shape[0]

        if combine_method.lower() == 'mean':
            # Combine assuming uncorrelated noise: ?_combined = sqrt(sum ?_i^2) / N
            combined_rms = np.sqrt(np.nansum(clipped_rms**2, axis=0)) / N

        elif combine_method.lower() in ['weight', 'weighted']:
            # Already computed weights in image combination block
            # combined_rms = sqrt(1 / sum w_i)
            combined_rms = np.sqrt(1.0 / np.where(weight_total > 0, weight_total, np.nan))

        elif combine_method.lower() == 'median':
            # Approximation for standard error of median: ?_combined ? 1.253 / sqrt(N) * median(?_i)
            combined_rms = 1 / np.sqrt(N) * np.nanmedian(clipped_rms, axis=0)

        else:
            combined_rms = None
    else:
        combined_rms = None

    return i_start, i_end, j_start, j_end, combined, combined_rms

def combine_patch_worker(args):
    (i0, i1, j0, j1,
     shm_image_name, img_shape, img_dtype,
     shm_bkgrms_name, bkgrms_shape, bkgrms_dtype,
     combine_method, clip_method, sigma, nlow, nhigh) = args

    shm_img = SharedMemory(name=shm_image_name)
    stack = np.ndarray(img_shape, dtype=np.dtype(img_dtype), buffer=shm_img.buf)

    shm_bkgrms = None
    if shm_bkgrms_name is not None:
        shm_bkgrms = SharedMemory(name=shm_bkgrms_name)
        bkgrms = np.ndarray(bkgrms_shape, dtype=np.dtype(bkgrms_dtype), buffer=shm_bkgrms.buf)
    else:
        bkgrms = None

    try:
        tile = stack[:, i0:i1, j0:j1]
        bk_tile = bkgrms[:, i0:i1, j0:j1] if bkgrms is not None else None

        return combine_patch(
            (i0, i1, j0, j1, tile, bk_tile),
            combine_method, clip_method, sigma, nlow, nhigh
        )
    finally:
        shm_img.close()
        if shm_bkgrms is not None:
            shm_bkgrms.close()



class Combiner:
    
    def __init__(self,
                 n_proc: int = 8):
        self.n_proc = cpu_count() if n_proc is None else n_proc

    def make_patches(self, H, W, patch_size=512):
        patches = []
        for i0 in range(0, H, patch_size):
            for j0 in range(0, W, patch_size):
                i1 = min(H, i0 + patch_size)
                j1 = min(W, j0 + patch_size)
                patches.append((i0, i1, j0, j1))
        return patches

    def combine_images_parallel(self, 
                                image_list, 
                                bkgrms_list=None,
                                combine_method='mean',
                                clip_method='sigma',
                                sigma=3.0, 
                                nlow=1,
                                nhigh=1,
                                verbose=True,
                                **kwargs):
        if verbose:
            print(f"[Combiner] Combining {len(image_list)} images with combine='{combine_method}', clip='{clip_method}', using {self.n_proc} processes")
            
        # Check image size first. If one of the image or bkgrms has different dimensions, raise an error.
        total_size_image = 0
        N = len(image_list)
        H = image_list[0].shape[0]
        W = image_list[0].shape[1]
        patches = self.make_patches(H, W, patch_size=512)

        if self.n_proc == 1:
            shared_image = np.stack(image_list)
            shared_bkgrms = np.stack(bkgrms_list) if bkgrms_list is not None else None
            image_out = np.zeros((H, W), dtype=np.float32)
            bkgrms_out = None
            if bkgrms_list is not None:
                bkgrms_out = np.zeros((H, W), dtype=np.float32)
            for (i0, i1, j0, j1) in tqdm(patches, desc="Combining (single CPU)..."):
                image_tile = shared_image[:, i0:i1, j0:j1]
                bkgrms_tile = None
                if bkgrms_list is not None:
                    bkgrms_tile = shared_bkgrms[:, i0:i1, j0:j1]

                _, _, _, _, patch_data, patch_bkgrms = combine_patch(
                    (i0, i1, j0, j1, image_tile, bkgrms_tile),
                    combine_method, clip_method, sigma, nlow, nhigh
                )
                image_out[i0:i1, j0:j1] = patch_data
                if bkgrms_out is not None:
                    bkgrms_out[i0:i1, j0:j1] = patch_bkgrms
        else:
            dtype_image = image_list[0].dtype
            dtype_bkgrms = bkgrms_list[0].dtype if bkgrms_list is not None else None
            for image in image_list:
                total_size_image += image.nbytes
                if image.shape[0] != H or image.shape[1] != W:
                    raise ValueError("All images must have the same dimensions.")
                if image.dtype != dtype_image:
                    raise ValueError("All images must have the same dtype.")
            shm = SharedMemory(create=True, size=total_size_image)
            shared_image = np.ndarray((N, H, W), dtype=dtype_image, buffer=shm.buf)
            for i in range(N):
                shared_image[i] = image_list[i]
            
            if bkgrms_list is not None:
                total_size_bkgrms = 0
                for bkgrms in bkgrms_list:
                    total_size_bkgrms += bkgrms.nbytes
                    if bkgrms.shape[0] != H or bkgrms.shape[1] != W:
                        raise ValueError("All bkgrms must have the same dimensions.")
                    if bkgrms.dtype != dtype_bkgrms:
                        raise ValueError("All bkgrms must have the same dtype.")
                shm_bkgrms = SharedMemory(create=True, size=total_size_bkgrms)
                shared_bkgrms = np.ndarray((N, H, W), dtype=dtype_bkgrms, buffer=shm_bkgrms.buf)
                for i in range(N):
                    shared_bkgrms[i] = bkgrms_list[i]
                    
            patch_args = []
            for patch in patches:
                i0, i1, j0, j1 = patch
                patch_args.append((
                    i0, i1, j0, j1,
                    shm.name, (N, H, W), dtype_image.str,
                    shm_bkgrms.name if bkgrms_list is not None else None,
                    (N, H, W), dtype_bkgrms.str if bkgrms_list is not None else None,
                    combine_method, clip_method, sigma, nlow, nhigh
                ))

            # Use persistent pool if available
            with Pool(processes=self.n_proc) as pool:
                results = list(
                    tqdm(pool.imap_unordered(combine_patch_worker, patch_args),
                            total=len(patch_args),
                            desc="Combining...",
                            ncols=80,
                            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]")
                )
            
            image_out = np.zeros((H, W), dtype=np.float32)
            bkgrms_out = np.zeros((H, W), dtype=np.float32) if bkgrms_list is not None else None
            for i_start, i_end, j_start, j_end, patch_result, patch_bkgrms in results:
                image_out[i_start:i_end, j_start:j_end] = patch_result
                if bkgrms_out is not None and patch_bkgrms is not None:
                    bkgrms_out[i_start:i_end, j_start:j_end] = patch_bkgrms
        

            shm.close()
            shm.unlink()
            if bkgrms_list is not None:
                shm_bkgrms.close()
                shm_bkgrms.unlink()
            
        del shared_image
        del image_list
        if bkgrms_list is not None:
            del shared_bkgrms
            del bkgrms_list
        gc.collect()

        return image_out, bkgrms_out

# Backgroud subtraction worker function
bkg_handler = BackgroundGenerator()
def _subtract_background_worker(args):
    (target_img, 
    target_bkg) = args

    subbkg_img = None
    subbkg_img = bkg_handler.subtract_background(
        target_img=target_img,
        target_bkg=target_bkg,
        save=False, 
        overwrite=False,
        visualize=False,
        verbose=False
    )
    target_img.clear(verbose = False)
    target_bkg.clear(verbose = False)
    return subbkg_img

projection_handler = Reproject()
def _reproject_worker(args):
    (target_img, 
     target_bkgrms, 
     resample_type, 
     center_ra, 
     center_dec, 
     x_size, 
     y_size, 
     pixel_scale,
     verbose) = args

    reprojected_img, reprojected_bkgrms, _ = projection_handler.reproject(
        target_img=target_img,
        target_errormap=target_bkgrms,
        swarp_params=None,
        resample_type=resample_type,
        center_ra=center_ra,
        center_dec=center_dec,
        x_size=x_size,
        y_size=y_size,
        pixelscale=pixel_scale,
        verbose=verbose,
        overwrite=False,
        save=False,
        return_ivpmask=False,
    )
    return reprojected_img, reprojected_bkgrms

def _scale_worker(args) -> Tuple:
    (target_img, 
     target_errormap, 
     ref_zp, 
     zp_key, 
     overwrite,
     verbose) = args

    zp = float(target_img.header[zp_key])
    delta_zp = ref_zp - zp
    scale_factor = 10 ** (0.4 * (delta_zp))
    if verbose:
        print(f"Scaling image {target_img.path} by {scale_factor}")

    if not overwrite:
        target_img_path = target_img.savepath.savedir / f"scaled_{target_img.savepath.savepath.name}"
        target_errormap_path = target_errormap.savepath.savedir / f"scaled_{target_errormap.savepath.savepath.name}" if target_errormap else None
    else:
        target_img_path = target_img.savepath.savepath
        target_errormap_path = target_errormap.savepath.savepath if target_errormap else None

    # Scale image
    scaled_img = type(target_img)(path=target_img_path, telinfo=target_img.telinfo, status=target_img.status.copy(), load=False)
    scaled_img.data = target_img.data * scale_factor
    scaled_img.header = target_img.header.copy()
    scaled_img.header[zp_key] = ref_zp
    scaled_img.header.update({
        'SCLEKEY': zp_key,
        'SCLEREF': ref_zp,
        'SCLEZP': delta_zp,
        'SCLEFACT': scale_factor,
    })
    for key in target_img.header.keys():
        if key.startswith('ZP_'):
            scaled_img.header[key] = target_img.header[key] + delta_zp
    scaled_img.update_status('ZPSCALE')

    # Scale error map
    scaled_errormap = None
    if target_errormap:
        emaptype = target_errormap.emaptype.lower()
        if emaptype == 'bkgrms':
            factor = scale_factor
        elif emaptype == 'bkgweight':
            factor = 1.0 / scale_factor**2
        else:
            raise ValueError(f"Unsupported emaptype '{emaptype}'")

        scaled_errormap = Errormap(path=target_errormap_path, emaptype=target_errormap.emaptype, status=target_errormap.status.copy(), load=False)
        scaled_errormap.data = target_errormap.data * factor
        scaled_errormap.header.update({
            'SCLEKEY': zp_key,
            'SCLEREF': ref_zp,
            'SCLEZP': delta_zp,
            'SCLEFACT': scale_factor,
        })
        scaled_errormap.add_status('zpscale', key=zp_key, ref_zp=ref_zp, scale_zp=delta_zp, scale_factor=scale_factor)
    target_img.clear(verbose = False)
    if target_errormap:
        target_errormap.clear(verbose = False)
    return scaled_img, scaled_errormap

def _convolve_worker(args):
    """
    Worker function for multiprocessing seeing-matching convolution.

    Parameters
    ----------
    args : tuple
        (target_img, target_errormap, current_seeing, ref_seeing,
         kernel_type, conv_kernel, seeing_key, overwrite)

        Notes:
        - conv_kernel = None for Gaussian kernel mode
        - conv_kernel = precomputed 2D kernel for 'image' mode

    Returns
    -------
    matched_img : ScienceImage | CalibrationImage
    matched_errormap : Errormap or None
    """
    
    (target_img,
     target_bkgrms,
     ref_seeing,
     seeing_key,
     verbose) = args

    current_seeing = float(target_img.header[seeing_key])
    if current_seeing is None:
        raise ValueError(f"SEEING key {seeing_key} not found in target_img header")

    convolved_data, updated_header = helper.img_convolve(
        target_img=target_img.data,
        input_type='image',
        kernel = 'gaussian',
        target_header=target_img.header,
        fwhm_target=current_seeing,
        fwhm_reference=ref_seeing,
        fwhm_key=seeing_key,
        verbose=verbose
    )

    convolved_img = target_img.copy()
    convolved_img.data = convolved_data
    convolved_img.header = updated_header
    
    convolved_bkgrms = None
    if target_bkgrms is not None:
        convolved_bkgrms_data, updated_bkgrms_header = helper.img_convolve(
        target_img=target_bkgrms.data,
        input_type='error',
        kernel='gaussian',
        target_header=target_bkgrms.header,
        fwhm_target=current_seeing,
        fwhm_reference=ref_seeing,
        fwhm_key=seeing_key,
        verbose=verbose
        )

        convolved_bkgrms = target_bkgrms.copy()
        convolved_bkgrms.data = convolved_bkgrms_data
        convolved_bkgrms.header = updated_bkgrms_header
    return convolved_img, convolved_bkgrms

def _prepare_image_worker(args):
    """
    1. Subtract background
    2. Scale image
    3. Convolve image
    4. Reproject image
    """
    (target_img, 
     target_bkg, 
     target_bkgrms, 
     scale, 
     ref_zp, 
     zp_key,
          
     convolve, 
     ref_seeing,
     seeing_key,

     reproject, 
     reproject_type, 
     center_ra,
     center_dec, 
     x_size,
     y_size,
     pixel_scale, 
     verbose,
     save,
     clear
     ) = args
    # time.sleep(3) # For delayed multiprocessing

    # Load data
    target_img.data
    target_bkg.data if target_bkg is not None else None
    target_bkgrms.data if target_bkgrms is not None else None
    
    # Subtract background
    if target_bkg is not None:
        args_bkg = (target_img, target_bkg)
        subbkg_img = _subtract_background_worker(args_bkg)
        if target_bkgrms is not None:
            target_bkgrms.path = target_bkgrms.path.parent / f"subbkg_{target_bkgrms.path.name}"
    else:
        subbkg_img = target_img

    # Scale image
    remove_scaled = False
    if scale:
        args_scale = (subbkg_img, target_bkgrms, ref_zp, zp_key, False, verbose)
        scaled_img, scaled_bkgrms = _scale_worker(args_scale)
        remove_scaled = True
        if (convolve == False) & (reproject == False):
            remove_scaled = False
    else:
        scaled_img = subbkg_img
        scaled_bkgrms = target_bkgrms

    # Convolve image
    if convolve:
        args_convolve = (scaled_img, scaled_bkgrms, ref_seeing, seeing_key, verbose)
        convolved_img, convolved_bkgrms = _convolve_worker(args_convolve)
    else:
        convolved_img = scaled_img
        convolved_bkgrms = scaled_bkgrms

    # Reproject image
    if reproject:
        args_reproject = (convolved_img, convolved_bkgrms, reproject_type, center_ra, center_dec, x_size, y_size, pixel_scale, verbose)
        reprojected_img, reprojected_bkgrms = _reproject_worker(args_reproject)
    else:
        reprojected_img = convolved_img
        reprojected_bkgrms = convolved_bkgrms

    if save:
        reprojected_img.write(verbose = verbose)
        if reprojected_bkgrms is not None:
            reprojected_bkgrms.write(verbose = verbose)

    target_img.clear(verbose = False)
    scaled_img.clear(verbose = False)
    convolved_img.clear(verbose = False)
    if target_bkg is not None:
        target_bkg.clear(verbose = False)
    if target_bkgrms is not None:
        target_bkgrms.clear(verbose = False)
        scaled_bkgrms.clear(verbose = False)
        convolved_bkgrms.clear(verbose = False)
    if remove_scaled:
        scaled_img.remove(verbose = False)
        if scaled_bkgrms is not None:
            scaled_bkgrms.remove(verbose = False)
    
    if clear:
        reprojected_img.clear(verbose = False)
        if reprojected_bkgrms is not None:
            reprojected_bkgrms.clear(verbose = False)

    return reprojected_img, reprojected_bkgrms

class Stack:
    
    def __init__(self):        
        self.helper = Helper()
        self.combiner = Combiner()
        self.background = BackgroundGenerator()
        self.psfphot    = PSFPhotometry()

    def __repr__(self):
        return f"Method class: {self.__class__.__name__}\n For help, use 'help(self)' or `self.help()`."

    def help(self):
        # Get all public methods from the class, excluding `help`
        methods = [
            (name, obj)
            for name, obj in inspect.getmembers(self.__class__, inspect.isfunction)
            if not name.startswith("_") and name != "help"
        ]

        # Build plain text list with parameters
        lines = []
        for name, func in methods:
            sig = inspect.signature(func)
            params = [str(p) for p in sig.parameters.values() if p.name != "self"]
            sig_str = f"({', '.join(params)})" if params else "()"
            lines.append(f"- {name}{sig_str}")

        # Final plain text output
        help_text = ""
        self.helper.print(f"Help for {self.__class__.__name__}\n{help_text}\n\nPublic methods:\n" + "\n".join(lines), True)

    def prepare_images(self,
                       target_imglist: Union[List[ScienceImage], List[CalibrationImage]],
                       target_bkglist: Optional[List[Background]] = None,
                       target_bkgrmslist: Optional[List[Errormap]] = None,
                       n_proc: int = 4,                       

                       # Scale parameters
                       scale: bool = True,
                       zp_key: str = 'ZP_APER_2',
                       
                       # Convolution parameters
                       convolve: bool = False,
                       seeing_key: str = 'SEEING', 

                       # Reproject parameters
                       reproject: bool = True,
                       reproject_type: str = 'LANCZOS3',
                       center_ra: float = None,
                       center_dec: float = None,
                       pixel_scale: float = None,
                       x_size: int = None,
                       y_size: int = None,     
                       
                       # Other parameters
                       verbose: bool = True,
                       save: bool = True,
                       clear: bool = True):
        """
        Prepare a list of images for stacking.
        """
        # Check whether 
        # If scale is True, ensure all images have the same ZP key
        ref_zp = None
        if scale:
            zp_values = []
            for target_img in target_imglist:
                if zp_key not in target_img.header:
                    raise ValueError(f"Missing ZP key '{zp_key}' in {target_img.path}")
                zp_values.append(float(target_img.header[zp_key]))
            
            ref_zp = np.min(zp_values)
        
        # If convolve is True, ensure all images have the same SEEING key
        ref_seeing = None
        if convolve:
            seeing_values = []
            for target_img in target_imglist:
                if seeing_key not in target_img.header:
                    raise ValueError(f"Missing SEEING key '{seeing_key}' in {target_img.path}")
                seeing_values.append(float(target_img.header[seeing_key]))
            
            ref_seeing = np.max(seeing_values)
        
        # Prepare images for multiprocessing
        if target_bkglist is None:
            target_bkglist = [None] * len(target_imglist)
        if target_bkgrmslist is None:
            target_bkgrmslist = [None] * len(target_imglist)
        
        input_list = []
        for target_img, target_bkg, target_bkgrms in zip(target_imglist, target_bkglist, target_bkgrmslist):
            args = (target_img, target_bkg, target_bkgrms, scale, ref_zp, zp_key, convolve, ref_seeing, seeing_key, reproject, reproject_type, center_ra, center_dec, x_size, y_size, pixel_scale, verbose, save, clear)
            input_list.append(args)
        
        if n_proc == 1:
            results = [_prepare_image_worker(args) for args in input_list]
        else:
            results = process_map(_prepare_image_worker, input_list, max_workers=n_proc, desc="Preparing images...", ncols=80, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]")

        prepared_imglist, prepared_bkgrmslist = zip(*results)

        return prepared_imglist, prepared_bkgrmslist  

    def stack_multiprocess(self,
                           target_imglist: Union[List[ScienceImage], List[CalibrationImage]],
                           target_bkgrmslist: Optional[List[Errormap]] = None,
                           target_outpath: str = None,
                           bkgrms_outpath: str = None,
                           n_proc=4,
                           
                           # Clip parameters
                           combine_type: str = 'mean',
                           clip_type: str = None,
                           sigma: float = 3.0,
                           nlow: int = 1,
                           nhigh: int = 1,
                           
                           # Other parameters
                           verbose: bool = True,
                           save: bool = True,
                           remove_intermediate: bool = False):
        """
        Stack a list of images.
        
        Parameters
        ----------
        
        target_imglist : List[ScienceImage] or List[CalibrationImage]
            The list of images to stack.
        target_bkgrmslist : List[Errormap], optional
            The list of background RMS maps to use for the stacking.
        target_outpath : str
            The path to save the stacked image.
        bkgrms_outpath : str
            The path to save the background RMS map.
        n_proc : int
            The number of processes to use for the stacking.
        combine_type : str
            The type of combination to use for the stacking.
        clip_type : str
            The type of clipping to use for the stacking. [sigma, extrema]
        sigma : float
            The sigma for the clipping.
        nlow : int
            The number of low values to clip.
        nhigh : int
            The number of high values to clip.
        verbose : bool
            Whether to print verbose output.
        save : bool
            Whether to save the stacked image.
        Returns
        -------
        (combined_image, combined_bkgrms) : Tuple[ScienceImage, Errormap]
            The stacked image and background RMS map.
        """
        
        if self.combiner.n_proc != n_proc:
            self.helper.print('[Combiner] Re-initializing Combiner with new n_proc', verbose)
            self.combiner = Combiner(n_proc=n_proc)

        # Define output paths if not provided
        bkgrms_exists = False
        if target_outpath is None:
            target_outpath = target_imglist[0].savepath.savepath.with_suffix('.com.fits')
        if (target_bkgrmslist is not None):
            bkgrms_exists = True
            if (bkgrms_outpath is None):
                suffix = '.com.fits' + target_bkgrmslist[0].savepath.savepath.suffix 
                bkgrms_outpath = target_imglist[0].savepath.savepath.with_suffix(suffix) 
        
        # Load target images and error maps ---
        image_datalist = []
        image_hdrlist = []
        
        iterator = tqdm(target_imglist, desc="Loading target images...", ncols=80, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]") if verbose else target_imglist
        for img in iterator:
            image_datalist.append(img.data)
            image_hdrlist.append(img.header)

        bkgrms_datalist = []
        if bkgrms_exists:
            bkgrms_datalist = []
            iterator = tqdm(target_bkgrmslist, desc="Loading target bkgrms maps...", ncols=80, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]") if verbose else target_bkgrmslist
            for target_bkgrms in iterator:
                bkgrms_datalist.append(target_bkgrms.data)

        # Combine the image stack
        if clip_type == 'extrema':
            if len(image_datalist) - nlow - nhigh < 3:
                self.helper.print(f"[Combiner] Not enough images to clip: ({len(image_datalist)}). Clip type is set as None", verbose)
                clip_type = None
                
        combined_data, combined_bkgrms = self.combiner.combine_images_parallel(
            image_list=image_datalist,
            bkgrms_list = bkgrms_datalist,
            combine_method=combine_type,
            clip_method=clip_type,
            sigma=sigma,
            nlow=nlow,
            nhigh=nhigh,
            verbose=verbose
        )

        # Initialize combined header 
        combined_header = image_hdrlist[0].copy()

        # Update header keywords with mean
        update_header_keywords_mean = ['ALTITUDE', 'AZIMUTH', 'CENTALT', 'CENTAZ', 'RA', 'DEC', 'AIRMASS', 'SEEING', 'PEEING', 'ELLIP', 'SKYVAL', 'JD', 'MJD', 'MJD-OBS']
        for key in update_header_keywords_mean:
            values = [hdr.get(key) for hdr in image_hdrlist if hdr.get(key) not in [None, '']]
            try:
                if values:
                    combined_header[key] = float(np.nanmean(values))
            except Exception:
                pass  # Handle non-numeric or incompatible values
        for i, target_img in enumerate(target_imglist):
            combined_header[f'COMBIM{i+1}'] = target_img.path.name
        
        values = [Time(hdr.get('DATE-OBS')).jd for hdr in image_hdrlist if hdr.get('DATE-OBS') not in [None, '']]
        combined_header['DATE-OBS'] = Time(np.nanmean(values), format='jd').isot if values else None
        values = [Time(hdr.get('DATE-LOC')).jd for hdr in image_hdrlist if hdr.get('DATE-LOC') not in [None, '']]
        combined_header['DATE-LOC'] = Time(np.nanmean(values), format='jd').iso if values else None

        update_header_keywords_sum = ['EXPTIME', 'EXPOSURE']
        for key in update_header_keywords_sum:
            values = [hdr.get(key) for hdr in image_hdrlist if hdr.get(key) not in [None, '']]
            try:
                if values:
                    combined_header[key] = float(np.nansum(values))
            except Exception:
                pass
            
        # Remove unwanted header keywords
        update_header_keywords_remove = ['IMAGEID', 'NOTE', 'MAG_*', 'ZP*', 'UL*', 'EZP*', 'APER*', 'SKYSIG']
        for pattern in update_header_keywords_remove:
            if '*' in pattern:
                regex = re.compile('^' + pattern.replace('*', '.*') + '$')
                keys_to_remove = [k for k in combined_header if regex.match(k)]
            else:
                keys_to_remove = [k for k in combined_header if k == pattern]
            for k in keys_to_remove:
                del combined_header[k]

        # Save combined image
        # If CalibrationImage is input, Save it as CalibrationImage. This will be saved in the master_frame directory.
        # Else, save it in the target_outpath.
        stack_instance =  type(target_imglist[0])(path = target_outpath, telinfo = target_imglist[0].telinfo, load = False)
        stack_instance.data = combined_data
        stack_instance.header = combined_header
        stack_instance.update_status(process_name = 'STACK')
        
        stack_bkgrms_instance = None
        if target_bkgrmslist is not None:
            stack_bkgrms_instance = Errormap(path=bkgrms_outpath, emaptype = 'bkgrms', load=False)
            stack_bkgrms_instance.data = combined_bkgrms
            stack_bkgrms_instance.header = combined_header
        
        if save:
            stack_instance.write(verbose = verbose)
            stack_bkgrms_instance.write(verbose = verbose) if stack_bkgrms_instance is not None else None

        if remove_intermediate:
            for target_img in target_imglist:
                target_img.remove(verbose = verbose)
            if target_bkgrmslist is not None:
                for target_bkgrms in target_bkgrmslist:
                    target_bkgrms.remove(verbose = verbose)
        return stack_instance, stack_bkgrms_instance
    
    def stack_swarp(self,
                    target_imglist : Union[List[ScienceImage], List[CalibrationImage]],
                    target_bkglist: Optional[List[Background]] = None,
                    target_errormaplist: Optional[List[Errormap]] = None,
                    target_outpath: str = None,
                    errormap_outpath: str = None,
                    combine_type: str = 'median', # median, weighted, mean, sum, min, max
                    
                    # Resample parameters
                    resample: bool = False,
                    resample_type: str = 'LANCZOS3',
                    center_ra: float = None,
                    center_dec: float = None,
                    x_size: int = None,
                    y_size: int = None,
                    
                    # Scale parameters
                    scale: bool = False,
                    scale_type: str = 'min',
                    zp_key : str = 'ZP_APER_1',
                    
                    # Convolution parameters
                    convolve: bool = False,
                    seeing_key: str = 'SEEING',
                    kernel: str = 'gaussian',
                    
                    # Other parameters
                    save: bool = True,
                    verbose: bool = True,
                    **kwargs
                    ):
        """
        Stack multiple images using SWArp.
        
        Parameters
        ----------
        target_imglist : List[Union[ScienceImage, CalibrationImage]]
            List of images to stack
        target_bkglist : Optional[List[Background]]
            Optional list of background maps to stack
        target_errormaplist : Optional[List[Errormap]]
            Optional list of error maps to stack
        target_outpath : str
            Path to save the stacked image
        errormap_outpath : str
            Path to save the stacked error map
        combine_type : str
            Method to combine images ('median', 'weighted', 'mean', 'sum', 'min', 'max')
        resample : bool
            Whether to resample the images
        resample_type : str
            Type of resampling to use ('NEAREST', 'BILINEAR', 'BICUBIC', 'LANCZOS3', etc.)
        center_ra : float
            RA of the center of the stacked image
        center_dec : float
            Dec of the center of the stacked image
        pixel_scale : float
            Pixel scale of the stacked image
        x_size : int
            Size of the stacked image in pixels
        y_size : int
            Size of the stacked image in pixels
        scale : bool
            Whether to scale the images
        scale_type : str
            Method to scale images ('min', 'mean', 'median', 'max')
        zp_key : str
            Header keyword for zero point
        convolve : bool
            Whether to convolve the images
        seeing_key : str
            Header keyword for seeing/FWHM in pixel units
        kernel : str
            Convolution kernel type ('gaussian')
        save : bool
            Whether to save the stacked image and error map
        verbose : bool
            Print progress messages
        **kwargs : dict, optional
            Additional keyword arguments.

        Returns
        -------
        (stack_instance, stack_weight_instance) : Tuple[Union[ScienceImage, CalibrationImage], Optional[Errormap]]
            Stacked image and optionally its weight map
        """ 
        # Set default output paths if not provided
        if target_outpath is None:
            target_outpath = target_imglist[0].savepath.combinepath
        
        errormap_outpath = target_outpath + '.weight'
        
        # Set temporary output paths
        target_outpath_tmp = str(target_outpath) + '.tmp'
        
        # Prepare images
        target_imglist, target_errormaplist = self.prepare_images(
            target_imglist = target_imglist,
            target_errormaplist = target_errormaplist,
            target_outpath = target_outpath,
            errormap_outpath = errormap_outpath,
            combine_type = combine_type,
            resample = False,
            resample_type = resample_type,
            center_ra = center_ra,
            center_dec = center_dec,
            x_size = x_size,
            y_size = y_size,
            scale = scale,
            scale_type = scale_type,
            zp_key = zp_key,
            convolve = convolve,
            seeing_key = seeing_key,
            kernel = kernel,
            save = True,
            verbose = verbose
        )

        # Loading the images
        image_pathlist = []
        image_hdrlist = []
        remove_image = []
        iterator = tqdm(target_imglist, desc="Loading target images...", ncols=80, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]") if verbose else target_imglist
        for target_img in iterator:
            if not target_img.is_exists:
                target_img.write(verbose = verbose)
                remove_image.append(True)
            else:
                remove_image.append(False)
            image_pathlist.append(target_img.path)
            image_hdrlist.append(target_img.header)
            
        weight_pathlist = None
        remove_errormap = []
        if target_errormaplist is not None:
            weight_pathlist = []
            iterator = tqdm(target_errormaplist, desc="Loading target weight maps...", ncols=80, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]") if verbose else target_errormaplist
            for target_errormap in iterator:
                if target_errormap.emaptype.lower() != 'weight':
                    target_errormap.to_weight()
                    if not target_errormap.is_exists:
                        target_errormap.write(verbose = verbose)
                        remove_errormap.append(True)
                    else:
                        remove_errormap.append(False)
                else:
                    if not target_errormap.is_exists:
                        target_errormap.write(verbose = verbose)
                        remove_errormap.append(True)
                    else:
                        remove_errormap.append(False)
                    
                weight_pathlist.append(target_errormap.path)
                
        # Header modification
        combined_header = image_hdrlist[0].copy()

        # --- Update header keywords with mean ---
        update_header_keywords_mean = ['ALTITUDE', 'AZIMUTH', 'RA', 'DEC', 'AIRMASS', 'SEEING', 'PEEING', 'ELLIP', 'ELONG', 'SKYVAL', 'JD', 'MJD', 'MJD-OBS']
        for key in update_header_keywords_mean:
            values = [hdr.get(key) for hdr in image_hdrlist if hdr.get(key) not in [None, '']]
            try:
                if values:
                    combined_header[key] = float(np.nanmean(values))
            except Exception:
                pass  # Handle non-numeric or incompatible values
        for i, target_img in enumerate(target_imglist):
            combined_header[f'COMBIM{i+1}'] = target_img.path.name
            
        values = [Time(hdr.get('DATE-OBS')).jd for hdr in image_hdrlist if hdr.get('DATE-OBS') not in [None, '']]
        combined_header['DATE-OBS'] = Time(np.nanmean(values), format='jd').isot if values else None
        values = [Time(hdr.get('DATE-LOC')).jd for hdr in image_hdrlist if hdr.get('DATE-LOC') not in [None, '']]
        combined_header['DATE-LOC'] = Time(np.nanmean(values), format='jd').iso if values else None

        # --- Remove unwanted header keywords ---
        update_header_keywords_remove = ['IMAGEID', 'NOTE', 'MAG_*', 'ZP*', 'UL*', 'EZP*', 'APER*', 'SKYSIG']
        for pattern in update_header_keywords_remove:
            if '*' in pattern:
                regex = re.compile('^' + pattern.replace('*', '.*') + '$')
                keys_to_remove = [k for k in combined_header if regex.match(k)]
            else:
                keys_to_remove = [k for k in combined_header if k == pattern]
            for k in keys_to_remove:
                del combined_header[k]

        # Image combine
        self.helper.print(f"Start image combining...", verbose)
        imagestack_path = None
        weightstack_path = None
        # Run swarp 
        stack_pathlist = self.helper.run_swarp(
            target_path = image_pathlist,
            swarp_configfile = target_imglist[0].config['SWARP_CONFIG'],
            swarp_params = None,
            target_outpath = target_outpath,
            weight_inpath = weight_pathlist,
            weight_outpath = errormap_outpath,
            weight_type = 'MAP_WEIGHT',
            resample = resample,
            resample_type = resample_type,
            center_ra = center_ra,
            center_dec = center_dec,
            x_size = x_size,
            y_size = y_size,
            pixelscale = np.mean(target_imglist[0].pixelscale),
            combine = True,
            combine_type = combine_type,
            subbkg = False
        )
        imagestack_path, weightstack_path = stack_pathlist
            
        # If errormaplist is provided, run swarp for error maps (NEAREST resampling)
        if target_errormaplist is not None:
            self.helper.print(f"Start weight combining...", verbose)
            stack_pathlist = self.helper.run_swarp(
                target_path = image_pathlist,
                swarp_configfile = target_imglist[0].config['SWARP_CONFIG'],
                swarp_params = None,
                target_outpath = target_outpath_tmp,
                weight_inpath = weight_pathlist,
                weight_outpath = errormap_outpath,
                weight_type = 'MAP_WEIGHT',
                resample = resample,
                resample_type = 'NEAREST',
                center_ra = center_ra,
                center_dec = center_dec,
                x_size = x_size,
                y_size = y_size,
                pixelscale = np.mean(target_imglist[0].pixelscale),
                combine = True,
                combine_type = combine_type,
                subbkg = False
            )
            imagestack_tmppath, weightstack_path = stack_pathlist
            os.remove(imagestack_tmppath)
        
        if type(target_imglist[0]) == CalibrationImage:
            stack_instance = CalibrationImage(path = target_outpath, telinfo = target_imglist[0].telinfo, load = True)
            stack_instance.header = self.helper.merge_header(stack_instance.header, combined_header, exclude_keys = ['PV*'])
        else:
            stack_instance = type(target_imglist[0])(path = imagestack_path, telinfo = target_imglist[0].telinfo, load = True)
            stack_instance.header = self.helper.merge_header(stack_instance.header, combined_header, exclude_keys = ['PV*'])
            stack_instance.update_status(process_name = 'STACK')

        stack_weight_instance = None
        stack_weight_instance = Errormap(path = weightstack_path, emaptype = 'bkgweight', status = None, load = True)
        stack_weight_instance.header = self.helper.merge_header(stack_weight_instance.header, combined_header, exclude_keys = ['PV*'])
        event_details_kwargs = dict(
            stack_type = 'SWARP',
            combine_type = combine_type,
            resample = resample,
            resample_type = resample_type,
            ncombine = len(target_imglist)
        )
        stack_weight_instance.add_status('stack_swarp', **event_details_kwargs)
        
        if save:
            stack_instance.write(verbose = verbose)
            stack_weight_instance.write(verbose = verbose) if stack_weight_instance is not None else None
            self.helper.print(f"Stacked image saved to {stack_instance.path}", verbose)
            self.helper.print(f"Stacked weight map saved to {stack_weight_instance.path}", verbose)
        else:
            stack_instance.load()
            stack_weight_instance.load()
            stack_instance.remove(verbose = verbose)
            stack_weight_instance.remove(verbose = verbose) if stack_weight_instance is not None else None
        
        if any(remove_errormap):
            for remove_key, target_errormap in zip(remove_errormap, target_errormaplist):
                if remove_key:
                    target_errormap.remove(verbose = verbose)
        if any(remove_image):
            for remove_key, target_img in zip(remove_image, target_imglist):
                if remove_key:
                    target_img.remove(verbose = verbose)
                
        return stack_instance, stack_weight_instance
    
    def select_quality_images(self, 
                              target_imglist: Union[List[ScienceImage], List[ReferenceImage]],
                              min_obsdate: Union[Time, str, float] = None,
                              max_obsdate: Union[Time, str, float] = None,
                              seeing_key: str = 'SEEING',
                              depth_key: str = 'UL5SKY_APER_1',
                              ellipticity_key: str = 'ELLIP',
                              obsdate_key: str = 'DATE-OBS',
                              weight_ellipticity: float = 3.0,
                              weight_seeing: float = 1.0,
                              weight_depth: float = 2.0,
                              max_numbers: int = None,
                              seeing_limit: float = 6.0,
                              depth_limit: float = 18.0,
                              ellipticity_limit: float = 0.3,
                              visualize: bool = False,
                              verbose: bool = True):
        """
        Select the best images based on seeing, depth, and ellipticity.
        
        Parameters
        ----------
        target_imglist : List[Union[ScienceImage, ReferenceImage]]
            List of images to select from.
        min_obsdate : Union[Time, str, float]
            Minimum observation date.
        max_obsdate : Union[Time, str, float]
            Maximum observation date.
        seeing_key : str
            Header keyword for seeing/FWHM in pixel units
        depth_key : str
            Header keyword for depth in AB magnitude
        ellipticity_key : str
            Header keyword for ellipticity
        obsdate_key : str
            Header keyword for observation date
        weight_ellipticity : float
            Weight for ellipticity
        weight_seeing : float
            Weight for seeing
        weight_depth : float
            Weight for depth
        max_numbers : int, optional
            Maximum number of images to select.
        seeing_limit : float
            Maximum seeing limit in arcseconds
        depth_limit : float
            Minimum depth limit in AB magnitude
        ellipticity_limit : float
            Maximum ellipticity limit
        visualize : bool
            Whether to visualize the selected images
        verbose : bool
            Whether to print verbose output.

        Returns
        -------
        (selected_imglist, selected_errormaplist) : Tuple[List[Union[ScienceImage, ReferenceImage]], Optional[List[Errormap]]]
            List of selected images and optionally their error maps
        """
        
        seeinglist = []
        depthlist = []
        ellipticitylist = []
        obsdatelist = []
        iterator = tqdm(target_imglist, desc="Querying images...", ncols=80, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]") if verbose else target_imglist
        for target_img in iterator:
            seeinglist.append(target_img.header.get(seeing_key, None))
            depthlist.append(target_img.header.get(depth_key, None))
            ellipticitylist.append(target_img.header.get(ellipticity_key, None))
            obsdatelist.append(target_img.header.get(obsdate_key, None))
        
        try:
            obsdate_time = Time(obsdatelist)
            min_obs_time = self.helper.flexible_time_parser(min_obsdate) if min_obsdate is not None else Time('1990-01-01')
            max_obs_time = self.helper.flexible_time_parser(max_obsdate) if max_obsdate is not None else Time.now()
        except Exception as e:
            raise ValueError(f"Invalid date format: {e}")          
        
        # Mask for images before max_obsdate
        valid_obs_mask = (obsdate_time < max_obs_time) & (obsdate_time > min_obs_time)
        
        # Also apply validity mask for seeing, depth, ellipticity
        seeinglist = np.array([v if v is not None else np.nan for v in seeinglist], dtype=float)
        depthlist = np.array([v if v is not None else np.nan for v in depthlist], dtype=float)
        ellipticitylist = np.array([v if v is not None else np.nan for v in ellipticitylist], dtype=float)
        valid_value_mask = (~np.isnan(seeinglist)) & (~np.isnan(depthlist)) & (~np.isnan(ellipticitylist))
        if not np.any(valid_value_mask):
            return []
                    
        # Apply limits mask
        valid_seeing_mask = seeinglist < seeing_limit
        valid_ellipticity_mask = ellipticitylist < ellipticity_limit
        valid_depth_mask = depthlist > depth_limit
        
        # Final combined mask (same length as target_imglist)
        combined_mask = (
            valid_obs_mask &
            valid_value_mask &
            valid_seeing_mask &
            valid_ellipticity_mask &
            valid_depth_mask
        )
        if not np.any(combined_mask):
            return []
        
        # Apply final mask
        ell_all = np.array(ellipticitylist)[valid_value_mask]
        see_all = np.array(seeinglist)[valid_value_mask]
        dep_all = np.array(depthlist)[valid_value_mask]
        obsdate_all = np.array(obsdatelist)[valid_value_mask]
        
        ell_filtered = np.array(ellipticitylist)[combined_mask]
        see_filtered = np.array(seeinglist)[combined_mask]
        dep_filtered = np.array(depthlist)[combined_mask]
        imgs_filtered = np.array(target_imglist)[combined_mask]
        obsdate_filtered = np.array(obsdate_time)[combined_mask]
        from sklearn.preprocessing import MinMaxScaler
        from matplotlib.gridspec import GridSpec

        # Normalize
        scaler = MinMaxScaler()
        ell_norm = scaler.fit_transform(ell_filtered.reshape(-1, 1)).flatten()
        see_norm = scaler.fit_transform(see_filtered.reshape(-1, 1)).flatten()
        dep_norm = scaler.fit_transform(dep_filtered.reshape(-1, 1)).flatten()

        # Compute combined score
        # You can adjust weights if needed
        score = (1 - ell_norm) * weight_ellipticity + (1 - see_norm) * weight_seeing + dep_norm * weight_depth

        # Rank and select best images
        sorted_idx = np.argsort(score)[::-1]  # descending
        best_images = imgs_filtered[sorted_idx]
        if max_numbers is None:
            num_select = max(1, int(len(sorted_idx)))  # select top 90%
        else:
            num_select = max_numbers
        selected_idx = sorted_idx[:num_select]

        # Top N or just best
        best_image = best_images[0]
        
        # Data for plotting
        x_all = np.array(see_all)
        y_all = np.array(dep_all)
        c_all = np.array(ell_all)
        x_valid = np.array(see_filtered)
        y_valid = np.array(dep_filtered)
        c_valid = np.array(ell_filtered)
        x_selected = x_valid[selected_idx]
        y_selected = y_valid[selected_idx]
        c_selected = c_valid[selected_idx]
        idx_best = sorted_idx[0]
        x_best = x_valid[idx_best]
        y_best = y_valid[idx_best]
        c_best = c_valid[idx_best]

        # Create marker masks with full length
        marker_sizes_full = np.where(combined_mask, 50, 10)
        marker_alphas_full = np.where(combined_mask, 0.8, 0.2)

        # Apply valid_value_mask to match x_all, y_all
        marker_sizes = marker_sizes_full[valid_value_mask]
        marker_alphas = marker_alphas_full[valid_value_mask]

        # Calculate percentiles (90%, 75%, and 50%)
        p90_x, p75_x, p50_x, p25_x, p10_x = np.percentile(x_all, [10, 25, 50, 75, 90])
        p90_y, p75_y, p50_y, p25_y, p10_y = np.percentile(y_all, [90, 75, 50, 25, 10])

        # Calculate the number of images for each percentile
        num_images_p90 = np.sum((x_all <= p90_x) & (y_all >= p90_y))  # Number of images below or equal to the 10th percentile
        num_images_p75 = np.sum((x_all <= p75_x) & (y_all >= p75_y))  # Number of images below or equal to the 25th percentile
        num_images_p50 = np.sum((x_all <= p50_x) & (y_all >= p50_y))  # Number of images below or equal to the 50th percentile
        num_images_p25 = np.sum((x_all <= p25_x) & (y_all >= p25_y))  # Number of images below or equal to the 75th percentile

        # Create figure with GridSpec layout
        if visualize:
            fig = plt.figure(figsize=(6, 6), dpi=300)
            gs = GridSpec(4, 4, fig, wspace=1.5, hspace=0.5)

            # Create scatter plot
            ax_main = fig.add_subplot(gs[1:, :-1])
            sc = ax_main.scatter(x_all, y_all,
                                c=c_all,
                                s=marker_sizes,
                                alpha=marker_alphas,
                                cmap='viridis', edgecolors='k', linewidths=0.5,
                                label = f'All images ({len(x_all)})')        
            ax_main.scatter(0,0, s = 10, alpha = 0.2, label = f'Filtered out images ({len(x_all) - len(x_selected)})')
            cbar = fig.colorbar(sc, ax=ax_main, pad=0.01)
            cbar.set_label('Ellipticity')
            ax_main.axvline(p90_x, color='r', linestyle='--')
            ax_main.axvline(p75_x, color='b', linestyle='--')
            ax_main.axvline(p50_x, color='g', linestyle='--')
            ax_main.axvline(p25_x, color='k', linestyle='--')
            ax_main.axhline(p90_y, color='r', linestyle='--')
            ax_main.axhline(p75_y, color='b', linestyle='--')
            ax_main.axhline(p50_y, color='g', linestyle='--')
            ax_main.axhline(p25_y, color='k', linestyle='--')
            ax_main.set_xlim(p90_x - 0.5, p10_x + 0.5)
            ax_main.set_ylim(p10_y - 1, p90_y + 1)
            ax_main.set_xlabel('Seeing [arcsec]')
            ax_main.set_ylabel('Depth [AB]')
            ax_main.scatter(x_selected, y_selected, marker='*', s=200, c='red', edgecolors='black', label=f'Selected ({len(selected_idx)}) images')
            ax_main.scatter(x_best, y_best, marker='*', s=200, c='red', edgecolors='black')
            ax_main.text(x_best, y_best + 0.3,
                        f"Best\nSeeing = {x_best:.2f} arcsec\nDepth = {y_best:.2f} AB\nEllipticity = {c_best:.2f}",
                        color='red', fontsize=8, ha='center', va='bottom',
                        bbox=dict(facecolor='white', edgecolor='red', boxstyle='round,pad=0.3'))
            ax_main.legend(loc='upper right', fontsize=8, frameon=True)


            # Create top histogram
            ax_histx = fig.add_subplot(gs[0, :-1], sharex=ax_main)
            ax_histx.hist(x_valid, bins=30, color='black', edgecolor='black', alpha=0.7)
            ax_histx.spines['top'].set_visible(False)  # Hide top spine
            ax_histx.spines['right'].set_visible(False)  # Hide right spine

            # Create right histogram
            ax_histy = fig.add_subplot(gs[1:, -1], sharey=ax_main)
            ax_histy.hist(y_valid, bins=30, color='black', edgecolor='black', alpha=0.7, orientation='horizontal')
            ax_histy.spines['top'].set_visible(False)  # Hide top spine
            ax_histy.spines['right'].set_visible(False)  # Hide right spine

            # Set limits for histograms to fit within the black box
            ax_histx.set_xlim(ax_main.get_xlim())
            ax_histy.set_ylim(ax_main.get_ylim())

            # Plot vertical regions for percentiles in histograms
            ax_histx.axvline(p90_x, color='r', linestyle='--', label='90%')
            ax_histx.axvline(p75_x, color='b', linestyle='--', label='75%')
            ax_histx.axvline(p50_x, color='g', linestyle='--', label='50%')
            ax_histx.axvline(p25_x, color='k', linestyle='--', label='25%')

            ax_histy.axhline(p90_y, color='r', linestyle='--', label='90%')
            ax_histy.axhline(p75_y, color='b', linestyle='--', label='75%')
            ax_histy.axhline(p50_y, color='g', linestyle='--', label='50%')
            ax_histy.axhline(p25_y, color='k', linestyle='--', label='25%')

            # Add text annotation in the upper right region of the scatter plot
            text = f'Percentile (# of images, Seeing, Depth):\n'
            text += f'90% ({num_images_p90}, {p90_x:.2f}, {p90_y:.2f})\n'
            text += f'75% ({num_images_p75}, {p75_x:.2f}, {p75_y:.2f})\n'
            text += f'50% ({num_images_p50}, {p50_x:.2f}, {p50_y:.2f})\n'
            text += f'25% ({num_images_p25}, {p25_x:.2f}, {p25_y:.2f})'
            ax_main.text(0.5, 0.15, text,
                        ha='center', va='center',
                        transform=ax_main.transAxes,
                        bbox=dict(facecolor='white', edgecolor='black', boxstyle='round,pad=0.5'))
            
            from matplotlib.lines import Line2D

            dashed_lines = [
                Line2D([0], [0], color='red', linestyle='--', label='90%'),
                Line2D([0], [0], color='blue', linestyle='--', label='75%'),
                Line2D([0], [0], color='green', linestyle='--', label='50%'),
                Line2D([0], [0], color='black', linestyle='--', label='25%')
            ]

            fig.legend(handles=dashed_lines,
                    loc='upper right',
                    bbox_to_anchor=(0.95, 0.95),
                    fontsize=10, frameon=True)
            plt.tight_layout()
            plt.show()
        
        selected_images = imgs_filtered[selected_idx]
        
        return selected_images

