"""
Atmosphere module for atmospheric modeling and radiative transfer calculations.

This module provides the Atmosphere class for handling atmospheric models,
including support for both petitRADTRANS and PyratBay radiative transfer codes.
"""

from pathlib import Path
from traceback import format_exc
import pyratbay.spectrum as ps
import numpy as np
from petitRADTRANS.plotlib import plot_opacity_contributions
from petitRADTRANS.radtrans import Radtrans
from petitRADTRANS import physical_constants as phys_const
from petitRADTRANS import physics as phys
from petitRADTRANS.stellar_spectra.phoenix import PhoenixStarTable
import pyratbay as pb
from scipy.integrate import trapezoid
from scipy.interpolate import splev
from scipy.stats import norm
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

from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.Pressure import Pressure
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.SpeciesMolec import SpeciesMolec
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.Resolution import Resolution
from GUIBRUSHR.Retrieval.ModelCalculation.Classes.Wavelength import Wavelength
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
    generate_chemcat,
    make_pyratbay_config, generate_star_model_hr, get_stellar_model
)


def cloud_opas(lambda0_micron, kappa_cond, kappa_opac, omega, xi, temperatures):
    """
    Create a cloud opacity function for atmospheric modeling.
    
    This function returns a closure that calculates cloud opacities based on
    wavelength-dependent scattering and absorption properties.
    
    Parameters
    ----------
    lambda0_micron : float
        Central wavelength in microns
    kappa_cond : float
        Condensation opacity coefficient
    kappa_opac : float or None
        Opacity coefficient (if None, calculated from kappa_cond)
    omega : float
        Width parameter for wavelength dependence
    xi : float
        Asymmetry parameter
    temperatures : array_like
        Temperature profile
        
    Returns
    -------
    function
        Opacity function that takes wavelengths and pressures as input
    """
    def get_opacity(wavelengths, pressures):
        """
        Calculate cloud opacities for given wavelengths and pressures.
        
        Parameters
        ----------
        wavelengths : array_like
            Wavelengths in microns
        pressures : array_like
            Pressures in bar
            
        Returns
        -------
        array_like
            Cloud opacities
        """
        # Convert pressure from bar to dyn/cm^2 for CGS units
        # The pressure must be in dyne because the kB constant is in erg/K (CGS)
        # Conversion: P_dyn = 1e6 * P_bar
        if kappa_opac is None:
            pressures_dyn_cm_2 = 1e6 * pressures  # dyn/cm^2
            # Calculate total number density in cm^-3
            ntot = pressures_dyn_cm_2 / (phys_const.kB * temperatures)
            kappa_opac_f = kappa_cond * ntot
        else:
            kappa_opac_f = kappa_opac
            
        # Calculate wavelength-dependent terms
        arg = (wavelengths - lambda0_micron) / omega  # Dimensionless
        phi_min = norm.pdf(arg)  # Probability density function (dimensionless)
        phi_maiusc = norm.cdf(xi * arg)  # Cumulative distribution function (dimensionless)
        
        # Calculate final opacities
        opacities = 2 * (phi_min * phi_maiusc)[:, np.newaxis] * kappa_opac_f
        return opacities
        
    return get_opacity

class Atmosphere:
    """
    Atmospheric model class for radiative transfer calculations.
    
    This class handles atmospheric modeling using either petitRADTRANS or PyratBay
    radiative transfer codes. It manages pressure profiles, chemical compositions,
    wavelength grids, and stellar spectra for both high and low resolution calculations.
    """
    
    def __init__(self):
        """Initialize the Atmosphere object with default None values."""
        # Stellar spectrum configuration
        self.temporary_force_new_binning = True
        self.hwhm_hr = None
        self.hwhm_lr = None
        self.stellar_spectrum_type_lr_mode = None
        self.use_hr_linelists_for_lr = None
        self.do_mean = False
        self.stellar_spline_model_hr = None
        self.stellar_spectrum_type_lr = None
        self.stellar_spectrum_type_hr = None
        self.stellar_spectrum = {}
        
        # Model configuration
        self.rv_sampling = None
        self.retrieval_model = None
        self.rad_mode = None
        self.additional_opac = False
        
        # Target and instrument configuration
        self.target = None
        self.resolution_obj = None
        
        # Atmospheric structure
        self.pressure_data = None
        self.temp_min = None
        self.temp_max = None
        self.temp_step = None
        
        # Chemical composition
        self.chemistry = None
        self.chemcat_obj = None
        self.species_compatible_with_prt = None
        self.species_obj = None
        
        # Wavelength grids
        self.wavelength = None
        self.wl_hr_nm = None
        self.wl_lr_nm = None
        
        # Radiative transfer objects
        self.atmosphere_hr = None  # High resolution atmosphere object
        self.atmosphere_lr = None  # Low resolution atmosphere object

    def update_params(
        self, line_species, line_species_isotope, 
        line_species_complete_name_hr, line_species_complete_name_lr, list_condensed_molecules,
        rayleigh_species, range_min_pressures, range_max_pressures, 
        layers, lbl_high_res, lbl_low_res, resolution=None, instruments_HR=None,
        instruments_LR=None, continum_opacity=None, mass_vector=None,
        target=None, chemistry=None, manual_model_obj=None,
        temp_min=300, temp_max=3000, temp_step=150,
        initial_params=None, path_default="", rad_mode="Transmission",
        manual_model_composition=None, retrieval_model="petitRADTRANS",
        rv_sampling=0.1, additional_opac=False, start_molecs=0
    ):
        """
        Update atmospheric model parameters.
        
        This method configures the atmospheric model with species, pressure grid,
        resolution settings, and other parameters needed for radiative transfer
        calculations.
        
        Parameters
        ----------
        line_species : list
            List of line species for opacity calculations
        line_species_isotope : list
            List of isotope-specific line species
        line_species_complete_name_hr : list
            Complete names of line species HR
        line_species_complete_name_lr : list
            Complete names of line species LR
        list_condensed_molecules : list
            List of condensed/cloud species
        rayleigh_species : list
            List of Rayleigh scattering species
        range_min_pressures : float
            Minimum pressure in the atmospheric grid
        range_max_pressures : float
            Maximum pressure in the atmospheric grid
        layers : int
            Number of atmospheric layers
        lbl_high_res : float
            Line-by-line high resolution sampling
        lbl_low_res: float
            Line-by-line low resolution sampling
        resolution : float, optional
            Spectral resolution
        instruments_HR : list, optional
            High resolution instruments
        instruments_LR : list, optional
            Low resolution instruments
        continum_opacity : list, optional
            Continuum opacity contributors
        mass_vector : array_like, optional
            Mass fractions for chemical species
        target : object, optional
            Target object containing system parameters
        chemistry : str, optional
            Chemistry model type
        manual_model_obj : object, optional
            Manual model configuration object
        temp_min : float, default=300
            Minimum temperature for calculations
        temp_max : float, default=3000
            Maximum temperature for calculations
        temp_step : float, default=150
            Temperature step for calculations
        initial_params : dict, optional
            Initial parameter values
        path_default : str, default=""
            Default path for data files
        rad_mode : str, default="Transmission"
            Radiative transfer mode
        manual_model_composition : dict, optional
            Manual composition specification
        retrieval_model : str, default="petitRADTRANS"
            Radiative transfer code to use
        rv_sampling : float, default=0.1
            Radial velocity sampling
        additional_opac : bool, default=False
            Whether to include additional opacity sources
        start_molecs : int, default=0
            Starting molecule index
        """
        # Initialize default parameters if not provided
        if initial_params is None:
            initial_params = {}
            
        # Unpack manual model object if provided
        # This allows overriding individual parameters with a complete model object
        if manual_model_obj is not None:
            resolution = manual_model_obj.resolution
            instruments_HR = manual_model_obj.instruments_HR
            instruments_LR = manual_model_obj.instruments_LR
            continum_opacity = manual_model_obj.continum_opacity
            mass_vector = manual_model_obj.mass_vector
            target = manual_model_obj.target
            temp_min = manual_model_obj.temp_min
            temp_max = manual_model_obj.temp_max
            temp_step = manual_model_obj.temp_step
            rad_mode = manual_model_obj.rad_mode
            retrieval_model = manual_model_obj.retrieval_model
            manual_model_composition = manual_model_obj.manual_model_composition
            rv_sampling = manual_model_obj.rv_sampling
            additional_opac = manual_model_obj.additional_opac
            
        # Set core atmospheric model parameters
        self.additional_opac = additional_opac
        print(f"Additional opac {self.additional_opac}")
        self.retrieval_model = retrieval_model
        self.target = target
        self.rad_mode = rad_mode
        
        # Initialize resolution object for wavelength grid management
        self.resolution_obj = Resolution(
            resolution, instruments_HR, instruments_LR, 
            target.systemic_velocity, lbl_high_res, lbl_low_res
        )
        
        # Initialize pressure grid
        self.pressure_data = Pressure(
            range_min_pressures, range_max_pressures, layers
        )

        # Set temperature range parameters
        self.temp_min = temp_min
        self.temp_max = temp_max
        self.temp_step = temp_step
        self.chemistry = chemistry
        self.rv_sampling = rv_sampling
        
        # Initialize chemistry catalog object
        self.chemcat_obj = None
        
        # Handle different chemistry models
        if chemistry == ConstantVariables.LIST_CHEMISTRY_TABLE[1]:
            # Use provided mass vector for manual chemistry
            mass_vector = np.array(mass_vector)
        else:
            # Generate chemistry catalog for equilibrium chemistry
            mass_vector, self.chemcat_obj, self.species_compatible_with_prt = generate_chemcat(
                self.pressure_data.range_min_pressures,
                self.pressure_data.range_max_pressures,
                self.pressure_data.n_layers,
                path_default,
            )
            
        # Initialize species and molecular composition object
        self.species_obj = SpeciesMolec(
            continum_opacity,
            line_species,
            line_species_isotope,
            line_species_complete_name_hr,
            line_species_complete_name_lr,
            list_condensed_molecules,
            rayleigh_species,
            mass_vector,
            start_molecs,
            initial_params,
            manual_model_composition,
        )

    def set_wl_range(self, minwlen_lr=None, maxwlen_lr=None,
                     minwlen_hr=None, maxwlen_hr=None):
        """
        Set wavelength ranges for high and low resolution calculations.
        
        This method determines the wavelength boundaries for both high and low
        resolution calculations based on instrument specifications and data ranges.
        If no explicit ranges are provided, they are automatically determined
        from the instrument configurations.

        Parameters
        ----------
        minwlen_lr : float, optional
            Minimum wavelength for low resolution calculations (microns)
        maxwlen_lr : float, optional
            Maximum wavelength for low resolution calculations (microns)
        minwlen_hr : float, optional
            Minimum wavelength for high resolution calculations (microns)
        maxwlen_hr : float, optional
            Maximum wavelength for high resolution calculations (microns)
        """
        # Determine LOW RESOLUTION wavelength range
        if minwlen_lr is None and self.resolution_obj.low_resolution():
            minwlen_lr = []
            maxwlen_lr = []
            
            for instrument in self.resolution_obj.instruments_LR:
                limit_ptr3 = 0.3  # Minimum wavelength limit in microns
                
                # Use instrumental wavelength limits if specified
                if instrument.wl_min is not None:
                    minwlen_lr.append(instrument.wl_min)
                if instrument.wl_max is not None:
                    maxwlen_lr.append(instrument.wl_max)
                    
                # Determine wavelength limits from actual data range
                lr_data = instrument.lr_data
                
                # Find minimum wavelength with buffer
                position = np.argmin(lr_data.ldata_LR)
                min_wl_temp = lr_data.ldata_LR[position]
                step_min_wl = lr_data.step_LR[position]
                # Apply lower limit or add buffer to data range
                min_wl_temp = (
                    limit_ptr3 if min_wl_temp < limit_ptr3 else ((min_wl_temp - 2 * step_min_wl) / 1000)
                )
                minwlen_lr.append(min_wl_temp)
                
                # Find maximum wavelength with buffer
                position = np.argmax(lr_data.ldata_LR)
                max_wl_temp = lr_data.ldata_LR[position]
                step_max_wl = lr_data.step_LR[position]
                max_wl_temp = ((max_wl_temp + 2 * step_max_wl) / 1000)
                maxwlen_lr.append(max_wl_temp)
                
            # Take the most restrictive (minimum) and extensive (maximum) ranges
            minwlen_lr = np.min(minwlen_lr)
            maxwlen_lr = np.max(maxwlen_lr)

        # Determine HIGH RESOLUTION wavelength range
        if minwlen_hr is None and self.resolution_obj.high_resolution():
            minwlen_hr = []
            maxwlen_hr = []
            
            for instrument in self.resolution_obj.instruments_HR:
                minwlen_night = []
                maxwlen_night = []
                
                # Collect wavelength ranges from all observation nights
                for icheck in range(instrument.ndate):
                    # Convert from nm to microns and add buffer
                    minwlen_night.append(
                        np.min(instrument.night_arr[icheck].lambdas) / 1000 - 0.1
                    )
                    maxwlen_night.append(
                        np.max(instrument.night_arr[icheck].lambdas) / 1000 + 0.1
                    )
                    
                # Choose between instrumental limits and night data limits
                # This balances completeness with computational efficiency
                minwlen_hr.append(np.maximum(instrument.wl_min, np.min(minwlen_night)))
                maxwlen_hr.append(np.minimum(instrument.wl_max, np.max(maxwlen_night)))
                
            # Take the most restrictive ranges across all instruments
            minwlen_hr = np.min(minwlen_hr)
            maxwlen_hr = np.max(maxwlen_hr)

        # Initialize wavelength object with determined ranges
        self.wavelength = Wavelength(
            minwlen_lr, maxwlen_lr, minwlen_hr, maxwlen_hr
        )

    def init_pyratbay(self, output_folder, initial_params):
        """
        Initialize PyratBay radiative transfer module.
        
        This method sets up the PyratBay atmospheric model with the specified
        configuration and identifies available opacity models for clouds and
        Rayleigh scattering.
        
        Parameters
        ----------
        output_folder : str or Path
            Directory for PyratBay output files
        initial_params : dict or None
            Initial parameter values for the model
        """
        # Generate PyratBay configuration file
        config_file = make_pyratbay_config(
            output_folder,
            self.target, self.pressure_data, self.species_obj,
            self.wavelength, self.resolution_obj, initial_params,
        )
        
        # Initialize PyratBay object
        self.atmosphere_hr = pyrat = pb.Pyrat(config_file)
        
        # Set logging verbosity to minimal
        pyrat.log.verb = 0
        
        # Identify and store indices for specific opacity models
        opacity_models = [model.name for model in pyrat.opacity.models]
        
        # Set cloud opacity model index if available
        if 'deck' in opacity_models:
            pyrat.icloud = opacity_models.index('deck')
            
        # Set Rayleigh scattering model index if available
        if 'lecavelier' in opacity_models:
            pyrat.irayleigh = opacity_models.index('lecavelier')

    def _read_opacities_low_resolution(self, path_target):
        """
        Initialize low-resolution opacity tables and stellar spectra.

        This method sets up the petitRADTRANS Radtrans object for low-resolution
        calculations. It creates the atmospheric model grid and prepares stellar
        spectra for emission calculations using one of three stellar spectrum types,
        each with different processing modes.

        Steps:
        1. Define wavelength boundaries from configured LR range
        2. Create Radtrans object:
           - If use_hr_linelists_for_lr: Use line-by-line (lbl) opacity mode
           - Otherwise: Use correlated-k opacity tables
        3. Extract wavelength grid from Radtrans and compute bin widths
        4. Load or compute stellar spectrum (for emission mode only):

           Stellar Spectrum Types:
           - Planck: Blackbody function at stellar effective temperature
             (Same processing for both use_hr_linelists_for_lr modes)

           - Phoenix HR: High-resolution Phoenix stellar atmosphere model
             * If use_hr_linelists_for_lr: Bin HR spectrum to LR grid using ps.bin_spectrum
             * Otherwise: Convolve HR spectrum with instrumental profile via spline

           - Phoenix LR: Low-resolution Phoenix stellar atmosphere model
             * If use_hr_linelists_for_lr: Bin Phoenix spectrum to LR grid using ps.bin_spectrum
             * Otherwise: Interpolate Phoenix spectrum to LR wavelength grid

        5. Store wavelength grid in nanometers

        Parameters
        ----------
        path_target : str
            Path to target directory containing stellar spectrum files (for Phoenix HR)

        Notes
        -----
        - For emission mode, stellar spectrum is required for planet-to-star flux ratio
        - Sets self.atmosphere_lr, self.stellar_spectrum["model_PRT_LR"], and self.wl_lr_nm
        - The use_hr_linelists_for_lr flag determines both opacity mode and stellar spectrum processing
        """
        # Define wavelength range for low-resolution calculations
        wl_min = self.wavelength.min_wl_lr
        wl_max = self.wavelength.max_wl_lr
        
        if self.use_hr_linelists_for_lr:
            self.atmosphere_lr = Radtrans(
                pressures=self.pressure_data.pressures,
                line_species=self.species_obj.line_species_complete_name_lr,
                rayleigh_species=self.species_obj.rayleigh_species,
                gas_continuum_contributors=self.species_obj.continum_opacity,
                wavelength_boundaries=np.array([wl_min, wl_max]),
                line_opacity_mode="lbl",  # Line-by-line mode for high resolution
                line_by_line_opacity_sampling=self.resolution_obj.lbl_low_res,
            )
        else:
            # Create low-resolution Radtrans object with correlated-k opacities
            self.atmosphere_lr = Radtrans(
                pressures=self.pressure_data.pressures,
                line_species=self.species_obj.line_species_complete_name_lr,
                cloud_species=self.species_obj.list_condensed_molecules,
                rayleigh_species=self.species_obj.rayleigh_species,
                gas_continuum_contributors=self.species_obj.continum_opacity,
                wavelength_boundaries=np.array([wl_min, wl_max]),
            )

        # Extract wavelength grid from Radtrans (in cm)
        wl_lr = self.atmosphere_lr.get_wavelengths()
        # Calcola le differenze tra punti consecutivi
        diffs = np.diff(wl_lr)

        # Bin width per ogni punto
        bin_width = np.zeros_like(wl_lr)
        bin_width[0] = diffs[0] / 2  # primo: metà della distanza successiva
        bin_width[1:-1] = (diffs[:-1] + diffs[1:]) / 2  # centrali: media delle due distanze
        bin_width[-1] = diffs[-1] / 2  # ultimo: metà della distanza precedente

        # Setup stellar spectrum for emission spectroscopy
        # (Not needed for transmission mode)
        if self.rad_mode == "Emission":
            # Step 1: Load/compute base stellar spectrum based on type
            if self.stellar_spectrum_type_lr == "Planck":
                # Planck blackbody - same for both HR/LR processing modes
                self.stellar_spectrum["model_PRT_LR"] = np.pi * phys.planck_function_cm(
                    self.target.stellar_effective_temperature, wl_lr
                )
                self.star_wavelength_cm = self.stellar_spectrum["model_PRT_LR"]
                self.stellar_spectrum_LR = wl_lr
                print("Used Planck function for stellar spectrum at LR")

            elif self.stellar_spectrum_type_lr == "Phoenix HR":
                # Load high-resolution Phoenix spectrum
                path_stellar_spectrum = str(Path(
                    path_target, "Star_Spectrum", "star_spectrum_default.pkl"
                ))

                if self.use_hr_linelists_for_lr or self.temporary_force_new_binning:
                    # Bin HR spectrum to LR grid
                    self.star_wavelength_cm, self.star_spectrum = get_stellar_model(
                        path_stellar_spectrum, np.min(wl_lr), np.max(wl_lr)
                    )
                    self.stellar_spectrum["model_PRT_LR"] = ps.bin_spectrum(
                        wl_lr * 1e4, self.star_wavelength_cm * 1e4, self.star_spectrum,
                        half_widths=bin_width * 1e4, gaps=None
                    )
                    print("Using high-resolution stellar spectrum for emission mode (binned)")
                else:
                    # Convolve HR spectrum with instrumental profile
                    wl_mid = (wl_lr[:-1] + wl_lr[1:]) / 2
                    delta_wl = np.diff(wl_lr)
                    resolution = int(np.round(np.mean(wl_mid / delta_wl)))
                    self.hwhm_lr = ConstantVariables.CLIGHT / (2 * resolution)

                    self.star_wavelength_cm, self.star_spectrum, self.stellar_spline_model_lr = generate_star_model_hr(
                        path_stellar_spectrum, self.hwhm_lr, np.min(wl_lr), np.max(wl_lr)
                    )
                    self.stellar_spectrum["model_PRT_LR"] = splev(wl_lr, self.stellar_spline_model_lr)
                    print(f"Using high-resolution stellar spectrum for emission mode at resolution = {resolution}")

            else:  # Phoenix LR
                # Load low-resolution Phoenix model
                star_table = PhoenixStarTable()
                stellar_spectrum, _ = star_table.compute_spectrum(
                    temperature=self.target.stellar_effective_temperature
                )
                self.star_wavelength_cm = stellar_spectrum[:, 0]
                flux_star_Fnu = stellar_spectrum[:, 1]

                # Convert from F_nu to F_lambda (erg/cm^2/s/cm)
                self.star_spectrum = flux_star_Fnu * ConstantVariables.CLIGHT * 1e5 / (self.star_wavelength_cm**2)

                # Process based on mode
                if self.use_hr_linelists_for_lr or self.temporary_force_new_binning:
                    self.stellar_spectrum["model_PRT_LR"] = ps.bin_spectrum(
                        wl_lr * 1e4, self.star_wavelength_cm * 1e4, self.star_spectrum,
                        half_widths=bin_width * 1e4, gaps=None
                    )
                    print("Used Phoenix LR function for stellar spectrum at LR (binned)")
                else:
                    self.stellar_spectrum["model_PRT_LR"] = np.interp(
                        wl_lr, self.star_wavelength_cm, self.star_spectrum
                    )
                    print("Used Phoenix LR function for stellar spectrum at LR (interpolated)")

        # Convert wavelength grid from cm to nm for convenience
        self.wl_lr_nm = wl_lr * 1e7

    def _read_opacities_high_resolution(self, path_target):
        """
        Initialize high-resolution opacity tables and stellar spectra.

        This method sets up the petitRADTRANS Radtrans object for high-resolution
        calculations using line-by-line (lbl) opacity mode. It creates the atmospheric
        model grid at high spectral resolution and prepares stellar spectra for
        emission calculations. Handles special case where HR linelists are used
        for LR calculations.

        Steps:
        1. Define wavelength boundaries from configured HR range
        2. Create Radtrans object with line-by-line opacity mode
        3. Extract high-resolution wavelength grid from Radtrans
        4. Load or compute stellar spectrum (for emission mode only):
           - Planck: Use blackbody function at stellar temperature
           - Phoenix/File: Load and spline-interpolate high-resolution stellar spectrum
        5. If HR linelists used for LR: bin HR stellar spectrum to LR instrument resolution
        6. Store wavelength grid in nanometers

        Parameters
        ----------
        path_target : str
            Path to target-specific data directory containing stellar spectra

        Notes
        -----
        - Only executed if high_resolution() is True or use_hr_linelists_for_lr is True
        - Line-by-line mode provides much higher spectral resolution than correlated-k
        - When use_hr_linelists_for_lr is True, HR opacities are binned to LR resolution
        - Spline interpolation used for stellar spectrum to match Radtrans grid
        - Sets self.atmosphere_hr, self.stellar_spectrum["model_PRT_HR"], and self.wl_hr_nm
        """
        # Define wavelength range for high-resolution calculations
        wl_min = self.wavelength.min_wl_hr
        wl_max = self.wavelength.max_wl_hr

        # Create high-resolution Radtrans object with line-by-line opacity mode
        # This provides much finer spectral resolution than low-resolution correlated-k tables
        self.atmosphere_hr = Radtrans(
            pressures=self.pressure_data.pressures,
            line_species=self.species_obj.line_species_complete_name_hr,
            rayleigh_species=self.species_obj.rayleigh_species,
            gas_continuum_contributors=self.species_obj.continum_opacity,
            wavelength_boundaries=np.array([wl_min, wl_max]),
            line_opacity_mode="lbl",  # Line-by-line mode for high resolution
            line_by_line_opacity_sampling=self.resolution_obj.lbl_high_res,
        )

        # Extract wavelength grid from Radtrans (in cm)
        wl_hr = self.atmosphere_hr.get_wavelengths()

        # Setup stellar spectrum for emission spectroscopy
        # (Not needed for transmission mode)
        if self.rad_mode == "Emission":
            if self.stellar_spectrum_type_hr == "Planck":
                # METHOD 1: Use Planck blackbody function
                # Compute flux at each wavelength using stellar effective temperature
                self.stellar_spectrum["model_PRT_HR"] = np.pi * phys.planck_function_cm(
                    self.target.stellar_effective_temperature, wl_hr
                )
                print("Used Planck function for stellar spectrum at HR")

            else:  # Phoenix or custom stellar spectrum from file
                # METHOD 2: Load and spline-interpolate stellar spectrum
                # Load high-resolution stellar spectrum from file
                path_stellar_spectrum = str(Path(
                    path_target, "Star_Spectrum", "star_spectrum_default.pkl"
                ))

                # Generate stellar model with spline interpolation
                # Returns wavelength grid, full-resolution spectrum, and spline coefficients
                _, _, self.stellar_spline_model_hr = generate_star_model_hr(
                    path_stellar_spectrum, self.hwhm_hr, np.min(wl_hr), np.max(wl_hr)
                )

                # Evaluate spline at Radtrans wavelength points
                self.stellar_spectrum["model_PRT_HR"] = splev(wl_hr, self.stellar_spline_model_hr)
                print(f"Used star spectrum {path_stellar_spectrum} for stellar spectrum at HR")

        # Convert wavelength grid from cm to nm for convenience
        self.wl_hr_nm = wl_hr * 1e7


    def _stellar_spectrum_instrument_LR(self, instrument_LR):
        """
        Bin stellar spectrum to low-resolution instrument wavelength bins.

        This method takes either high-resolution or low-resolution stellar spectra
        and bins them to match each instrument's wavelength resolution. The binning
        approach depends on the use_hr_linelists_for_lr flag.

        Two processing modes:

        1. **use_hr_linelists_for_lr = True**:
           - Uses HR stellar spectrum (self.star_spectrum, self.star_wavelength_cm)
           - Bins using petitRADTRANS ps.bin_spectrum for all instruments
           - Ignores stellar_spectrum_type_lr_mode setting

        2. **use_hr_linelists_for_lr = False**:
           - Uses LR stellar spectrum (self.stellar_spectrum["model_PRT_LR"])
           - Bins using method specified by stellar_spectrum_type_lr_mode:
             * "Integral": Integrates over each bin using trapezoidal rule
             * "Point by Point": Preserves spectrum points within each bin

        Process:
        1. Skip if no low-resolution calculations needed
        2. For each instrument:
           - Bin stellar spectrum to instrument's wavelength grid
           - Store result in self.stellar_spectrum[instrument.name]

        Parameters
        ----------
        instrument_LR : list
            List of low-resolution instrument objects, each with:
            - lr_data.ldata_LR_micron: wavelength centers (microns)
            - lr_data.step_LR_micron: wavelength bin half-widths (microns)
            - lr_data.npix: number of pixels
            - lr_data.lbottom, lr_data.ltop: wavelength bin edges (nanometers)

        Notes
        -----
        - Only applies to emission spectroscopy (not needed for transmission)
        - Stores binned spectra in self.stellar_spectrum[instrument.name]
        - For use_hr_linelists_for_lr=True, uses consistent binning with atmosphere
        """
        # Skip if no low-resolution calculations needed
        if not self.resolution_obj.low_resolution():
            return

        if self.use_hr_linelists_for_lr or self.temporary_force_new_binning:
            # Use HR stellar spectrum and bin with ps.bin_spectrum (same as atmosphere binning)
            for instrument in instrument_LR:
                self.stellar_spectrum[instrument.name] = ps.bin_spectrum(
                    instrument.lr_data.ldata_LR_micron,
                    self.star_wavelength_cm * 1e4,
                    self.star_spectrum,
                    half_widths=instrument.lr_data.step_LR_micron,
                    gaps=None
                )
        else:
            # Use LR stellar spectrum and bin according to stellar_spectrum_type_lr_mode
            stellar_spc_reference = self.stellar_spectrum["model_PRT_LR"]
            stellar_wavelength_nm = self.wl_lr_nm

            for instrument in instrument_LR:
                binned_spectrum = []

                # Bin to each wavelength bin of the instrument
                for i in range(instrument.lr_data.npix):
                    # Find stellar spectrum points within this instrument bin
                    mask = (
                        (stellar_wavelength_nm > instrument.lr_data.lbottom[i]) &
                        (stellar_wavelength_nm <= instrument.lr_data.ltop[i])
                    )

                    # Apply binning method
                    if self.stellar_spectrum_type_lr_mode == "Integral":
                        # Integrate stellar spectrum over bin using trapezoidal rule
                        flux_value = trapezoid(
                            stellar_spc_reference[mask],
                            stellar_wavelength_nm[mask]
                        )
                    else:  # Point by Point
                        # Use stellar spectrum points directly (no integration)
                        flux_value = stellar_spc_reference[mask]

                    binned_spectrum.append(flux_value)

                # Store binned stellar spectrum for this instrument
                self.stellar_spectrum[instrument.name] = binned_spectrum


    def read_opacities(
        self, table_output_file=None, path_target="", stellar_spectrum_type_lr="Phoenix",
        stellar_spectrum_type_hr="Phoenix", min_hwhm_hr=1.5, instrument_LR=None,
        stellar_spectrum_type_lr_mode="Point by Point", use_hr_linelists_for_lr=False
    ):
        """
        Initialize opacity tables and stellar spectra for radiative transfer.

        This method orchestrates the setup of petitRADTRANS Radtrans objects for both
        high and low resolution calculations. It delegates to specialized methods for
        LR and HR initialization, loads opacity data, and prepares stellar spectra
        for emission calculations.

        The method handles three resolution scenarios:
        1. Low resolution only: Uses correlated-k opacity tables
        2. High resolution only: Uses line-by-line opacity mode
        3. Mixed mode: HR linelists binned to LR resolution (use_hr_linelists_for_lr)

        Parameters
        ----------
        table_output_file : str, optional
            Path to output file for error messages
        path_target : str, default=""
            Path to target-specific data directory
        stellar_spectrum_type_lr : str, default="Phoenix"
            Type of stellar spectrum for LR ("Planck", "Phoenix", or "Integral")
        stellar_spectrum_type_hr : str, default="Phoenix"
            Type of stellar spectrum for HR ("Planck" or path to spectrum file)
        min_hwhm_hr : float, default=1.5
            Minimum pixel velocity width for high resolution (km/s)
        instrument_LR : list, optional
            List of instruments for low resolution calculations
        stellar_spectrum_type_lr_mode : str, default="Point by Point"
            Type of stellar spectrum for LR ("Point by Point" or "Integral")
        use_hr_linelists_for_lr : bool, default=False
            Use high resolution line lists for low resolution calculations

        Warnings
        --------
        If opacity initialization fails, the error is caught internally,
        a diagnostic message is written to table_output_file (if provided),
        and the process exits via ``exit()``. No exception is raised to the caller.

        Notes
        -----
        - Stellar spectra only loaded for emission mode (not transmission)
        - Sets attributes: atmosphere_lr, atmosphere_hr, stellar_spectrum (dict with
          "model_PRT_LR" and "model_PRT_HR" keys), wl_lr_nm, wl_hr_nm
        """
        self.use_hr_linelists_for_lr = use_hr_linelists_for_lr
        # Store stellar spectrum type configurations
        self.stellar_spectrum_type_lr = stellar_spectrum_type_lr
        self.stellar_spectrum_type_hr = stellar_spectrum_type_hr
        self.stellar_spectrum_type_lr_mode = stellar_spectrum_type_lr_mode
        self.stellar_spectrum = {}
        self.hwhm_hr = min_hwhm_hr

        try:
            # Initialize LOW RESOLUTION radiative transfer
            if self.resolution_obj.low_resolution():
                self._read_opacities_low_resolution(
                    path_target=path_target
                )

            # Initialize HIGH RESOLUTION radiative transfer
            if self.resolution_obj.high_resolution():
                self._read_opacities_high_resolution(
                    path_target=path_target
                )

            # Bin stellar spectrum to instrument resolution (for emission mode)
            if self.rad_mode != "Transmission" and instrument_LR is not None:
                self._stellar_spectrum_instrument_LR(instrument_LR=instrument_LR)

        except Exception as e:
            # Handle errors in opacity initialization
            print("Error initializing atmospheric opacities: " + str(e))
            print(format_exc())
            
            # Write error message to output file if specified
            if table_output_file is not None:
                with open(table_output_file, "w") as f:
                    f.write("Generate Atmosphere: Some molecules cannot be retrieved")
                exit()

    def calc_model(self, temperature, mass_fraction, vmr, MMW, dict_calc_model, 
                   HR=True):
        """
        Compute the transmission or emission spectrum for given atmospheric model.
        
        This method calculates either transmission or emission spectra using
        either petitRADTRANS or PyratBay radiative transfer codes, depending
        on the configuration.

        Parameters
        ----------
        temperature : array_like
            Temperature profile
        mass_fraction : dict
            Mass fractions of atmospheric species
        vmr : dict
            Volume mixing ratios of atmospheric species
        MMW : array_like
            Mean molecular weight profile
        dict_calc_model : dict
            Dictionary containing model parameters
        HR : bool, default=True
            Whether to use high resolution calculation

        Returns
        -------
        wl_full_resolution : array_like
            Wavelengths (nm)
        model : array_like or None
            Raw model output (transit radii or flux). None when using PyratBay.
        depth_full_resolution : array_like
            Final spectrum (transmission depth or emission)
        """
        # Extract model parameters from dictionary
        Pc = dict_calc_model["Pc"]  # Cloud top pressure
        rp = dict_calc_model["rp"]  # Planet radius
        k0 = dict_calc_model["k0"]  # Haze opacity at reference wavelength
        gamma = dict_calc_model["gamma"]  # Haze opacity power law index
        gravity = dict_calc_model["gravity"]  # Surface gravity
        haze_factor = dict_calc_model["haze_factor"]  # Haze enhancement factor
        cloud_fraction = dict_calc_model["cloud_fraction"]  # Cloud coverage fraction
        k_cond = dict_calc_model["k_cond"]  # Condensation opacity
        k_opac = dict_calc_model["k_opac"]  # Additional opacity
        lambda0_micron = dict_calc_model["lambda0_micron"]  # Reference wavelength
        P_ref = dict_calc_model["P_ref"]  # Reference pressure
        xi = dict_calc_model["xi"]  # Cloud asymmetry parameter
        std_radius_distribution = dict_calc_model["std_radius_distribution"]  # Cloud particle size distribution
        cloud_fsed = dict_calc_model["cloud_fsed"]  # Cloud sedimentation efficiency
        eddy_diff_coeff = dict_calc_model["eddy_diff_coeff"]  # Eddy diffusion coefficient
        omega_scale_micron = dict_calc_model["omega_scale_micron"]  # Cloud wavelength scale
        
        # Initialize output variables
        wl_full_resolution = None
        depth_full_resolution = None
        model = None

        # Calculate spectrum using petitRADTRANS
        if self.retrieval_model == 'petitRADTRANS':
            # Setup additional cloud opacity function if enabled
            additional_absorption_opacities_function = (
                cloud_opas(lambda0_micron, k_cond, k_opac, omega_scale_micron, xi, temperature) 
                if self.additional_opac else None
            )
            
            # Select appropriate resolution model
            rt_model = self.atmosphere_hr if HR else self.atmosphere_lr
            
            if self.rad_mode == "Transmission":
                # Calculate transmission spectrum
                _, model, _ = rt_model.calculate_transit_radii(
                    temperatures=temperature,
                    mass_fractions=mass_fraction,
                    mean_molar_masses=MMW,
                    planet_radius=rp,
                    reference_pressure=P_ref,
                    reference_gravity=gravity,
                    opaque_cloud_top_pressure=Pc,
                    haze_factor=haze_factor,
                    cloud_fraction=cloud_fraction,
                    power_law_opacity_350nm=k0,
                    power_law_opacity_coefficient=gamma,
                    frequencies_to_wavelengths=True,
                    additional_absorption_opacities_function=additional_absorption_opacities_function,
                    cloud_f_sed=cloud_fsed,
                    eddy_diffusion_coefficients=eddy_diff_coeff,
                    cloud_particle_radius_distribution_std=std_radius_distribution,
                )
                
                # Convert transit radii to transmission depth
                mol_modf = (
                       (model / phys_const.r_jup_mean) /
                       (self.target.stellar_radius * ConstantVariables.RATIO_RSUN_RJUP_MEAN)
                ) ** 2
                depth_full_resolution = 1 - mol_modf  # Transmission depth
            else:
                # Calculate emission spectrum
                _, model, _ = rt_model.calculate_flux(
                    temperatures=temperature,
                    mass_fractions=mass_fraction,
                    mean_molar_masses=MMW,
                    planet_radius=rp,
                    reference_gravity=gravity,
                    opaque_cloud_top_pressure=Pc,
                    haze_factor=haze_factor,
                    cloud_fraction=cloud_fraction,
                    power_law_opacity_350nm=k0,
                    power_law_opacity_coefficient=gamma,
                    frequencies_to_wavelengths=True,
                    additional_absorption_opacities_function=additional_absorption_opacities_function,
                    cloud_f_sed=cloud_fsed,
                    eddy_diffusion_coefficients=eddy_diff_coeff,
                    cloud_particle_radius_distribution_std=std_radius_distribution,
                )

                stellar_spectrum = self.stellar_spectrum["model_PRT_HR"] if HR else self.stellar_spectrum["model_PRT_LR"]

                # Calculate scaled flux ratio (planet/star)
                # Stellar spc is already multiplied by np.pi to go from erg/cm^2/cm/s/sr to erg/cm^2/cm/s
                Fscaled = (
                    (model / stellar_spectrum) *
                    (rp / (self.target.stellar_radius * phys_const.r_sun)) ** 2
                )
                depth_full_resolution = 1 + Fscaled  # Relative flux
                
            # Set wavelength np.array
            wl_full_resolution = self.wl_hr_nm if HR else self.wl_lr_nm
        # Calculate spectrum using PyratBay
        elif self.retrieval_model == 'PyratBay':
            pyrat = self.atmosphere_hr
            
            # Update atmospheric model parameters
            pyrat.atm.rplanet = rp
            pyrat.atm.refpressure = np.log10(P_ref)
            
            # Update cloud parameters if specified
            if Pc is not None:
                cloud = pyrat.opacity.models[pyrat.icloud]
                cloud.pars[0] = np.log10(Pc)
                
            # Update Rayleigh scattering parameters if specified
            if k0 is not None:
                rayleigh = pyrat.opacity.models[pyrat.irayleigh]
                rayleigh.pars[0] = np.log10(k0)
            if gamma is not None:
                rayleigh = pyrat.opacity.models[pyrat.irayleigh]
                rayleigh.pars[1] = gamma
                
            # Run radiative transfer calculation
            pyrat.run(temperature, vmr)
            
            # Extract results and convert units
            wl_full_resolution = 1.0 / np.flip(pyrat.spec.wn) * 1e7  # Convert wavenumber to wavelength (nm)
            depth = np.flip(pyrat.spec.spectrum)  # Transit depth
            depth_full_resolution = 1.0 - depth  # Transmission spectrum
            
        return wl_full_resolution, model, depth_full_resolution

    def calc_model_contribution(self, temperature, mass_fraction, MMW, 
                                dict_calc_model, HR=True):
        """
        Calculate transmission contribution function for atmospheric layers.
        
        This method computes the contribution of different atmospheric layers
        to the total transmission spectrum, useful for understanding which
        pressure levels dominate the observed signal.

        Parameters
        ----------
        temperature : array_like
            Temperature profile
        mass_fraction : dict
            Mass fractions of atmospheric species
        MMW : array_like
            Mean molecular weight profile
        dict_calc_model : dict
            Dictionary containing model parameters
        HR : bool, default=True
            Whether to use high resolution calculation

        Returns
        -------
        wl : array_like
            Wavelengths (nm)
        contribution : array_like
            Transmission contribution function
        """
        # Extract model parameters
        Pc = dict_calc_model["Pc"]
        rp = dict_calc_model["rp"]
        k0 = dict_calc_model["k0"]
        gamma = dict_calc_model["gamma"]
        gravity = dict_calc_model["gravity"]
        haze_factor = dict_calc_model["haze_factor"]
        cloud_fraction = dict_calc_model["cloud_fraction"]
        k_cond = dict_calc_model["k_cond"]
        k_opac = dict_calc_model["k_opac"]
        lambda0_micron = dict_calc_model["lambda0_micron"]
        P_ref = dict_calc_model["P_ref"]
        xi = dict_calc_model["xi"]
        omega_scale_micron = dict_calc_model["omega_scale_micron"]
        std_radius_distribution = dict_calc_model["std_radius_distribution"]
        cloud_fsed = dict_calc_model["cloud_fsed"]
        eddy_diff_coeff = dict_calc_model["eddy_diff_coeff"]
        
        # Select appropriate atmosphere model
        atmo = self.atmosphere_hr if HR else self.atmosphere_lr
        
        # Setup additional opacity function if enabled
        additional_absorption_opacities_function = (
            cloud_opas(lambda0_micron, k_cond, k_opac, omega_scale_micron, xi, temperature) 
            if self.additional_opac else None
        )
        
        # Calculate transit radii with contribution function
        wl, transit_radii, additional = atmo.calculate_transit_radii(
            temperatures=temperature,
            mass_fractions=mass_fraction,
            mean_molar_masses=MMW,
            planet_radius=rp,
            reference_pressure=P_ref,
            reference_gravity=gravity,
            opaque_cloud_top_pressure=Pc,
            haze_factor=haze_factor,
            cloud_fraction=cloud_fraction,
            power_law_opacity_350nm=k0,
            power_law_opacity_coefficient=gamma,
            frequencies_to_wavelengths=True,
            return_contribution=True,  # Request contribution function
            additional_absorption_opacities_function=additional_absorption_opacities_function,
            cloud_f_sed=cloud_fsed,
            eddy_diffusion_coefficients=eddy_diff_coeff,
            cloud_particle_radius_distribution_std=std_radius_distribution,
        )
        
        # Convert wavelengths to nanometers
        wl *= 1e7
        
        return wl, additional["transmission_contribution"]

    def opacity_contribution(self, temperatures, mass_fractions, MMW, 
                             dict_calc_model, HR=True):
        """
        Generate opacity contribution plots for atmospheric species.
        
        This method creates plots showing the relative contribution of different
        opacity sources (molecular species, clouds, hazes) to the total spectrum.
        
        Parameters
        ----------
        temperatures : array_like
            Temperature profile
        mass_fractions : dict
            Mass fractions of atmospheric species
        MMW : array_like
            Mean molecular weight profile
        dict_calc_model : dict
            Dictionary containing model parameters
        HR : bool, default=True
            Whether to use high resolution calculation
            
        Returns
        -------
        opacity_contributions : matplotlib.figure.Figure or None
            Figure object with opacity contribution plots, or None if failed
        """
        opacity_contributions = None
        
        try:
            # Extract model parameters
            Pc = dict_calc_model["Pc"]
            rp = dict_calc_model["rp"]
            k0 = dict_calc_model["k0"]
            gamma = dict_calc_model["gamma"]
            gravity = dict_calc_model["gravity"]
            haze_factor = dict_calc_model["haze_factor"]
            cloud_fraction = dict_calc_model["cloud_fraction"]
            k_cond = dict_calc_model["k_cond"]
            k_opac = dict_calc_model["k_opac"]
            lambda0_micron = dict_calc_model["lambda0_micron"]
            P_ref = dict_calc_model["P_ref"]
            xi = dict_calc_model["xi"]
            omega_scale_micron = dict_calc_model["omega_scale_micron"]
            std_radius_distribution = dict_calc_model["std_radius_distribution"]
            cloud_fsed = dict_calc_model["cloud_fsed"]
            eddy_diff_coeff = dict_calc_model["eddy_diff_coeff"]
            
            # Select appropriate atmosphere model
            rt_model = self.atmosphere_hr if HR else self.atmosphere_lr
            
            # Setup additional opacity function if enabled
            additional_absorption_opacities_function = (
                cloud_opas(lambda0_micron, k_cond, k_opac, omega_scale_micron, xi, temperatures) 
                if self.additional_opac else None
            )
            # Generate opacity contribution plots based on radiative transfer mode
            if self.rad_mode == "Transmission":
                # Plot transmission opacity contributions
                opacity_contributions = plot_opacity_contributions(
                    rt_model,
                    mode=self.rad_mode.lower(),
                    fill_below=False,  # Use curves instead of filled areas
                    x_axis_scale='log',  # Logarithmic pressure scale
                    temperatures=temperatures,
                    mass_fractions=mass_fractions,
                    mean_molar_masses=MMW,
                    reference_gravity=gravity,
                    planet_radius=rp,
                    reference_pressure=P_ref,
                    opaque_cloud_top_pressure=Pc,
                    haze_factor=haze_factor,
                    cloud_fraction=cloud_fraction,
                    power_law_opacity_350nm=k0,
                    power_law_opacity_coefficient=gamma,
                    additional_absorption_opacities_function=additional_absorption_opacities_function,
                    cloud_f_sed=cloud_fsed,
                    eddy_diffusion_coefficients=eddy_diff_coeff,
                    cloud_particle_radius_distribution_std=std_radius_distribution,
                )
            else:
                # TODO adjust and insert again plot
                print("Emission opacity contributions not yet fully implemented")
                return None
                # Plot emission opacity contributions
                # opacity_contributions = plot_opacity_contributions(
                #     rt_model,
                #     mode=self.rad_mode.lower(),
                #     fill_below=False,  # Use curves instead of filled areas
                #     x_axis_scale='log',  # Logarithmic pressure scale
                #     temperatures=temperatures,
                #     mass_fractions=mass_fractions,
                #     mean_molar_masses=MMW,
                #     reference_gravity=gravity,
                #     planet_radius=rp,
                #     opaque_cloud_top_pressure=Pc,
                #     haze_factor=haze_factor,
                #     cloud_fraction=cloud_fraction,
                #     power_law_opacity_350nm=k0,
                #     power_law_opacity_coefficient=gamma,
                #     additional_absorption_opacities_function=additional_absorption_opacities_function,
                #     cloud_f_sed=cloud_fsed,
                #     eddy_diffusion_coefficients=eddy_diff_coeff,
                #     cloud_particle_radius_distribution_std=std_radius_distribution,
                # )
                
            # Close the plot to free memory
            plt.close()
        except Exception as e:
            # Handle errors in opacity contribution calculation
            print("Opacity contributions failed: " + str(e) + "\n" + format_exc())
            
        return opacity_contributions

