Source code for GUIBRUSHR.Retrieval.ModelCalculation.Classes.Resolution

"""
Resolution class for handling high and low resolution spectroscopic data.

This module provides functionality to manage different resolution modes
for spectroscopic instruments and process low-resolution data.
"""

import pickle
from pathlib import Path

from GUIBRUSHR.Retrieval.ModelCalculation.Classes.LowResolutionData import LowResolutionData
from GUIBRUSHR.General_Constants.FunctionsAndConstants.general_functions import read_fits




class Resolution:
    """
    Manages resolution settings and data processing for spectroscopic instruments.
    
    This class handles both high-resolution (HR) and low-resolution (LR) modes,
    processing instrument data accordingly and creating appropriate data objects.
    
    Attributes:
        instruments_HR: High-resolution instruments collection
        instruments_LR: Low-resolution instruments collection
        lbl_high_res: High-resolution label data (optional)
        LR: Boolean flag indicating if low-resolution mode is active
        HR: Boolean flag indicating if high-resolution mode is active
    """
    


    def __init__(
        self,
        resolution,
        instruments_HR,
        instruments_LR,
        target_vsys,
        lbl_high_res=None,
        lbl_low_res=None
    ):
        """
        Initialize the Resolution object with specified parameters.
        
        Args:
            resolution: String indicating resolution mode ("Low", "High", or both)
            instruments_HR: Collection of high-resolution instruments
            instruments_LR: Collection of low-resolution instruments
            target_vsys: Target system velocity parameter
            lbl_high_res: Optional high-resolution label data
            lbl_low_res: Optional low-resolution label data
        """
        # Store instrument collections
        self.instruments_HR = instruments_HR
        self.instruments_LR = instruments_LR
        self.lbl_high_res = lbl_high_res
        self.lbl_low_res = lbl_low_res
        if lbl_high_res is not None:
            self.HR_resolution = 1e6 / self.lbl_high_res
        else:
            self.HR_resolution = None
        if lbl_low_res is not None:
            self.LR_resolution = 1e6 / self.lbl_low_res
        else:
            self.LR_resolution = None
        
        # Initialize resolution flags
        self.LR = False
        self.HR = False
        
        # Check resolution mode and set appropriate flags
        if "Low" in resolution:
            self.LR = True
            # Create low-resolution data objects for all LR instruments
            self.create_low_resolution_obj(target_vsys)
        
        if "High" in resolution:
            self.HR = True




    def create_low_resolution_obj(self, target_vsys):
        """
        Create low-resolution data objects for all low-resolution instruments.

        This method loads pickled low-resolution data for each instrument
        and creates LowResolutionData objects with the loaded parameters.

        Args:
            target_vsys: Target system velocity parameter for data processing
        """
        for element in self.instruments_LR:
            path_LR = element.path_instrument

            path_LR_classic = Path(path_LR, "lr_data.pkl")
            path_LR_norm = Path(path_LR, "lr_data_norm.pkl")
            path_LR = path_LR_classic if path_LR_classic.exists() else path_LR_norm

            with open(path_LR, "rb") as fo:
                lr_data = pickle.load(fo)

            spectr = lr_data.get("lr_spectr", None)
            lambd = lr_data.get("lr_lambd", None)
            lambdr = lr_data.get("lr_lambdr", None)
            flux_norm = lr_data.get("flux_norm", False)
            npix = len(spectr)

            # Load sigma: prefer separate plus/minus arrays,
            # fall back to legacy lr_sigma for both if not present
            sigma_legacy = lr_data.get("lr_sigma", None)
            sigma_plus = lr_data.get("lr_sigma_plus", sigma_legacy)
            sigma_minus = lr_data.get("lr_sigma_minus", sigma_legacy)

            element.lr_data = LowResolutionData(
                element, spectr, lambd, lambdr, sigma_plus, sigma_minus, npix, target_vsys, flux_norm
            )




    def high_resolution(self):
        """
        Check if high-resolution mode is active.
        
        Returns:
            Boolean indicating if high-resolution mode is enabled
        """
        return self.HR




    def low_resolution(self):
        """
        Check if low-resolution mode is active.
        
        Returns:
            Boolean indicating if low-resolution mode is enabled
        """
        return self.LR




    def only_high_resolution(self):
        """
        Check if only high-resolution mode is active (excluding low-resolution).
        
        Returns:
            Boolean indicating if only high-resolution mode is enabled
        """
        return self.HR and not self.LR




    def only_low_resolution(self):
        """
        Check if only low-resolution mode is active (excluding high-resolution).
        
        Returns:
            Boolean indicating if only low-resolution mode is enabled
        """
        return self.LR and not self.HR