using Plots, Plots.PlotMeasures, Interpolations, LsqFit, DataFrames, Distributions 
using DifferentialMobilityAnalyzers, SpecialFunctions, Random
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 = 6                                            # Upper number of charges
Λ = DMAconfig(t,p,qsa,qsa/β,r₁,r₂,l,leff,:-,m,:cylindrical)   # Specify DMA with negative polarity
bins,z₁,z₂ = 256, dtoz(Λ,1000e-9), dtoz(Λ,10e-9) # bins, upper, lower mobility limit
δ = setupDMA(Λ, z₁, z₂, bins);                   # Compute matrices

# CCN efficiency function 
#  - 𝕟 is a size distribution from which the diameter vector is used
#  - μ is the activation diameter
#  - σ is the spread
Taf = (𝕟,μ,σ) -> 0.5 .* (1.0 .+ erf.((𝕟.Dp .- μ)./(sqrt(2σ))))
μ,σ = 60,8

# Generate a lognormal size distribution that serves as input. 
# The size distribution is computed on the DMA size grid defined by DMA δ1 and δ2
#   - Nt is in [# m-3] for calculations with k12
#   - Dg is in nm in accordance with the grid uings
#   - σ is the geometeric standard deviation (σ > 1)
# Multiple modes can be defined by concatenation. A single mode is permissible A = [[Nt1,Dg1,σ1]]
𝕟ᶜⁿ = DMALognormalDistribution([[1800, 80, 1.4]], δ)
𝕟ᶜᶜⁿ = Taf(𝕟ᶜⁿ,μ,σ)*𝕟ᶜⁿ  # CCN distribution function
𝕣ᶜⁿ = δ.𝐀*𝕟ᶜⁿ               # CN response function. δ.A is the forward matrix for DMA δ
𝕣ᶜᶜⁿ = δ.𝐀*𝕟ᶜᶜⁿ             # CCN response function. δ.A is the forward matrix for DMA δ  
𝕒𝕗₁ =  𝕟ᶜᶜⁿ/𝕟ᶜⁿ            # ratio of CCN to CN based on real size distribution. This recovers g
𝕒𝕗₂ =  𝕣ᶜᶜⁿ/𝕣ᶜⁿ;           # ratio of CCN to CN based on response function (measured)

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

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

# Panel (a): Mobility size and CCN distribution
p1 = plot(𝕟ᶜⁿ.Dp, 𝕟ᶜⁿ.S, xaxis = :log10, xticks = [10, 100, 1000], ylabel = "dN/dlnD (cm⁻³)", 
    label = "𝕟ᶜⁿ = 𝓛𝓝 (N,Dₘ,σ)", color = p[1],  left_margin = 20px)
p1 = plot!(𝕟ᶜᶜⁿ.Dp, 𝕟ᶜᶜⁿ.S, label = "𝕟ᶜᶜⁿ = Taf.*𝕟ᶜⁿ", color = p[2])

# Panel (b): Size and CCN response function 
p2 = plot(𝕣ᶜⁿ.Dp, 𝕣ᶜⁿ.N , xaxis = :log10, xticks = [10, 100, 1000], label = "𝕣ᶜⁿ = 𝐀*𝕟ᶜⁿ", 
    ylabel = "Concentration (cm⁻³)", color = p[3])
p2 = plot!(𝕣ᶜᶜⁿ.Dp, 𝕣ᶜᶜⁿ.N, label = "𝕣ᶜᶜⁿ = 𝐀*𝕟ᶜᶜⁿ", color = p[4], left_margin = 36px)

# Panel (c): Activated fraction. Note that g and 𝕒𝕗₁ are identical
p3 = plot(𝕒𝕗₁.Dp, 𝕒𝕗₁.S, xaxis = :log10, xticks = [10, 100, 1000], xlabel = "Mobility Diameter (nm)", 
    ylabel = "Activated fraction", label = "𝕟ᶜᶜⁿ ./ 𝕟ᶜⁿ", ylim = (-0.05,1.05), lw = 1.5, 
    yticks = [0, 0.2, 0.4, 0.6, 0.8, 1], color = :black,  left_margin = 20px)
p3 = plot!(𝕟ᶜⁿ.Dp, Taf(𝕟ᶜⁿ,μ,σ), label = "Taf = ½[1+erf((Dₚ-μ)/(√2σ)]", ls = :dash, 
    color = :darkorange, lw = 1.5)

# Panel (d): Activated fraction from response function ("measured ratio" from raw signal)
p4 = plot(𝕒𝕗₂.Dp, 𝕒𝕗₂.S, xaxis = :log10, xticks = [10, 100, 1000],  
    xlabel = "Apparent +1 Mobility Diam. (nm)", ylim = (-0.05,1.05), label = "𝕣ᶜᶜⁿ ./ 𝕣ᶜⁿ", 
    color = p[7], ylabel = "Activated fraction", yticks = [0, 0.2, 0.4, 0.6, 0.8, 1], left_margin = 20px)

plot(p1,p2,p3,p4, layout = grid(2,2), lw = 1.5, right_margin = 30px, 
    xlim = (20,200), legend=:right, top_margin = 15px, fmt = :svg)

t,p = 295.15, 1e5                                # Temperature [K], Pressure [Pa]
qsa,β = 1.66e-5, 1/5                             # 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 = 6                                            # Upper number of charges
Λ = DMAconfig(t,p,qsa,qsa/β,r₁,r₂,l,leff,:-,m,:cylindrical)   # Specify DMA with negative polarity
bins,z₁,z₂ = 60, dtoz(Λ,300e-9), dtoz(Λ,30e-9)   # bins, upper, lower mobility limit
δ = setupDMA(Λ, z₁, z₂, bins);                   # Compute matrices

𝕟ᶜⁿ = DMALognormalDistribution([[7000.0, 80, 1.4]], δ)
Taf = (x,μ,σ) -> @. 0.5 .* (1.0 .+ erf.((x .- μ)/(sqrt(2σ))))
tscan = 180        # SMPS scan time [s] 
Qcpc = 16.66       # CPC flow rate [cm³ s⁻¹]. 16.66 cm³ s⁻¹ = 1 L min⁻¹
Qccn = 0.05*16.66  # CCN sample flow rate [cm³ s⁻¹]. 16.66 cm³ s⁻¹ = 1 L min⁻¹
t = tscan./bins    # time in each bin
Random.seed!(705)
μ,σ = 60,8
𝕣ᶜⁿ = δ.𝐀*𝕟ᶜⁿ;                     # DMA response function 
𝕣ᶜᶜⁿ = δ.𝐀*(Taf(δ.Dp,μ,σ)*𝕟ᶜⁿ);   # DMA response function 
cᶜⁿ = 𝕣ᶜⁿ.N*Qcpc*t;                # number of counts in each bin
cᶜᶜⁿ = 𝕣ᶜᶜⁿ.N*Qccn*t;              # number of CCN counts in each bin

Rcn = Float64[]    # Construct a noisy response function 𝕣+ϵ, where ϵ is counting uncertainty
Rccn = Float64[]   # Construct a noisy response function 𝕣+ϵ, where ϵ is counting uncertainty
for i = cᶜⁿ
    f = rand(Poisson(i),1)     # draw Poisson random number with mean = counts
    push!(Rcn,f[1]/(Qcpc*t))   # convert back to concentration
end

for i = cᶜᶜⁿ
    f = rand(Poisson(i),1)      # draw Poisson random number with mean = counts
    push!(Rccn,f[1]/(Qccn*t))   # convert back to concentration
end

𝕣ᶜⁿ.N, 𝕣ᶜᶜⁿ.N  = Rcn, Rccn;

function threshold(𝕟::SizeDistribution, c::Float64, n1::Float64, n2::Float64)
   N = 𝕟.N  
   N[(N .<= c) .& (𝕟.Dp .> 150)] .= n1
   N[(N .<= c) .& (𝕟.Dp .> 150)] .= n2
   𝕟.N = N
end

threshold(𝕣ᶜⁿ, 6.0, 0.1, 0.1);
threshold(𝕣ᶜᶜⁿ, 6.0, 0.0, 0.1);

𝕟ᶜⁿ = rinv(𝕣ᶜⁿ.N, δ, λ₁= 1e-3, λ₂=1e1)
𝕟ᶜᶜⁿ = rinv(𝕣ᶜᶜⁿ.N, δ, λ₁= 1e-3, λ₂=1e1)
threshold(𝕟ᶜⁿ, 60.0, 0.1, 0.1)
threshold(𝕟ᶜᶜⁿ, 60.0, 0.0, 0.1)
𝕒𝕗₁ =  𝕟ᶜᶜⁿ/𝕟ᶜⁿ           # ratio of CCN to CN based on inverted data

model(x,p) = (δ.𝐀*(𝕟ᶜⁿ.N.*Taf(δ.Dp, p[1], p[2])))./(δ.𝐀*𝕟ᶜⁿ.N)
fit = curve_fit(model, δ.Dp, 𝕣ᶜᶜⁿ.N./𝕣ᶜⁿ.N, [65.0, 3.0])
Ax = fit.param;
DataFrame(μ=round(Ax[1],digits=1), σ=round(Ax[2],digits=1))

figure("Nimbus Sans L", 2, 5.5, 3, 8)
pre = 1e-3

p1 = plot(δ.Dp, 𝕣ᶜⁿ.N, xaxis = :log10, xticks = [10, 100, 1000], ylabel = "Concentration (cm⁻³)", 
    label = "𝕣ᶜⁿ",  color = :black, lt = :steppre)
p1 = plot!(δ.Dp, 𝕣ᶜᶜⁿ.N, label = "𝕣ᶜᶜⁿ", color = RGBA(163/255,0,0,1), lt = :steppre,
        left_margin = 40px)

p2 = plot(𝕟ᶜⁿ.Dp, pre.*𝕟ᶜⁿ.S, xaxis = :log10, xticks = [10, 100, 1000], 
     ylabel = "dN/dlnD × 10⁻³ (cm⁻³)",  label = "𝕟ᶜⁿ (Eq. 10)", color = :black, 
     ls = :dashdot, lt = :steppre)
p2 = plot!(𝕟ᶜᶜⁿ.Dp, pre.*𝕟ᶜᶜⁿ.S, label = "𝕟ᶜᶜⁿ (Eq. 10)", color = RGBA(163/255,0,0,1), 
    ls = :dashdot, lt = :steppre, left_margin = 52px)

p3 = scatter(δ.Dp, 𝕣ᶜᶜⁿ.N./𝕣ᶜⁿ.N, xaxis = :log10, xticks = [10, 100],  
    xlabel = "Apparent +1 Mobility Diameter (nm)", ms = 3, 
    ylabel = "Activated fraction", ylim = (-0.05,1.25), label = "𝕣ᶜᶜⁿ./𝕣ᶜⁿ", markercolor = :darkorange,
     left_margin = 30px)

p3 = plot!(δ.Dp, model(δ.Dp, Ax), label = "Eq. (12)", color = :black , lw = 2)

p4 = scatter(𝕒𝕗₁.Dp, 𝕒𝕗₁.N, xaxis = :log10, xticks = [10, 100, 1000], ms = 3,
    xlabel = "Mobility Diameter (nm)", ylabel = "Activated fraction", label = "𝕟ᶜᶜⁿ/𝕟ᶜⁿ", 
    markercolor = :white, ylim = (-0.05,1.25), left_margin = 25px)

p4 = plot!(δ.Dp, Taf(δ.Dp,Ax[1],Ax[2]), label = "Eq. (11)", color = :black, lw = 2, ls = :dashdot)

plot(p1,p2,p3,p4, layout = grid(2,2),  bottom_margin = 10px, xlim =(30,130), 
    legend = :topright, lw = 2, right_margin = 20px, fmt = :svg)