"""
Frame Corner Plot Module

This module provides a GUI panel for creating and displaying corner plots
for atmospheric retrieval analysis. It handles plotting, data visualization,
and statistical analysis of MCMC chains.

Classes:
    FrameCornerPlot: Main panel class for corner plot functionality
"""

import os
import traceback
from pathlib import Path
import numpy as np
import pickle
import pandas as pd
from tkinter import messagebox
import subprocess
import re

from GUIBRUSHR.GUI.LAYOUT.MyPanel import MyPanel
from GUIBRUSHR.GUI.WIDGET.MyCheckBox import MyCheckBox
from GUIBRUSHR.GUI.WIDGET.MyDropDown import MyDropdown
from GUIBRUSHR.GUI.WIDGET.MyFigure import MyFigure
from GUIBRUSHR.GUI.WIDGET.MyLabel import MyLabel
from GUIBRUSHR.GUI.WIDGET.MyButton import MyButton
from GUIBRUSHR.GUI.WIDGET.MyTextField import MyTextField
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.General_Constants.Classes.HelpButton import HelpButton
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import (
    get_csv_value, pre_plot, plot_tp_profile, plot_corner_metallicity
)
from petitRADTRANS import physical_constants as pc


def calculate_solar_metallicity(names_molecules=None, dataset="Lodders09"):
    """
    Calculate Z/H (metals to hydrogen ratio) based on numerical abundances.

    Parameters:
        dataset: str, either "AGSS09" (Asplund et al. 2009) or "Lodders09" (Lodders et al. 2009)

    Returns:
        float: Z/H ratio (sum of all metal abundances relative to hydrogen)
    """

    def parse_molecule_formula(formula):
        """
        Parse a molecular formula and return element counts.

        Parameters:
            formula: str, molecular formula (e.g., "H2O", "CO2", "CH4")

        Returns:
            dict: element symbols as keys, counts as values
        """
        import re
        pattern = r'([A-Z][a-z]?)(\d*)'
        matches = re.findall(pattern, formula)
        elements = {}
        for element, count in matches:
            if element:
                count = int(count) if count else 1
                elements[element] = elements.get(element, 0) + count
        return elements
    if dataset == "AGSS09":
        # Asplund et al. (2009) - Photospheric abundances (Ni / NH)
        abundances = {
            'C': 2.69e-4,
            'N': 6.76e-5,
            'O': 4.90e-4,
            'Ne': 8.51e-5,
            'Mg': 3.98e-5,
            'Si': 3.24e-5,
            'S': 1.32e-5,
            'Fe': 3.16e-5,
            'Na': 1.74e-6,
            'Al': 2.82e-6,
            'P': 2.57e-7,
            'Cl': 3.16e-7,
            'Ar': 2.51e-6,
            'K': 1.07e-7,
            'Ca': 2.19e-6,
            'Ti': 8.91e-8,
            'Cr': 4.37e-7,
            'Mn': 2.45e-7,
            'Ni': 1.66e-7
        }
    elif dataset == "Lodders09":
        # Lodders et al. (2009) - Proto-Sun abundances (Ni / NH)
        abundances = {
            'C': 2.77e-4,
            'N': 8.19e-5,
            'O': 6.07e-4,
            'Ne': 8.71e-5,
            'Mg': 4.07e-5,
            'Si': 3.47e-5,
            'S': 1.55e-5,
            'Fe': 3.24e-5,
            'Na': 2.14e-6,
            'Al': 2.82e-6,
            'P': 2.82e-7,
            'K': 1.35e-7,
            'Ca': 2.19e-6,
            'Ti': 8.91e-8,
            'V': 1.00e-8,
            'Cr': 4.47e-7,
            'Ni': 1.78e-6
        }
    else:
        raise ValueError(f"Dataset '{dataset}' not recognized. Use 'AGSS09' or 'Lodders09'.")
    if names_molecules is None:
        # Z/H = sum of all metal abundances (Ni/NH)
        z_over_h = sum(abundances.values())
    else:
        elements_in_molecules = set()
        for mol in names_molecules:
            elements_in_molecules.update(parse_molecule_formula(mol).keys())
        elements_in_molecules.discard('H')
        elements_in_molecules.discard('He')
        z_over_h = sum(abundances.get(el, 0) for el in elements_in_molecules)
    return z_over_h

# noinspection PyTypeChecker
class FrameCornerPlot(MyPanel):
    """
    A GUI panel for creating and displaying corner plots for atmospheric retrieval analysis.
    
    This class provides functionality to:
    - Generate corner plots from MCMC chains
    - Display statistical information (likelihood, AIC, etc.)
    - Show reduced parameter plots
    - Display component details and comments
    - Plot temperature-pressure profiles
    """

    def __init__(self, parent, color, row, column, path_default, frame_input, 
                 width, height, window, **kwargs):
        """
        Initialize the FrameCornerPlot panel.
        
        Args:
            parent: Parent widget
            color: Background color for the panel
            row: Grid row position
            column: Grid column position
            path_default: Default path for the application
            frame_input: Input frame containing filter and table data
            width: Panel width
            height: Panel height
            window: Main window reference
            **kwargs: Additional keyword arguments for parent class
        """
        super().__init__(parent, color, row, column, **kwargs)
        
        # Initialize data attributes
        self.stringa_par = None
        self.aic_var = None
        self.current_max_likelihood = None
        self.chi2_conc = None
        self.last_max_likelihood = None
        self.stringa_med = None
        self.number_bad_chains = None
        self._initialize_data_attributes()
        
        # Set up main attributes
        self.path_folder_targets = frame_input.frame_input.frame_path.path_folder_targets.get_text()
        self.window = window
        self.path_default = path_default
        self.frame_input = frame_input
        
        # Initialize processing attributes
        self.folder_retrieval = None
        self.p = None  # Subprocess reference
        self.filepng = None  # PNG file path
        self.lista_parametri = None
        self.path_imm = None  # Image path
        self.last_plot_reduced = None
        
        # Create GUI components
        self._create_gui_components(color, width, height, path_default)
        
        # Initialize parameter attributes
        self._initialize_parameter_attributes()

    def _initialize_data_attributes(self):
        """Initialize data-related attributes to None."""
        self.stringa_par = None
        self.aic_var = None
        self.current_max_likelihood = None
        self.chi2_conc = None
        self.last_max_likelihood = None
        self.stringa_med = None
        self.number_bad_chains = None

    def _create_gui_components(self, color, width, height, path_default):
        """Create all GUI components for the panel."""
        # Single top row: controls + buttons together
        self.row1 = MyPanel(self, color, 0, 1)
        self._create_top_row(color)

        # Information panel below
        self.panel_info = MyPanel(self, color, 1, 1)
        self._create_info_panel(color)

    def _create_top_row(self, color):
        """Create single top row with all controls and buttons."""
        # Help button
        self._create_help_button()

        # Corner plot type dropdown
        self.dropdown_kind_corner = MyDropdown(
            self.row1, 0, 1, ["corner", "pygtc"], "Plot\nkind:",
            color=color, width=15, columnspan=2
        )

        # Corner plot style dropdown
        self.dropdown_color_corner = MyDropdown(
            self.row1, 0, 3, ["black", "orange", "blue", "purple", "green", "brown", "blue_cmap", "orange_cmap", "gnuplot2"],
            "corner.py\nstyle:", color=color, width=20, columnspan=2
        )

        # Remove bad chains checkbox
        self.checkbox_remove_chain = MyCheckBox(
            self.row1, 0, 5, "RM bad\nchain(s)", initial_value=1
        )

        # Best parameters extraction dropdown
        self.dropdown_best_params = MyDropdown(
            self.row1, 0, 6, ["Max LH", "Median", "KDE", "Max from param"],
            "Best param.\nextraction:", color=color, width=15, columnspan=2
        )

        # Metallicity selection textfield
        self.textfield_metallicity = MyTextField(
            self.row1, 0, 8, "c,o",
            "Metallicity:",
            color=color, width=10
        )

        # Buttons
        self.button_plot = MyButton(
            self.row1, 0, 10, 'PLOT', '#a366ff', self.plot, color_panel=color
        )
        self.button_show = MyButton(
            self.row1, 0, 11, 'SHOW LAST PLOT', '#ffcc00', self.show, color_panel=color
        )
        self.button_display_component = MyButton(
            self.row1, 0, 12, 'DISPLAY COMPONENTS', '#FF0000',
            self.display_comp, color_panel=color
        )
        self.button_display_comments = MyButton(
            self.row1, 0, 13, 'SHOW COMMENTS', '#66b3ff',
            self.show_comments, color_panel=color
        )
        self.checkbox_use_taskset = MyCheckBox(
            self.row1, 0, 14, "Use taskset", initial_value=1
        )

    def _create_info_panel(self, color):
        """Create information display panel with compact layout."""
        # Row 0: Likelihood stats in a single horizontal row
        self.textfield_gravity = MyTextField(
            self.panel_info, 0, 1, "0", "Gravity:", color=color, width=15
        )
        self.textfield_last_max_likelihood = MyTextField(
            self.panel_info, 0, 3, "0", "Last Max LH:", color=color, width=15
        )
        self.textfield_current_max_likelihood = MyTextField(
            self.panel_info, 0, 5, "0", "Current Max LH:", color=color, width=15
        )
        self.textfield_aic = MyTextField(
            self.panel_info, 0, 7, "0", "Aic:", color=color, width=15
        )
        self.textfield_bad_chains = MyTextField(
            self.panel_info, 0, 9, "0", "Bad Chains:", color=color, width=15
        )

        # Row 1: String reduced field + reduced plot buttons
        self.textfield_string_reduced = MyTextField(
            self.panel_info, 1, 1, "0", "String reduced:", color=color,
            columnspan=6, height=3, width=60
        )
        self.button_plot_reduced = MyButton(
            self.panel_info, 1, 7, 'PLOT REDUCED', '#FF3366',
            self.plot_reduce, color_panel=color
        )
        self.button_show_reduced = MyButton(
            self.panel_info, 1, 9, 'SHOW LAST REDUCED', '#B8002E',
            self.show_reduced, color_panel=color
        )

        # Row 2: Medians text field
        self.textfield_medians_errors = MyTextField(
            self.panel_info, 2, 1, "", "Medians (M.) / Errors (E.) / Upper Limits (U.L.):",
            color=color, columnspan=10, rowspan=8, height=10
        )

        self.label_comment = MyLabel(
            self.panel_info, 10, 1,
            label_text="The posterior distribution always includes the median value (in red). "
                       "Values shown above correspond to the selected likelihood representation.",
            columnspan=10, color=color
        )

    def _initialize_parameter_attributes(self):
        """Initialize parameter-related attributes."""
        self.chemistry = None
        self.format_temperature = None
        self.scattering = None
        self.ecc = None
        self.resolution = None
        self.global_exc = 1

    def _create_help_button(self):
        """Create help button with widget descriptions."""
        help_text = {
            "Plot kind dropdown": "Select the corner plot library to use. 'corner' uses the corner.py library for standard corner plots. 'pygtc' uses the pyGTC library for publication-quality corner plots with different styling.",

            "corner.py style dropdown": "Color scheme for corner plots when using corner.py library. Options include monochrome (black) and various colormaps (blue_cmap, orange_cmap, gnuplot2) for gradient-style plots, and solid colors (orange, blue, purple, green, brown) for single-color histograms.",

            "RM bad chain(s) checkbox": "When enabled, removes MCMC chains identified as poorly converged before generating the corner plot. This improves plot quality by excluding chains that didn't properly explore parameter space. Bad chains are identified using Gelman-Rubin convergence diagnostics.",

            "Best param. extraction dropdown": "Method for extracting best-fit parameter values from the posterior distribution. 'Max LH': Maximum likelihood point. 'Median': Median of the posterior (50th percentile, most robust). 'KDE': Peak of kernel density estimate. 'Max from param': Maximum of individual parameter distributions.",

            "Metallicity textfield": "Specifies which elements to include in metallicity [M/H] calculation. Format: comma-separated element symbols. Examples: 'c,o' calculates [(C+O)/H], 'full' includes all available heavy elements, 'c,o,n' calculates [(C+O+N)/H]. The calculation follows: [M/H] = log₁₀[(ΣNᵢ)/(2×nH₂×(ΣNᵢ/NH)☉)] where Nᵢ are the selected element number densities and (ΣNᵢ/NH)☉ is the solar abundance sum for those elements. Uses AGSS09 solar abundances.",

            "Use taskset checkbox": "When enabled, uses Linux taskset to run the corner plot generation on a dedicated CPU core (ncore+1) with low priority (nice=19). This prevents corner plot computation from interfering with other processes. Disable if taskset is not available or if you want to use all available cores.",

            "PLOT button": "Generate a new corner plot from the selected retrieval's MCMC chains. This processes the raw chains, computes statistics, creates the corner plot showing parameter correlations and posteriors, calculates C/O ratio and metallicity if molecules are present, generates TP profile plots, and saves all results.",

            "SHOW LAST PLOT button": "Display the most recently generated corner plot without recomputing. Loads previously saved plot and statistics from pickle files.",

            "DISPLAY COMPONENTS button": "Show detailed component information from the retrieval, including the full list of fitted parameters, fixed parameters, priors, molecular species, cloud properties, and complete model configuration from string_res.txt file.",

            "SHOW COMMENTS button": "Display user-defined comments associated with this retrieval run. Comments are stored in df_general_info.csv and can be used to document retrieval settings, assumptions, or results.",

            "PLOT REDUCED PLOT button": "Generate a corner plot showing only a subset of parameters. The parameters to include are specified in the 'String reduced' field as a comma-separated list (e.g., 'log_CH4,log_H2O,T_int'). Useful for focusing on parameters of interest or creating publication-quality figures.",

            "SHOW LAST REDUCED PLOT button": "Display the most recently generated reduced corner plot without recomputing. Loads the reduced plot from saved pickle file.",

            "Gravity field": "Computed surface gravity of the planet in m/s². Calculated from the best-fit planet mass and radius: g = G×M_planet/R_planet². Used in atmospheric modeling and scale height calculations.",

            "Last Max LH field": "Maximum log-likelihood value from the last MCMC run before any chain removal. This is the absolute best fit found by any chain during the entire sampling process.",

            "Current Max LH field": "Maximum log-likelihood value after removing bad chains (if RM bad chain(s) is enabled). This represents the best fit from well-converged chains only and is more reliable than Last Max LH.",

            "Aic field": "Akaike Information Criterion (AIC) value for this retrieval. AIC = -2×ln(L_max) + 2×k where k is the number of free parameters. Lower AIC values indicate better model fit accounting for model complexity. Use for comparing different retrieval models (different chemistry, TP profiles, cloud models).",

            "Bad Chains field": "Number of MCMC chains identified as poorly converged using Gelman-Rubin diagnostics (R̂ > 1.1). These chains are excluded from the corner plot if RM bad chain(s) is enabled. High numbers indicate sampling issues.",

            "String reduced field": "Comma-separated list of parameter names to include in a reduced corner plot. Parameter names must match those from the retrieval exactly (case-sensitive). Use parameter names as they appear in the full corner plot labels. Example: 'log_H2O,log_CO,T_int,log_g'.",

            "Medians (M.) / Errors (E.) / Upper Limits (U.L.) field": "Statistical summary of all fitted parameters extracted from posterior distributions. Shows median values (M.) as best estimates, asymmetric 1-sigma error bars (E.) from 16th and 84th percentiles (68% credible interval), and 2-sigma upper limits (U.L.) at 97.5th percentile for poorly constrained parameters. Format: parameter_name M: value E: -lower_err/+upper_err U.L.: upper_limit."
        }

        self.help_button = HelpButton(
            self.row1, 0, 0,
            "Corner Plot Panel Help",
            help_text,
            columnspan=1
        )

    def _get_selected_target_info(self):
        """
        Get selected target and retrieval information from the GUI.
        
        Returns:
            tuple: (target, selected_name, arr_values) or raises exception
            
        Raises:
            Exception: If no row is selected in the table
        """
        try:
            target = self.frame_input.frame_filter_table.frame_filter.dropdown_target.get_value()
            arr_values = self.frame_input.frame_filter_table.frame_table.table.get_arr_values()
            selected_name = arr_values[ConstantVariables.DICT_FILTER_TABLE["ID"]["index"]]
            return target, selected_name, arr_values
        except Exception:
            messagebox.showerror("Choose a row", "Choose a row")
            raise

    def _set_retrieval_parameters(self, arr_values):
        """
        Set retrieval parameters from array values.
        
        Args:
            arr_values: Array containing retrieval parameter values
        """
        self.chemistry = arr_values[ConstantVariables.DICT_FILTER_TABLE["#modes"]["index"]]
        self.format_temperature = arr_values[ConstantVariables.DICT_FILTER_TABLE["#atmos"]["index"]]
        self.scattering = arr_values[ConstantVariables.DICT_FILTER_TABLE["#scatterings"]["index"]]
        self.ecc = arr_values[ConstantVariables.DICT_FILTER_TABLE["#eccs"]["index"]]
        self.resolution = arr_values[ConstantVariables.DICT_FILTER_TABLE["#resolutions"]["index"]]

    def _disable_buttons(self):
        """Disable all control buttons during processing."""
        self.button_plot.button.configure(state="disabled")
        self.button_show.button.configure(state="disabled")
        self.button_plot_reduced.button.configure(state="disabled")
        self.button_show_reduced.button.configure(state="disabled")

    def _enable_buttons(self):
        """Enable all control buttons after processing."""
        self.button_plot.button.configure(state="normal")
        self.button_show.button.configure(state="normal")
        self.button_plot_reduced.button.configure(state="normal")
        self.button_show_reduced.button.configure(state="normal")

    def _save_output_data(self, arr1, lista_par, stringa_med, last_max_likelihood,
                         best_fit_parameters, value_fix, lista_par_fix, error_list,
                         current_max_likelihood):
        """
        Save output data to pickle file.
        
        Args:
            arr1: Array data
            lista_par: Parameter list
            stringa_med: Median string
            last_max_likelihood: Last maximum likelihood
            best_fit_parameters: Median values
            value_fix: Fixed values
            lista_par_fix: Fixed parameter list
            error_list: Error list
            current_max_likelihood: Current maximum likelihood
        """
        output_dictionary = {
            "arr1": arr1,
            "lista_par": lista_par,
            "stringa_med": stringa_med,
            "last_max_likelihood": last_max_likelihood,
            "lista_parametri": self.lista_parametri,
            "aic_var": self.aic_var,
            "number_bad_chains": self.number_bad_chains,
            "best_fit_parameters": best_fit_parameters,
            "filename" : "",
            "empty_string": "",
            "value_fix": value_fix,
            "lista_par_fix": lista_par_fix,
            "error_list": error_list,
            "current_max_likelihood": current_max_likelihood,
            "kind_corner": self.dropdown_kind_corner.get_value(),
            "style_corner": self.dropdown_color_corner.get_value(),
            "mode_LH": self.dropdown_best_params.get_value()
        }
        with open(str(Path(self.folder_retrieval, "output_data.pkl")), "wb") as fo:
            pickle.dump(output_dictionary, fo)

    def _start_plot_subprocess(self, target, suffix=""):
        """
        Start the plotting subprocess with optional taskset on dedicated core.

        Args:
            target: Target name
            suffix: File suffix for reduced plots
        """
        import sys

        plot_script_path = str(Path(
            self.path_default, "GUIBRUSHR", "GUI", "Input_Output_Panels", "Input_Panels",
            "TabPanels", "Frame_Retrieval_Analysis", "Tab_Analysis", "FrameCornerPlot", "plot_pmcmcg.py"
        ))

        # Build command based on checkbox value
        if self.checkbox_use_taskset.get_value():
            # Use taskset to run corner plot on core ncore+1 with priority 19
            corner_core = ConstantVariables.ncore + 1
            command = [
                "taskset",
                "-c",
                str(corner_core),
                "nice",
                "-n",
                "19",
                sys.executable,
                plot_script_path,
                self.folder_retrieval,
                suffix,
                target,
                self.path_default
            ]
        else:
            # Run without taskset
            command = [
                sys.executable,
                plot_script_path,
                self.folder_retrieval,
                suffix,
                target,
                self.path_default
            ]

        self.p = subprocess.Popen(
            command,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
        )
        self.window.after(500, self.check_proc)

    def _update_text_fields_with_results(self):
        """Update text fields with analysis results."""
        self.textfield_bad_chains.insert_text(str(self.number_bad_chains))
        self.textfield_last_max_likelihood.insert_text(str(self.last_max_likelihood))
        self.textfield_current_max_likelihood.insert_text(str(self.current_max_likelihood))
        self.textfield_aic.insert_text(str(self.aic_var))
        self.textfield_medians_errors.insert_text(self.stringa_med)
        self.textfield_string_reduced.insert_text(self.stringa_par)

    def _create_likelihood_distribution_plot(self):
        """Create and save the log likelihood distribution plot."""
        root = MyFigure("Log Likelihood distribution", type_data="Plot")
        ax = root.return_ax(["Log Likelihood distribution"], ["Log Likelihood"], ["#"])
        
        # Calculate bin width for histogram
        binwidth = (np.max(self.chi2_conc) - np.max(
            np.median(self.chi2_conc) - np.abs(np.median(self.chi2_conc) - np.max(self.chi2_conc))
        )) / 100
        
        # Create histogram
        _, _, _ = ax.hist(
            self.chi2_conc,
            bins=np.arange(
                np.median(self.chi2_conc) - np.abs(
                    np.median(self.chi2_conc) - np.max(self.chi2_conc)) - binwidth,
                np.max(self.chi2_conc) + binwidth,
                binwidth,
            ),
            alpha=0.5,
        )
        root.create_figure(str(Path(self.path_imm, "lh_distribution.jpg")), 600)

        # Save plot data to pkl for later reconstruction
        pkl_data = {"chi2_conc": self.chi2_conc}
        with open(str(Path(self.path_imm, "lh_distribution.pkl")), "wb") as _f:
            pickle.dump(pkl_data, _f)

    def _save_aic_likelihood_medians(self):
        """Save AIC, likelihood, and medians data to pickle file."""
        aic_likelihood_dictionary = {
            "stringa_med": self.stringa_med,
            "last_max_likelihood": self.last_max_likelihood,
            "aic_var": self.aic_var,
            "number_bad_chains": self.number_bad_chains,
            "current_max_likelihood": self.current_max_likelihood
        }
        with open(str(Path(self.folder_retrieval, "aic_likelihood_medians.pkl")), "wb") as fo:
            pickle.dump(aic_likelihood_dictionary, fo)

    @staticmethod
    def _parse_molecule_formula(formula):
        """
        Parse a molecular formula to count atoms of each element.

        Args:
            formula: Molecular formula string

        Returns:
            dict: Dictionary with element symbols as keys and counts as values
        """
        element_counts = {}

        # Pattern to match element symbols followed by optional numbers
        # Matches: Capital letter, optional lowercase letter, optional digit(s)
        pattern = r'([A-Z][a-z]?)(\d*)'

        matches = re.findall(pattern, formula)

        for element, count in matches:
            if element:  # Skip empty matches
                count = int(count) if count else 1
                element_counts[element] = element_counts.get(element, 0) + count

        return element_counts

    @staticmethod
    def _process_metallicity_string(metallicity_string):
        """
        Process metallicity string according to specific formatting rules.

        Args:
            metallicity_string: String containing element symbols separated by commas

        Returns:
            list: List of properly formatted element symbols
        """
        # Split by comma and strip whitespace
        elements = [elem.strip() for elem in metallicity_string.split(',')]

        processed_elements = []
        for elem in elements:
            if not elem:  # Skip empty strings
                continue

            # Check if contains "full" (case insensitive)
            if "full" in elem.lower():
                processed_elements.append("full")
            elif len(elem) == 1:
                # Single letter -> uppercase
                processed_elements.append(elem.upper())
            elif len(elem) == 2:
                # Two letters -> first uppercase, second lowercase
                processed_elements.append(elem[0].upper() + elem[1].lower())
            else:
                # For longer strings, just capitalize first letter
                processed_elements.append(elem[0].upper() + elem[1:].lower())

        return processed_elements

    def _load_molecules_data(self):
        """
        Load molecules data from pickle file and calculate metallicity and C/O ratio.
        """
        with open(str(Path(self.folder_retrieval, "c_o_ratio_met.pkl")), "rb") as fo:
            molecules_data = pickle.load(fo)

        array_molecules = molecules_data.get("array_molecules", None)
        names_molecules = molecules_data.get("names_molecules", None)
        names_molecules_corner = molecules_data.get("names_molecules_corner", None)
        H2 = molecules_data.get("H2", None)
        He = molecules_data.get("He", None)

        if array_molecules is not None:
            n_samples = array_molecules.shape[0]
            array_molecules_pow = np.power(10, array_molecules)
            H2_He_ratio = H2 / He
            total_vmr_metals = np.sum(array_molecules_pow, axis=1)
            total_he = (1.0 - total_vmr_metals) / (1.0 + H2_He_ratio)
            n_H2 = H2_He_ratio * total_he

            # Initialize arrays
            total_carbon = np.zeros(n_samples)
            total_oxygen = np.zeros(n_samples)
            total_metals = np.zeros(n_samples)  # Selected heavy atoms for metallicity calculation

            carbon_found = False
            oxygen_found = False

            # Get and process metallicity string from textfield
            # This determines which elements to include in the metallicity calculation
            # Example: "c,o" -> calculate [(C+O)/H], "full" -> calculate all metals
            metallicity_string = self.textfield_metallicity.get_text()
            processed_elements = self._process_metallicity_string(metallicity_string)

            # Parse each molecule and accumulate atomic counts
            # Following the generalized formula: [(M₁+M₂+...)/H]ₛ
            # where the numerator coefficient for each molecule is the sum of the counts
            # of selected elements in that molecule
            # Example for C+O: H₂O→1, CO→2, CH₄→1, HCN→1, CO₂→3
            for i, molecule_name in enumerate(names_molecules):
                element_counts = self._parse_molecule_formula(molecule_name)
                n_molecule = array_molecules_pow[:, i]

                # Accumulate carbon if present (for C/O ratio calculation)
                if 'C' in element_counts:
                    total_carbon += element_counts['C'] * n_molecule
                    carbon_found = True

                # Accumulate oxygen if present (for C/O ratio calculation)
                if 'O' in element_counts:
                    total_oxygen += element_counts['O'] * n_molecule
                    oxygen_found = True

                # Calculate the coefficient for this molecule in the metallicity numerator
                # This is the sum of atomic counts for all selected elements
                # Example: for CO₂ with elements=[C,O], coefficient = 1*C + 2*O = 3
                molecule_metal_coefficient = 0
                for element, count in element_counts.items():
                    # Skip H and He as they are not metals
                    if element not in ['H', 'He']:
                        # Include this element if "full" is selected or if it's in the selected list
                        if "full" in processed_elements or element in processed_elements:
                            molecule_metal_coefficient += count

                # Add contribution from this molecule: coefficient × abundance
                total_metals += molecule_metal_coefficient * n_molecule

            if carbon_found and oxygen_found:
                # Calculate C/O ratio
                c_o_ratio = np.divide(
                    total_carbon,
                    total_oxygen,
                    out=np.zeros_like(total_carbon),
                    where=total_oxygen != 0
                )
                array_molecules = np.column_stack((array_molecules, c_o_ratio))
                names_molecules_corner.append(r"$C/O$")

            # Calculate the solar metallicity ratio (denominator)
            # Use fictitious "molecules" (just element symbols) to pass to calculate_solar_metallicity
            # This allows us to get the solar abundance sum for only the selected elements
            # Example: for "c,o" -> pass ["C", "O"] to get solar (C+O)/H ratio
            if "full" in processed_elements:
                # Use all molecules to extract all elements
                solar_metals_to_h = calculate_solar_metallicity(dataset="AGSS09", names_molecules=names_molecules)
            else:
                # Pass element symbols as "molecules" to get solar abundance for selected elements only
                # This works because calculate_solar_metallicity extracts elements from molecule formulas
                solar_metals_to_h = calculate_solar_metallicity(dataset="AGSS09", names_molecules=processed_elements)
            denominator = 2 * n_H2 * solar_metals_to_h
            ratio = np.divide(
                total_metals,
                denominator,
                out=np.ones_like(total_metals) * 1e-10,
                where=denominator != 0
            )
            print(f"n_H2 <= 0: {np.sum(n_H2 <= 0)}")
            print(f"n_H2 min: {np.min(n_H2)}")
            print(f"solar_metals_to_h: {solar_metals_to_h}")
            print(f"denominator <= 0: {np.sum(denominator <= 0)}")
            print(f"denominator min: {np.min(denominator)}")
            print(f"total_metals <= 0: {np.sum(total_metals <= 0)}")
            print(f"total_metals min: {np.min(total_metals)}")
            print(f"ratio <= 0: {np.sum(ratio <= 0)}")
            print(f"ratio min: {np.min(ratio)}")
            print(f"total_metals <= 0: {np.sum(total_metals <= 0)}")
            # metallicity = np.log10(ratio)
            metallicity = np.log10(np.clip(ratio, a_min=1e-30, a_max=None))
            array_molecules = np.column_stack((array_molecules, metallicity))
            names_molecules_corner.append(r"$[M/H]$")
            plot_corner_metallicity(array_molecules, names_molecules_corner, self.folder_retrieval,
                                    "co_ratio_met_molecules", "C/O, met and molecules")
            plot_corner_metallicity(array_molecules[:, -2:], names_molecules_corner[-2:], self.folder_retrieval,
                                    "co_ratio_met", "C/O and met")

    def _load_temperature_pressure_data(self):
        """
        Load temperature-pressure profile data from pickle file.
        
        Returns:
            dict: Dictionary containing all temperature-pressure parameters
        """
        with open(str(Path(self.folder_retrieval, "temperature_pression.pkl")), "rb") as fo:
            tp_data = pickle.load(fo)
            
        return {
            "T0": tp_data["T0"], 
            "kappa_IR": tp_data["kappa_IR"], 
            "gamma_g": tp_data["gamma_g"], 
            "T_int": tp_data["T_int"],
            "rp": tp_data["rp"], 
            "target_mass": tp_data["target_mass"], 
            "T_low": tp_data["T_low"], 
            "T_high": tp_data["T_high"],
            "P_low": tp_data["P_low"], 
            "P_high": tp_data["P_high"], 
            "p1": tp_data["p1"], 
            "p2": tp_data["p2"], 
            "p3": tp_data["p3"],
            "alpha1": tp_data["alpha1"], 
            "alpha2": tp_data["alpha2"], 
            "T0_node": tp_data["T0_node"], 
            "T1_node": tp_data["T1_node"],
            "T2_node": tp_data["T2_node"], 
            "T3_node": tp_data["T3_node"], 
            "P1_node": tp_data["P1_node"], 
            "P2_node": tp_data["P2_node"]
        }

    def _calculate_gravity_and_plot_tp_profile(self, tp_data, df_general_info):
        """
        Calculate gravity and plot temperature-pressure profile.
        
        Args:
            tp_data: Temperature-pressure data dictionary
            df_general_info: General information DataFrame
        """
        # Calculate gravity
        gravity = (
            pc.G * (tp_data["target_mass"] * pc.m_jup) / 
            np.power(tp_data["rp"].value * pc.r_jup_mean, 2)
        )
        self.textfield_gravity.insert_text(str(gravity))
        
        # Set up random number generator
        rng = np.random.default_rng(seed=17)
        
        # Convert logarithmic values to linear if needed
        if tp_data["kappa_IR"].value is not None:
            tp_data["kappa_IR"].value = 10**tp_data["kappa_IR"].value
        if tp_data["gamma_g"].value is not None:
            tp_data["gamma_g"].value = 10**tp_data["gamma_g"].value
        
        # Get pressure range and layers from general info
        range_min_pressures = float(get_csv_value(df_general_info, "range_min_pressures"))
        range_max_pressures = float(get_csv_value(df_general_info, "range_max_pressures"))
        n_layers = int(get_csv_value(df_general_info, "layers"))
        
        # Create parameters dictionary for TP profile
        parameters_tp_profile = {
            "T0": tp_data["T0"], "kappa_IR": tp_data["kappa_IR"], "gamma_g": tp_data["gamma_g"],
            "T_int": tp_data["T_int"], "T_low": tp_data["T_low"], "T_high": tp_data["T_high"],
            "P_low": tp_data["P_low"], "P_high": tp_data["P_high"], "p1": tp_data["p1"],
            "p2": tp_data["p2"], "p3": tp_data["p3"], "alpha1": tp_data["alpha1"],
            "alpha2": tp_data["alpha2"], "T0_node": tp_data["T0_node"], "T1_node": tp_data["T1_node"],
            "T2_node": tp_data["T2_node"], "T3_node": tp_data["T3_node"], 
            "P1_node": tp_data["P1_node"], "P2_node": tp_data["P2_node"],
        }
        
        # Plot temperature-pressure profile
        plot_tp_profile(
            range_min_pressures, range_max_pressures, n_layers, parameters_tp_profile, 
            gravity, self.format_temperature, self.path_imm, True, rng
        )

    def display_comp(self):
        """
        Display component details from the retrieval results.
        
        Shows the string_res.txt file content in a popup window.
        """
        try:
            target, selected_name, _ = self._get_selected_target_info()
        except Exception:
            return
            
        try:
            self.folder_retrieval = str(
                Path(self.path_folder_targets, target, ConstantVariables.RESULTS_FOLDER, selected_name)
            )
            with open(str(Path(self.folder_retrieval, "string_res.txt")), "r") as fo:
                str_disp = fo.read()
                str_disp = str_disp.replace("\t", " ")
        except Exception as e:
            print("Generic error: " + str(e))
            print(traceback.format_exc())
            messagebox.showerror("File not ok", "File not ok")
            return
            
        try:
            MyFigure("Detail components", message=str_disp)
        except Exception:
            print("")

    def show_comments(self):
        """
        Display comments associated with the selected retrieval.
        
        Reads comments from df_general_info.csv and displays them in a popup window.
        """
        try:
            target, selected_name, _ = self._get_selected_target_info()
        except Exception:
            return
            
        try:
            self.folder_retrieval = str(
                Path(self.path_folder_targets, target, ConstantVariables.RESULTS_FOLDER, selected_name)
            )
            df_general_info = pd.read_csv(str(Path(self.folder_retrieval, "df_general_info.csv")))
        except Exception as e:
            print("Generic error: " + str(e))
            print(traceback.format_exc())
            messagebox.showerror("File not ok", "File not ok")
            return
            
        try:
            comments = get_csv_value(df_general_info, "comments")
        except Exception:
            messagebox.showerror(
                "There are no comments for this retrieval",
                "There are no comments for this retrieval",
            )
            return
            
        try:
            MyFigure("Comments details", message=comments)
        except Exception:
            print("")

    def plot(self):
        """
        Generate and display the main corner plot.
        
        This method:
        1. Extracts data from the selected retrieval
        2. Processes MCMC chains using pre_plot function
        3. Saves processed data to pickle files
        4. Starts subprocess to generate corner plot
        5. Creates additional plots (likelihood distribution, TP profile)
        """
        self.last_plot_reduced = False
        
        try:
            # Get selected target information
            target, selected_name, arr_values = self._get_selected_target_info()
            
            # Set retrieval parameters
            self._set_retrieval_parameters(arr_values)
            
            # Set folder path for retrieval results
            self.folder_retrieval = str(
                Path(self.path_folder_targets, target, ConstantVariables.RESULTS_FOLDER, selected_name)
            )
            
            # Process plot data using pre_plot function
            (
                number_bad_chains, aic_var, last_max_likelihood, current_max_likelihood,
                chi2_conc, lista_par, arr1, best_fit_parameters, value_fix, lista_par_fix,
                stringa_med, stringa_par, self.lista_parametri, error_list,
            ) = pre_plot(
                self.folder_retrieval, arr_values, self.checkbox_remove_chain.get_value(), True,
                self.dropdown_best_params.get_value()
            )
            
            # Create image directory
            self.path_imm = str(Path(self.folder_retrieval, "pmcmcg"))
            os.system("mkdir -p " + self.path_imm)

            # Store results in instance variables
            self.stringa_par = stringa_par
            self.aic_var = aic_var
            self.current_max_likelihood = current_max_likelihood
            self.chi2_conc = chi2_conc
            self.last_max_likelihood = last_max_likelihood
            self.stringa_med = stringa_med
            self.number_bad_chains = number_bad_chains
            
            # Save output data to pickle file
            self._save_output_data(
                arr1, lista_par, stringa_med, last_max_likelihood, best_fit_parameters,
                value_fix, lista_par_fix, error_list, current_max_likelihood
            )

            # Set up plot file path and disable buttons
            self.filepng = str(Path(self.folder_retrieval, "pmcmcg.png"))
            self._disable_buttons()
            
            # Start plotting subprocess
            self._start_plot_subprocess(target)
            
        except Exception as e:
            MyFigure("Error in corner plot", message=str(e) + "\n" + str(traceback.format_exc()))
            print(str(e) + "\n" + str(traceback.format_exc()))

    def check_proc(self):
        """
        Check if the plotting subprocess has completed.
        
        This method is called periodically to monitor the subprocess status.
        When complete, it updates the GUI with results and creates additional plots.
        """
        if self.p.poll() is not None:
            # Subprocess completed - re-enable buttons
            self._enable_buttons()
            
            # Display the generated plot
            MyFigure(title="Corner Plot", type_data="Image", image_path=self.filepng)
            
            # Update text fields and create additional plots for main plot (not reduced)
            if not self.last_plot_reduced:
                self._update_text_fields_with_results()

                self._load_molecules_data()
                
                # Create log likelihood distribution plot
                self._create_likelihood_distribution_plot()
                
                # Save AIC, likelihood, and medians data
                self._save_aic_likelihood_medians()
                
                # Load general information and create temperature-pressure profile
                df_general_info = pd.read_csv(str(Path(self.folder_retrieval, "df_general_info.csv")))
                tp_data = self._load_temperature_pressure_data()
                self._calculate_gravity_and_plot_tp_profile(tp_data, df_general_info)
        else:
            # Subprocess still running - check again later
            self.window.after(500, self.check_proc)

    def show(self):
        """
        Display the last generated corner plot.
        
        Loads data from saved pickle file and displays the plot without regenerating it.
        """
        try:
            # Get selected target information
            _, selected_name, arr_values = self._get_selected_target_info()
            target = self.frame_input.frame_filter_table.frame_filter.dropdown_target.get_value()
            
            # Set folder path and file path
            self.folder_retrieval = str(
                Path(self.path_folder_targets, target, ConstantVariables.RESULTS_FOLDER, selected_name)
            )
            self.filepng = str(Path(self.folder_retrieval, "pmcmcg.png"))
            
            # Load data from pickle file
            with open(str(Path(self.folder_retrieval, "output_data.pkl")), "rb") as fo:
                output_data = pickle.load(fo)
                # Skip arr1 - not used
                lista_par = output_data["lista_par"]
                stringa_med = output_data["stringa_med"]
                likelihood_var = output_data["last_max_likelihood"]
                # Skip lista_parametri - not used
                aic_var = output_data["aic_var"]
                number_bad_chains = output_data["number_bad_chains"]
                # Skip best_fit_parameters - not used
                # Skip empty string - not used
                # Skip value_fix - not used
                # Skip lista_par_fix - not used
                # Skip error_list - not used
                current_max_likelihood = output_data["current_max_likelihood"]
                # Skip dropdown value - not used
            
            # Create parameter string
            stringa_par = ""
            for elem in lista_par:
                stringa_par += elem + ","
            stringa_par = stringa_par[:-1]
            tp_data = self._load_temperature_pressure_data()
            gravity = (
                    pc.G * (tp_data["target_mass"] * pc.m_jup) /
                    np.power(tp_data["rp"].value * pc.r_jup_mean, 2)
            )
            self.textfield_gravity.insert_text(str(gravity))

            # Update text fields with loaded data
            self.textfield_bad_chains.insert_text(str(number_bad_chains))
            self.textfield_last_max_likelihood.insert_text(str(likelihood_var))
            self.textfield_current_max_likelihood.insert_text(str(current_max_likelihood))
            self.textfield_aic.insert_text(str(aic_var))
            self.textfield_medians_errors.insert_text(stringa_med)
            self.textfield_string_reduced.insert_text(stringa_par)
            
            # Display the plot
            MyFigure(title="Corner Plot", type_data="Image", image_path=self.filepng)
            
        except Exception as e:
            print("Generic error: " + str(e))
            print(traceback.format_exc())
            messagebox.showerror(
                "Error during the files import", "No previous plot generated"
            )

    def plot_reduce(self):
        """
        Generate a reduced corner plot with selected parameters only.
        
        This method creates a corner plot using only the parameters specified
        in the reduced string text field.
        """
        self.last_plot_reduced = True
        
        try:
            # Get selected target information
            target, selected_name, arr_values = self._get_selected_target_info()
            
            # Set retrieval parameters
            self._set_retrieval_parameters(arr_values)
            
            # Set folder path
            self.folder_retrieval = str(
                Path(self.path_folder_targets, target, ConstantVariables.RESULTS_FOLDER, selected_name)
            )
            
            # Process plot data
            (
                number_bad_chains, aic_var, last_max_likelihood, current_max_likelihood,
                chi2_conc, lista_par, arr1, best_fit_parameters, value_fix, lista_par_fix,
                stringa_med, stringa_par, self.lista_parametri, error_list,
            ) = pre_plot(
                self.folder_retrieval, arr_values, self.checkbox_remove_chain.get_value(), True,
                self.dropdown_best_params.get_value()
            )
            
            # Update text fields with basic information
            self.textfield_bad_chains.insert_text(str(number_bad_chains))
            self.textfield_current_max_likelihood.insert_text(str(current_max_likelihood))
            self.textfield_aic.insert_text(str(aic_var))
            self.textfield_medians_errors.insert_text(stringa_med)
            
            # Process parameter reduction
            self.lista_parametri = np.array(self.lista_parametri)
            reduced_string = self.textfield_string_reduced.get_text()
            parameter_reduced_list = reduced_string.split(",")
            self.textfield_string_reduced.insert_text(stringa_par)
            
            # Find column indices for reduced parameters
            column_reduced = np.array([], dtype=np.int16)
            for elem in parameter_reduced_list:
                column_reduced = np.append(column_reduced, lista_par.index(elem))
                
            # Reduce data arrays to selected parameters
            self.lista_parametri = self.lista_parametri[column_reduced]
            lista_par = np.array(lista_par)[column_reduced]
            lista_par = lista_par.tolist()
            
            # Create filename from parameter names
            filename = ""
            for elem in lista_par:
                filename += elem + "_"
            filename = filename[:-1]
            filename = filename.replace("/", "_over_")
            
            # Reduce array data and calculate new medians
            arr1 = np.array(arr1[:, column_reduced])
            best_fit_parameters = best_fit_parameters[column_reduced]
            
            # Create median string for reduced parameters
            stringa = ""
            for i in range(len(parameter_reduced_list)):
                stringa += (
                    parameter_reduced_list[i] + " = " + 
                    str("{:.3f}".format(best_fit_parameters[i])) + "\n"
                )
            self.textfield_medians_errors.insert_text(stringa)
            
            # Save reduced data to pickle file
            output_data_reduced = {
                "arr1": arr1,
                "lista_par": lista_par,
                "stringa_med": stringa,
                "last_max_likelihood": last_max_likelihood,
                "lista_parametri": self.lista_parametri,
                "aic_var": aic_var,
                "number_bad_chains": number_bad_chains,
                "best_fit_parameters": best_fit_parameters,
                "filename": filename,
                "value_fix": value_fix,
                "lista_par_fix": lista_par_fix,
                "error_list": error_list,
                "current_max_likelihood": current_max_likelihood,
                "kind_corner": self.dropdown_kind_corner.get_value(),
                "style_corner": self.dropdown_color_corner.get_value(),
                "mode_LH": self.dropdown_best_params.get_value()
            }
            with open(str(Path(self.folder_retrieval, "output_data_reduced.pkl")), "wb") as fo:
                pickle.dump(output_data_reduced, fo)
                
            # Set up reduced plot file path and disable buttons
            self.filepng = str(Path(self.folder_retrieval, filename + "_reducedpmcmcg.png"))
            self._disable_buttons()
            
            # Start plotting subprocess for reduced plot
            self._start_plot_subprocess(target, "_reduced")
            
        except Exception as e:
            print("Generic error: " + str(e))
            print(traceback.format_exc())
            messagebox.showerror(
                "Error during the files import", "Error during the files import"
            )

    def show_reduced(self):
        """
        Display the last generated reduced corner plot.
        
        Loads data from saved reduced pickle file and displays the plot.
        """
        try:
            # Get selected target information
            _, selected_name, arr_values = self._get_selected_target_info()
            target = self.frame_input.frame_filter_table.frame_filter.dropdown_target.get_value()
            
            # Set folder path
            self.folder_retrieval = str(
                Path(self.path_folder_targets, target, ConstantVariables.RESULTS_FOLDER, selected_name)
            )
            
            # Load reduced data from pickle file
            with open(str(Path(self.folder_retrieval, "output_data_reduced.pkl")), "rb") as fo:
                output_data_reduced = pickle.load(fo)
                
            # Extract data from dictionary
            lista_par = output_data_reduced["lista_par"]
            stringa_med = output_data_reduced["stringa_med"]
            likelihood_var = output_data_reduced["last_max_likelihood"]
            aic_var = output_data_reduced["aic_var"]
            number_bad_chains = output_data_reduced["number_bad_chains"]
            filename = output_data_reduced["filename"]
            current_max_likelihood = output_data_reduced["current_max_likelihood"]
            
            # Set reduced plot file path
            self.filepng = str(Path(self.folder_retrieval, filename + "_reducedpmcmcg.png"))
            
            # Create parameter string
            stringa_par = ""
            for elem in lista_par:
                stringa_par += elem + ","
            stringa_par = stringa_par[:-1]
            
            # Update text fields with loaded data
            self.textfield_bad_chains.insert_text(str(number_bad_chains))
            self.textfield_last_max_likelihood.insert_text(str(likelihood_var))
            self.textfield_current_max_likelihood.insert_text(str(current_max_likelihood))
            self.textfield_aic.insert_text(str(aic_var))
            self.textfield_medians_errors.insert_text(stringa_med)
            self.textfield_string_reduced.insert_text(stringa_par)
            
            # Display the reduced plot
            MyFigure(title="Corner Plot", type_data="Image", image_path=self.filepng)
            
        except Exception as e:
            print("Generic error: " + str(e))
            print(traceback.format_exc())
            messagebox.showerror(
                "Error during the files import", "No previous plot generated"
            )
