"""
Telluric removal module using PCA (Principal Component Analysis) in the wavelength domain.
This module processes spectroscopic data to remove telluric features using PCA decomposition.
"""

import traceback
import os
import sys
import pickle
import yaml
from pathlib import Path
from scipy import ndimage
from astropy.io import fits
from sklearn.decomposition import PCA
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
import copy
import matplotlib

from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import estimate_continuum

if not os.environ.get('SPHINX_BUILD'):
    matplotlib.use("TkAgg")
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
import numpy as np

# Module-level imports - path setup moved to main()


def combine(path_night_completef, name, transpose_read, repeatf=None, sort_indices=None):
    """
    Combine A and B channel data into a single dataset.

    Args:
        path_night_completef: Path to the night's data
        name: Base name for the files
        transpose_read: Whether to transpose the data during reading
        repeatf: Optional file name for repeating dimensions
        sort_indices: Optional indices to sort the combined data (if None, interleave A and B)
    """
    from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import read_fits
    path_A = str(Path(path_night_completef, "A_" + name + ".fits"))
    path_B = str(Path(path_night_completef, "B_" + name + ".fits"))
    if os.path.exists(path_A) and os.path.exists(path_B):
        # Read A and B channel data
        arr_A = read_fits(path_A, transpose=transpose_read)
        arr_B = read_fits(path_B, transpose=transpose_read)

        # Handle dimension repetition if specified
        if repeatf is not None:
            arr_dimension = read_fits(str(Path(path_night_completef, "A_" + repeatf + ".fits")), transpose=transpose_read)
            arr_A = np.repeat(arr_A[:, np.newaxis, :], arr_dimension.shape[1], axis=1)
            arr_B = np.repeat(arr_B[:, np.newaxis, :], arr_dimension.shape[1], axis=1)

        if sort_indices is not None:
            # Concatenate and sort using provided indices
            combined = np.concatenate([arr_A, arr_B], axis=1)
            combined = combined[:, sort_indices, :]
        else:
            # Create combined array with interleaving (original behavior)
            combined = np.empty(
                (arr_A.shape[0], arr_A.shape[1] + arr_B.shape[1], arr_A.shape[2]),
                dtype=arr_A.dtype
            )

            # Interleave A and B data
            combined[:, 0::2, :] = arr_A
            combined[:, 1::2, :] = arr_B

        # Save combined data
        hduf = fits.PrimaryHDU(combined.T)
        hdulistf = fits.HDUList([hduf])
        hdulistf.writeto(str(Path(path_night_completef, name + ".fits")), overwrite=True)


def combine_nfc(path_night_completef):
    """
    Copy A_nfc.fits to nfc.fits.

    Args:
        path_night_completef: Path to the night's data
    """
    source = str(Path(path_night_completef, "A_nfc.fits"))
    dest = str(Path(path_night_completef, "nfc.fits"))
    os.system(f"cp {source} {dest}")

def load_fase_data(file_path):
    """Load phase data from file, handling both dictionary and legacy formats."""
    with open(file_path, "rb") as fo:
        # Try to load as dictionary first (new format)
        phase_data = pickle.load(fo)
        return (
            np.array(phase_data["barycentric_velocity"]),
            np.array(phase_data["transit_phase"]),
            np.array(phase_data["hjd"]),
            np.array(phase_data["BJD_TDB"])
        )

def combine_fase(path_night_completef):
    """
    Combine phase data from A and B channels into a single file.

    Handles both new dictionary format and legacy sequential format for backward compatibility.

    Args:
        path_night_completef: Path to the night directory

    Returns:
        sort_indices: Indices used to sort the combined data based on hjd
    """


    # Load A channel data
    barycentric_velocity_A, transit_phase_A, hjd_A, BJD_TDB_A = load_fase_data(
        str(Path(path_night_completef, "A_phase.pkl"))
    )

    # Load B channel data
    barycentric_velocity_B, transit_phase_B, hjd_B, BJD_TDB_B = load_fase_data(
        str(Path(path_night_completef, "B_phase.pkl"))
    )

    # Concatenate all arrays
    combined_hjd = np.concatenate([hjd_A, hjd_B])
    combined_barycentric_velocity = np.concatenate([barycentric_velocity_A, barycentric_velocity_B])
    combined_transit_phase = np.concatenate([transit_phase_A, transit_phase_B])
    combined_BJD_TDB = np.concatenate([BJD_TDB_A, BJD_TDB_B])

    # Sort all arrays using the same indices from hjd
    sort_indices = np.argsort(combined_hjd)
    combined_hjd = combined_hjd[sort_indices]
    combined_barycentric_velocity = combined_barycentric_velocity[sort_indices]
    combined_transit_phase = combined_transit_phase[sort_indices]
    combined_BJD_TDB = combined_BJD_TDB[sort_indices]

    # Save combined data as dictionary (new format)
    phase_dictionary = {
        "barycentric_velocity": combined_barycentric_velocity,
        "transit_phase": combined_transit_phase,
        "hjd": combined_hjd,
        "BJD_TDB": combined_BJD_TDB
    }
    with open(str(Path(path_night_completef, "phase.pkl")), "wb") as fof:
        pickle.dump(phase_dictionary, fof)

    return sort_indices


def combine_maskinv_err_pca(path_night_completef):
    """
    Combine mask_good and error information from A and B channels.

    Args:
        path_night_completef: Path to the night's data
    """
    # Load masks from A and B channels
    with open(str(Path(path_night_completef, "A_mask_inv_complete.pkl")), "rb") as fof:
        maskinvm_A = pickle.load(fof)
    with open(str(Path(path_night_completef, "B_mask_inv_complete.pkl")), "rb") as fof:
        maskinvm_B = pickle.load(fof)

    # Combine masks based on operation type
    maskinvm_func = np.logical_or(maskinvm_A, maskinvm_B)

    # Save combined mask_good
    with open(str(Path(path_night_completef, "mask_inv_complete.pkl")), "wb") as fof:
        pickle.dump(maskinvm_func, fof)


def median_idl(arr):
    """
    Calculate median using IDL-like tell_rm_method.

    Args:
        arr: Input array

    Returns:
        Median value
    """
    arr_s = np.sort(arr)
    return arr_s[round((len(arr_s) + 1) / 2)]


def linear_solver(topca_l, res_l):
    """
    Solve linear system using SVD decomposition.

    Args:
        topca_l: Input data matrix
        res_l: Result matrix

    Returns:
        Reconstructed array
    """
    # Perform SVD decomposition
    u, s, vh = np.linalg.svd(res_l, full_matrices=False)
    uT = np.transpose(u)
    nD = np.diag(1 / s)
    vHcT = np.transpose(vh.conj())
    # U diag(s) Vh x = b <=> diag(s) Vh x = U.T b = c
    # Solve linear system
    c = np.dot(uT, topca_l)
    # diag(s) Vh x = c <=> Vh x = diag(1/s) c = w (trivial inversion of a diagonal matrix)
    w = np.dot(nD, c)
    # Vh x = w <=> x = Vh.H w (where .H stands for hermitian = conjugate transpose)
    x = np.dot(vHcT, w)
    array_reconstructf = np.dot(res_l, x)

    return array_reconstructf


def plot_tangent(pdf_tangentf, current_order, tan_sncf, der_secf, tangent_thresholdf, n_components_pcaf):
    """
    Plot tangent information for PCA component selection.

    Args:
        pdf_tangentf: PDF file to save plot
        current_order: Order 
        tan_sncf: Tangent signal-to-noise curve
        der_secf: Second derivative
        tangent_thresholdf: Threshold for tangent
        n_components_pcaf: Number of PCA components
    """
    fig_tangent, (ax0, ax1) = plt.subplots(2, 1, sharex=True, figsize=(8, 6), constrained_layout=True)
    fig_tangent.suptitle(f"Order n.{current_order}")

    # Configure first subplot
    ticksx = np.arange(1, tan_sncf.size + 1, 2)
    ax0.set_ylabel(r'$\Delta\sigma$ (%)', fontsize=20)
    ax0.set_xticks(ticksx)
    ax0.tick_params(labelbottom=True)
    ax0.tick_params(axis='x', labelsize=0, length=5., width=2., top=True, bottom=True,
                    direction='in')
    ax0.tick_params(axis='y', labelsize=20, length=6., width=2., left=True, right=True,
                    direction='in')
    ax0.minorticks_on()
    ax0.tick_params(axis='y', which='minor', direction='in', length=4.0, left=True,
                    right=True)
    ax0.tick_params(axis='x', which='minor', direction='in', length=3.0, top=False,
                    bottom=False)
    ax0.set_yscale('log')

    # Configure second subplot
    ax1.set_ylabel(r'$\delta$', fontsize=20)
    ax1.set_xlabel('Number of principal components', fontsize=20)
    ax1.tick_params(axis='both', which='both')
    ax1.tick_params(labelbottom=True)
    ax1.tick_params(axis='x', labelsize=20)
    ax1.tick_params(axis='y', labelsize=20)
    ax1.tick_params(direction="in", bottom=True, top=True, left=True, right=True,
                    width=2)
    ax1.set_ylim(-1.01, 0.1)
    ax0.set_ylim(0.80, 123.)

    # Plot data
    ax0.plot(np.arange(1, tan_sncf.size + 1), 100. * tan_sncf, "-*", label="Variance first derivative")
    ax1.plot(np.arange(2, der_secf.size + 2), der_secf, "-*", label="Variance second derivative")
    ax1.axhline(y=-tangent_thresholdf, color='black', label=f'Threshold = {-tangent_thresholdf}')
    ax1.axvline(x=n_components_pcaf, color='red', label='Chosen component')
    ax1.set_xticks(ticksx)
    ax1.minorticks_off()
    ax0.legend()
    ax1.legend()

    # Save plot
    pdf_tangentf.savefig(fig_tangent)
    fig_tangent.clear()
    plt.close(fig_tangent)


def plot_process(
    initial_spectrum, spectrum_after_mask, log_spectrum_after_mask, log_normalized_spectrum_after_mask, log_residuals_after_reconstruction,
    log_residuals_after_reconstruction_and_eventual_smooth, wl_order, pdf_processf, n_pixels, oi, phases
):
    """
    Plot processing steps for telluric removal.

    Args:
        initial_spectrum: Initial data
        spectrum_after_mask: Data after masking
        log_spectrum_after_mask: Logarithmic of data after masking
        log_normalized_spectrum_after_mask: Logarithmic of normalized data after masking
        log_residuals_after_reconstruction: Residuals after reconstruction
        log_residuals_after_reconstruction_and_eventual_smooth: Residuals after reconstruction and eventual smoothing
        wl_order: Wavelength array for order
        pdf_processf: PDF file to save plots
        n_pixels: Number of pixels
        oi: Order index
        phases: Planetary orbital phases
    """
    # Set font sizes
    fontsizes = {
        'suptitle': 26,
        'title': 22,
        'label': 22,
        'tick': 18,
        'colorbar': 16
    }

    fig, (ax_a, ax_b, ax_c, ax_d, ax_e, ax_f) = plt.subplots(6, 1, figsize=(18, 32), sharex=True,
                                                             constrained_layout=True)
    fig.suptitle(f"Order: {oi}", fontsize=fontsizes['suptitle'], fontweight='bold')
    cmap = matplotlib.colormaps['inferno']

    tick_positions = np.linspace(0, n_pixels - 1, 5)
    tick_labels = np.linspace(wl_order.min(), wl_order.max(), 5)
    phases = np.linspace(phases.min(), phases.max(), 5)

    # Plot initial data
    im_a = ax_a.imshow(initial_spectrum, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(initial_spectrum), vmax=np.nanmax(initial_spectrum))
    ax_a.set_facecolor('white')
    ax_a.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_a.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_a.set_title('Initial data', fontsize=fontsizes['title'])
    ax_a.set_xticks(tick_positions)
    ax_a.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_a = plt.colorbar(im_a, ax=ax_a, shrink=0.8)
    cbar_a.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after threshold
    im_b = ax_b.imshow(spectrum_after_mask, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(spectrum_after_mask), vmax=np.nanmax(spectrum_after_mask))
    ax_b.set_facecolor('white')
    ax_b.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_b.set_title('Data after masking', fontsize=fontsizes['title'])
    ax_b.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_b.set_xticks(tick_positions)
    ax_b.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_b = plt.colorbar(im_b, ax=ax_b, shrink=0.8)
    cbar_b.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after logarithm
    im_c = ax_c.imshow(log_spectrum_after_mask, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(log_spectrum_after_mask), vmax=np.nanmax(log_spectrum_after_mask))
    ax_c.set_facecolor('white')
    ax_c.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_c.set_title('Logarithm of data after masking', fontsize=fontsizes['title'])
    ax_c.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_c.set_xticks(tick_positions)
    ax_c.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_c = plt.colorbar(im_c, ax=ax_c, shrink=0.8)
    cbar_c.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after normalization
    im_d = ax_d.imshow(log_normalized_spectrum_after_mask, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(log_normalized_spectrum_after_mask),
                       vmax=np.nanmax(log_normalized_spectrum_after_mask))
    ax_d.set_facecolor('white')
    ax_d.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_d.set_title('Logarithm of data after masking and normalization', fontsize=fontsizes['title'])
    ax_d.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_d.set_xticks(tick_positions)
    ax_d.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_d = plt.colorbar(im_d, ax=ax_d, shrink=0.8)
    cbar_d.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after PCA
    im_e = ax_e.imshow(log_residuals_after_reconstruction, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(log_residuals_after_reconstruction),
                       vmax=np.nanmax(log_residuals_after_reconstruction))
    ax_e.set_facecolor('white')
    ax_e.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_e.set_title('Residuals after reconstruction', fontsize=fontsizes['title'])
    ax_e.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_e.set_xticks(tick_positions)
    ax_e.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_e = plt.colorbar(im_e, ax=ax_e, shrink=0.8)
    cbar_e.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot final data
    im_f = ax_f.imshow(log_residuals_after_reconstruction_and_eventual_smooth, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(log_residuals_after_reconstruction_and_eventual_smooth),
                       vmax=np.nanmax(log_residuals_after_reconstruction_and_eventual_smooth))
    ax_f.set_facecolor('white')
    ax_f.set_xlabel('Wavelength [nm]', fontsize=fontsizes['label'])
    ax_f.set_title('Residuals after reconstruction and smoothing', fontsize=fontsizes['title'])
    ax_f.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_f.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_f.set_xticks(tick_positions)
    ax_f.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    ax_f.set_xticklabels([f'{x:.1f}' for x in tick_labels], fontsize=fontsizes['tick'])
    cbar_f = plt.colorbar(im_f, ax=ax_f, shrink=0.8)
    cbar_f.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Save plot
    pdf_processf.savefig(fig)
    fig.clear()
    plt.close(fig)

def plot_process2(
    initial_spectrum, log_spectrum_after_mask, log_residuals_after_reconstruction, 
        array_reconstructed_from_PCA_components, wl_order, pdf_processf, n_pixels, oi, phases
):
    """
    Plot processing steps for telluric removal.

    Args:
        initial_spectrum: Initial data
        log_spectrum_after_mask: Logarithmic of data after masking
        log_residuals_after_reconstruction: Residuals after reconstruction
        array_reconstructed_from_PCA_components: array reconstructed from PCA components
        wl_order: Wavelength array for order
        pdf_processf: PDF file to save plots
        n_pixels: Number of pixels
        oi: Order index
        phases: Planetary orbital phases
    """
    # Set font sizes
    fontsizes = {
        'suptitle': 26,
        'title': 22,
        'label': 22,
        'tick': 18,
        'colorbar': 16
    }

    fig, (ax_a, ax_b, ax_c, ax_d) = plt.subplots(4, 1, figsize=(12, 16), sharex=True,
                                                             constrained_layout=True)
    fig.suptitle(f"Order: {oi}", fontsize=fontsizes['suptitle'], fontweight='bold')
    cmap = matplotlib.colormaps['inferno']

    tick_positions = np.linspace(0, n_pixels - 1, 5)
    tick_labels = np.linspace(wl_order.min(), wl_order.max(), 5)
    phases = np.linspace(phases.min(), phases.max(), 5)

    # Plot initial data
    im_a = ax_a.imshow(initial_spectrum, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(initial_spectrum), vmax=np.nanmax(initial_spectrum))
    ax_a.set_facecolor('white')
    ax_a.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_a.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_a.set_title('Initial data', fontsize=fontsizes['title'])
    ax_a.set_xticks(tick_positions)
    ax_a.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_a = plt.colorbar(im_a, ax=ax_a, shrink=0.8)
    cbar_a.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after logarithm
    im_b = ax_b.imshow(log_spectrum_after_mask, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(log_spectrum_after_mask), vmax=np.nanmax(log_spectrum_after_mask))
    ax_b.set_facecolor('white')
    ax_b.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_b.set_title('Logarithm of data after masking and normalization', fontsize=fontsizes['title'])
    ax_b.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_b.set_xticks(tick_positions)
    ax_b.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_b = plt.colorbar(im_b, ax=ax_b, shrink=0.8)
    cbar_b.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after PCA
    im_c = ax_c.imshow(array_reconstructed_from_PCA_components, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(array_reconstructed_from_PCA_components),
                       vmax=np.nanmax(array_reconstructed_from_PCA_components))
    ax_c.set_facecolor('white')
    ax_c.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_c.set_title('Array Reconstructed from PCA components', fontsize=fontsizes['title'])
    ax_c.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_c.set_xticks(tick_positions)
    ax_c.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    cbar_c = plt.colorbar(im_c, ax=ax_c, shrink=0.8)
    cbar_c.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Plot data after PCA
    im_d = ax_d.imshow(log_residuals_after_reconstruction, aspect='auto', origin='lower', cmap=cmap,
                       vmin=np.nanmin(log_residuals_after_reconstruction),
                       vmax=np.nanmax(log_residuals_after_reconstruction))
    ax_d.set_facecolor('white')
    ax_d.set_ylabel('Planetary Orbital Phase', fontsize=fontsizes['label'])
    ax_d.set_title('Residuals after reconstruction', fontsize=fontsizes['title'])
    ax_d.tick_params(axis='both', which='major', labelsize=fontsizes['tick'])
    ax_d.set_xticks(tick_positions)
    ax_d.set_yticklabels([f'{x:.2f}' for x in phases], fontsize=fontsizes['tick'])
    ax_d.set_xticklabels([f'{x:.1f}' for x in tick_labels], fontsize=fontsizes['tick'])
    ax_d.set_xlabel('Wavelength [nm]', fontsize=fontsizes['label'])
    cbar_d = plt.colorbar(im_d, ax=ax_d, shrink=0.8)
    cbar_d.ax.tick_params(labelsize=fontsizes['colorbar'])

    # Save plot
    pdf_processf.savefig(fig)
    fig.clear()
    plt.close(fig)


def plot_order(current_order, x_arrf, maskf, resf, pdff):
    """
    Plot order information.

    Args:
        current_order: Order index
        x_arrf: X-axis array
        maskf: Mask array
        resf: Result array
        pdff: PDF file to save plot
    """
    fig = plt.figure()
    plt.title(f"Order n.{current_order}")
    plt.xlabel("Pixel")
    plt.ylabel("Flux")
    for iii in range(len(resf[0, :])):
        plt.plot(x_arrf[maskf], resf[:, iii] + 2 * iii, alpha=0.6)
    pdff.savefig(fig)
    fig.clear()
    plt.close(fig)


def plot_mean_sp(
    current_order, meanspf, normmeansp, continuumf, sigma_coefff, maskinv1f, maskinv2f, cutofff,
    threshold_mean, pdf_mean_spf, wlen
):
    """
    Plot mean spectrum with masks.

    Args:
        current_order: Order
        meanspf: Mean spectrum
        normmeansp: Normalized Mean spectrum
        continuumf: Continuum
        sigma_coefff: Sigma coefficient for variance-based masking
        maskinv1f: First mask inverse (cutoff + threshold)
        maskinv2f: Second mask inverse (sigma-based)
        cutofff: Cutoff value for telluric masking
        threshold_mean: Threshold for mean spectrum masking
        pdf_mean_spf: PDF file to save plot
        wlen: Wavelength array for order
    """
    # Set global font sizes
    plt.rcParams.update({
        'font.size': 16,
        'axes.titlesize': 20,
        'axes.labelsize': 18,
        'xtick.labelsize': 16,
        'ytick.labelsize': 16,
        'legend.fontsize': 14,
        'figure.titlesize': 22
    })

    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 16), dpi=300, sharex=True, constrained_layout=True)
    fig.suptitle(f"Order n.{current_order}", fontsize=22, fontweight='bold')

    ax1.plot(wlen, meanspf, label="Mean spectrum", linewidth=1.5)
    ax1.plot(wlen, continuumf, label="Continuum", linewidth=1.5)
    ax1.axhline(y=threshold_mean, color='black', linestyle='--', linewidth=1.5,
                label=f"Threshold mean spectrum = {threshold_mean}")
    ax1.set_title("Mean spectrum after throughput correction", fontsize=16)
    ax1.set_ylabel("Flux [normalized at brightest pixels]", fontsize=14)
    ax1.legend(fontsize=12)
    ax1.tick_params(axis='both', which='major', labelsize=12)

    ax2.plot(wlen, normmeansp, label="Normalized Mean spectrum", linewidth=1.5)
    ax2.axhline(y=cutofff, color='black', linestyle='--', linewidth=1.5, label=f"Cutoff telluric = {cutofff}")
    ax2.set_ylabel("Flux [continuum normalized]", fontsize=14)
    ax2.legend(fontsize=12)
    ax2.tick_params(axis='both', which='major', labelsize=12)

    ax3.plot(wlen, meanspf, label="Mean spectrum", linewidth=1.5)
    ax3.plot(wlen, continuumf, label="Continuum", linewidth=1.5)
    ax3.scatter(wlen[np.where(maskinv2f)[0]], meanspf[maskinv2f], marker="x", color="g", s=50,
                label=r"$\sigma$ >" + str(sigma_coefff) + r" * $\sigma_{mean}$")
    ax3.scatter(wlen[np.where(maskinv1f)[0]], meanspf[maskinv1f], marker="x", color="r", s=50,
                label="Cutoff + Threshold mean spectrum")
    ax3.set_ylabel("Flux [normalized at brightest pixels]", fontsize=14)
    ax3.set_xlabel('Wavelength [nm]', fontsize=14)
    ax3.legend(fontsize=12)
    ax3.tick_params(axis='both', which='major', labelsize=12)

    pdf_mean_spf.savefig(fig)
    fig.clear()
    plt.close(fig)

def divide_by_brightest(spectra_order, n_images, n_brightest=50):
    """
    Normalize spectra by dividing each exposure by the median of its brightest pixels.

    Args:
        spectra_order: 2D array of spectral data (pixels x images)
        n_images: Number of images (exposures)
        n_brightest: Number of brightest pixels to use for normalization

    Returns:
        tuple: (spectra_order, spectra_order_before_mask) - cleaned input and normalized spectra
    """
    spectra_order_before_mask = np.zeros_like(spectra_order)

    # Normalize data based on mean value. In this case, for example,
    # the exposures that have higher flux are normalized by a higher mean, while the ones that have less flux are normalized
    # by a lower value. This will lead to a matrix that has similar values per each channel.
    if np.mean(spectra_order) > 100:
        spectra_order[np.where(spectra_order <= 1)] = 1
        spectra_order[np.where(np.isfinite(spectra_order) == 0)] = 1
        # divide the single spectrum by the median of the first 50 brightest spectras
        for img in range(n_images):
            sorted_data = np.sort(spectra_order[:, img])[::-1]
            first_brightest = sorted_data[:n_brightest]
            spectra_order_before_mask[:, img] = np.array(spectra_order[:, img]) / median_idl(first_brightest)
    else:
        spectra_order[np.where(spectra_order <= 0)] = 0
        spectra_order_before_mask = np.array(spectra_order)

    return spectra_order, spectra_order_before_mask


def masking_function(
    masking, spectra_order_before_mask, cutoff, threshold_mean_spec, sigma_coeff, checkbox_save_plot,
    good_order_oi, pdf_mean_sp, n_pixels, do_not_fit_continuum, wlen
):
    """
    Apply spectral masking based on telluric absorption and variance criteria.

    Args:
        masking: Whether to apply masking
        spectra_order_before_mask: Normalized spectral data (pixels x images)
        cutoff: Cutoff threshold for telluric line detection
        threshold_mean_spec: Threshold for mean spectrum signal level
        sigma_coeff: Coefficient for variance-based masking
        checkbox_save_plot: Whether to save diagnostic plots
        good_order_oi: Order index for plot labeling
        pdf_mean_sp: PDF file for saving mean spectrum plots
        n_pixels: Number of spectral pixels
        do_not_fit_continuum: Whether to skip continuum fitting
        wlen: Wavelength array for order

    Returns:
        tuple: (std_after_full_norm, spectra_order_after_mask, mask_inv,
                count_mask_inv, mask_good, amount_good_pixels)
    """
    # Apply masking if enabled
    mean_spectra_images = np.mean(spectra_order_before_mask, axis=1)
    if masking:
        # Check if continuum fitting should be skipped
        if do_not_fit_continuum:
            # Use a flat continuum (all ones) when fitting is disabled
            max_continuum = np.max(mean_spectra_images)
            continuum = np.ones(len(mean_spectra_images)) * max_continuum
            norm_mean_spectra_images = mean_spectra_images / continuum
            print(f"Continuum fitting disabled - using flat continuum at max {max_continuum}")
        else:
            continuum, norm_mean_spectra_images = estimate_continuum(mean_spectra_images, num_windows=40, top_window=10, check_gaussian=True)

        # Normalize by wavelength (divide each column by its mean)
        # Normalize each spectral channel (e.g. 2048 pixels) of each spectrum by the mean of the spectrum itself (e.g. 60 values)
        # So each cell of the matrix (i.2. 2048 x 60) is normalized by the mean of the exposure.
        # This operation will lead to exposures with mean 1 (np.mean(spec_wave_norm, axis=0) = [1 .... 1]
        # While having different std, the spectra divided by the std is equal:
        # np.std(spec_wave_norm)
        # 0.37127575062412854
        # np.std(spectra_order_before_mask)
        # 0.2526833738723915
        # np.mean(spec_wave_norm / np.std(spec_wave_norm))
        # 2.6934158730241933
        # np.mean(spectra_order_before_mask / np.std(spectra_order_before_mask))
        # 2.6933605602288786
        spec_wave_norm = spectra_order_before_mask / np.mean(spectra_order_before_mask, axis=0)
        # Full normalization: normalize also by mean of spectral channels (e.g. 2048 values)
        # This is done to calculate properly the standard deviation for the masking, by correcting the continuum to flat the spectra.
        spec_full_norm = spec_wave_norm / np.expand_dims(np.mean(spec_wave_norm, axis=1), axis=1)
        # Standard deviation along each spectral channel, to see how the value changes along the night because the telluric lines
        # vary due to the airmass and PWV.
        std_after_full_norm = np.std(spec_full_norm, axis=1, ddof=1)

        # Apply masks
        # If the value is greater than the cutoff (suppose 0.8), it means that it is ok because there is not much absorption (in
        # our case 1 means no absorption, 0 means line saturated). So in this case, having a value greater than
        # the cutoff means that the line is probably not telluric, so we keep it. Same for the mean: if the mean
        # is greater than the threshold (suppose 0.3) we keep it because it means that there is a signal.
        #
        # Mask_inv are the pixels masked fpr the telluric, mask_good the one we keep
        mask_good_1 = ((norm_mean_spectra_images > cutoff) & (mean_spectra_images > threshold_mean_spec))
        mask_inv_1 = ~mask_good_1
        mask_good_2 = (std_after_full_norm <= median_idl(std_after_full_norm[mask_good_1]) * sigma_coeff)
        mask_inv_2 = ~mask_good_2

        # Plot mean spectrum if enabled
        if checkbox_save_plot:
            plot_mean_sp(
                good_order_oi, mean_spectra_images, norm_mean_spectra_images, continuum, sigma_coeff,
                mask_inv_1, mask_inv_2, cutoff, threshold_mean_spec, pdf_mean_sp, wlen
            )

        # Combine masks
        mask_good = np.where(mask_good_1 & mask_good_2)[0]
        mask_inv = np.where((mask_inv_1 | mask_inv_2))[0]
        count_mask_inv = np.size(mask_inv)
        amount_good_pixels = n_pixels - count_mask_inv

        # Apply mask_good to data
        # TODO spectra_order_before_mask might be wrong, spec_full_norm
        spectra_order_after_mask = copy.copy(spectra_order_before_mask)
        if count_mask_inv > 0:
            spectra_order_after_mask[mask_inv, :] = 1
    else:
        # Handle case without masking
        count_mask_inv = 0
        std_after_full_norm = np.std(spectra_order_before_mask, axis=1, ddof=1)
        spectra_order_after_mask = np.array(spectra_order_before_mask)
        mask_good = np.arange(n_pixels)
        amount_good_pixels = n_pixels
        mask_inv = []

    return std_after_full_norm, spectra_order_after_mask, mask_inv, count_mask_inv, mask_good, amount_good_pixels

def tangent_method_for_PCA(
    pca_obj, log_spectra_order_after_mask, max_n_components_per_order, spectra_order_after_mask, result_first_pca,
    mask_good, epsilon, tangent_threshold, substitute_for_comp, current_order, checkbox_save_plot, pdf_tangent, pca_mode
):
    """
    Determine optimal number of PCA components using the tangent method.

    Analyzes the variance reduction curve to find where adding more PCA
    components provides diminishing returns, using derivatives of the
    standard deviation.

    Args:
        pca_obj: Fitted PCA object with ``explained_variance_`` attribute
        log_spectra_order_after_mask: Log-transformed masked spectral data
        max_n_components_per_order: Range of component numbers to test
        spectra_order_after_mask: Masked spectral data (used for shape)
        result_first_pca: PCA components from initial fit
        mask_good: Indices of good (unmasked) pixels
        epsilon: Small value to avoid division by zero
        tangent_threshold: Threshold for second derivative test
        substitute_for_comp: Fallback number of components
        current_order: Current order index for logging
        checkbox_save_plot: Whether to save tangent diagnostic plot
        pdf_tangent: PDF file for saving tangent plot
        pca_mode: PCA decomposition mode ('spatial' or 'temporal')

    Returns:
        int: Optimal number of PCA components
    """
    rerr = []
    variance = []
    snc = [np.std(log_spectra_order_after_mask, ddof=1)]
    # eigenvectors = pca_obj.components_
    eigenvalues = pca_obj.explained_variance_

    # Calculate statistics for each number of components
    for n in max_n_components_per_order:
        lcrm_temp = np.zeros_like(spectra_order_after_mask)
        res_temp = np.array(result_first_pca)
        res_temp[:, n:] = 0
        ones_column = np.ones((res_temp.shape[0], 1))
        array_of_components_plus_column_of_ones = np.hstack((res_temp, ones_column))
        array_reconstruct = linear_solver(log_spectra_order_after_mask, array_of_components_plus_column_of_ones)

        if pca_mode == "spatial":
            lcrm_temp[mask_good, :] = (log_spectra_order_after_mask - array_reconstruct)
        else:
            lcrm_temp[mask_good, :] = (log_spectra_order_after_mask - array_reconstruct).T
        rerr.append(np.sum((array_reconstruct - log_spectra_order_after_mask) ** 2))
        variance.append(np.sum(eigenvalues[:n]))
        snc.append(np.std(lcrm_temp[mask_good, :], ddof=1))

    # Calculate tangent and second derivative
    snc = np.array(snc)
    tan_snc = -(snc[1:] - snc[:-1]) / (snc[:-1] + epsilon)
    der_sec = (tan_snc[1:] - tan_snc[:-1]) / (tan_snc[:-1] + epsilon)

    # Determine number of components based on threshold
    above_thresolds = np.where(der_sec[1:] > -tangent_threshold)[0] + 3
    if len(above_thresolds) < 1 or above_thresolds[0] < substitute_for_comp:
        n_components_pca = substitute_for_comp
        if len(above_thresolds) < 1:
            print(
                f"Order {current_order}: no components found under {-tangent_threshold}. Taking {n_components_pca} components."
            )
        else:
            print(
                f"Order {current_order}: components found under {-tangent_threshold}: {above_thresolds[0]}. Taking minimum n. components = {n_components_pca}"
            )
    else:
        n_components_pca = above_thresolds[0]
        print(
            f"Order {current_order}: components found under {-tangent_threshold}. Taking {n_components_pca} components."
        )

    # Plot tangent information if enabled
    if checkbox_save_plot:
        plot_tangent(
            pdf_tangent, current_order, tan_snc, der_sec, tangent_threshold, n_components_pca
        )
    return n_components_pca

def cross_val_score_components(pca_obj, pca_scores, max_n_components_per_order, log_spectra_order_after_mask):
    for n in max_n_components_per_order:
        pca_obj.n_components = n
        pca_scores.append(np.mean(cross_val_score(pca_obj, log_spectra_order_after_mask)))
    n_components_pca = max_n_components_per_order[np.argmax(pca_scores)]
    return n_components_pca

def get_number_of_components_PCA(
    auto_components, log_spectra_order_after_mask, max_n_components_per_order, spectra_order_after_mask,
    mask_good, epsilon, tangent_threshold, substitute_for_comp, current_order,
    checkbox_save_plot, pdf_tangent, entry_min_comp, entry_max_comp, entry_fixed_comp, tangent_method, apply_standardize,
    pca_mode
):
    # Create pipeline with optional standardization
    if apply_standardize:
        pipeline = Pipeline([
            ('scaler', StandardScaler()),
            ('pca', PCA(svd_solver="full"))
        ])
    else:
        pipeline = Pipeline([
            ('pca', PCA(svd_solver="full"))
        ])

    result_first_pca = pipeline.fit_transform(log_spectra_order_after_mask)
    pca_obj = pipeline.named_steps['pca']
    pca_scores = []
    if auto_components:  # true only if fixed is not 0 or tangent_method is true
        if tangent_method:
            n_components_pca = tangent_method_for_PCA(
                pca_obj, log_spectra_order_after_mask, max_n_components_per_order,
                spectra_order_after_mask, result_first_pca,
                mask_good, epsilon, tangent_threshold, substitute_for_comp, current_order,
                checkbox_save_plot, pdf_tangent, pca_mode
            )
        else:
            # Use cross-validation to determine number of components
            n_components_pca = cross_val_score_components(
                pca_obj, pca_scores, max_n_components_per_order, log_spectra_order_after_mask
            )

        # Apply component limits
        # NB a < b != c is equal to say (a < b) ad (b != c)
        if n_components_pca < int(entry_min_comp) != 0:
            n_components_pca = int(entry_min_comp)
        if n_components_pca > int(entry_max_comp) != 0:
            n_components_pca = int(entry_max_comp)
    else:
        n_components_pca = int(entry_fixed_comp)
    return n_components_pca

def apply_smooth(log_final_pca, smooth_size, mask_good):
    B = np.array(log_final_pca)
    mask_final = np.zeros_like(B)
    mask_final[mask_good, :] = 1

    # Apply uniform filter
    filtered_sum = ndimage.uniform_filter(B, size=[smooth_size, 0], mode="nearest")
    filtered_count = ndimage.uniform_filter(mask_final, size=[smooth_size, 0], mode="nearest")

    # Calculate smoothed result
    with np.errstate(divide='ignore', invalid='ignore'):
        ynew = np.true_divide(filtered_sum, filtered_count)
        ynew[mask_final == 0] = 0
    log_final_pca = B - ynew

    return log_final_pca

def tell_rem_wl_log2(exchange_pathf):
    """
    Main function for telluric removal using PCA with wavelength calibration.

    Args:
        exchange_pathf: Path to exchange directory containing parameters
    """
    from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
    try:
        # Load parameters from exchange file
        filename_pkl = str(Path(exchange_pathf, "exchange.pkl"))
        option_list_method = ConstantVariables.LIST_TELL_TABLE
        with open(filename_pkl, "rb") as fo:
            exchange_data = pickle.load(fo)
            path_instrument = exchange_data["path_instrument"]
            molec = exchange_data["order_selection"]
            nights = exchange_data["nights"]
            tell_rm_method = exchange_data["tell_rm_method"]
            nfcoo = exchange_data["use_last_component_selection"]
            sigma_coeff = exchange_data["sigma_coeff"]
            smooth_size = int(exchange_data["smooth_size"])
            yesmask = exchange_data["yesmask"]
            entry_min_comp = int(exchange_data["min_comp"])
            entry_max_comp = int(exchange_data["max_comp"])
            entry_fixed_comp = int(exchange_data["fixed_comp"])
            checkbox_save_plot = exchange_data["save_plot"]
            cutoff = exchange_data["cutoff"]
            tangent_method = exchange_data["tangent_components_method"]
            tangent_threshold = exchange_data["tangent_value"]
            substitute_for_comp = exchange_data["substitute_for_comp"]
            threshold_mean_spec = exchange_data["threshold_mean_spec"]
            subtraction_of_the_average_spectrum_checkbox = exchange_data["subtraction_of_the_average_spectrum"]
            do_not_fit_continuum = exchange_data["do_not_fit_continuum"]
            apply_standardize = exchange_data["apply_standardize"]
            # PCA mode: "spatial" (default, loop over orders) or "temporal" (loop over spectral channels)
            pca_mode = exchange_data.get("pca_mode", "spatial")

        # Initialize processing parameters
        auto_components = (entry_min_comp != 0) | tangent_method
        smooth_on = smooth_size > 0
        # Array of 0 or 1 depending if applying the mask_good or not
        yesmask = np.array([int(elem) for elem in yesmask])
        # string where to save the removed components to be shown in the GUI
        string_tell_rem = ""

        # Process each order and night
        for order_sel in molec:
            for i in range(len(nights)):
                current_night = nights[i]
                print("Night: ", current_night)
                # Set masking parameters
                masking = bool(yesmask[i])

                # Determine prefix based on night type
                prefixes = ["A_", "B_"] if "A" in nights[i] else [""]
                path_order = Path(path_instrument, current_night, "orders_selection", order_sel)
                # TODO Remove the mv in future updates
                if not path_order.exists():
                    path_order_tmp = Path(path_instrument, current_night, order_sel)
                    if not path_order_tmp.exists():
                        print("Order selection does not exist")
                        with open(str(Path(exchange_pathf, "error_tell_rm.pkl")), "wb") as fo:
                            pickle.dump("Order selection does not exist", fo)
                        return
                    path_order_without_sel = Path(path_instrument, current_night, "orders_selection")
                    os.system(f"mkdir -p {path_order_without_sel}")
                    os.system(f"mv {path_order_tmp} {path_order_without_sel}")
                path_order_and_telluric_method = str(Path(path_order, tell_rm_method))
                os.system("mkdir -p " + str(path_order))
                os.system("mkdir -p " + path_order_and_telluric_method)
                os.system("mkdir -p " + str(Path(path_order_and_telluric_method, "test")))

                # Create exchange.yaml with the same information
                filename_yaml = str(Path(path_order_and_telluric_method, "exchange.yaml"))
                # Create a copy of exchange_data and convert numpy arrays to lists for proper YAML serialization
                exchange_data_for_yaml = {}
                for key, value in exchange_data.items():
                    if isinstance(value, np.ndarray):
                        exchange_data_for_yaml[key] = value.tolist()
                    else:
                        exchange_data_for_yaml[key] = value
                exchange_data_for_yaml["index_night_for_mask"] = i

                with open(filename_yaml, "w") as yaml_file:
                    yaml.dump(exchange_data_for_yaml, yaml_file, default_flow_style=False, sort_keys=False)

                # Load good orders
                with open(str(Path(path_order, "good_order_" + order_sel + ".pkl")), "rb") as fo:
                    good_order = pickle.load(fo)
                n_good_orders = len(good_order)

                # create path for the final night to merge A and B slits
                path_night_complete = None if prefixes == [""] else str(Path(path_instrument, current_night))

                # Process each prefix (A/B or single)
                for prefix in prefixes:
                    # Initialize arrays
                    mask_good = None

                    # Load and prepare data
                    dataf = fits.open(str(Path(path_instrument, current_night, prefix + "aligned.fits")))
                    data = np.transpose(dataf[0].data)
                    wlenf = fits.open(str(Path(path_instrument, current_night, prefix + "wlen_refined.fits")))
                    wlen = np.transpose(wlenf[0].data)
                    n_pixels, n_images, n_total_orders = np.shape(data)
                    barycentric_velocity, phases, hjd, BJD_TDB = load_fase_data(
                        str(Path(path_instrument, current_night, prefix + "phase.pkl"))
                    )
                    
                    # Select only the good_orders spectra
                    spectra_good_orders = data[:, :, good_order]

                    # Initialize result arrays
                    signal_residuals_after_PCA_subtraction = np.zeros_like(spectra_good_orders)
                    matrix_pre_PCA = np.zeros_like(spectra_good_orders)
                    spAb_tr_cor = np.zeros_like(spectra_good_orders)
                    array_reconstructed_from_PCA_components = np.zeros_like(spectra_good_orders)

                    # Initialize processing variables
                    rmsorder = []
                    nfc = []
                    nfc_old = None

                    # Load NFC data if available
                    if nfcoo:
                        ncompnfcf = fits.open(str(Path(path_order_and_telluric_method, "test", prefix + "nfc.fits")))
                        nfc_old = ncompnfcf[0].data
                        ncompnfcf.close()

                    # Initialize PCA processing arrays
                    x_arr = np.arange(n_pixels)
                    mask_inv_complete = np.zeros([n_pixels, n_good_orders])
                    err_pca = np.zeros([n_pixels, n_good_orders])
                    list_amount_good_pixels = np.zeros(n_good_orders, dtype=int)
                    # Initialize PDF files for plotting
                    pdf = None
                    pdf_mean_sp = None
                    pdf_process = None
                    pdf_process2 = None
                    pdf_tangent = None
                    if checkbox_save_plot:
                        pdf = PdfPages(str(Path(path_order_and_telluric_method, prefix + "tell_rm.pdf")))
                        pdf_mean_sp = PdfPages(str(Path(path_order_and_telluric_method, prefix + "mean_sp_mask.pdf")))
                        pdf_process = PdfPages(str(Path(path_order_and_telluric_method, prefix + "pdf_process.pdf")))
                        pdf_process2 = PdfPages(str(Path(path_order_and_telluric_method, prefix + "pdf_process2.pdf")))
                        if tangent_method:
                            pdf_tangent = PdfPages(str(Path(path_order_and_telluric_method, prefix + "pdf_tangent.pdf")))

                    final_average_spectrum_orders = []
                    # Process each order
                    for oi in range(n_good_orders):
                        # Extract and prepare data for current order
                        spectra_order = np.array(spectra_good_orders[:, :, oi])
                        full_log_spectrum = np.zeros_like(spectra_good_orders[:, :, oi])
                        full_normalized_log_spectrum = np.zeros_like(spectra_good_orders[:, :, oi])
                        full_diff_masked_spc_with_reconstruction = np.zeros_like(spectra_good_orders[:, :, oi])
                        spectra_order, spectra_order_before_mask = divide_by_brightest(spectra_order, n_images)
                        spAb_tr_cor[:, :, oi] = np.array(spectra_order_before_mask)

                        (
                            err_pca[:, oi], spectra_order_after_mask, mask_inv, count_mask_inv,
                            mask_good, list_amount_good_pixels[oi]
                        ) = masking_function(
                            masking, spectra_order_before_mask, cutoff, threshold_mean_spec, sigma_coeff,
                            checkbox_save_plot, good_order[oi], pdf_mean_sp, n_pixels, do_not_fit_continuum, wlen[:, oi]
                        )

                        # Apply mask_good to data
                        if count_mask_inv > 0:
                            mask_inv_complete[mask_inv, oi] = 1
                        #
                        matrix_pre_PCA[:, :, oi] = spectra_order_after_mask
                        spAb_tr_cor[mask_inv, :, oi] = 1

                        # Prepare data for PCA, transforming the data in log (ex topca)
                        log_spectra_order_after_mask = np.log10(spectra_order_after_mask[mask_good, :])
                        # Normalize components by mean of spectra
                        means = np.mean(log_spectra_order_after_mask, axis=0)
                        log_normalized_spectra_order_after_mask = log_spectra_order_after_mask - np.expand_dims(means, axis=0)
                        final_average_spectrum = None
                        if subtraction_of_the_average_spectrum_checkbox:
                            average_spectrum = np.mean(log_normalized_spectra_order_after_mask, axis=1) # to divide each spectrum by a mean one to remove possinle tellar contribution
                            final_average_spectrum = np.expand_dims(average_spectrum, axis=1)
                            log_normalized_spectra_order_after_mask = log_normalized_spectra_order_after_mask - final_average_spectrum   # subtract the spectra mean to each spectrum (in log) to remove common spectral features
                        final_average_spectrum_orders.append(final_average_spectrum)
                        # The previous operations were made to obtain log_normalized_spectra_order_after_mask at zero mean for the PCA
                        # Determine number of PCA components
                        n_features = int(n_images / 3)  # number of spectra
                        max_n_components_per_order = np.arange(0, n_features)
                        epsilon = 1e-9
                        n_good_components = None
                        array_reconstruct_final = None
                        array_of_components = None
                        diff_masked_spc_with_reconstruction = None
                        
                        # Perform PCA
                        if tell_rm_method == option_list_method[0]:  # PCA
                            # Prepare data based on PCA mode
                            # Spatial mode (default): samples=images, features=pixels -> matrix shape (n_pixels, n_images)
                            # Temporal mode: samples=pixels, features=images -> matrix shape (n_images, n_pixels)
                            if pca_mode == "temporal":
                                # Transpose for temporal PCA: now rows are images, columns are pixels
                                pca_input_matrix = log_normalized_spectra_order_after_mask.T
                                print(f"Order {good_order[oi]}: PCA mode = temporal, matrix shape = {pca_input_matrix.shape}")
                            else:
                                # Spatial PCA (default): rows are pixels, columns are images
                                pca_input_matrix = log_normalized_spectra_order_after_mask
                                print(f"Order {good_order[oi]}: PCA mode = spatial, matrix shape = {pca_input_matrix.shape}")

                            if nfc_old is None:
                                # Determine the number of components
                                n_components_pca = get_number_of_components_PCA(
                                    auto_components, pca_input_matrix, max_n_components_per_order,
                                    spectra_order_after_mask,
                                    mask_good, epsilon, tangent_threshold, substitute_for_comp, good_order[oi],
                                    checkbox_save_plot, pdf_tangent, entry_min_comp, entry_max_comp, entry_fixed_comp,
                                    tangent_method, apply_standardize, pca_mode
                                )
                            else:
                                n_components_pca = nfc_old[oi]

                            # Perform final PCA with a selected number of components using pipeline
                            # The standardize procedure is to obtain a dataset with std at 1.
                            # The mean 0 is obtsined in the previous step. Here, all spectra are at mean 0 and std 1
                            if apply_standardize:
                                pca_pipeline = Pipeline([
                                    ('scaler', StandardScaler()),
                                    ('pca', PCA(n_components=n_components_pca, svd_solver="full"))
                                ])
                            else:
                                pca_pipeline = Pipeline([
                                    ('pca', PCA(n_components=n_components_pca, svd_solver="full"))
                                ])

                            array_of_components = pca_pipeline.fit_transform(pca_input_matrix) # array_of_components is a matrix with dimensione [n_samples, n_components]. Each array is an eigenvector that, in the linear solver, will be multiplied by the eigenvalue found by the solver. So The reconstrude spectrum will be a1*v1 + a2*v2 + a3*v3 etc
                            ones_column = np.ones((array_of_components.shape[0], 1)) # This is an extra array for a n+1 component. This is a set of "1" that will be multiplied by the corresponding coefficient obtained by the Linear solver.
                            array_of_components_plus_column_of_ones = np.hstack((array_of_components, ones_column))
                            n_good_components = n_components_pca
                            array_reconstruct_pca = linear_solver(pca_input_matrix, array_of_components_plus_column_of_ones)
                            diff_pca = pca_input_matrix - array_reconstruct_pca

                            # Transform back to original orientation if temporal mode was used
                            if pca_mode == "temporal":
                                # Transpose back: from (n_images, n_pixels) to (n_pixels, n_images)
                                array_reconstruct_final = array_reconstruct_pca.T
                                diff_masked_spc_with_reconstruction = diff_pca.T
                            else:
                                array_reconstruct_final = array_reconstruct_pca
                                diff_masked_spc_with_reconstruction = diff_pca

                        # Plot order information if enabled
                        if checkbox_save_plot:
                            plot_order(good_order[oi], x_arr, mask_good, array_of_components, pdf)
                        
                        # Store number of components
                        nfc.append(n_good_components)

                        # Store test data if enabled
                        array_reconstructed_from_PCA_components[mask_good, :, oi] = array_reconstruct_final

                        log_final_pca = signal_residuals_after_PCA_subtraction[:, :, oi]
                        log_final_pca[mask_good, :] = diff_masked_spc_with_reconstruction # only planetary spectrum in principle

                        # Apply smoothing if enabled
                        if smooth_on:
                            log_final_pca = apply_smooth(log_final_pca, smooth_size, mask_good)
                        
                        # Store final result
                        signal_residuals_after_PCA_subtraction[:, :, oi] = log_final_pca
                        rmsorder.append(np.std(log_final_pca, ddof=1))

                        full_log_spectrum[mask_good, :] = log_spectra_order_after_mask
                        full_normalized_log_spectrum[mask_good, :] = log_normalized_spectra_order_after_mask
                        full_diff_masked_spc_with_reconstruction[mask_good, :] = diff_masked_spc_with_reconstruction

                        # Plot process steps if enabled
                        if checkbox_save_plot:
                            plot_process(
                                np.transpose(spectra_good_orders[:, :, oi]), np.transpose(spectra_order_after_mask),
                                np.transpose(full_log_spectrum), np.transpose(full_normalized_log_spectrum), np.transpose(full_diff_masked_spc_with_reconstruction),
                                np.transpose(log_final_pca), wlen[:, oi], pdf_process, n_pixels, oi, phases
                            )
                            plot_process2(
                                np.transpose(spectra_good_orders[:, :, oi]), np.transpose(full_normalized_log_spectrum), np.transpose(full_diff_masked_spc_with_reconstruction),
                                np.transpose(array_reconstructed_from_PCA_components[:, :, oi]), wlen[:, oi], pdf_process2, n_pixels, oi, phases
                            )

                    # Close PDF files
                    if checkbox_save_plot:
                        pdf.close()
                        pdf_mean_sp.close()
                        pdf_process.close()
                        pdf_process2.close()
                        if pdf_tangent is not None:
                            pdf_tangent.close()

                    # Handle smoothing parameters
                    smooth_file = str(Path(path_order_and_telluric_method, "smooth_file.pkl"))
                    smooth_data_dictionary = {
                        "smooth_on": smooth_on,
                        "smooth_size": smooth_size,
                        "subtraction_of_the_average_spectrum_checkbox": subtraction_of_the_average_spectrum_checkbox,
                        "standardize": apply_standardize,
                        "average_spectrum_to_be_added": final_average_spectrum_orders,
                        "pca_mode": pca_mode  # "spatial" or "temporal"
                    }
                    with open(smooth_file, "wb") as fo:
                        pickle.dump(smooth_data_dictionary, fo)

                    # Save mask_good and error information
                    file = open(str(Path(path_order_and_telluric_method, prefix + "mask_inv_complete.pkl")), "wb")
                    pickle.dump(mask_inv_complete, file)
                    file.close()

                    file = open(str(Path(path_order_and_telluric_method, prefix + "err_pca.pkl")), "wb")
                    pickle.dump(err_pca, file)
                    file.close()

                    # Save RMS information
                    hdu = fits.PrimaryHDU(rmsorder)
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, prefix + "rms_pca.fits")), overwrite=True)

                    # Save final result
                    hdu = fits.PrimaryHDU(np.transpose(signal_residuals_after_PCA_subtraction))
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, prefix + "signal_residuals_after_PCA_subtraction.fits")), overwrite=True)

                    hdu = fits.PrimaryHDU(np.transpose(matrix_pre_PCA))
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, "test", prefix + "matrix_pre_PCA.fits")),
                                 overwrite=True)

                    hdu = fits.PrimaryHDU(np.transpose(spAb_tr_cor))
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, "test", prefix + "spAb_tr_cor.fits")),
                                 overwrite=True)

                    hdu = fits.PrimaryHDU(np.transpose(array_reconstructed_from_PCA_components))
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, "test", prefix + "array_reconstructed_from_PCA_components.fits")), overwrite=True)

                    hdu = fits.PrimaryHDU(nfc)
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, "test", prefix + "nfc.fits")), overwrite=True)

                    # Calculate and save RMS per image
                    rms_img_pca = np.array([])
                    for img in range(n_images):
                        rms_img_pca = np.append(rms_img_pca, np.std(signal_residuals_after_PCA_subtraction[mask_good, img, :], ddof=1))
                    hdu = fits.PrimaryHDU(rms_img_pca)
                    hdul = fits.HDUList([hdu])
                    hdul.writeto(str(Path(path_order_and_telluric_method, "test", prefix + "rms_img_pca.fits")), overwrite=True)

                    # Generate and save results string
                    pca_mode_desc = "spatial (PCA applied on wavelenghts)" if pca_mode == "spatial" else "temporal (PCA applied on spectral channels)"
                    string_res = f"Night: {nights[i]}\nPCA mode: {pca_mode_desc}\nOrder:\t\tRMS:\t\t\tN.Comp:\t\tGood pixels (no telluric lines)\n"
                    for oi in range(n_good_orders):
                        string_res += (
                                str(good_order[oi])
                                + "\t\t"
                                + str(rmsorder[oi])
                                + "\t\t\t"
                                + str(nfc[oi])
                                + "\t\t"
                                + str(list_amount_good_pixels[oi]) + "/" + str(n_pixels)
                                + "\n"
                        )
                    string_res += "\n\n"
                    string_tell_rem += string_res

                    # Save results string
                    with open(str(Path(path_order_and_telluric_method, "string_res.txt")), "w") as fo:
                        fo.write(string_res)

                # Save telluric removal string
                with open(str(Path(exchange_pathf, "string_tell_rem.pkl")), "wb") as fo:
                    pickle.dump(string_tell_rem, fo)

                # Combine results if night is separated
                if prefixes != [""]:  # meaning the night is separated
                    # Combine phase information and get sort indices
                    sort_indices = combine_fase(path_night_complete)
                    # Combine aligned data using sort_indices
                    combine(path_night_complete, "aligned", True, sort_indices=sort_indices)
                    # Combine snr_total data using sort_indices
                    combine(path_night_complete, "snr_total", True, sort_indices=sort_indices)
                    # Combine wavelength data using sort_indices
                    combine(path_night_complete, "wlen_refined", True, repeatf="aligned", sort_indices=sort_indices)
                    # Combine mask_good and error information
                    combine_maskinv_err_pca(path_order_and_telluric_method)
                    # Combine test data using sort_indices
                    combine(str(Path(path_order_and_telluric_method, "test")), "matrix_pre_PCA", True, sort_indices=sort_indices)
                    combine(str(Path(path_order_and_telluric_method, "test")), "spAb_tr_cor", True, sort_indices=sort_indices)
                    combine(str(Path(path_order_and_telluric_method, "test")), "array_reconstructed_from_PCA_components", True, sort_indices=sort_indices)
                    combine_nfc(str(Path(path_order_and_telluric_method, "test")))
                    # Combine final result using sort_indices
                    combine(path_order_and_telluric_method, "signal_residuals_after_PCA_subtraction", True, sort_indices=sort_indices)

                # Remove error file if it exists
                if os.path.exists(str(Path(exchange_pathf, "error_tell_rm.pkl"))):
                    os.remove(str(Path(exchange_pathf, "error_tell_rm.pkl")))

    except Exception as e:
        # Save error information
        with open(str(Path(exchange_pathf, "error_tell_rm.pkl")), "wb") as fo:
            pickle.dump(str(e) + "\n" + traceback.format_exc(), fo)
        print(str(e) + "\n" + str(traceback.format_exc()))


def main():
    """
    Main function for telluric removal using PCA in wavelength domain.

    Command line usage:
        python tell_rem_pca_wl_log_final.py <exchange_path> <path_default>

    Args (from command line):
        exchange_path: Path to exchange directory with input data
        path_default: Default path for GUIBRUSHR modules
    """
    # Setup paths from command line arguments
    path_default = sys.argv[2]
    sys.path.append(path_default)
    os.chdir(path_default)

    # Execute telluric removal
    exchange_path = sys.argv[1]
    tell_rem_wl_log2(exchange_path)


if __name__ == '__main__':
    main()

