import pandas as pd
import numpy as np
import os 

from .aerosol_optics import aerosol_optics_single as brc_absorption_optics
from .aerosol_optics import aerosol_optics as ao
from .aerosol_absorption_calc import total_absorption_in_layer #calc_absorption_by_layer as total_absorption_in_layer

def extract_layer_data_by_number(rh_df, thickness_df, brc_conc_df, layer_number):
    """
    Extrae los datos de una capa específica basada en el número (1, 2, 3, ...).
    
    Args:
    rh_df (pd.DataFrame): DataFrame con columnas de capas para RH.
    thickness_df (pd.DataFrame): DataFrame con columnas de capas para Thickness.
    brc_conc_df (pd.DataFrame): DataFrame con columnas de capas para BRC Concentration.
    layer_number (int): Número de la capa a extraer (1 para "layer_1", 2 para "layer_2", etc.).
    
    Returns:
    tuple: Tres Series renombradas ('RH', 'Thickness', 'BRC') correspondientes a la capa especificada.
    """
    # Convertir el número de capa en el nombre correspondiente
    layer = f"layer_{layer_number}"
    
    # Validar que la capa exista en los DataFrames
    if layer not in rh_df.columns or layer not in thickness_df.columns or layer not in brc_conc_df.columns:
        raise ValueError(f"La capa {layer} no está presente en todos los dataframes.")
    
    # Extraer y renombrar las columnas, y redondear a dos decimales
    rh = rh_df[layer].rename('RH')
    thickness = thickness_df[layer].rename('Thickness')
    brc = brc_conc_df[layer].rename('BRC')
    
    return rh, thickness, brc


def calculate_total_absorption_by_layer(rh, layer_thickness, brc_conc, layers, **kwargs):
    """
    Calcula los valores de absorción para todas las capas y devuelve un DataFrame con el resultado.

    Args:
        rh (pd.DataFrame): DataFrame con columnas de capas para RH.
        layer_thickness (pd.DataFrame): DataFrame con columnas de capas para Thickness.
        brc_conc (pd.DataFrame): DataFrame con columnas de capas para BRC Concentration.
        layers (int): Número de capas (por ejemplo, 24).
        ri_brc (float): Índice de refracción para el BRC.
        wavelenght (float): Longitud de onda de la luz incidente (nm).

    Returns:
        pd.DataFrame: DataFrame con una columna por capa y el índice igual al de rh.
    """
    # DataFrame para almacenar resultados
    absorption_results = pd.DataFrame(index=rh.index)

    for layer in range(1, layers + 1):
        # Extraer los datos para la capa actual
        rh_l, layer_thickness_l, brc_conc_l = extract_layer_data_by_number(rh, layer_thickness, brc_conc, layer)

        # Calcular absorción para cada fila
        absc_brc = []
        for i in range(len(rh_l)):
            # Calcular parámetros ópticos para el valor actual de rh
            optical_parameters = brc_absorption_optics(
                np.array([rh_l[i]]), ['brc'], wvln=kwargs['wavelenght'], ri_brc=kwargs['ri_brc']
            )
            # Calcular absorción para el valor actual de brc y thickness
            absorption = total_absorption_in_layer(
                [brc_conc_l[i]], optical_parameters, layer_thickness_l[i], 'brc'
            )
            absc_brc.append(absorption)

        # Agregar los resultados de la capa actual al DataFrame
        absorption_results[f"brc_{layer}"] = absc_brc

    return absorption_results


#doing for other aerosol types
def aerosol_extraction(rh_df, thickness_df, conc, layer_number, aero_name):
    """
    Extrae los datos de una capa específica basada en el número (1, 2, 3, ...).
    
    Args:
    rh_df (pd.DataFrame): DataFrame con columnas de capas para RH.
    thickness_df (pd.DataFrame): DataFrame con columnas de capas para Thickness.
    conc (pd.DataFrame): DataFrame con columnas de capas para Concentration.
    layer_number (int): Número de la capa a extraer (1 para "layer_1", 2 para "layer_2", etc.).
    aero_name (str): Nombre del aerosol a extraer (por ejemplo, 'SO4').
    
    Returns:
    tuple: Tres Series renombradas ('RH', 'Thickness', 'Concentration') correspondientes al aerosol especificado.
    """
    # Convertir el número de capa en el nombre correspondiente
    layer = f"layer_{layer_number}"
    
    # Validar que la capa exista en los DataFrames
    if layer not in rh_df.columns or layer not in thickness_df.columns or layer not in conc.columns:
        raise ValueError(f"Layer {layer} is not in all dataframes.")
    
    # Extraer y renombrar las columnas, y redondear a dos decimales
    rh = rh_df[layer].rename('RH')
    thickness = thickness_df[layer].rename('Thickness')
    concentration = conc[layer].rename(aero_name)
    
    return rh, thickness, concentration

def aerosol_total_abs(rh_df, thickness_df, aero_conc_df, layers, aero_name, **kwargs):
    """
    Calcula los valores de absorción para todas las capas y devuelve un DataFrame con el resultado.

    Args:
        rh (pd.DataFrame): DataFrame con columnas de capas para RH.
        layer_thickness (pd.DataFrame): DataFrame con columnas de capas para Thickness.
        brc_conc (pd.DataFrame): DataFrame con columnas de capas para aero Concentration.
        layers (int): Número de capas (por ejemplo, 24).

    Returns:
        pd.DataFrame: DataFrame con una columna por capa y el índice igual al de rh.
    """
    # DataFrame para almacenar resultados
    absorption_results = pd.DataFrame(index=rh_df.index)

    for layer in range(1, layers + 1):
        # Extraer los datos para la capa actual
        rh_l, layer_thickness_l, aero_conc_l = aerosol_extraction(rh_df, thickness_df, aero_conc_df, layer, aero_name)

        # Calcular absorción para cada fila
        absc_aero = []
        for i in range(len(rh_l)):
            # Calcular parámetros ópticos para el valor actual de rh
            optical_parameters = ao(
                np.array([rh_l[i]]), [aero_name], wvln=kwargs['wavelenght'])
            # Calcular absorción para el valor actual de brc y thickness
            absorption = total_absorption_in_layer(
                                                    [aero_conc_l[i]], 
                                                     optical_parameters,
                                                     layer_thickness_l[i], 
                                                     aero_name)
            absc_aero.append(absorption)

        # Agregar los resultados de la capa actual al DataFrame
        absorption_results[f"{aero_name}_{layer}"] = absc_aero

    return absorption_results
