"""
Module for plotting atmospheric model contributions and spectra.

This module provides functions to generate high-resolution and low-resolution
plots of atmospheric models, including contribution functions and spectral data
visualization with observational data comparison.

The module supports two main workflows:
1. Single model plotting - plots from best-fit parameters
2. Multi-model plotting - plots with confidence intervals from MCMC samples

Main Functions:
    - high_resolution_plot: Standard HR spectrum plot from single model
    - plot_envelope_model_HR: HR spectrum plot with confidence intervals
    - low_resolution_plot: Standard LR spectrum plot from single model
    - plot_envelope_model_LR: LR spectrum plot with confidence intervals
    - plot_contribution: Contribution function contour plots
    - main: Command-line entry point for plotting operations

Helper Functions:
    - _save_fits_file: Save numpy array to FITS file
    - _add_file_path_annotations: Add file path text to plot
    - _configure_wavelength_scale: Set wavelength axis scale (log/linear)
    - load_contribution_data: Load analysis parameters from pickle
    - save_model_outputs: Save model results to pickle files
"""

import traceback
from pathlib import Path
import matplotlib
matplotlib.rcParams['agg.path.chunksize'] = 10000
import os
if not os.environ.get('SPHINX_BUILD'):
    matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import numpy as np
import pickle
import os
import sys
from astropy.io import fits
import pandas as pd
from typing import List

from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import estimate_continuum
from GUIBRUSHR.GUI.LAYOUT.GraphicsConfig import GraphicsConfig


# =====================
# Helper Functions
# =====================

def _save_fits_file(data: np.ndarray, filepath: str) -> None:
    """
    Save a numpy array to a FITS file.

    Args:
        data: Numpy array to save
        filepath: Full path where to save the FITS file
    """
    hdu = fits.PrimaryHDU(data)
    hdul = fits.HDUList([hdu])
    hdul.writeto(filepath, overwrite=True)


def _add_file_path_annotations(fig: plt.Figure, filepaths: List[str]) -> None:
    """
    Add file path annotations to a matplotlib figure.

    Annotations are placed outside the plot area (xy=(1.6, 1.0)) and are
    typically used to document output file locations.

    Args:
        fig: Matplotlib figure object
        filepaths: List of file paths to annotate
    """
    for filepath in filepaths:
        fig.gca().annotate(
            filepath,
            xy=(1.6, 1.0),
            xycoords="figure fraction",
            ha="right",
            va="top",
            color="black",
            fontsize=15,
        )


def _configure_wavelength_scale(ax: plt.Axes, wl_scale: str) -> None:
    """
    Configure the wavelength axis scale.

    Args:
        ax: Matplotlib axis object
        wl_scale: Scale type - 'log' for logarithmic, otherwise linear
    """
    if wl_scale == "log":
        ax.set_xscale("log")


def _configure_plot_ticks(ax: plt.Axes, direction: str = 'in',
                         all_sides: bool = True) -> None:
    """
    Configure plot tick parameters for better appearance.

    Args:
        ax: Matplotlib axis object
        direction: Tick direction ('in', 'out', 'inout')
        all_sides: If True, show ticks on all sides (top, right, bottom, left)
    """
    ax.minorticks_on()
    tick_params = {'which': 'both', 'direction': direction}
    if all_sides:
        tick_params.update({'top': True, 'right': True})
    ax.tick_params(**tick_params)



# =====================
# High-Resolution Plotting Functions
# =====================

def high_resolution_plot(folder_retrieval, overplot, ylabel, wlen_list_overplot,
                        not_retrieval_objf, rad_modef, wl_scalef):
    """
    Generate high-resolution model spectrum plot from a single model.

    Creates plots showing the full resolution atmospheric model spectrum,
    with optional overplotting of selected wavelength orders. This function
    is used when plotting a single best-fit model (not MCMC samples).

    The function performs three main tasks:
    1. Saves wavelength and spectrum data to FITS files
    2. Creates a high-resolution plot with optional order highlighting
    3. Saves the plot in multiple formats (PNG, PDF)

    Args:
        folder_retrieval (str): Path to the retrieval folder
        overplot (bool): Flag indicating whether to overplot selected wavelength orders
        ylabel (str): Y-axis label for the spectrum
        wlen_list_overplot (list): List of wavelength arrays for highlighting specific orders
        not_retrieval_objf: Model object containing spectrum data with attributes:
            - wl_full_resolution_HR: High-resolution wavelength array (nm)
            - depth_full_resolution_HR: High-resolution spectrum values
        rad_modef (str): Radiation mode (e.g., 'Transmission', 'Emission')
        wl_scalef (str): Wavelength scale - 'log' for logarithmic, otherwise linear

    Output Files:
        - wavelength_HR.fits: High-resolution wavelength data
        - transit_depth_HR.fits: High-resolution spectrum data
        - high_res.pdf: Plot in PDF format
        - high_res.png: Plot in PNG format
        - model_high_low_res.png: Temporary plot file
    """
    # Define output file paths
    model_folder = Path(folder_retrieval, "model_high_low_res")
    filepng_temp = str(Path(folder_retrieval, "model_high_low_res.png"))
    filepdf = str(model_folder / "high_res.pdf")
    filepng = str(model_folder / "high_res.png")

    # Save high-resolution data to FITS files
    _save_fits_file(
        not_retrieval_objf.wl_full_resolution_HR,
        str(model_folder / "wavelength_HR.fits")
    )
    _save_fits_file(
        not_retrieval_objf.depth_full_resolution_HR,
        str(model_folder / "transit_depth_HR.fits")
    )

    # Create the main plot
    fig, ax = plt.subplots(1, 1, figsize=(12, 10), dpi=600)

    # Plot the full high-resolution spectrum
    ax.plot(not_retrieval_objf.wl_full_resolution_HR,
            not_retrieval_objf.depth_full_resolution_HR,
            "purple", alpha=0.8, label="Model")

    # Add overplot of selected wavelength orders if requested
    if overplot:
        _plot_selected_orders(ax, not_retrieval_objf, wlen_list_overplot)

    # Configure plot appearance
    ax.set_xlabel("Wavelength (nm)")
    ax.set_ylabel(ylabel)
    ax.set_title(f"{rad_modef} spectrum")
    ax.legend()

    # Set wavelength scale (log or linear)
    _configure_wavelength_scale(ax, wl_scalef)

    # Save plots in multiple formats
    plt.savefig(fname=filepdf, bbox_inches="tight", pad_inches=0.1)
    plt.savefig(fname=filepng, bbox_inches="tight", pad_inches=0.1)
    plt.savefig(fname=filepng_temp, bbox_inches="tight", pad_inches=0.1)

    # Add file path annotations
    _add_file_path_annotations(fig, [filepdf, filepng, filepng_temp])

    plt.close(fig)

    # Save plot data to pkl for later reconstruction
    pkl_data = {
        "wl_full_resolution_HR": not_retrieval_objf.wl_full_resolution_HR,
        "depth_full_resolution_HR": not_retrieval_objf.depth_full_resolution_HR,
        "overplot": overplot,
        "wlen_list_overplot": wlen_list_overplot,
        "ylabel": ylabel,
        "rad_mode": rad_modef,
        "wl_scalef": wl_scalef,
    }
    with open(str(model_folder / "high_res.pkl"), "wb") as _f:
        pickle.dump(pkl_data, _f)
    with open(str(Path(folder_retrieval, "model_high_low_res.pkl")), "wb") as _f:
        pickle.dump(pkl_data, _f)


def _plot_selected_orders(ax, not_retrieval_obj, wlen_list_overplot):
    """
    Plot selected wavelength orders on an existing axis.

    Highlights specific wavelength regions (orders) on the spectrum plot
    by overplotting them in a different color (orange).

    Args:
        ax: Matplotlib axis object
        not_retrieval_obj: Model object with wl_full_resolution_HR and depth_full_resolution_HR
        wlen_list_overplot: List of wavelength arrays defining the orders to highlight
    """
    # Process wavelength intervals for overplotting
    interval_list = []
    for elem in wlen_list_overplot:
        mins = np.min(elem, axis=1)
        maxs = np.max(elem, axis=1)
        interval_list.append(np.array([mins, maxs]))

    # Plot each selected order
    first_plot = True  # Flag to track the first plotted segment for legend
    for elem in interval_list:
        for i in range(len(elem[0])):
            # Find indices corresponding to wavelength range
            min_idx = (np.abs(not_retrieval_obj.wl_full_resolution_HR - elem[0][i])).argmin()
            max_idx = (np.abs(not_retrieval_obj.wl_full_resolution_HR - elem[1][i])).argmin()

            # Plot selected wavelength order in orange
            ax.plot(
                not_retrieval_obj.wl_full_resolution_HR[min_idx:max_idx],
                not_retrieval_obj.depth_full_resolution_HR[min_idx:max_idx],
                "orange",
                alpha=0.5,
                label="Selected orders" if first_plot else "_nolegend_",
            )
            first_plot = False  # Disable legend for subsequent plots


def _calculate_percentiles(data):
    """
    Calculate percentiles for confidence interval plotting.

    Computes median and percentiles for 68% (1-sigma) and 95% (2-sigma)
    credible regions from an array of model realizations.

    Args:
        data: 2D numpy array (n_models, n_wavelengths) containing multiple spectra

    Returns:
        dict: Dictionary with keys:
            - 'median': 50th percentile
            - 'lower_1sigma': 16th percentile
            - 'upper_1sigma': 84th percentile
            - 'lower_2sigma': 2.5th percentile
            - 'upper_2sigma': 97.5th percentile
    """
    return {
        'median': np.percentile(data, 50, axis=0),
        'lower_1sigma': np.percentile(data, 16, axis=0),
        'upper_1sigma': np.percentile(data, 84, axis=0),
        'lower_2sigma': np.percentile(data, 2.5, axis=0),
        'upper_2sigma': np.percentile(data, 97.5, axis=0),
    }


def _plot_confidence_intervals(ax, wavelength, percentiles, color_scheme=None):
    """
    Plot confidence intervals on an axis.

    Adds filled regions for 68% and 95% credible intervals plus median line.

    Args:
        ax: Matplotlib axis object
        wavelength: Wavelength array
        percentiles: Dictionary with percentile data (from _calculate_percentiles)
        color_scheme: Optional dict with keys 'color_2sigma', 'color_1sigma', 'color_model'
    """
    # Default color scheme (orange theme)
    if color_scheme is None:
        color_scheme = {
            'color_2sigma': "#fdc086",  # Light orange
            'color_1sigma': "#ff7f00",  # Orange
            'color_model': "#8c2d04",   # Dark orange
        }

    # Plot 95% credible region (2-sigma)
    ax.fill_between(wavelength,
                   percentiles['lower_2sigma'],
                   percentiles['upper_2sigma'],
                   color=color_scheme['color_2sigma'],
                   alpha=0.5,
                   label='95% credible region')

    # Plot 68% credible region (1-sigma)
    ax.fill_between(wavelength,
                   percentiles['lower_1sigma'],
                   percentiles['upper_1sigma'],
                   color=color_scheme['color_1sigma'],
                   alpha=0.5,
                   label='68% credible region')

    # Plot median model
    ax.plot(wavelength,
           percentiles['median'],
           color=color_scheme['color_model'],
           lw=1.5,
           label='Median model',
           zorder=3)


def plot_envelope_model_HR(folder_retrieval, wlen_hr, all_spectra_hr, ylabel,
                          rad_mode, wl_scalef):
    """
    Generate high-resolution spectrum plot with confidence intervals from MCMC samples.

    Creates a plot showing the median model spectrum with 68% and 95% credible regions
    from multiple model realizations. This function is used when plotting results from
    MCMC sampling (multiple models) rather than a single best-fit model.

    The function performs three main tasks:
    1. Calculates median and percentiles from all model spectra
    2. Saves median wavelength and spectrum data to FITS files
    3. Creates a plot with confidence interval bands

    Args:
        folder_retrieval (str): Path to the retrieval folder
        wlen_hr (np.ndarray): High-resolution wavelength array (nm), shape (n_wavelengths,)
        all_spectra_hr (np.ndarray): 2D array of spectra, shape (n_models, n_wavelengths)
        ylabel (str): Y-axis label for the spectrum
        rad_mode (str): Radiation mode (e.g., 'Transmission', 'Emission')
        wl_scalef (str): Wavelength scale - 'log' for logarithmic, otherwise linear

    Output Files:
        - wavelength_HR.fits: High-resolution wavelength data
        - transit_depth_HR.fits: Median high-resolution spectrum
        - high_res.pdf: Plot in PDF format
        - high_res.png: Plot in PNG format
        - model_high_low_res.png: Temporary plot file

    Note:
        The FITS files contain the median spectrum (50th percentile) rather than
        a single model realization.
    """
    # Define output file paths
    model_folder = Path(folder_retrieval, "model_high_low_res")
    filepng_temp = str(Path(folder_retrieval, "model_high_low_res.png"))
    filepdf = str(model_folder / "high_res.pdf")
    filepng = str(model_folder / "high_res.png")

    # Calculate percentiles for confidence intervals
    percentiles = _calculate_percentiles(all_spectra_hr)

    # Save high-resolution wavelength and median spectrum to FITS files
    _save_fits_file(wlen_hr, str(model_folder / "wavelength_HR.fits"))
    _save_fits_file(percentiles['median'], str(model_folder / "transit_depth_HR.fits"))

    # Create the main plot
    fig, ax = plt.subplots(1, 1, figsize=(12, 10), dpi=600)

    # Plot confidence intervals and median
    _plot_confidence_intervals(ax, wlen_hr, percentiles)

    # Configure plot appearance
    ax.set_xlabel("Wavelength (nm)")
    ax.set_ylabel(ylabel)
    ax.set_title(f"{rad_mode} spectrum")
    ax.legend(loc='best', frameon=True, fancybox=True, fontsize=10)

    # Configure tick parameters for professional appearance
    _configure_plot_ticks(ax)

    # Set wavelength scale (log or linear)
    _configure_wavelength_scale(ax, wl_scalef)

    # Save plots in multiple formats
    plt.savefig(fname=filepdf, bbox_inches="tight", pad_inches=0.1)
    plt.savefig(fname=filepng, bbox_inches="tight", pad_inches=0.1)
    plt.savefig(fname=filepng_temp, bbox_inches="tight", pad_inches=0.1)

    # Add file path annotations
    _add_file_path_annotations(fig, [filepdf, filepng, filepng_temp])

    plt.close(fig)

    # Save plot data to pkl for later reconstruction
    pkl_data = {
        "wlen_hr": wlen_hr,
        "all_spectra_hr": all_spectra_hr,
        "ylabel": ylabel,
        "rad_mode": rad_mode,
        "wl_scalef": wl_scalef,
    }
    with open(str(model_folder / "high_res.pkl"), "wb") as _f:
        pickle.dump(pkl_data, _f)
    with open(str(Path(folder_retrieval, "model_high_low_res.pkl")), "wb") as _f:
        pickle.dump(pkl_data, _f)



# =====================
# Low-Resolution Plotting Functions
# =====================

def low_resolution_plot(path_imm, not_retrieval_obj_f, instruments_LR, wl_scalef, ylabel):
    """
    Generate low-resolution spectrum plot from a single model with observational data.

    Creates two separate plots (regular and normalized flux) showing:
    1. Upper panel: Model spectrum with observational data points and error bars
    2. Lower panel: Residuals (data - binned model)

    The function automatically separates instruments into normalized and non-normalized
    flux datasets, creating separate plots for each category.

    Args:
        path_imm (str): Output path for the main plot (PNG format)
        not_retrieval_obj_f: Model object containing attributes:
            - wl_full_resolution_LR: Low-resolution model wavelengths (nm)
            - depth_full_resolution_LR: Low-resolution model spectrum values
            - wl_binned_LR: List of binned wavelengths for each instrument
            - spectrum_binned_LR: List of binned spectrum for each instrument
            - offsetLRarr: Array of vertical offsets for instrument alignment
        instruments_LR (list): Collection of instrument objects, each with 'lr_data' containing:
            - ldata_LR: Observational wavelengths (nm)
            - rdata_LR: Observational flux data
            - errdata_LR: Flux uncertainties
            - step_LR: Wavelength bin width (used as x-error bars)
            - flux_norm: Boolean flag - True if flux is normalized
        wl_scalef (str): Wavelength scale - 'log' for logarithmic, otherwise linear
        ylabel (str): Y-axis label for the spectrum

    Output Files:
        - {path_imm}: Main plot (regular flux) in PNG format
        - {path_imm}.pdf: Main plot in PDF format
        - {path_imm}_normalized.png: Normalized flux plot (if applicable)
        - {path_imm}_normalized.pdf: Normalized flux plot PDF (if applicable)

    Notes:
        - Normalized flux is computed using estimate_continuum() function
        - Each instrument can be vertically offset for better visibility
        - Residuals help assess model-data agreement
    """
    # Define color scheme
    color_model = "#5782ff"
    colors = ["#DC143C", "#FFA500", "#6A5ACD", "#32CD32", "#008080",
              "#00CED1", "#FFD700", "#800080", "#FF00FF", "#A52A2A"]

    # Create two subplot structures: one for regular flux, one for normalized flux
    figure_regular, ax_regular = plt.subplots(
        2, 1, figsize=(12, 10), dpi=600,
        constrained_layout=True,
        gridspec_kw={'hspace': 0.01, 'wspace': 0.1, 'height_ratios': [7, 3]},
        sharex=True,
    )

    figure_norm, ax_norm = plt.subplots(
        2, 1, figsize=(12, 10), dpi=600,
        constrained_layout=True,
        gridspec_kw={'hspace': 0.01, 'wspace': 0.1, 'height_ratios': [7, 3]},
        sharex=True,
    )

    # Configure plot labels and titles for regular flux
    ax_regular[0].set_title("Low-resolution model", fontsize=GraphicsConfig.PLOT_FONT_SIZE_TITLE)
    ax_regular[0].set_ylabel(ylabel, fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_regular[1].set_ylabel("Residuals = data - binned", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_regular[1].set_xlabel("Wavelength (nm)", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    for a in ax_regular:
        a.tick_params(axis='both', labelsize=GraphicsConfig.PLOT_FONT_SIZE_TICK)

    # Configure plot labels and titles for normalized flux
    ax_norm[0].set_title("Low-resolution model (Normalized Flux)", fontsize=GraphicsConfig.PLOT_FONT_SIZE_TITLE)
    ax_norm[0].set_ylabel("Flux normalized", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_norm[1].set_ylabel("Residuals = data - binned", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_norm[1].set_xlabel("Wavelength (nm)", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    for a in ax_norm:
        a.tick_params(axis='both', labelsize=GraphicsConfig.PLOT_FONT_SIZE_TICK)

    # Plot the low-resolution model data on both figures
    ax_regular[0].plot(
        not_retrieval_obj_f.wl_full_resolution_LR,
        not_retrieval_obj_f.depth_full_resolution_LR,
        alpha=0.5, label="Model", color=color_model
    )

    _, spectrum_normalized = estimate_continuum(not_retrieval_obj_f.depth_full_resolution_LR, num_windows=40, top_window=10, check_gaussian=False)
    ax_norm[0].plot(
        not_retrieval_obj_f.wl_full_resolution_LR,
        spectrum_normalized,
        alpha=0.5, label="Model", color=color_model
    )

    # Get offset array for instrument data alignment
    offsetLR_arr = not_retrieval_obj_f.offsetLRarr

    # Counters to track how many instruments are plotted on each figure
    count_regular = 0
    count_norm = 0
    instruments_data = []

    # Loop over each instrument and plot its low-resolution data with error bars
    for index_LR, instrument in enumerate(instruments_LR):
        # Calculate offset for current instrument (first instrument has no offset)
        offsetLR = 0 if (index_LR == 0 or offsetLR_arr is None) else offsetLR_arr[index_LR - 1]

        # Extract instrument data
        ldata = instrument.lr_data.ldata_LR
        rdata = instrument.lr_data.rdata_LR - offsetLR  # Shift data to model level
        errdata = instrument.lr_data.errdata_LR
        step = instrument.lr_data.step_LR

        # Get binned spectrum for this instrument (adjusted for offset)
        spectrum_binned_LR = (not_retrieval_obj_f.spectrum_binned_LR[index_LR] - offsetLR)  # Return to model level

        # Check if this instrument has normalized flux
        is_normalized = instrument.lr_data.flux_norm
        instruments_data.append({
            "name": instrument.name,
            "ldata_LR": ldata,
            "rdata_LR": instrument.lr_data.rdata_LR,
            "errdata_LR": errdata,
            "step_LR": step,
            "flux_norm": is_normalized,
        })

        # Select the appropriate axes based on flux normalization and increment counter
        if is_normalized:
            ax = ax_norm
            count_norm += 1
        else:
            ax = ax_regular
            count_regular += 1

        # Plot observational data with error bars
        ax[0].errorbar(
            ldata, rdata, yerr=errdata, xerr=step, ls='none',
            label=f"{instrument.name} + {-offsetLR:.5f}",
            alpha=0.5, color=colors[index_LR], fmt='o', markersize=4,
            capsize=2, elinewidth=1, zorder=2
        )

        # Plot binned model data
        ax[0].plot(
            not_retrieval_obj_f.wl_binned_LR[index_LR],
            spectrum_binned_LR, alpha=0.8,
            label=f"{instrument.name} binned", color=colors[index_LR]
        )

        # Plot residuals (data - model)
        ax[1].errorbar(
            ldata, rdata - spectrum_binned_LR, yerr=errdata, xerr=step,
            alpha=0.8, ls='none',
            label=f"{instrument.name} + {-offsetLR:.5f}", color=colors[index_LR], fmt='o', markersize=4,
            capsize=2, elinewidth=1, zorder=2
        )

    # Set logarithmic scale if requested for both figures
    if wl_scalef == "log":
        ax_regular[0].set_xscale('log')
        ax_regular[1].set_xscale('log')
        ax_norm[0].set_xscale('log')
        ax_norm[1].set_xscale('log')

    # Configure and save regular flux figure only if it has instruments
    if count_regular > 0:
        ax_regular[0].legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
        ax_regular[1].set_xlim(ax_regular[0].get_xlim())
        ax_regular[1].axhline(0, color='black')  # Add zero line for residuals
        ax_regular[1].legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)

        # Save the regular flux plot in multiple formats
        figure_regular.savefig(path_imm, dpi=600)
        figure_regular.savefig(str(path_imm).replace("png", "pdf"), dpi=600)

    # Configure and save normalized flux figure only if it has instruments
    if count_norm > 0:
        ax_norm[0].legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
        ax_norm[1].set_xlim(ax_norm[0].get_xlim())
        ax_norm[1].axhline(0, color='black')  # Add zero line for residuals
        ax_norm[1].legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)

        # Save the normalized flux plot in multiple formats
        path_imm_norm = str(path_imm).replace(".png", "_normalized.png")
        figure_norm.savefig(path_imm_norm, dpi=600)
        figure_norm.savefig(str(path_imm_norm).replace("png", "pdf"), dpi=600)

    # Save plot data to pkl for later reconstruction
    pkl_data = {
        "wl_full_resolution_LR": not_retrieval_obj_f.wl_full_resolution_LR,
        "depth_full_resolution_LR": not_retrieval_obj_f.depth_full_resolution_LR,
        "offsetLRarr": offsetLR_arr,
        "wl_binned_LR": [np.array(s) for s in not_retrieval_obj_f.wl_binned_LR],
        "spectrum_binned_LR": [np.array(s) for s in not_retrieval_obj_f.spectrum_binned_LR],
        "instruments_data": instruments_data,
        "ylabel": ylabel,
        "wl_scalef": wl_scalef,
    }
    if count_regular > 0:
        with open(str(path_imm).replace(".png", ".pkl"), "wb") as _f:
            pickle.dump(pkl_data, _f)
    if count_norm > 0:
        with open(str(path_imm).replace(".png", "_normalized.pkl"), "wb") as _f:
            pickle.dump(pkl_data, _f)

    # Close both figures
    plt.close(figure_regular)
    plt.close(figure_norm)



# =====================
# Contribution Function Plotting
# =====================

def plot_contribution(folder_retrieval, rad_mode, model_obj, wlen_mu_contribution,
                     contribution_data, pressures, range_min_pressures,
                     range_max_pressures, resolution, wl_scalef):
    """
    Generate atmospheric contribution function contour plots.

    Creates 2D contour plots showing how different atmospheric pressure layers
    contribute to the observed spectrum as a function of wavelength. This is
    a diagnostic tool to understand which atmospheric layers are being probed
    at different wavelengths.

    The contribution function shows where in the atmosphere (at which pressure)
    photons are being absorbed or emitted for a given wavelength. High values
    indicate strong contribution from that layer.

    Args:
        folder_retrieval (str): Path to the retrieval folder
        rad_mode (str): Radiative transfer mode ('Transmission' or 'Emission')
        model_obj: Main model object containing atmosphere data with attributes:
            - atmosphere.pressure_data.pressures: Pressure grid (bar)
        wlen_mu_contribution (np.ndarray): Wavelength array in microns, shape (n_wavelengths,)
        contribution_data (np.ndarray): 2D contribution values, shape (n_pressures, n_wavelengths)
        pressures (np.ndarray): Pressure array in bars, shape (n_pressures,)
        range_min_pressures (float): Minimum pressure range in log10(bar)
        range_max_pressures (float): Maximum pressure range in log10(bar)
        resolution (str): Resolution identifier - 'HR' (high-res) or 'LR' (low-res)
        wl_scalef (str): Wavelength scale - 'log' for logarithmic, otherwise linear

    Output Files:
        In folder {folder_retrieval}/contribution/:
        - contribution_{resolution}.pdf/png: Contour plots
        - X_meshgrid_contribution_{resolution}.fits: Wavelength meshgrid
        - Y_meshgrid_contribution_{resolution}.fits: Pressure meshgrid
        - wlen_mu_contribution_{resolution}.fits: Wavelength array
        - pressures_contribution_{resolution}.fits: Pressure array
        - atmosphere.contr_tr_{resolution}.fits: Contribution function data

    Notes:
        - Y-axis (pressure) is always logarithmic
        - Colormap: 'gnuplot2_r' (reversed gnuplot2)
        - 30 contour levels are used for smooth visualization
    """
    # Define output file paths
    filepng_temp_c = str(Path(folder_retrieval, f"contribution_{resolution}.png"))
    filepdf_c = str(Path(folder_retrieval, "contribution", f"contribution_{resolution}.pdf"))
    filepng_c = str(Path(folder_retrieval, "contribution", f"contribution_{resolution}.png"))
    
    # Create meshgrid for contour plotting
    X, Y = np.meshgrid(wlen_mu_contribution, pressures)
    
    # Create the contribution function plot
    fig, ax = plt.subplots(1, 1, figsize=(12, 10), dpi=600)
    cf = ax.contourf(X, Y, contribution_data, 30, cmap="gnuplot2_r")
    cbar = plt.colorbar(cf)
    cbar.set_label("Contribution")
    
    # Set plot limits and scales
    ax.set_ylim([10 ** range_max_pressures, 10 ** range_min_pressures])
    ax.set_xlim([np.min(wlen_mu_contribution), np.max(wlen_mu_contribution)])
    ax.set_yscale("log")
    ax.set_xlabel("Wavelength (microns)")
    ax.set_ylabel("P (bar)")
    
    # Set logarithmic wavelength scale if requested
    if wl_scalef == "log":
        plt.xscale("log")
    
    ax.set_title(rad_mode + " contribution function")
    
    # Save plots in multiple formats
    plt.savefig(fname=filepdf_c, bbox_inches="tight", pad_inches=0.1)
    plt.savefig(fname=filepng_c, bbox_inches="tight", pad_inches=0.1)
    plt.savefig(fname=filepng_temp_c, bbox_inches="tight", pad_inches=0.1)
    
    # Add file path annotations
    plt.gca().annotate(
        filepdf_c,
        xy=(1.6, 1.0),
        xycoords="figure fraction",
        ha="right",
        va="top",
        color="black",
        fontsize=15,
    )
    plt.gca().annotate(
        filepng_c,
        xy=(1.6, 1.0),
        xycoords="figure fraction",
        ha="right",
        va="top",
        color="black",
        fontsize=15,
    )
    plt.gca().annotate(
        filepng_temp_c,
        xy=(1.6, 1.0),
        xycoords="figure fraction",
        ha="right",
        va="top",
        color="black",
        fontsize=15,
    )
    plt.close(fig)
    
    # Save data arrays to FITS files for further analysis
    
    # Save X meshgrid (wavelength grid)
    hdu = fits.PrimaryHDU(X)
    hdul = fits.HDUList([hdu])
    hdul.writeto(
        str(Path(folder_retrieval, "contribution", f"X_meshgrid_contribution_{resolution}.fits")), 
        overwrite=True
    )
    
    # Save Y meshgrid (pressure grid)
    hdu = fits.PrimaryHDU(Y)
    hdul = fits.HDUList([hdu])
    hdul.writeto(
        str(Path(folder_retrieval, "contribution", f"Y_meshgrid_contribution_{resolution}.fits")), 
        overwrite=True
    )
    
    # Save wavelength array
    hdu = fits.PrimaryHDU(wlen_mu_contribution)
    hdul = fits.HDUList([hdu])
    hdul.writeto(
        str(Path(folder_retrieval, "contribution", f"wlen_mu_contribution_{resolution}.fits")), 
        overwrite=True
    )
    
    # Save pressure array
    hdu = fits.PrimaryHDU(model_obj.atmosphere.pressure_data.pressures)
    hdul = fits.HDUList([hdu])
    hdul.writeto(
        str(Path(folder_retrieval, "contribution", f"pressures_contribution_{resolution}.fits")), 
        overwrite=True
    )
    
    # Save contribution function data
    hdu = fits.PrimaryHDU(contribution_data)
    hdul = fits.HDUList([hdu])
    hdul.writeto(
        str(Path(folder_retrieval, "contribution", f"atmosphere.contr_tr_{resolution}.fits")),
        overwrite=True
    )

    # Save plot data to pkl for later reconstruction
    pkl_data = {
        "wlen_mu_contribution": wlen_mu_contribution,
        "contribution_data": contribution_data,
        "pressures": pressures,
        "pressures_model": model_obj.atmosphere.pressure_data.pressures,
        "range_min_pressures": range_min_pressures,
        "range_max_pressures": range_max_pressures,
        "rad_mode": rad_mode,
        "wl_scalef": wl_scalef,
    }
    with open(str(Path(folder_retrieval, "contribution", f"contribution_{resolution}.pkl")), "wb") as _f:
        pickle.dump(pkl_data, _f)
    with open(str(Path(folder_retrieval, f"contribution_{resolution}.pkl")), "wb") as _f:
        pickle.dump(pkl_data, _f)



# =====================
# Data Loading and Saving Functions
# =====================

def load_contribution_data(folder_retrieval):
    """
    Load contribution analysis configuration from pickle file.

    Reads the analysis parameters that were saved by the GUI before launching
    the plotting script. These parameters control the model calculation and
    plotting behavior.

    Args:
        folder_retrieval (str): Path to the retrieval folder containing
            'output_data_contribution.pkl'

    Returns:
        tuple: Configuration parameters in the following order:
            - best_fit_parameters (np.ndarray): Best-fit or MCMC sample parameters
            - contribution (str): Contribution mode flag
            - lbl_sampling_hr (int): Line-by-line sampling parameter for HR
            - lbl_sampling_lr (int): Line-by-line sampling parameter for LR
            - range_min (float): Minimum pressure range (log scale)
            - range_max (float): Maximum pressure range (log scale)
            - n_layers (int): Number of atmospheric layers
            - aic_var: AIC variance parameter
            - overplot (bool): Whether to overplot selected orders
            - min_wl_lr (float): Minimum wavelength for low-res (nm)
            - max_wl_lr (float): Maximum wavelength for low-res (nm)
            - min_wl_hr (float): Minimum wavelength for high-res (nm)
            - max_wl_hr (float): Maximum wavelength for high-res (nm)
            - wl_scale_hr (str): Wavelength scale for HR plots ('log' or 'linear')
            - wl_scale_lr (str): Wavelength scale for LR plots ('log' or 'linear')
            - plot_opacity (bool): Whether to plot opacity data
    """
    with open(str(Path(folder_retrieval, "output_data_contribution.pkl")), "rb") as fo:
        contribution_data = pickle.load(fo)
        
    best_fit_parameters = contribution_data["best_fit_parameters"]
    contribution = contribution_data["contribution"]
    lbl_sampling_hr = contribution_data.get("lbl_sampling_hr", None)
    if lbl_sampling_hr is None: lbl_sampling_hr = contribution_data.get("lbl_sampling", 4)
    lbl_sampling_lr = contribution_data.get("lbl_sampling_lr", 4)
    range_min = contribution_data["min_pressures"]
    range_max = contribution_data["max_pressures"]
    n_layers = contribution_data["layers"]
    aic_var = contribution_data["aic_var"]
    overplot = contribution_data["overplot_orders"]
    min_wl_lr = contribution_data["min_wl_lr"]
    max_wl_lr = contribution_data["max_wl_lr"]
    min_wl_hr = contribution_data["min_wl_hr"]
    max_wl_hr = contribution_data["max_wl_hr"]
    wl_scale_hr = contribution_data["wl_scale_hr"]
    wl_scale_lr = contribution_data["wl_scale_lr"]
    plot_opacity = contribution_data.get("plot_opacity", False)
    
    return (best_fit_parameters, contribution, lbl_sampling_hr, lbl_sampling_lr, range_min, range_max,
            n_layers, aic_var, overplot, min_wl_lr, max_wl_lr, 
            min_wl_hr, max_wl_hr, wl_scale_hr, wl_scale_lr, plot_opacity)


def save_model_outputs(folder_retrieval, not_retrieval_obj, model_obj):
    """
    Save model calculation results to pickle files for later analysis.

    Exports three types of data for post-processing and analysis:
    1. VMR profiles - Volume mixing ratios as a function of pressure
    2. VMR summary - Mean values and mass fractions
    3. Distribution data - Full model object for further analysis

    These files are used by the GUI for displaying results and by other
    analysis scripts for post-processing.

    Args:
        folder_retrieval (str): Path to the retrieval folder
        not_retrieval_obj: Model calculation object containing:
            - vmr_dict: Dictionary of volume mixing ratios by species
            - MMW: Mean molecular weight
            - vmr: Volume mixing ratio array
            - mass_fraction_names_lr/hr: Species names for mass fractions
            - mean_VMR_and_MF_string/dict: Summary statistics
        model_obj: Main model object with atmosphere data containing:
            - atmosphere.pressure_data.pressures: Pressure grid
            - atmosphere.resolution_obj.instruments_LR: Low-res instruments

    Output Files:
        - vmr_profile.pkl: VMR profiles vs pressure
        - vmr.pkl: Mean VMR values and mass fractions
        - distribution.pkl: Complete model object and instrument data
    """
    # Save VMR (Volume Mixing Ratio) profile
    file_vmr_prof = str(Path(folder_retrieval, "vmr_profile.pkl"))
    vmr_profile_dictionary = {
        "vmr_dict": not_retrieval_obj.vmr_dict,
        "pressures": model_obj.atmosphere.pressure_data.pressures
    }
    with open(file_vmr_prof, "wb") as fo:
        pickle.dump(vmr_profile_dictionary, fo)
    
    # Save VMR data
    file_vmr = str(Path(folder_retrieval, "vmr.pkl"))
    dictionary_file_vmr = {
        "MMW": not_retrieval_obj.MMW,
        "vmr": not_retrieval_obj.vmr,
        "mass_fraction_names_lr": not_retrieval_obj.mass_fraction_names_lr,
        "mass_fraction_names_hr": not_retrieval_obj.mass_fraction_names_hr,
        "string": not_retrieval_obj.mean_VMR_and_MF_string,
        "dictinary": not_retrieval_obj.mean_VMR_and_MF_dict
    }
    with open(file_vmr, "wb") as fo:
        pickle.dump(dictionary_file_vmr, fo)
    
    # Save distribution data
    file_distribution = str(Path(folder_retrieval, "distribution.pkl"))
    distribution_dictionary = {
        "not_retrieval_obj": not_retrieval_obj,
        "instruments_LR": model_obj.atmosphere.resolution_obj.instruments_LR
    }
    with open(file_distribution, "wb") as fo:
        pickle.dump(distribution_dictionary, fo)

def plot_envelope_model_LR(wlen_nm, all_spectra, all_binned_spectra, datasets_lr,
                          ylabel, wl_scalef, path_imm):
    """
    Generate low-resolution spectrum plot with confidence intervals from MCMC samples.

    Creates plots showing the median model spectrum with 68% and 95% credible regions
    from multiple model realizations, along with observational data and residuals.
    Similar to plot_envelope_model_HR but includes data comparison and residual panels.

    The function creates two separate plots:
    1. Regular flux plot (for non-normalized instruments)
    2. Normalized flux plot (for normalized instruments)

    Each plot has two panels:
    - Upper: Model confidence intervals with observational data points
    - Lower: Residuals (data - median binned model)

    Args:
        wlen_nm (np.ndarray): Wavelength array in nm, shape (n_wavelengths,)
        all_spectra (np.ndarray): 2D array of full-resolution spectra, shape (n_models, n_wavelengths)
        all_binned_spectra (list): List of dictionaries, one per model, each containing
            binned spectra for each instrument (key: instrument name, value: binned spectrum)
        datasets_lr (dict): Dictionary of observational data, keyed by instrument name, each containing:
            - ldata: Wavelength data points
            - rdata: Flux data (offset-corrected)
            - errdata: Flux uncertainties
            - step: Wavelength bin width
            - flux_norm: Boolean indicating if flux is normalized
        ylabel (str): Y-axis label for the spectrum
        wl_scalef (str): Wavelength scale - 'log' for logarithmic, otherwise linear
        path_imm (str): Output path for the main plot (PNG format)

    Output Files:
        - {path_imm}: Main plot (regular flux) in PNG format
        - {path_imm}.pdf: Main plot in PDF format
        - {path_imm}_normalized.png: Normalized flux plot (if applicable)
        - {path_imm}_normalized.pdf: Normalized flux plot PDF (if applicable)

    Color Scheme:
        - Light orange (#fdc086): 95% credible region
        - Orange (#ff7f00): 68% credible region
        - Dark orange (#8c2d04): Median model line
        - Various colors for different instruments' data points
    """
    color_2sigma = "#fdc086"  # Light orange
    color_1sigma = "#ff7f00"  # Orange
    color_model = "#8c2d04"  # Dark orange
    colors = ["#DC143C", "#FFA500", "#6A5ACD", "#32CD32", "#008080",
              "#00CED1", "#FFD700", "#800080", "#FF00FF", "#A52A2A"]

    # Create two subplot structures: one for regular flux, one for normalized flux
    figure_regular, ax_regular = plt.subplots(
        2, 1, figsize=(12, 10), dpi=600,
        constrained_layout=True,
        gridspec_kw={'hspace': 0.01, 'wspace': 0.1, 'height_ratios': [7, 3]},
        sharex=True,
    )

    figure_norm, ax_norm = plt.subplots(
        2, 1, figsize=(12, 10), dpi=600,
        constrained_layout=True,
        gridspec_kw={'hspace': 0.01, 'wspace': 0.1, 'height_ratios': [7, 3]},
        sharex=True,
    )

    # Configure plot labels and titles for regular flux
    ax_regular[0].set_title("Low-resolution model", fontsize=GraphicsConfig.PLOT_FONT_SIZE_TITLE)
    ax_regular[0].set_ylabel(ylabel, fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_regular[1].set_ylabel("Residuals = data - binned", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_regular[1].set_xlabel("Wavelength (nm)", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    for a in ax_regular:
        a.tick_params(axis='both', labelsize=GraphicsConfig.PLOT_FONT_SIZE_TICK)

    # Configure plot labels and titles for normalized flux
    ax_norm[0].set_title("Low-resolution model (Normalized Flux)", fontsize=GraphicsConfig.PLOT_FONT_SIZE_TITLE)
    ax_norm[0].set_ylabel("Flux normalized", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_norm[1].set_ylabel("Residuals = data - binned", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    ax_norm[1].set_xlabel("Wavelength (nm)", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
    for a in ax_norm:
        a.tick_params(axis='both', labelsize=GraphicsConfig.PLOT_FONT_SIZE_TICK)

    count_norm = 0
    count_regular = 0

    median_spectrum = np.percentile(all_spectra, 50, axis=0)
    lower_1sigma = np.percentile(all_spectra, 16, axis=0)
    upper_1sigma = np.percentile(all_spectra, 84, axis=0)
    lower_2sigma = np.percentile(all_spectra, 2.5, axis=0)
    upper_2sigma = np.percentile(all_spectra, 97.5, axis=0)

    counter_instruments = -1
    for instrument_name in datasets_lr.keys():
        counter_instruments += 1
        current_instrument = datasets_lr[instrument_name]
        # Check if this instrument has normalized flux
        is_normalized = current_instrument["flux_norm"]
        flux_data = current_instrument["rdata"]
        wave_data = current_instrument["ldata"]
        bin_width_data = current_instrument["step"]
        err_data = current_instrument["errdata"]

        # Select the appropriate axes based on flux normalization and increment counter
        if is_normalized:
            ax = ax_norm
            count_norm += 1
        else:
            ax = ax_regular
            count_regular += 1

        all_spectra_binned_single_instrument = []
        for index_model in range(all_spectra.shape[0]):
            current_model = all_binned_spectra[index_model]
            all_spectra_binned_single_instrument.append(current_model[instrument_name])
        all_spectra_binned_single_instrument = np.array(all_spectra_binned_single_instrument)
        median_binned = np.percentile(all_spectra_binned_single_instrument, 50, axis=0)
        residuals = flux_data - median_binned

        if counter_instruments == 0:
        # Upper panel: spectrum with envelope
            ax[0].fill_between(wlen_nm, lower_2sigma, upper_2sigma,
                               color=color_2sigma, alpha=0.5, label='95% credible region')
            ax[0].fill_between(wlen_nm, lower_1sigma, upper_1sigma,
                               color=color_1sigma, alpha=0.5, label='68% credible region')
            ax[0].plot(wlen_nm, median_spectrum, color=color_model, lw=1.5,
                       label='Median model', zorder=3)
        else:
            ax[0].fill_between(wlen_nm, lower_2sigma, upper_2sigma, color=color_2sigma, alpha=0.5)
            ax[0].fill_between(wlen_nm, lower_1sigma, upper_1sigma, color=color_1sigma, alpha=0.5)
            ax[0].plot(wlen_nm, median_spectrum, color=color_model, lw=1.5, zorder=3)

        ax[0].errorbar(wave_data, flux_data, yerr=err_data, xerr=bin_width_data,
                       fmt='o', color=colors[counter_instruments], alpha=0.8, markersize=4,
                       label='Data', capsize=2, elinewidth=1, zorder=4)

        # Lower panel: residuals
        ax[1].axhline(0, color='gray', ls='--', lw=1, zorder=0)
        ax[1].errorbar(wave_data, residuals, yerr=err_data, xerr=bin_width_data,
                       fmt='o', color=colors[counter_instruments], alpha=0.8, markersize=4,
                       capsize=2, elinewidth=1)
        ax[1].axhspan(-np.median(err_data), np.median(err_data), color='gray', alpha=0.2)

    # Set logarithmic scale if requested for both figures
    if wl_scalef == "log":
        ax_regular[0].set_xscale('log')
        ax_regular[1].set_xscale('log')
        ax_norm[0].set_xscale('log')
        ax_norm[1].set_xscale('log')

    # Configure and save regular flux figure only if it has instruments
    if count_regular > 0:
        ax_regular[0].legend(loc='best', frameon=True, fancybox=True, fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
        ax_regular[0].tick_params(labelbottom=False)
        ax_regular[1].set_xlim(ax_regular[0].get_xlim())
        ax_regular[1].axhline(0, color='black')  # Add zero line for residuals
        ax_regular[1].legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
        for a in ax_regular:
            a.minorticks_on()
            a.tick_params(which='both', direction='in', top=True, right=True)

        # Save the regular flux plot in multiple formats
        figure_regular.savefig(path_imm, dpi=600)
        figure_regular.savefig(str(path_imm).replace("png", "pdf"), dpi=600)

    # Configure and save normalized flux figure only if it has instruments
    if count_norm > 0:
        ax_norm[0].legend(loc='best', frameon=True, fancybox=True, fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
        ax_norm[0].tick_params(labelbottom=False)
        ax_norm[1].set_xlim(ax_norm[0].get_xlim())
        ax_norm[1].axhline(0, color='black')  # Add zero line for residuals
        ax_norm[1].legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
        for a in ax_norm:
            a.minorticks_on()
            a.tick_params(which='both', direction='in', top=True, right=True)

        # Save the normalized flux plot in multiple formats
        path_imm_norm = str(path_imm).replace(".png", "_normalized.png")
        figure_norm.savefig(path_imm_norm, dpi=600)
        figure_norm.savefig(str(path_imm_norm).replace("png", "pdf"), dpi=600)

    # Save plot data to pkl for later reconstruction
    pkl_data = {
        "wlen_nm": wlen_nm,
        "all_spectra": all_spectra,
        "all_binned_spectra": all_binned_spectra,
        "datasets_lr": datasets_lr,
        "ylabel": ylabel,
        "wl_scalef": wl_scalef,
    }
    if count_regular > 0:
        with open(str(path_imm).replace(".png", ".pkl"), "wb") as _f:
            pickle.dump(pkl_data, _f)
    if count_norm > 0:
        with open(str(path_imm).replace(".png", "_normalized.pkl"), "wb") as _f:
            pickle.dump(pkl_data, _f)

    # Close both figures
    plt.close(figure_regular)
    plt.close(figure_norm)



# =====================
# Main Execution Block
# =====================

def main():
    """
    Main entry point for atmospheric model plotting script.

    This script is called by the GUIBRUSHR GUI after a retrieval analysis has
    completed. It generates various types of plots depending on the analysis
    configuration:

    1. **Single Model Mode** (best-fit parameters):
       - High-resolution spectrum plot (if HR data available)
       - Low-resolution spectrum plot with data comparison (if LR data available)
       - Contribution function plots (if requested)

    2. **Multi-Model Mode** (MCMC samples):
       - High-resolution spectrum with confidence intervals
       - Low-resolution spectrum with confidence intervals and data comparison

    The script automatically determines which plots to generate based on:
    - Resolution mode (HR, LR, or both)
    - Contribution function flag
    - Number of parameter sets (single vs multiple models)

    Command Line Usage:
        python plot_contribution.py <folder_retrieval> <target> <path_default>

    Args (from command line):
        sys.argv[1] - folder_retrieval: Path to retrieval results folder
        sys.argv[2] - target: Target name (currently unused in code)
        sys.argv[3] - path_default: Default path for GUIBRUSHR modules

    Input Files (in folder_retrieval):
        - output_data_contribution.pkl: Analysis configuration
        - df_general_info.csv: General retrieval information
        - df_parameters.csv: Parameter definitions

    Output:
        Various plots and FITS files (see individual plotting functions)
        error_log_model.pkl: Error log (None if successful)

    Workflow:
        1. Load configuration from pickle file
        2. Initialize ModelData object with appropriate settings
        3. Determine if single model or multi-model mode
        4. Generate appropriate plots for each resolution
        5. Save results and error log
    """
    # ========== Setup and Configuration ==========
    # Add GUIBRUSHR path to sys.path for imports
    path_default = sys.argv[3]
    sys.path.append(path_default)
    os.chdir(path_default)

    # Import GUIBRUSHR modules (must be done after path setup)
    from GUIBRUSHR.Retrieval.ModelCalculation import ModelData
    from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import get_csv_value

    # Extract command line arguments
    folder_retrieval = sys.argv[1]

    try:
        # ========== Load Analysis Configuration ==========
        # Load all analysis parameters from pickle file
        (best_fit_parameters, contribution, lbl_sampling_hr, lbl_sampling_lr, range_min, range_max,
         n_layers, aic_var, overplot, min_wl_lr, max_wl_lr,
         min_wl_hr, max_wl_hr, wl_scale_hr, wl_scale_lr, plot_opacity) = load_contribution_data(folder_retrieval)

        # ========== Read General Retrieval Information ==========
        df_general_info = pd.read_csv(str(Path(folder_retrieval, "df_general_info.csv")))

        # Determine radiation mode (Transmission or Emission) and set y-axis label
        rad_mode = get_csv_value(df_general_info, "rad_mode")
        if rad_mode == "Transmission":
            ylabel = r"1 - Transit depth ($\rm 1 - (R_{P}/R_{S})^2$)"
        else:
            ylabel = r"1 + Emission lines ($\rm 1 + (R_{P}/R_{S})^2$)"

        # Check if using HR line list for LR calculations (for consistency)
        use_hr_linelist_for_lr = get_csv_value(df_general_info, "use_hr_linelist_for_lr")
        use_hr_linelist_for_lr = use_hr_linelist_for_lr == "1"
        resolution = get_csv_value(df_general_info, "resolution")

        # Adjust wavelength ranges if using HR line list for LR
        if use_hr_linelist_for_lr:
            if "High" in resolution:
                # Expand HR range to include both HR and LR wavelengths
                min_wl_hr = min(min_wl_lr, min_wl_hr)
                max_wl_hr = max(max_wl_lr, max_wl_hr)
            else:
                # Use LR wavelengths for HR calculation
                min_wl_hr = min_wl_lr
                max_wl_hr = max_wl_lr

        # ========== Determine Plotting Mode ==========
        # Check if we have multiple models (MCMC samples) or single best-fit
        n_dim_bestfit = best_fit_parameters.shape
        model_from_samples = len(n_dim_bestfit) == 2  # 2D array means multiple samples

        # Add opacity flag to model type if requested (single model only)
        plot_opacity = ", Opacity" if plot_opacity and not model_from_samples else ""

        # ========== Initialize Atmospheric Model Object ==========
        model_obj = ModelData.ModelData(
            path_params=str(Path(folder_retrieval, "df_parameters.csv")),
            path_df=str(Path(folder_retrieval, "df_general_info.csv")),
            id_process=None,
            table_output_file=None,
            minwlen_lr=min_wl_lr,
            maxwlen_lr=max_wl_lr,
            minwlen_hr=min_wl_hr,
            maxwlen_hr=max_wl_hr,
            model_type="Model" + contribution + plot_opacity,
            lbl_sampling_hr=lbl_sampling_hr,
            lbl_sampling_lr=lbl_sampling_lr,
            range_min=range_min,
            range_max=range_max,
            nlayers=n_layers,
            manual_model_obj=None,
            load_new_opacities=True,
            path_default=path_default,
        )

        # ========== Create Output Directories ==========
        os.system("mkdir -p " + str(Path(folder_retrieval, "model_high_low_res")))
        os.system("mkdir -p " + str(Path(folder_retrieval, "contribution")))

        # ========== MULTI-MODEL MODE: MCMC Samples ==========
        # When we have multiple parameter sets from MCMC sampling,
        # we calculate spectra for each sample and plot confidence intervals
        if model_from_samples:
            # IMPORTANT: Force LR-only mode temporarily to avoid computing HR for all samples
            # (HR calculations are expensive, so we compute median HR separately below)
            was_HR = model_obj.atmosphere.resolution_obj.HR
            model_obj.atmosphere.resolution_obj.HR = False

            # --- Low-Resolution Envelope Plot ---
            if model_obj.atmosphere.resolution_obj.low_resolution():
                # Initialize storage for all model spectra
                list_spc_lr = []  # Full-resolution spectra
                list_spectrum_binned_LR = []  # Binned spectra for each instrument
                datasets_lr = {}  # Observational data (extracted once)
                wlen_lr = None  # Wavelength array (same for all models)

                # Loop through all MCMC samples
                for i in range(n_dim_bestfit[0]):
                    # Progress indicator (print every 50 models)
                    if i % 50 == 0:
                        print(f"Processing model {i+1}/{n_dim_bestfit[0]}")

                    # Calculate model spectrum for this parameter set
                    param_full = model_obj.create_param_full(best_fit_parameters[i, :])
                    _, _, not_retrieval_obj = model_obj.lh_function_gib(param_full)

                    # Store full-resolution spectrum
                    wlen_lr = not_retrieval_obj.wl_full_resolution_LR
                    list_spc_lr.append(not_retrieval_obj.depth_full_resolution_LR)

                    # Extract binned spectra for each instrument
                    offsetLR_arr = not_retrieval_obj.offsetLRarr
                    spectrum_binned_LR_current_model = {}

                    for index_LR, instrument in enumerate(model_obj.atmosphere.resolution_obj.instruments_LR):
                        # Calculate vertical offset for this instrument
                        offsetLR = 0 if (index_LR == 0 or offsetLR_arr is None) else offsetLR_arr[index_LR - 1]

                        # Extract observational data on first iteration only
                        if i == 0:
                            datasets_lr[instrument.name] = {
                                "ldata": instrument.lr_data.ldata_LR,
                                "rdata": instrument.lr_data.rdata_LR - offsetLR,
                                "errdata": instrument.lr_data.errdata_LR,
                                "step": instrument.lr_data.step_LR,
                                "flux_norm": instrument.lr_data.flux_norm,
                            }

                        # Store binned spectrum (offset-corrected)
                        spectrum_binned_LR_current_model[instrument.name] = np.array(
                            not_retrieval_obj.spectrum_binned_LR[index_LR] - offsetLR
                        )

                    list_spectrum_binned_LR.append(spectrum_binned_LR_current_model)

                # Generate LR envelope plot with confidence intervals
                path_LR = str(Path(folder_retrieval, "model_high_low_res", "low_res.png"))
                plot_envelope_model_LR(
                    wlen_lr, np.array(list_spc_lr), list_spectrum_binned_LR,
                    datasets_lr, ylabel, wl_scale_lr, path_LR
                )

            # --- High-Resolution Plot (Median Model Only) ---
            # Restore HR flag and compute median model if HR is enabled
            model_obj.atmosphere.resolution_obj.HR = was_HR
            if model_obj.atmosphere.resolution_obj.high_resolution():
                # Use median of all parameters (more stable than single sample)
                param_full = model_obj.create_param_full(np.median(best_fit_parameters, axis=0))
                _, _, not_retrieval_obj_hr = model_obj.lh_function_gib(param_full)

                # Generate standard HR plot (not envelope, too expensive for HR)
                # NOTE: plot_envelope_model_HR is commented out - HR envelope would require
                # computing HR spectra for all samples, which is computationally prohibitive
                # plot_envelope_model_HR(folder_retrieval, wlen_hr, np.array(list_spc_hr),
                #                       ylabel, rad_mode, wl_scale_hr)
                high_resolution_plot(
                    folder_retrieval, overplot, ylabel, model_obj.wlen_list_overplot,
                    not_retrieval_obj_hr, rad_mode, wl_scale_hr
                )

        # ========== SINGLE-MODEL MODE: Best-Fit Parameters ==========
        # When we have a single parameter set (e.g., maximum likelihood estimate),
        # we calculate one model and optionally plot contribution functions
        else:
            # Calculate single model spectrum
            param_full = model_obj.create_param_full(best_fit_parameters)
            _, _, not_retrieval_obj = model_obj.lh_function_gib(param_full)

            # Save detailed model outputs (VMR, mass fractions, etc.)
            save_model_outputs(folder_retrieval, not_retrieval_obj, model_obj)

            # --- High-Resolution Plots ---
            if model_obj.atmosphere.resolution_obj.high_resolution():
                if not contribution:
                    # Standard spectrum plot
                    high_resolution_plot(
                        folder_retrieval, overplot, ylabel,
                        model_obj.wlen_list_overplot, not_retrieval_obj,
                        rad_mode, wl_scale_hr
                    )
                else:
                    # Contribution function plot
                    plot_contribution(
                        folder_retrieval, rad_mode, model_obj,
                        not_retrieval_obj.wlen_mu_contribution_HR,
                        not_retrieval_obj.contribution_HR,
                        model_obj.atmosphere.pressure_data.pressures,
                        range_min, range_max, "HR", wl_scale_hr
                    )

            # --- Low-Resolution Plots ---
            if model_obj.atmosphere.resolution_obj.low_resolution():
                if not contribution:
                    # Standard spectrum plot with data comparison
                    path_LR = str(Path(folder_retrieval, "model_high_low_res", "low_res.png"))
                    low_resolution_plot(
                        path_LR, not_retrieval_obj,
                        model_obj.atmosphere.resolution_obj.instruments_LR,
                        wl_scale_lr, ylabel
                    )
                else:
                    # Contribution function plot
                    plot_contribution(
                        folder_retrieval, rad_mode, model_obj,
                        not_retrieval_obj.wlen_mu_contribution_LR,
                        not_retrieval_obj.contribution_LR,
                        model_obj.atmosphere.pressure_data.pressures,
                        range_min, range_max, "LR", wl_scale_lr
                    )

        # ========== Cleanup ==========
        # Delete model object to free memory
        del model_obj
        error = None  # No errors occurred

    except Exception as e:
        # ========== Error Handling ==========
        # Capture full error traceback for debugging
        error = str(e) + "\n" + traceback.format_exc()

    # ========== Save Error Log ==========
    # Always save error log (None if successful, traceback if failed)
    with open(str(Path(folder_retrieval, "error_log_model.pkl")), "wb") as fo:
        pickle.dump(error, fo)
        if error:
            # Print error to console for immediate visibility
            print(error)


if __name__ == '__main__':
    main()


# =====================
# Module Summary
# =====================
#
# This module provides comprehensive plotting capabilities for atmospheric
# retrieval analysis results from GUIBRUSHR. It supports both single best-fit
# models and MCMC sample ensembles, generating publication-quality plots with
# proper error propagation.
#
# Key Features:
# - High-resolution spectrum plots (with optional order highlighting)
# - Low-resolution spectrum plots (with data comparison and residuals)
# - Confidence interval plots from MCMC samples (68% and 95% credible regions)
# - Contribution function contour plots (wavelength vs pressure)
# - Automatic FITS file export for further analysis
# - Support for both transmission and emission spectroscopy
# - Handles multiple instruments with automatic normalization detection
#
# Performance Considerations:
# - HR envelope plots are computationally expensive (disabled for multi-model mode)
# - LR calculations are optimized by temporarily disabling HR computations
# - Progress indicators show processing status for long calculations
#
# Error Handling:
# - All errors are caught and saved to error_log_model.pkl
# - Traceback information is preserved for debugging
# - GUI can check error log to report failures to user
#
