# This file is a part of PyPartMC licensed under the GNU General Public License v3
# Copyright (C) 2022 University of Illinois Urbana-Champaign
# Authors: https://github.com/open-atmos/PyPartMC-examples/graphs/contributors

import sys
import os
if 'google.colab' in sys.modules:
    !pip --quiet install open-atmos-jupyter-utils
    from open_atmos_jupyter_utils import pip_install_on_colab
    pip_install_on_colab('PyPartMC', 'PySDM')
elif 'JUPYTER_IMAGE' in os.environ and '.arm.gov' in os.environ['JUPYTER_IMAGE']:
    !pip --quiet install PyPartMC PySDM open-atmos-jupyter-utils
    _pypartmc_path = !pip show PyPartMC | fgrep Location | cut -f2 -d' '
    sys.path.extend(_pypartmc_path if _pypartmc_path[0] not in sys.path else [])

from collections import namedtuple
import ipywidgets as widgets
import numpy as np
from PySDM.environments import Box
from PySDM import Builder
from PySDM.backends import CPU
from PySDM.initialisation import equilibrate_wet_radii
from matplotlib import pyplot
from IPython.display import display, clear_output
from open_atmos_jupyter_utils import show_plot
from PyPartMC import si
import PyPartMC as ppmc

sliders = {
    'n_part': {
        'widget': widgets.IntSlider(min=8, max=128, value=64),
        'label': 'Number of Computational Particles [#]'
    },
    'temp': {
        'widget': widgets.FloatSlider(min=250, max=350),
        'label': 'Temperature [K]'
    },
    'RH_percent': {
        'widget': widgets.FloatSlider(min=0, max=100.0, value=55, readout_format='.1f'),
        'label': 'Relative Humidity [%]'
    },
    'kappa': {
        'widget': widgets.FloatSlider(min=0, max=2, value=1),
        'label': 'Kappa []'
    },
    'mode_1_n_per_cc': {
        'widget': widgets.IntSlider(min=0, max=100000, value=50000),
        'label': 'Mode 1 Number [#/cc]'
    },
    'mode_1_gsd': {
        'widget': widgets.FloatSlider(min=1.1, max=5, value=1.3),
        'label' : 'Mode 1 Geometric Standard Deviation'
    },
    'mode_1_gm_microns': {
        'widget': widgets.FloatSlider(min=0.001, max=10, value=0.9, readout_format='.3f'),
        'label': 'Mode 1 Geometric Mean Diameter [microns]'
    },
    'mode_2_n_per_cc' : {
        'widget': widgets.IntSlider(min=0, max=100000, value=80000),
        'label': 'Mode 2 Number [#/cc]'
    },
    'mode_2_gsd': {
        'widget': widgets.FloatSlider(min=1.1, max=5, value=2),
        'label': 'Mode 2 Geometric Standard Deviation'
    },
    'mode_2_gm_microns': {
        'widget': widgets.FloatSlider(min=0.001, max=10, value=5.8, readout_format='.3f'),
        'label':'Mode 2 Geometric Mean Diameter [microns]'
    },
    'log_base': {
        'widget': widgets.RadioButtons(options=['10','e','none'], value='10'),
        'label': 'Distribution function logarithm base'
    }
}

def dn_ddp(diam, num, geom_mean, geom_stdev):
    return ((num / (np.sqrt(2*np.pi)*diam*np.log(geom_stdev))) *
        np.exp(-(np.log(diam) - np.log(geom_mean))** 2 / (2*np.log(geom_stdev)** 2)))

def dn_dlogdp(diam, num, geom_mean, geom_stdev):
    return np.log(10) * diam * dn_ddp(diam, num, geom_mean, geom_stdev)

def dn_dlndp(diam, num, geom_mean, geom_stdev):
    return diam * dn_ddp(diam, num, geom_mean, geom_stdev)

log_bases = {
    '10': {
        'func': dn_dlogdp,
        'y_unit': 1/si.cm**3,
        'y_label': '$dN/dlogD_p$ [$cm^{-3}$]'
    },
    'e': {
        'func': dn_dlndp,
        'y_unit': 1/si.cm**3,
        'y_label': '$dN/dlnD_p$ [$cm^{-3}$]'
    },
    'none': {
        'func': dn_ddp,
        'y_unit': (1/si.um)*(1/(si.cm**3)),
        'y_label': r'$dN/dD_p$ [$\mu m^{-1} cm^{-3}$]'
    }
}

linestyles = {
    'PyPartMC': 'dashed',
    'PySDM': 'solid'
}

def pypartmc(dry_diameters, temp, rel_humid, kappa):
    env_state = ppmc.EnvState({
        'rel_humidity': rel_humid,
        'latitude': 0.,
        'longitude': 0.,
        'altitude': 0.,
        'start_time': 0.,
        'start_day': 0
    })
    
    env_state.set_temperature(temp)
    
    composition = (
        {"H2O": [1000 * si.kg / si.m**3, 0, 18e-3 * si.kg / si.mol, 0]},
        {"XXX": [np.nan * si.kg / si.m**3, 0, np.nan * si.kg / si.mol, kappa]}
    )
    
    aero_data = ppmc.AeroData(composition)
    
    dry_volumes = (np.pi / 6) * dry_diameters**3
    aero_particles = [
        ppmc.AeroParticle(aero_data, np.array([0, 1])*volume) for volume in dry_volumes
    ]

    for aero_particle in aero_particles:
        ppmc.condense_equilib_particle(env_state, aero_data, aero_particle)

    wet_volumes = [np.sum(particle.volumes) for particle in aero_particles]
    wet_diameters = ((6 / np.pi) * np.asarray(wet_volumes))**(1/3)

    return wet_diameters

def pysdm(dry_diameters, temp, rel_humid, kappa):
    r_dry = dry_diameters / 2
    builder = Builder(n_sd=0, backend=CPU())
    environment = Box(dt=np.nan, dv=np.nan)
    environment.register(builder)
    environment['T'] = temp
    environment['RH'] = rel_humid
    kappa_times_dry_volume = kappa * (np.pi / 6) * dry_diameters**3
    return 2 * equilibrate_wet_radii(
        r_dry=r_dry,
        environment=environment,
        kappa_times_dry_volume=kappa_times_dry_volume
    )

models = {
    'PyPartMC': pypartmc,
    'PySDM': pysdm
}