"""
NotRetrieval class module for handling atmospheric retrieval calculations.

This module contains the NotRetrieval class which processes atmospheric data
including volume mixing ratios (VMR), spectra, and opacity contributions for
both high resolution (HR) and low resolution (LR) data.
"""

import numpy as np
from numpy import mean
from petitRADTRANS.spectral_model import SpectralModel

from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
    get_line_lists,
    get_keys_by_value,
    get_condensed_line_list
)


class NotRetrieval:
    """
    A class for handling atmospheric retrieval calculations and data processing.
    
    This class manages atmospheric data including volume mixing ratios (VMR),
    spectral data at different resolutions, and opacity contributions. It processes
    molecular species data and creates dictionaries for easy access to VMR values.
    """
    
    def __init__(
        self,
        mass_fraction_names_hr,
        mass_fraction_names_lr,
        MMW,
        vmr,
        mean_VMR_and_MF_string,
        mean_VMR_and_MF_dict,
        wl_full_resolution_HR,
        depth_full_resolution_HR,
        wlen_mu_contribution_HR,
        contribution_HR,
        wl_full_resolution_LR,
        depth_full_resolution_LR,
        wl_binned_LR,
        spectrum_binned_LR,
        offsetLRarr,
        wlen_mu_contribution_LR,
        contribution_LR,
        chemistry,
        params_only_molec,
        species,
        opacity_contribution_HR,
        opacity_contribution_LR,
        instruments_LR,
        rp,
        rad_mode,
        stellar_radius,
        stellar_spectrum,
        lh_HR,
        lh_LR,
        HR_res_present,
        LR_res_present,
        use_hr_linelists_for_lr
    ):
        """
        Initialize the NotRetrieval class with atmospheric and spectral data.
        
        Args:
            mass_fraction_names_hr: High resolution mass fraction names
            mass_fraction_names_lr: Low resolution mass fraction names
            MMW: Mean molecular weight data
            vmr: Volume mixing ratio data
            mean_VMR_and_MF_string: Mean VMR and Molecular Fraction string
            mean_VMR_and_MF_dict: Mean VMR and Molecular Fraction dictionary
            wl_full_resolution_HR: Full resolution wavelength data
            depth_full_resolution_HR: Full resolution spectrum data
            wlen_mu_contribution_HR: High resolution wavelength contribution
            contribution_HR: High resolution contribution data
            wl_full_resolution_LR: Low resolution full wavelength data
            depth_full_resolution_LR: Low resolution full spectrum data
            wl_binned_LR: Low resolution binned wavelength data
            spectrum_binned_LR: Low resolution binned spectrum data
            offsetLRarr: Low resolution offset array
            wlen_mu_contribution_LR: Low resolution wavelength contribution
            contribution_LR: Low resolution contribution data
            chemistry: Chemistry table selection
            params_only_molec: Molecular parameters
            species: List of chemical species
            opacity_contribution_HR: High resolution opacity contribution
            opacity_contribution_LR: Low resolution opacity contribution
            instruments_LR: List of instruments for low resolution data
            rp: Planet radius parameter
            rad_mode: Radiative transfer mode selection
            stellar_radius: Stellar radius parameter
            stellar_spectrum: Stellar spectrum for resolution data
            lh_HR: Likelihood dictionary for high resolution data
            lh_LR: Likelihood dictionary for low resolution data
            HR_res_present: Boolean indicating if high resolution data is present
            LR_res_present: Boolean indicating if low resolution data is present
            use_hr_linelists_for_lr: Use high resolution linelists for low resolution data
        """
        # Store all input parameters as instance attributes
        self.mass_fraction_names_hr = mass_fraction_names_hr
        self.mass_fraction_names_lr = mass_fraction_names_lr
        self.MMW = MMW
        self.vmr = vmr
        self.mean_VMR_and_MF_string = mean_VMR_and_MF_string
        self.mean_VMR_and_MF_dict = mean_VMR_and_MF_dict
        self.wl_full_resolution_HR = wl_full_resolution_HR
        self.depth_full_resolution_HR = depth_full_resolution_HR
        self.wlen_mu_contribution_HR = wlen_mu_contribution_HR
        self.contribution_HR = contribution_HR
        self.wlen_mu_contribution_LR = wlen_mu_contribution_LR
        self.contribution_LR = contribution_LR
        self.wl_full_resolution_LR = wl_full_resolution_LR
        if use_hr_linelists_for_lr:
            self.depth_full_resolution_LR = SpectralModel.convolve(
                wl_full_resolution_LR, depth_full_resolution_LR, 5000,
            )
        else:
            self.depth_full_resolution_LR = depth_full_resolution_LR
        self.wl_binned_LR = wl_binned_LR
        self.spectrum_binned_LR = spectrum_binned_LR
        self.offsetLRarr = offsetLRarr
        self.instruments_LR = instruments_LR
        self.rp = rp
        self.rad_mode = rad_mode
        self.stellar_radius = stellar_radius
        self.lh_HR = lh_HR
        self.lh_LR = lh_LR

        # Initialize VMR-related attributes
        self.vmr_dict = dict()
        self.vmr_mean = []
        self.opacity_contribution_HR = None
        self.opacity_contribution_LR = None
        
        # Process molecular species and create VMR dictionary
        self._process_molecular_species(chemistry, params_only_molec, species)
        if self.rad_mode == "Emission":
            # TODO opacities contribution for Emission"
            pass
        else:
            if HR_res_present and opacity_contribution_HR is not None:
                self.opacity_contribution_HR = self._process_opacities(opacity_contribution_HR, stellar_spectrum.get("model_PRT_HR", None))
            if LR_res_present and opacity_contribution_LR is not None:
                self.opacity_contribution_LR = self._process_opacities(opacity_contribution_LR, stellar_spectrum.get("model_PRT_LR", None))

    def _process_opacities(self, opacities_contributions, stellar_spectrum):
        """
        Process opacity contributions by converting units and applying stellar correction.

        Args:
            opacities_contributions: Dictionary of opacity contribution data, where values
                can be tuples or nested dicts of species-level contributions
            stellar_spectrum: Stellar spectrum array used for emission mode correction

        Returns:
            The processed opacity contributions dictionary with converted units,
            or None if input is None.
        """
        if opacities_contributions is None:
            return None
        for key, value in opacities_contributions.copy().items():
            if value is None:
                continue

            if isinstance(value, dict):
                specie_dict = {}
                for species_name, species_data in value.copy().items():
                    specie_dict[species_name] = self.adapt_opacities(species_data, stellar_spectrum)
                opacities_contributions[key] = specie_dict
            else:
                opacities_contributions[key] = self.adapt_opacities(value, stellar_spectrum)

        return opacities_contributions

    def adapt_opacities(self, value, stellar_spectrum=None):
        """
        Convert a single opacity contribution tuple to display-ready units.

        Wavelength is converted from cm to nm. In transmission mode, the contribution
        is converted from cm to Jupiter radii.

        Args:
            value: Tuple of (wavelength, contribution, extra_data) arrays in cgs units
            stellar_spectrum: Stellar spectrum array (reserved for emission mode)

        Returns:
            Tuple of (wavelength_nm, contribution_converted, extra_data) arrays.
        """
        temp_value1 = value[0].copy() * 1e7  # wl in nanometers
        temp_value2 = value[1].copy()  # contribution
        temp_value3 = value[2].copy()
        if self.rad_mode == "Transmission":
            temp_value2 = temp_value2 / ConstantVariables.R_JUP_MEAN  # from cm to Rjup
        else:
            # TODO opacities contribution for Emission
            pass
            # temp_value2 = (
            #     (temp_value2 / (np.pi * stellar_spectrum)) *
            #     (self.rp * ConstantVariables.R_JUP_MEAN / (self.stellar_radius * ConstantVariables.R_SUN)) ** 2
            # ) + 1
        return temp_value1, temp_value2, temp_value3

    def _process_molecular_species(self, chemistry, params_only_molec, species):
        """
        Process molecular species data and populate VMR dictionary.
        
        This method handles the creation of VMR dictionaries based on the chemistry
        type and molecular parameters. It processes species names and creates
        mean VMR values for display.
        
        Args:
            chemistry: Chemistry table selection
            params_only_molec: Molecular parameters
            species: List of chemical species
        """
        # Get molecular line lists and condensed data
        molecs, isotopes, opacity_line_lists, _, isotope_masses = get_line_lists()
        (
            condensed_molecs_list, condensed_names_list,
            condensed_masses, condensed_visualization_name
        ) = get_condensed_line_list()
        
        # Process species based on chemistry type
        if chemistry == ConstantVariables.LIST_CHEMISTRY_TABLE[1]:
            self._process_chemistry_table_species(
                species, opacity_line_lists, isotopes, 
                condensed_molecs_list, condensed_names_list
            )
        else:
            self._process_molecular_parameters(
                params_only_molec, species, opacity_line_lists, isotopes
            )
        
        # Convert vmr mean list to formatted string
        self.vmr_mean = '\n'.join(self.vmr_mean)
    
    def _process_chemistry_table_species(self, species, opacity_line_lists, isotopes, 
                                       condensed_molecs_list, condensed_names_list):
        """
        Process species from chemistry table.
        
        Args:
            species: List of chemical species
            opacity_line_lists: opacity_line_list mapping
            isotopes: Isotope mapping
            condensed_molecs_list: List of condensed molecules
            condensed_names_list: List of condensed names
        """
        # List of species to exclude from processing
        excluded_species = ["H2", "He", "H", "H-", "e-"]
        
        for index, key in enumerate(species):
            if key not in excluded_species:
                # Try to find isotope mapping
                isotop = get_keys_by_value(opacity_line_lists, key)
                if isotop is not None:
                    key = get_keys_by_value(isotopes, isotop)
                else:
                    # Use condensed molecule list if isotope not found
                    key = condensed_molecs_list[np.where(condensed_names_list == key)][0]
            
            # Store VMR data and create mean VMR string
            self.vmr_dict[key] = self.vmr[:, index]
            self.vmr_mean.append(f"VMR({key}) = {mean(self.vmr[:, index]):.7e}")
    
    def _process_molecular_parameters(self, params_only_molec, species, 
                                    opacity_line_lists, isotopes):
        """
        Process molecular parameters for VMR calculation.
        
        Args:
            params_only_molec: Molecular parameters
            species: List of chemical species
            opacity_line_lists: opacity_line_list mapping
            isotopes: Isotope mapping
        """
        # List of species to exclude from processing
        excluded_species = ["H2", "He", "H", "H-", "e-"]
        
        for mol in params_only_molec:
            if mol is not None:
                name = mol.name_for_list_molec
                
                if name not in excluded_species:
                    # Find isotope mapping for the molecule
                    isotop = get_keys_by_value(opacity_line_lists, name)
                    name = get_keys_by_value(isotopes, isotop)
                
                # Find species index and store VMR data
                i = species.index(mol.molec_formula)
                self.vmr_dict[name] = self.vmr[:, i]
                self.vmr_mean.append(f"VMR({name}) = {mean(self.vmr[:, i]):.7e}")
