using Plots, Plots.PlotMeasures, DifferentialMobilityAnalyzers, LinearAlgebra, Printf
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 = 2                                            # Upper number of charges
β₁ = 1.0/3.0      # Sample-to-sheath ratio DMA1
β₂ = 1.0/3.0      # Sample-to-sheath ratio DMA2
β₃ = 1.0/10.0     # Sample-to-sheath ratio DMA3
Λ₁ = 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 positive polarity
Λ₃ = DMAconfig(t,p,qsa,qsa/β₃,r₁,r₂,l,leff,:-,m,:cylindrical)  # Specify DMA with negative polarity
bins,z₁,z₂ = 512, dtoz(Λ₁,500e-9),dtoz(Λ₁,20e-9)  # bins, upper, lower mobility limit
δ₁ = setupDMA(Λ₁, z₁, z₂, bins)                   # Compute matrices
δ₂ = setupDMA(Λ₂, z₁, z₂, bins)                   # Compute matrices
δ₃ = setupDMA(Λ₃, z₁, z₂, bins);                  # Compute matrices

Dₘ = 50.0                       # Size select 60 nm mobility diameter
zˢ = dtoz(Λ₁, Dₘ*1e-9);         # Compute corresponding electrical mobility
β₁₂ = map(k->β12(ztod(Λ₁,k,zˢ)*1e-9, ztod(Λ₂,k,zˢ)*1e-9, k,-k),1:m) # rate for +i/-i charges
βᵈ = 3e-5                       # kd is the rate of spontaneous decharge [s-1], number is a guess
f = 2.0^(1.0/3.0)               # diameter shift of dimers (Dₘ^3+Dₘ^3)^(1/3)/Dₘ
t = 8.0                         # time to coagualate 
𝕟ᶜⁿ = DMALognormalDistribution([[1e12, 30., 1.4], [4e12, 80., 1.5]], δ₁);   # Returns a size distribution 

T = (zˢ,k,Λ,δ) -> δ.Ω(Λ,δ.Z,zˢ/k).*δ.Tc(k,δ.Dp).*δ.Tl(Λ,δ.Dp)

𝕟ᵖ = map(k -> T(zˢ,k,Λ₁,δ₁)*𝕟ᶜⁿ,1:m)
𝕟ᵐ = map(k -> T(zˢ,k,Λ₂,δ₂)*𝕟ᶜⁿ,1:m)
𝕟ᵗ = sum(𝕟ᵐ + 𝕟ᵖ);

𝕟ᵈ = βᵈ*t*(𝕟ᵐ[1] + 𝕟ᵖ[1]);
ℂ = map(k -> f⋅((β₁₂[k]*t)*(𝕟ᵐ[k]*𝕟ᵖ[k])), 1:m) 
𝕔ᵗ = sum(ℂ) + 𝕟ᵈ

𝕊  =  map(k -> δ₃.𝐀*ℂ[k],1:m)
𝕤ᵈ =  δ₃.𝐀*𝕟ᵈ                   
𝕤ᵗ =  δ₃.𝐀*𝕔ᵗ;                  

figure("Nimbus Sans L", 2, 4, 3.5, 8)

pre = 1e-10
# Get a list of colors from the theme palette
p = Plots.get_color_palette(:auto, default(:bgcolor), 100)

# Panel 1
p1 = plot(𝕟ᵖ[1].Dp,pre*𝕟ᵐ[1].S, xaxis = :log10, xlim = (40,400), ylabel = "dN/dlnD (m⁻³)", 
    label = "𝕟ᵖ₁ = T₁(δ₁).*𝕟ᶜⁿ", color = p[1])
p1 = plot!(𝕟ᵐ[1].Dp,pre*𝕟ᵖ[1].S, ls = :dot,  label = "𝕟ᵐ₁ = T₁(δ₂)*𝕟ᶜⁿ", color = p[1])
p1 = plot!(𝕟ᵖ[2].Dp,pre*𝕟ᵖ[2].S, ls = :dot, label = "𝕟ᵖ₂ = T₂(δ₁)*𝕟ᶜⁿ", color = p[2])   
p1 = plot!(𝕟ᵐ[2].Dp,pre*𝕟ᵐ[2].S, label = "𝕟ᵐ₂ = T₂(δ₂)*𝕟ᶜⁿ", color = p[2])   
p1 = plot!(𝕟ᵗ.Dp,pre*𝕟ᵗ.S, color = :black, label = "𝕟ᵗ = ∑(Tᵢ * 𝕟ᶜⁿ)", ls = :dashdot)  

# Panel 2
p2 = plot(ℂ[1].Dp,pre*ℂ[1].S, xaxis = :log10, xlim = (40,400), 
    ylabel = "dN/dlnD (m⁻³)", label = "ℂ₁ = f⋅[(β₁₂ᶻ⁽¹ ⁻¹⁾t)*(𝕟ᵖ₁*𝕟ᵐ₁)]")
p2 = plot!(ℂ[2].Dp,pre*ℂ[2].S, label = "ℂ₂ = f⋅[(β₁₂ᶻ⁽² ⁻²⁾t)*(𝕟ᵖ₂*𝕟ᵐ₂)]")   
p2 = plot!(𝕟ᵈ.Dp,pre*𝕟ᵈ.S, label = "𝕟ᵈ = (βᵈ*t).*(𝕟ᵖ₁+𝕟ᵐ₁)", color = p[4])   
p2 = plot!(𝕔ᵗ.Dp,pre*𝕔ᵗ.S, color = :black, label = "∑𝕔ᵗ + 𝕟ᵈ", ls = :dashdot)   
 
# Panel 3
pre = 1e-6     # Note change in units from m-3 to cm-3
p3 = plot(𝕊[1].Dp,pre*𝕊[1].N, xaxis = :log10, xlim = (40, 400), ylabel = "Conc. (cm⁻³)", ylim = (0,25),
    xlabel = "Diameter (nm)", label = "𝕊₁ = 𝐀₃*ℂ₁")
p3 = plot!(𝕊[2].Dp,pre*𝕊[2].N,  label = "𝕊₂ = 𝐀₃*ℂ₂")   
p3 = plot!(𝕤ᵈ.Dp,pre*𝕤ᵈ.N,  label = "𝕤ᵈ = 𝐀₃*𝕟ᵈ", color = p[4])   
p3 = plot!(𝕤ᵗ.Dp,pre*𝕤ᵗ.N, color = :black, label = "𝕤ᵗ = 𝐀₃*(∑ℂₖ + 𝕟ᵈ)", ls = :dashdot)   

plot(p1,p2,p3, layout = grid(3,1), left_margin = 40px,right_margin = 25px,legend=:right, 
    top_margin = 15px, xlim = (30,120), fmt = :svg)

figure("Nimbus Sans L", 2, 6.5, 2.2, 8)

pre = 1e-10
# Get a list of colors from the theme palette
p = Plots.get_color_palette(:auto, default(:bgcolor), 100)

# Panel 1
p1 = plot(𝕟ᵖ[1].Dp,pre*𝕟ᵐ[1].S, xaxis = :log10, xlim = (40,400), ylabel = "dN/dlnD × 10⁻¹⁰ (m⁻³)", 
    label = "ℕ₁(-)", color = p[1], xlabel = "Mobility Diameter (nm)", left_margin = 20px)
p1 = plot!(𝕟ᵐ[1].Dp,pre*𝕟ᵖ[1].S, ls = :dot,  label = "ℕ₁(+)", color = p[1])
p1 = plot!(𝕟ᵐ[2].Dp,pre*𝕟ᵐ[2].S, label = "ℕ₂(-)", color = p[2])   
p1 = plot!(𝕟ᵖ[2].Dp,pre*𝕟ᵖ[2].S, ls = :dot, label = "ℕ₂(+)", color = p[2])   
p1 = plot!(𝕟ᵗ.Dp,pre*𝕟ᵗ.S, color = :black, label = "∑ℕₖ", ls = :dashdot)  

# Panel 2
p2 = plot(ℂ[1].Dp,pre*ℂ[1].S, xaxis = :log10, xlim = (40,400), 
    ylabel = "dN/dlnD × 10⁻¹⁰ (m⁻³)", label = "ℂ₁", xlabel = "Mobility Diameter (nm)")
p2 = plot!(ℂ[2].Dp,pre*ℂ[2].S, label = "ℂ₂")   
p2 = plot!(𝕟ᵈ.Dp,pre*𝕟ᵈ.S, label = "𝕟ᵈ", color = p[4])   
p2 = plot!(𝕔ᵗ.Dp,pre*𝕔ᵗ.S, color = :black, label = "∑ℂₖ + 𝕟ᵈ", ls = :dashdot)   
 
# Panel 3
pre = 1e-6     # Note change in units from m-3 to cm-3
p3 = plot(𝕊[1].Dp,pre*𝕊[1].N, xaxis = :log10, xlim = (40, 400), ylabel = "Concentration (cm⁻³)", 
    xlabel = "Diameter (nm)", label = "𝕊₁",left_margin = 30px, right_margin = 10px)
p3 = plot!(𝕊[2].Dp,pre*𝕊[2].N,  label = "𝕊₂")   
p3 = plot!(𝕤ᵈ.Dp,pre*𝕤ᵈ.N,  label = "𝕤ᵈ", color = p[4])   
p3 = plot!(𝕤ᵗ.Dp,pre*𝕤ᵗ.N, color = :black, label = "𝕤ᵗ=∑𝕊ₖ + 𝕤ᵈ", ls = :dashdot, 
    xlabel = "Apparent +1 Mobility Diameter (nm)")   

plot(p1,p2,p3, layout = grid(1,3), legend=:right, top_margin = 30px, xlim = (40,125), lw = 2, fmt = :svg)

@printf("Number concentration after DMA 1: %i cm-3\n", sum((sum(𝕟ᵐ)).N)*1e-6)
@printf("Number concentration after DMA 2: %i cm-3\n", sum((sum(𝕟ᵖ)).N)*1e-6)
@printf("Number concentration after electrostatic filter: %i cm-3\n", sum(𝕤ᵗ.N)*1e-6)
@printf("Number concentration at peak during DMA3 scan: %i cm-3\n", maximum(𝕤ᵗ.N)*1e-6)