"""
Aerosol Optical Properties Analysis Module

This module is dedicated to the computation of optical properties of atmospheric aerosols. 
It leverages detailed physical properties of various aerosol types to calculate key optical 
metrics like extinction, scattering, and asymmetry parameters. The module is designed to 
support atmospheric science research, particularly in the fields of climate modeling and 
environmental studies where understanding aerosol behavior is crucial.

The module includes functions to calculate optical properties for a broad range of aerosol types,
including specific computations for brown carbon aerosols. It utilizes data from NetCDF files 
and integrates with other modules, such as `physics`, for comprehensive property analysis. 
This approach allows for flexible, accurate calculations that can be tailored to specific research needs.

"""

import numpy as np
import netCDF4 as nc
from .physics import physical_properties_single, physical_properties_multi, physical_properties
import os 
import sys 

def aerosol_optics_single(hamb, tids, wvln, ri_brc):
    """
    Calculate the optical properties of aerosols for a single aerosol type.

    Parameters:
    hamb (numpy.ndarray): The relative humidity values for the aerosol.
    tids (list): A list of aerosol type identifiers.
    wvln (float): The wavelength of the light in nm.
    ri_brc (float): The real part of the refractive index of BrC.

    Returns:
    dict: A dictionary containing the optical properties of the aerosols.
    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).
    """
    #should be one hamb value (relative humidity) for each hour

    #if tids list is in lower case convert to upper
    tids = [t.upper() for t in tids]
    wvln =  wvln/1e3
    phys = physical_properties_single(ri_brc)

    with nc.Dataset('../data/nkxrnel.nc') as ds:
        ngrd = ds.variables['NGRD'][:]
        kgrd = ds.variables['KGRD'][:]
        xvnb = ds.variables['XVNB'][:]
        qext = ds.variables['QEXT'][:,:,:]
        qsca = ds.variables['QSCA'][:,:,:]
        asym = ds.variables['ASYM'][:,:,:]

    rvnb = xvnb * wvln / (2.0 * np.pi)
    hlev = np.array([0.0, 0.50, 0.70, 0.80, 0.90, 0.95, 0.99])

    nwat = phys['WATE'][0]
    kwat = phys['WATE'][1]

    hamb = hamb.flatten()
    H = hamb.shape[0]
    I = len(tids)

    eout, sout, gout, dout, rout, vout = [], [], [], [], [], []
    for h in np.arange(H):
        eens, sens, gens, dens, rens, vens = [], [], [], [], [], []
        for i in np.arange(I):
            rgeo = phys[tids[i]][0]
            sdev = phys[tids[i]][1]
            rmin = phys[tids[i]][2]
            rmax = phys[tids[i]][3]
            real = phys[tids[i]][4]
            imag = phys[tids[i]][5]
            dens.append(phys[tids[i]][6])
            alfa = phys[tids[i]][7]
            beta = phys[tids[i]][8]

            # dry geometric integration
            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)

            # hygroscopic growth factors
            h0 = np.searchsorted(hlev, hamb[h]) - 1
            h1 = np.searchsorted(hlev, hamb[h])
            MA = (alfa[h1] - alfa[h0])/(hlev[h1] - hlev[h0])
            alfh = alfa[h0] + MA*(hamb[h] - hlev[h0])
            MB = (beta[h1] - beta[h0])/(hlev[h1] - hlev[h0])
            beth = beta[h0] + MB*(hamb[h] - hlev[h0])

            # 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)])

            # 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))

        eout.append(eens)
        sout.append(sens)
        gout.append(gens)
        dout.append(dens)
        rout.append(rens)
        vout.append(vens)
        
    optical_properties = {}
    for t in range(len(tids)):
        optical_properties[tids[t].lower()] = {
            'qext': eout[0][t],
            'qsca': sout[0][t],
            'asym': gout[0][t],
            'dens': dout[0][t],
            'reff': rout[0][t],
            'vphi': vout[0][t]
        }
    return optical_properties

def aerosol_optics(hamb, tids, wvln):
    """
    Calculate the optical properties of aerosols for a single aerosol type.

    Parameters:
    hamb (numpy.ndarray): The relative humidity values for the aerosol.
    tids (list): A list of aerosol type identifiers.
    wvln (float): The wavelength of the light in nm.
    ri_brc (float): The real part of the refractive index of BrC.

    Returns:
    dict: A dictionary containing the optical properties of the aerosols.
    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).
    """
    #should be one hamb value (relative humidity) for each hour

    #if tids list is in lower case convert to upper
    tids = [t.upper() for t in tids]
    wvln =  wvln/1e3 #convert nm to um
    phys = physical_properties()

    with nc.Dataset('../data/nkxrnel.nc') as ds:
        ngrd = ds.variables['NGRD'][:]
        kgrd = ds.variables['KGRD'][:]
        xvnb = ds.variables['XVNB'][:]
        qext = ds.variables['QEXT'][:,:,:]
        qsca = ds.variables['QSCA'][:,:,:]
        asym = ds.variables['ASYM'][:,:,:]

    rvnb = xvnb * wvln / (2.0 * np.pi)
    hlev = np.array([0.0, 0.50, 0.70, 0.80, 0.90, 0.95, 0.99])

    nwat = phys['WATE'][0]
    kwat = phys['WATE'][1]

    hamb = hamb.flatten()
    H = hamb.shape[0]
    I = len(tids)

    eout, sout, gout, dout, rout, vout = [], [], [], [], [], []
    for h in np.arange(H):
        eens, sens, gens, dens, rens, vens = [], [], [], [], [], []
        for i in np.arange(I):
            rgeo = phys[tids[i]][0]
            sdev = phys[tids[i]][1]
            rmin = phys[tids[i]][2]
            rmax = phys[tids[i]][3]
            real = phys[tids[i]][4]
            imag = phys[tids[i]][5]
            dens.append(phys[tids[i]][6])
            alfa = phys[tids[i]][7]
            beta = phys[tids[i]][8]

            # dry geometric integration
            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)

            # hygroscopic growth factors
            h0 = np.searchsorted(hlev, hamb[h]) - 1
            h1 = np.searchsorted(hlev, hamb[h])
            MA = (alfa[h1] - alfa[h0])/(hlev[h1] - hlev[h0])
            alfh = alfa[h0] + MA*(hamb[h] - hlev[h0])
            MB = (beta[h1] - beta[h0])/(hlev[h1] - hlev[h0])
            beth = beta[h0] + MB*(hamb[h] - hlev[h0])

            # 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)])

            # 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))

        eout.append(eens)
        sout.append(sens)
        gout.append(gens)
        dout.append(dens)
        rout.append(rens)
        vout.append(vens)
        
    optical_properties = {}
    for t in range(len(tids)):
        optical_properties[tids[t].lower()] = {
            'qext': eout[0][t],
            'qsca': sout[0][t],
            'asym': gout[0][t],
            'dens': dout[0][t],
            'reff': rout[0][t],
            'vphi': vout[0][t]
        }
    return optical_properties
    