import pandas as pd
import numpy as np

from typing import Dict, Union
from typing import Tuple
from typing import Optional

from .physics_aerosol_optical_properties import calculate_aerosol_optical_properties
from .extract_layer_data_by_number import extract_layer_data, extract_target_layer_data
from .aerosol_optical_functions import (
    total_absorption_in_layer,
    absorption_coefficient,
    aerosol_optical_depth,
)

def calculate_aaod_by_layer(
    rh_df: pd.DataFrame,
    thickness_df: pd.DataFrame,
    model_conc_df: pd.DataFrame,
    layers: int,
    wavelength: float,
    aerosol: str,
    time_resolution: Optional[str
    ] = None,
    **brc_refractive_indices: float,
    ) -> pd.DataFrame:
    """
    Calculate the Absorption Aerosol Optical Depth (AAOD) for each layer.

    Args:
        rh_df (pd.DataFrame): DataFrame with columns for relative humidity.
        thickness_df (pd.DataFrame): DataFrame with columns for thickness.
        model_conc_df (pd.DataFrame): DataFrame with columns for aerosol concentration.
        layers (int): Number of layers (e.g., 24).
        wavelength (float): Wavelength of the incident light (nm).
        aerosol (str): Name of the aerosol to calculate the AAOD for.
        ri_brc (Optional[float]): Refractive index for BRC. Defaults to None.
        ri_brc_strong (Optional[float]): Refractive index for strong BRC. Defaults to None.
        ri_brc_bleached (Optional[float]): Refractive index for bleached BRC. Defaults to None.

    Returns:
        pd.DataFrame: DataFrame with one column per layer and the index matching that of rh_df.
    """
    if aerosol.lower() == 'brc':
        required_ri = 'ri_brc'
    elif aerosol.lower() == 'poabrc':
        required_ri = 'ri_brc_strong'
    elif aerosol.lower() == 'soabrc':
        required_ri = 'ri_brc_bleached'
    else:
        required_ri = None

    # Validate that the required refractive index is provided
    if required_ri and required_ri not in brc_refractive_indices:
        raise ValueError(f"The refractive index '{required_ri}' is required for aerosol '{aerosol}'.")


    # DataFrame to store results
    aaod_results = pd.DataFrame()

    for layer in range(1, layers + 1):
        # Extract data for the current layer
        rh_l, thickness_l, aero_conc_l = extract_layer_data(rh_df, thickness_df, model_conc_df, layer, aerosol, time_resolution)
        # Calculate absorption for each row
        aero_abs= []
        for i in range(len(rh_l)):
            # Calculate optical parameters for the current rh value
            optical_parameters = calculate_aerosol_optical_properties(
                        relative_humidity=np.array([rh_l[i]]), 
                        aerosol_type_ids=[aerosol], 
                        wavelength_nm=wavelength, 
                        **brc_refractive_indices
                    )

            # Calculate absorption for the current aerosol and thickness value
            aaod = total_absorption_in_layer(
                model_conc=aero_conc_l[i], 
                optical_parameters=optical_parameters, 
                thickness_layer=thickness_l[i], 
                aerosol=aerosol
            )

            aero_abs.append(aaod)
        # Store the results in the DataFrame
        #set index as rh_l
        aaod_results[f"{aerosol}_{layer}"] = aero_abs
        #set index as rh_l
        aaod_results.index = rh_l.index

    return aaod_results

def calculate_absorption_coefficient(
    rh_df: pd.DataFrame,
    model_conc_df: pd.DataFrame,
    layers: int,
    wavelength: float,
    aerosol: str,
    time_resolution: Optional[str] = None,
    **brc_refractive_indices: float,
    ) -> pd.DataFrame:
    """
    Calculate the Absorption Coefficient for each layer.

    Args:
        rh_df (pd.DataFrame): DataFrame with columns for relative humidity.
        model_conc_df (pd.DataFrame): DataFrame with columns for aerosol concentration.
        layers (int): Number of layers (e.g., 24).
        wavelength (float): Wavelength of the incident light (nm).
        aerosol (str): Name of the aerosol to calculate the AAOD for.
        ri_brc (Optional[float]): Refractive index for BRC. Defaults to None.
        ri_brc_strong (Optional[float]): Refractive index for strong BRC. Defaults to None.
        ri_brc_bleached (Optional[float]): Refractive index for bleached BRC. Defaults to None.

    Returns:
        pd.DataFrame: DataFrame with one column per layer and the index matching that of rh_df.
    """
    if aerosol.lower() == 'brc':
        required_ri = 'ri_brc'
    elif aerosol.lower() == 'poabrc':
        required_ri = 'ri_brc_strong'
    elif aerosol.lower() == 'soabrc':
        required_ri = 'ri_brc_bleached'
    else:
        required_ri = None

    # Validate that the required refractive index is provided
    if required_ri and required_ri not in brc_refractive_indices:
        raise ValueError(f"The refractive index '{required_ri}' is required for aerosol '{aerosol}'.")

    # Determine the layer to calculate (last layer or the only layer)
    target_layer = layers if layers > 1 else 1

    # Extract data for the target layer
    rh_l, aero_conc_l = extract_target_layer_data(rh_df, model_conc_df, target_layer, aerosol, time_resolution)

    # Calculate absorption for each row
    aero_abs = []
    for i in range(len(rh_l)):
        # Calculate optical parameters for the current rh value
        optical_parameters = calculate_aerosol_optical_properties(
            relative_humidity=np.array([rh_l[i]]),
            aerosol_type_ids=[aerosol],
            wavelength_nm=wavelength,
            **brc_refractive_indices
        )

        # Calculate absorption for the current aerosol and thickness value
        abs_coeff = absorption_coefficient(
            model_conc=aero_conc_l[i],
            optical_parameters=optical_parameters,
            aerosol=aerosol
        )

        aero_abs.append(abs_coeff)

    # Store the results in a DataFrame
    abs_coeff_results = pd.DataFrame({f"{aerosol}_{target_layer}": aero_abs}, index=rh_l.index)

    return abs_coeff_results