using Plots, Plots.PlotMeasures, DifferentialMobilityAnalyzers, LinearAlgebra
plotlyjs();

t,p = 295.15, 1e5                                # Temperature [K], Pressure [Pa]
qsa,β = 1.66e-5, 1/10                            # Qsample [m3 s-1], Sample-to-sheath ratio,
r₁,r₂,l = 9.37e-3,1.961e-2,0.44369               # DMA geometry [m]
leff = 13.0                                      # DMA effective diffusion length [m]
m = 3                                            # Upper number of charges
Λ₁ = DMAconfig(t,p,qsa,qsa/β,r₁,r₂,l,leff,:-,m,:cylindrical)  # Specify DMA with negative polarity
Λ₂ = DMAconfig(t,p,qsa,qsa/β,r₁,r₂,l,leff,:+,m,:cylindrical)  # Specify DMA with negative polarity
bins,z₁,z₂ = 512, dtoz(Λ₁,250e-9),dtoz(Λ₁,20e-9) # bins, upper, lower mobility limit
δ₁ = setupDMA(Λ₁, z₁, z₂, bins);                 # Compute matrices
δ₂ = setupDMA(Λ₂, z₁, z₂, bins);                 # Compute matrices

Dm = 100.0                      # Size select 100 nm mobility diameter
zˢ = dtoz(Λ₁, Dm*1e-9);         # Compute corresponding electrical mobility
T = (zˢ,k,Λ,δ) -> δ.Ω(Λ,δ.Z,zˢ/k).*δ.Tc(k,δ.Dp).*δ.Tl(Λ,δ.Dp); # Transfer model through DMA

Ax = [[400., 60., 1.5], [3000., 200., 1.6]]   # Assumed lognormal size distribution
𝕟ᶜⁿ = DMALognormalDistribution(Ax, δ₁)      

# Evaporation function
t = 10.0                     # Time for particles to evaporate [s]
dt = 1e-4                    # Timestep for evaporation - needs to be small for small sizes
g = -3.0e2
D = deepcopy(𝕟ᶜⁿ.Dp)
for i in 1:length(D)
    clock = 0                # Start the clock
    while clock < t
        D[i] = (D[i] < 10) ? D[i]-1e-11 : D[i] + g/D[i]*dt  # evaporate
        clock = clock + dt
    end
end
l  = length(D[D .< 10])      # Check for complete evaporation and fill the arry with dummy values 

D[end-l+1:end] = reverse(exp10.(range(log10(9.8), stop=log10(10), length=l)))
EF = D./𝕟ᶜⁿ.Dp;              # Compute the size depenent evaporation factor

T = (zˢ,k,Λ,δ) -> δ.Ω(Λ,δ.Z,zˢ/k).*δ.Tc(k,δ.Dp).*δ.Tl(Λ,δ.Dp) # Transfer model through DMA
Nkz = (zˢ,k,Λ,δ) -> EF⋅(𝕥=T(zˢ,k,Λ,δ)*𝕟ᶜⁿ)  # Size distribution for particles with charge k
Nz = (zˢ,Λ,δ) -> Σ(k->EF⋅(𝕥=T(zˢ,k,Λ,δ)*𝕟ᶜⁿ),Λ₁.m) # Total size distribution for all charges 1..n
Nkm = (zˢ,k,Λ,δ) -> ((ztod(Λ,1,zˢ)/ztod(Λ,k,zˢ))⋅(EF⋅(𝕥=T(zˢ,k,Λ,δ)*𝕟ᶜⁿ))) # App. Mobility dist k charges
Nm =(zˢ,Λ,δ) -> Σ(k->((ztod(Λ,1,zˢ)/ztod(Λ,k,zˢ))⋅(EF⋅(𝕥=T(zˢ,k,Λ,δ)*𝕟ᶜⁿ))),Λ₁.m); # Total all charges

# Setup the special case
Dm = 100                 # Dry size 100 nm
zˢ = dtoz(Λ₁, Dm*1e-9);  # Compute corresponding electrical mobility
n = 3                    # Charges 1..3

figure("Nimbus Sans L", 2, 5.5, 3.5, 8)
p = Plots.get_color_palette(:auto, default(:bgcolor), 100) # Get a list of colors from the theme palette

# Panel 1: m\Mobility diameter
𝕩 = map(k -> Nkz(zˢ,k,Λ₁,δ₁), 1:3)
𝕪 = Nz(zˢ,Λ₁,δ₁)
p1 = plot([ztod(Λ₁,1,zˢ), ztod(Λ₁,1,zˢ)], [0,200], color = p[1], ls = :dot, label = "D₁(zˢ)", lw = 2)
p1 = plot!([ztod(Λ₁,2,zˢ), ztod(Λ₁,2,zˢ)], [0,200], color = p[2], ls = :dot, label = "D₂(zˢ)", lw = 2)
p1 = plot!([ztod(Λ₁,3,zˢ), ztod(Λ₁,3,zˢ)], [0,200], color = p[3], ls = :dot, label = "D₃(zˢ)", lw = 2)
p1 = plot!(𝕩[1].Dp, 𝕩[1].S, label = "ℕ₁<sup>δ₁</sup>", xaxis = :log10, xlim = (50,500), color = p[1])
p1 = plot!(𝕩[2].Dp, 𝕩[2].S, label = "ℕ₂<sup>δ₁</sup>", color = p[2])
p1 = plot!(𝕩[3].Dp, 𝕩[3].S, label = "ℕ₃<sup>δ₁</sup>", color = p[3])
p1 = plot!(𝕪.Dp, 𝕪.S, color = :black, label = "𝕟ₜ<sup>δ₁</sup>", ls = :dashdot, lw = 1.25, 
    ylabel = "dN/dlnD (cm⁻³)", left_margin = 20px, xlabel = "Mobility Diameter (nm)")

# Panel 2: Apparent +1 Mobility diameter
𝕩 = map(k -> Nkm(zˢ,k,Λ₁,δ₁), 1:3)
𝕪 = Nm(zˢ,Λ₁,δ₁)
p2 = plot([ztod(Λ₁,1,zˢ), ztod(Λ₁,1,zˢ)], [0,200], color = :black, ls = :dot, label = "D₁(zˢ)", lw = 2)
p2 = plot!(𝕩[1].Dp, 𝕩[1].S, xlim = (120, 220), label = "𝕄₁<sup>δ₁</sup>", color = p[1])
p2 = plot!(𝕩[2].Dp, 𝕩[2].S, label = "𝕄₂<sup>δ₁</sup>", color = p[2])
p2 = plot!(𝕩[3].Dp, 𝕩[3].S, label = "𝕄₃<sup>δ₁</sup>", color = p[3])
p2 = plot!(𝕪.Dp, 𝕪.S, color = :black, lw = 1.25, ls = :dashdot,
     label = "𝕞ₜ<sup>δ₁</sup>", xlabel = "Apparent +1 Mobility Diameter (nm)")

# Panel 3: Passage through DMA
𝕩 = map(k -> δ₂.𝐀*Nkz(zˢ,k,Λ₁,δ₁), 1:3)
𝕪 = δ₂.𝐀*Nz(zˢ,Λ₁,δ₁)
p3 = plot(𝕩[1].Dp, 𝕩[1].N, label = "ℕ₁<sup>δ₂</sup>")
p3 = plot!(𝕩[2].Dp, 𝕩[2].N, label = "ℕ₂<sup>δ₂</sup>")
p3 = plot!(𝕩[3].Dp, 𝕩[3].N, label = "ℕ₃<sup>δ₂</sup>")
p3 = plot!(𝕪.Dp, 𝕪.N, xaxis = :log10,  ylim = (-0.1,2.5), xlim = (50,500), color = :black, ls = :dashdot,
            left_margin = 34px, lw = 1.25, label = "𝕟ₜ<sup>δ₂</sup>",
            ylabel = "Concentration (cm⁻³)", xlabel = "Apparent +1 Mobility Diameter (nm)")

# Panel 4: Passage through DMA
𝕩 = map(k -> δ₂.𝐎*Nkm(zˢ,k,Λ₁,δ₁), 1:3)
𝕪 = δ₂.𝐎*Nm(zˢ,Λ₁,δ₁)
p4 = plot(𝕩[1].Dp, 𝕩[1].N, label = "𝕄₁<sup>δ₂</sup>")
p4 = plot!(𝕩[2].Dp, 𝕩[2].N, label = "𝕄₂<sup>δ₂</sup>")
p4 = plot!(𝕩[3].Dp, 𝕩[3].N, label = "𝕄₃<sup>δ₂</sup>")
p4 = plot!(𝕪.Dp, 𝕪.N, xlim = (120,220), xlabel = "Apparent +1 Mobility Diameter (nm)", lw = 1.25, 
    ls = :dashdot, color = :black, left_margin = 30px, ylim = (0,9.5), label = "𝕞ₜ<sup>δ₂</sup>")

plot(p1,p2,p3,p4, layout = grid(2,2), legend=:right, xlim = (23,230),
     xaxis = :log10, top_margin = 35px, right_margin = 20px, lw = 2,
     bottom_margin = 20px, fmt = :svg)