"""
Aerosol Optical Properties Module.

This module computes optical properties of atmospheric aerosols, including extinction, scattering,
and asymmetry parameters. It supports calculations for various aerosol types, with specific
functionality for brown carbon (BrC) aerosols. The module is designed for atmospheric science
research, particularly in climate modeling and environmental studies.

Key Features:
- Calculation of optical properties for multiple aerosol types.
- Support for brown carbon (BrC) refractive index variants.
- Integration with NetCDF datasets for optical and physical property data.
- Compatibility with the `physics` module for comprehensive aerosol analysis.

Functions:
- `calculate_aerosol_optical_properties`: Computes optical properties for given aerosol types
  and conditions.
"""
from typing import Dict, List, Optional, Tuple
import warnings
import os 
import sys 
import numpy as np
import netCDF4 as nc
from .physics import get_aerosol_physical_properties, physical_properties_single, physical_properties_multi, physical_properties 

def calculate_aerosol_optical_properties(
    relative_humidity: np.ndarray,
    wavelength_nm: float,
    aerosol_type_ids: List[str],
    dataset_path: str = '../data/nkxrnel.nc',
    **brc_refractive_indices
    ) -> Dict[str, Dict[str, float]]:

    """
    Calculate optical properties for different aerosols, and for BrC with single/multiple refractive index variants.

    Parameters: 
        relative_humidity (np.ndarray): Array of relative humidity values (0-1)
        aerosol_type_ids (List[str]): List of aerosol type identifiers as is in physics module.
        wavelength_nm (float): Wavelength in nanometers.
        dataset_path (str): Path to NetCDF dataset with optical data.
        brc_refractive_indices: Key-value pairs for BrC refractive index variants
            (e.g., ri_brc=0.03, ri_brc_strong=0.12, ri_brc_bleached=0.005).

    Returns: 
        Dict[str, Dict[str, float]]: Optical properties for each aerosol type.
            Keys are aerosol type IDs, and values are dictionaries containing:
                - qext: Extinction efficiency (unitless).
                - qsca: Scattering efficiency (unitless).
                - asym: Asymmetry factor (unitless).
                - dens: Mass density (dry species) [g/cm^3].
                - reff: Effective radius (wet species) [um].
                - vphi: Wet-to-dry volume ratio (unitless).
    """
    # Validate BrC parameters
    brc_param = validate_brc_parameters(aerosol_type_ids, brc_refractive_indices)

    # Log which BrC refractive index is being used
    if brc_param and brc_refractive_indices:
        warnings.warn(
            f"Using refractive index with value {brc_refractive_indices}",
            UserWarning
        )

    #convert wavelength from nm to um
    wavelength_um = wavelength_nm / 1e3

    #load optical dataset
    optical_data = load_optical_dataset(dataset_path)

    #load physical properties
    physical_properties = load_physical_properties(brc_refractive_indices)

    # Convert aerosol type IDs to uppercase
    aerosol_type_ids = [t.upper() for t in aerosol_type_ids]

     # Calculate optical properties
    optical_properties = compute_optical_properties(
        relative_humidity, wavelength_um, aerosol_type_ids, optical_data, physical_properties
    )

    return optical_properties



def validate_brc_parameters(
    aerosol_type_ids: List[str],
    brc_refractive_indices: Dict[str, float]
    ) -> Optional[str]:
    """
    Validate BrC refractive indices based on aerosol types.

    Parameters:
        aerosol_type_ids (List[str]): List of aerosol type identifiers
        brc_refractive_indices (Dict[str, float]): Provided BrC refractive indices

    Returns:
        Optional[str]: Name of the required BrC parameter if needed, otherwise None

    Raises:
        ValueError: If required BrC parameters are missing, incorrect, or conflicting
    """
    # Normalize aerosol types to uppercase
    aerosol_types = _normalize_aerosol_types(aerosol_type_ids)

    # Define BrC-related aerosol types and their required parameters
    brc_aerosols = {
        'BRC': 'ri_brc',
        'POABRC': 'ri_brc_strong',
        'SOABRC': 'ri_brc_bleached'
    }

    # Find which BrC-related aerosols are present
    present_brc_types = aerosol_types & brc_aerosols.keys()

    # If no BrC-related aerosols, no parameters are needed
    if not present_brc_types:
        if brc_refractive_indices:
            raise ValueError(
                "BrC refractive indices provided but no BrC-related aerosols found. "
                f"Aerosol types: {', '.join(aerosol_types)}"
            )
        return None

    # Ensure only one BrC-related aerosol type is present
    if len(present_brc_types) > 1:
        raise ValueError(
            "Only one BrC-related aerosol type can be processed at a time. "
            f"Found: {', '.join(present_brc_types)}"
        )

    # Get the single BrC-related aerosol type
    brc_type = present_brc_types.pop()
    required_param = brc_aerosols[brc_type]

    # Validate that the required parameter is provided
    if required_param not in brc_refractive_indices:
        raise ValueError(
            f"Missing required BrC refractive index: {required_param}. "
            f"Required for aerosol type: {brc_type}"
        )

    # Ensure no extra BrC parameters are provided
    extra_params = set(brc_refractive_indices.keys()) - {required_param}
    if extra_params:
        raise ValueError(
            f"Unexpected BrC refractive indices provided: {', '.join(extra_params)}. "
            f"Only {required_param} is required for aerosol type: {brc_type}"
        )

    return required_param


def load_optical_dataset(dataset_path: str) -> Dict[str, np.ndarray]:
    """
    Load optical data from a NetCDF file.

    Parameters:
        dataset_path (str): Path to the NetCDF file.

    Returns:
        Dict[str, np.ndarray]: Dictionary containing optical data.
    """
    with nc.Dataset(dataset_path) as ds:
        optical_data = {
            'ngrd': ds.variables['NGRD'][:],
            'kgrd': ds.variables['KGRD'][:],
            'xvnb': ds.variables['XVNB'][:],
            'qext': ds.variables['QEXT'][:, :, :],
            'qsca': ds.variables['QSCA'][:, :, :],
            'asym': ds.variables['ASYM'][:, :, :],
        }
    return optical_data

def load_physical_properties(brc_refractive_indices: Dict[str, float]) -> Dict[str, np.ndarray]:
    """
    Load physical properties based on BrC refractive index variants.

    Parameters:
        brc_refractive_indices (Dict[str, float]): Key-value pairs for BrC refractive index variants.

    Returns:
        Dict[str, np.ndarray]: Physical properties of aerosols.
    """
    if brc_refractive_indices:
        brc_param = list(brc_refractive_indices.keys())[0]
        brc_value = brc_refractive_indices[brc_param]
        return get_aerosol_physical_properties(**{brc_param: brc_value})
    else:
        return get_aerosol_physical_properties()

def validate_brc_parameters(aerosol_type_ids: List[str], brc_refractive_indices: Dict[str, float]) -> Optional[str]:
    """
    Validate if BrC refractive indices are provided for non-BrC aerosols.

    Parameters:
        aerosol_type_ids (List[str]): List of aerosol type identifiers.
        brc_refractive_indices (Dict[str, float]): Key-value pairs for BrC refractive index variants.

    Returns:
        Optional[str]: The BrC parameter key if BrC aerosols are present, otherwise None.

    Raises:
        Warning: If BrC refractive indices are provided but no BrC aerosols are requested.
    """
    # Define BrC aerosol types (case-insensitive)
    brc_aerosol_types = {"brc", "poabrc", "soabrc"}

    # Check if any requested aerosol type is BrC
    has_brc = any(t.lower() in brc_aerosol_types for t in aerosol_type_ids)

    # If BrC refractive indices are provided but no BrC aerosols are requested, raise a warning
    if brc_refractive_indices and not has_brc:
        warnings.warn(
            "BrC refractive indices were provided, but no BrC aerosols are requested. "
            "The provided BrC refractive indices will not be used.",
            UserWarning
        )
        return None

    # Return the BrC parameter key if BrC aerosols are present
    if has_brc and brc_refractive_indices:
        return list(brc_refractive_indices.keys())[0]

    return None

def compute_optical_properties(
    relative_humidity: np.ndarray,
    wavelength_um: float,
    aerosol_type_ids: List[str],
    optical_data: Dict[str, np.ndarray],
    physical_properties: Dict[str, np.ndarray]
    ) -> Dict[str, Dict[str, float]]:
    """
    Compute optical properties for each aerosol type.

    Parameters:
        relative_humidity (np.ndarray): Array of relative humidity values.
        wavelength_um (float): Wavelength in micrometers.
        aerosol_type_ids (List[str]): List of aerosol type identifiers.
        optical_data (Dict[str, np.ndarray]): Optical data from the NetCDF file.
        physical_properties (Dict[str, np.ndarray]): Physical properties of aerosols.

    Returns:
        Dict[str, Dict[str, float]]: Optical properties for each aerosol type.
    """
    
    # Extract optical data
    ngrd = optical_data['ngrd']
    kgrd = optical_data['kgrd']
    xvnb = optical_data['xvnb']
    qext = optical_data['qext']
    qsca = optical_data['qsca']
    asym = optical_data['asym']

    # Compute effective radius
    rvnb = xvnb * wavelength_um / (2.0 * np.pi)

    # Relative humidity levels
    hlev = np.array([0.0, 0.50, 0.70, 0.80, 0.90, 0.95, 0.99])

    nwat = physical_properties['WATE'][0]
    kwat = physical_properties['WATE'][1]
    # Initialize output dictionaries
    optical_properties = {}

    eout, sout, gout, dout, rout, vout = [], [], [], [], [], []
    for i, h in enumerate(relative_humidity):
        eens, sens, gens, dens, rens, vens = [], [], [], [], [], []
        for aero in aerosol_type_ids:
            try:
                rgeo, sdev, rmin, rmax, real, imag = physical_properties[aero][:6]
                dens.append(physical_properties[aero][6])
                alfa, beta = physical_properties[aero][7:9]

                r0, r1, rgrd, logn, drrr = dry_geometric_integration(rvnb, rmin, rmax, rgeo, sdev)
                alfh, beth = hygroscopic_growth(h.flatten(),hlev, alfa, beta, i)

                # wet geometric integration
                enqe, enqs, enqg, rens, vens, n0, n1, k0, k1 = wet_optical_properties(rgeo, sdev, rmin, rmax, rens, vens, drrr,\
                    real, imag, nwat, kwat, qext, qsca, asym, alfh, beth, rvnb, kgrd, ngrd)

                # kernel grid interpolation
                ddnk = (ngrd[n1] - ngrd[n0])*(kgrd[k1] - kgrd[k0])
                wt00 = (ngrd[n1] - real)*(kgrd[k1] - imag)/ddnk
                wt01 = (ngrd[n1] - real)*(imag - kgrd[k0])/ddnk
                wt10 = (real - ngrd[n0])*(kgrd[k1] - imag)/ddnk
                wt11 = (real - ngrd[n0])*(imag - kgrd[k0])/ddnk
                wgth = np.array([wt00, wt01, wt10, wt11])

                eens.append(np.average(enqe, weights=wgth))
                sens.append(np.average(enqs, weights=wgth))
                gens.append(np.average(enqg, weights=wgth))

            except Exception as e:
                print(f"Error in {h}, aero={aero}: {e}")
        
        eout.append(eens)
        sout.append(sens)
        gout.append(gens)
        dout.append(dens)
        rout.append(rens)
        vout.append(vens)

    for j in range(len(aerosol_type_ids)):
        optical_properties[aerosol_type_ids[j].lower()] = {
            'qext': eout[0][j],
            'qsca': sout[0][j],
            'asym': gout[0][j],
            'dens': dout[0][j],
            'reff': rout[0][j],
            'vphi': vout[0][j]
        }
    return optical_properties


def dry_geometric_integration(
    rvnb: np.ndarray,
    rmin: float,
    rmax: float,
    rgeo: float,
    sdev: float
    ) -> Tuple[int, int, np.ndarray, np.ndarray, float]:
    """
    Perform geometric integration over a range of radii.

    Parameters:
        rvnb (np.ndarray): Effective radius grid.
        rmin (float): Minimum radius.
        rmax (float): Maximum radius.
        rgeo (float): Geometric mean radius.
        sdev (float): Geometric standard deviation.

    Returns:
        Tuple[int, int, np.ndarray, np.ndarray, float]:
            - r0: Index of the minimum radius.
            - r1: Index of the maximum radius.
            - rgrd: Radius grid within the range [rmin, rmax].
            - logn: Logarithmic normal distribution values.
            - drrr: Integrated value over the radius grid.
    """
    r0 = np.abs(rvnb - rmin).argmin()
    r1 = np.abs(rvnb - rmax).argmin() + 1
    rgrd = rvnb[r0:r1]
    rrri = rvnb[r0:r1] ** 3
    logn = (1.0/(np.sqrt(2.0*np.pi)*np.log(sdev)))*(1.0/rvnb[r0:r1])*np.exp(-0.5*(np.log(rvnb[r0:r1]/rgeo)/np.log(sdev))**2)
    drrr = np.trapz(rrri * logn, x=rgrd)

    return r0, r1, rgrd, logn, drrr

def hygroscopic_growth(
    hamb: np.ndarray,
    hlev: np.ndarray,
    alfa: np.ndarray,
    beta: np.ndarray,
    h: int
    ) -> Tuple[float, float]:
    """
    Compute hygroscopic growth factors for a given relative humidity value.

    Parameters:
        hamb (np.ndarray): Array of relative humidity values.
        hlev (np.ndarray): Array of relative humidity levels.
        alfa (np.ndarray): Array of alpha values.
        beta (np.ndarray): Array of beta values.
        h (int): Index of the current relative humidity value.

    Returns:
        Tuple[float, float]: Hygroscopic growth factors (alfh, beth).
    """
    h0 = np.searchsorted(hlev, hamb[h]) - 1
    h1 = np.searchsorted(hlev, hamb[h])
    ma = (alfa[h1] - alfa[h0]) / (hlev[h1] - hlev[h0])
    mb = (beta[h1] - beta[h0]) / (hlev[h1] - hlev[h0])
    alfh = alfa[h0] + ma * (hamb[h] - hlev[h0])
    beth = beta[h0] + mb * (hamb[h] - hlev[h0])

    return alfh, beth

def wet_optical_properties(
    rgeo: float,
    sdev: float,
    rmin: float,
    rmax: float,
    rens: List[float],
    vens: List[float],
    drrr: float,
    real: float,
    imag: float,
    nwat: float,
    kwat: float,
    qext: np.ndarray,
    qsca: np.ndarray,
    asym: np.ndarray,
    alfh: float,
    beth: float,
    rvnb: np.ndarray,
    kgrd: np.ndarray,
    ngrd: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[float], List[float], int, int, int, int]:
    """
    Compute wet optical properties for a given aerosol type.

    Parameters:
        rgeo (float): Geometric mean radius.
        sdev (float): Geometric standard deviation.
        rmin (float): Minimum radius.
        rmax (float): Maximum radius.
        rens (List[float]): List to store effective radius values.
        vens (List[float]): List to store wet-to-dry volume ratios.
        drrr (float): Dry geometric integration result.
        real (float): Real part of the refractive index.
        imag (float): Imaginary part of the refractive index.
        nwat (float): Real part of the refractive index of water.
        kwat (float): Imaginary part of the refractive index of water.
        qext (np.ndarray): Extinction efficiency grid.
        qsca (np.ndarray): Scattering efficiency grid.
        asym (np.ndarray): Asymmetry factor grid.
        alfh (float): Alpha growth factor.
        beth (float): Beta growth factor.
        rvnb (np.ndarray): Effective radius grid.
        kgrd (np.ndarray): Imaginary refractive index grid.
        ngrd (np.ndarray): Real refractive index grid.

    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray, List[float], List[float], int, int, int, int]:
            - enqe: Extinction efficiency values.
            - enqs: Scattering efficiency values.
            - enqg: Asymmetry factor values.
            - rens: Updated list of effective radius values.
            - vens: Updated list of wet-to-dry volume ratios.
            - n0: Lower index for real refractive index grid.
            - n1: Upper index for real refractive index grid.
            - k0: Lower index for imaginary refractive index grid.
            - k1: Upper index for imaginary refractive index grid.
    """
    # wet geometric integration
    rgeo = alfh*rgeo**beth
    sdev = sdev**beth
    rmin = alfh*rmin**beth
    rmax = alfh*rmax**beth

    r0 = np.abs(rvnb - rmin).argmin()
    r1 = np.abs(rvnb - rmax).argmin() + 1
    rgrd = rvnb[r0:r1]
    rrii = rvnb[r0:r1]**2
    prri = np.pi*rvnb[r0:r1]**2
    rrri = rvnb[r0:r1]**3
    logn = (1.0/(np.sqrt(2.0*np.pi)*np.log(sdev)))*(1.0/rvnb[r0:r1])*np.exp(-0.5*(np.log(rvnb[r0:r1]/rgeo)/np.log(sdev))**2)
    wwrr = np.trapz(rrii*logn, x=rgrd)
    area = np.trapz(prri*logn, x=rgrd)
    wrrr = np.trapz(rrri*logn, x=rgrd)

    # wet geometric ratios
    rens.append(wrrr/wwrr)
    vens.append(wrrr/drrr)
    vfra = drrr/wrrr

    # wet optical integration
    real = real*vfra + nwat*(1.0 - vfra)
    imag = imag*vfra + kwat*(1.0 - vfra)

    n0 = np.searchsorted(ngrd, real) - 1
    n1 = np.searchsorted(ngrd, real)
    k0 = np.searchsorted(kgrd, imag) - 1
    k1 = np.searchsorted(kgrd, imag)

    qe00, qe01, qe10, qe11 = qext[n0,k0,r0:r1], qext[n0,k1,r0:r1], qext[n1,k0,r0:r1], qext[n1,k1,r0:r1]
    qs00, qs01, qs10, qs11 = qsca[n0,k0,r0:r1], qsca[n0,k1,r0:r1], qsca[n1,k0,r0:r1], qsca[n1,k1,r0:r1]
    qg00, qg01, qg10, qg11 = asym[n0,k0,r0:r1], asym[n0,k1,r0:r1], asym[n1,k0,r0:r1], asym[n1,k1,r0:r1]

    enqe = np.array([np.trapz(qe00*prri*logn, x=rgrd)/area, np.trapz(qe01*prri*logn, x=rgrd)/area,
                        np.trapz(qe10*prri*logn, x=rgrd)/area, np.trapz(qe11*prri*logn, x=rgrd)/area])
    enqs = np.array([np.trapz(qs00*prri*logn, x=rgrd)/area, np.trapz(qs01*prri*logn, x=rgrd)/area,
                        np.trapz(qs10*prri*logn, x=rgrd)/area, np.trapz(qs11*prri*logn, x=rgrd)/area])
    enqg = np.array([np.trapz(qg00*qs00*prri*logn, x=rgrd)/(enqs[0]*area), np.trapz(qg01*qs01*prri*logn, x=rgrd)/(enqs[1]*area),
                        np.trapz(qg10*qs10*prri*logn, x=rgrd)/(enqs[2]*area), np.trapz(qg11*qs11*prri*logn, x=rgrd)/(enqs[3]*area)])
    return enqe, enqs, enqg, rens, vens, n0, n1, k0, k1
