

using Juriba.Particles
const PT = Particles
using QuadGK

@Base.kwdef mutable struct Powerlaw <: AbstractParticles
    pspec::Float64 = 0
    plnorm::Float64 = 0
    pmin::Float64 = 0
    pmax::Float64 = 0

    size::Int
    mass_gr::Float64   # particle mass in grams
    mass_kev::Float64  # same as above but in keV, using electrons as "reference"
    p::Vector{Float64} # array of particle momenta
    ndens::Vector{Float64}			# array of number density per unit volume, per unit momentum
    gamma::Vector{Float64}			# array of particle kinetic energies for each momentum
    gdens::Vector{Float64}		    # array of number density per unit volume, per unit gamma
    gdens_diff::Vector{Float64}     # array with differential of number density for radiation calculation
    
end

function powerlaw(s)
    size = s

    plnorm = 1.0
    
    mass_gr = emgm
    mass_kev = emgm*gr_to_kev

    p = zeros(Float64, size)
    ndens = similar(p)
    gamma = similar(p)
    gdens = similar(p)
    gdens_diff = similar(p)

    return Powerlaw(plnorm=plnorm, size = size, mass_gr = mass_gr, mass_kev = mass_kev, p=p, ndens=ndens,
                    gamma=gamma, gdens=gdens, gdens_diff=gdens_diff)
end

function set_p(pl::Powerlaw, min, ucom, bfield, betaeff, r, fsc)

    pl.pmin = min
    pl.pmax = max_p(pl, ucom, bfield, betaeff, r, fsc)

    pinc = (log10(pmax) - log10(pmin))/(pl.size - 1)

    for i in eachindex(pl.p)
        pl.p[i] = 10^(log10(pl.pmin)+i*pinc)
        pl.gamma[i] = sqrt((pl.p[i]/(pl.mass_gr*cee))^2 + 1)
    end

end

function set_p(pl::Powerlaw, min, gmax)
    pl.pmin = min
    pl.pmax = sqrt(gmax^2 - 1)*pl.mass_gr*cee

    pinc = (log10(pmax) - log10(pmin))/(size - 1)

    for i in eachindex(pl.p)
        pl.p[i] = 10^(log10(pl.pmin) + i*pinc)
        pl.gamma[i] = sqrt((pl.p[i]/(pl.mass_gr*cee))^2 + 1)
    end

end

function set_ndens(pl::Powerlaw)
    
    for i = 1:pl.size
        pl.ndens[i] = pl.plnorm*pl.p[i]^(-pl.psec)*exp(-pl.p[i]/pl.pmax)
    end

    PT.initialize_gdens(pl)
    PT.gdens_differentiate(pl)
end

function set_pspec(pl::Powerlaw, s1)
    pl.spec = s1
end

function set_norm(pl::Powerlaw, n)
    pl.plnorm = n*(1.0 - pl.pspec)/(pl.p[end]^(1 - pl.spec) - pl.p[1]^(1.0 - pl.pspec))
end

""" Injection function to be integrated in cooling """
function injection_pl_int(x, p::PT.injection_pl_params)
    params = p

    s = params.s
    n = params.n
    m = params.m
    max = params.max

    mom_int = sqrt((x^2) - 1)*m*cee

    return n*mom_int^(-s)*exp(-mom_int/max)
end

function cooling_steadystate(pl::Powerlaw, ucom, n0, bfield, r, betaeff)

    Urad = bfield^2/8π + ucom
    pdot_ad = betaeff*cee/r
    pdot_rad = (4.0*sigtom*cee*Urad)/(3.0*pl.mass_gr*cee^2)
    tinj = r/cee

    params = PT.injection_pl_params(pl.psec, pl.plnorm, pl.mass_gr, pl.pmax)
    
    func(x) = PT.injection_pl_int(x, params)
    
    for i = 1:pl.size-1
        integral, error = quadgk(func, pl.gamma[i], pl.gamma[i+1], rtol=1e1, atol=1e1, maxevals=100)

        pl.ndens[i] = (integral/tinj)/(pdot_ad*pl.p[i]/(pl.mass_gr*cee)+pl.pdot_rad*(pl.gamma[i]*pl.p[i]/(pl.mass_gr*cee)))
    end

    pl.ndens[end] = pl.ndens[end-1]*(pl.p[end]/pl-p[end-2])^(-pl.pspec - 1)*exp(-1)

    # the last bin is set by arbitrarily assuming cooled distribution; this is necessary because the integral 
    # above is undefined for the last bin

    # The last step requires a renormalization. The reason is that the result of gsl_integration_qag strongly
    # depends on the value of "size". Without doing anything fancy, this can be fixed simply by ensuring that
    # the total integrated number of density equals n0 (which we know), and rescaling the array ndens[i] by
    # the appropriate constant.

    renorm = PT.count_particles(pl)/n0

    ndens[i] ./= renorm

    PT.initialize_gdens(pl)
    PT.gdens_differentiate(pl)
    
end

"""
Method to calculate maximum momentum of non thermal particles based on acceleration and cooling timescales
The estimate is identical to the old agnjet but in momentum space; see Lucchini et al. 2019 for the math of
the old version
"""
function max_p(pl::Powerlaw, ucom, bfield, betaeff, r, fsc)

    Urad = bfield^2/8π + ucom
    escom = betaeff*cee/r
    syncon = (4.0*sigtom*Urad)/(3.0*pl.mass_gr*cee)

    accon = (3.0*fsc*pl.charg*bfield)/(4.0*pl.mass_gr*cee)

    b = escom/syncon
    c = accon/syncon

    gmax = (-b + sqrt(b^2 + 4.0*c))/2.0
    return sqrt(gmax^2 - 1)*pl.mass_gr*cee

end

function test(pl::Powerlaw)

    println("Power-law distribution")
    println("pspec: $(pl.pspec)")
    println("Array size: $(pl.size)")
    println("Default normalization: $(pl.plnorm)")
    println("Particle mass in grams: $(pl.mass_gr)")

end

