"""Functions for computing CIB models."""

abstract type AbstractCIBModel{T<:Real} <: AbstractForegroundModel end

"""
    CIB_Planck2013{T}(; kwargs...)

Define CIB model parameters. Defaults are from Viero et al. 2013. All numbers
not typed are converted to type T. This model has the following parameters and default values:

* `nside::Int64 = 4096`
* `hod::String = "shang"`
* `Inu_norm = 0.3180384`
* `min_redshift = 0.0`
* `max_redshift = 5.0`
* `min_mass = 1e12`
* `box_size = 40000`
* `shang_zplat = 2.0`
* `shang_Td = 20.7`
* `shang_betan = 1.6`
* `shang_eta = 2.4`
* `shang_alpha = 0.2`
* `shang_Mpeak = 10^12.3`
* `shang_sigmaM = 0.3`
* `shang_Msmin = 1e11`
* `shang_Mmin = 1e10`
* `shang_I0 = 46`
* `jiang_gamma_1 = 0.13`
* `jiang_alpha_1 = -0.83`
* `jiang_gamma_2 = 1.33`
* `jiang_alpha_2 = -0.02`
* `jiang_beta_2 = 5.67`
* `jiang_zeta = 1.19`

"""
@with_kw struct CIB_Planck2013{T<:Real} <: AbstractCIBModel{T}
    @deftype T
    nside::Int64 = 4096
    hod::String = "shang"
    Inu_norm = 0.3180384
    min_redshift = 0.0
    max_redshift = 6.0
    min_mass = 2e10
    box_size = 40000

    # shang HOD
    shang_zplat = 2.0
    shang_Td = 20.7
    shang_beta = 1.6
    shang_eta = 2.4
    shang_alpha = 0.2
    shang_Mpeak = 10^12.3
    shang_sigmaM = 0.3
    shang_Msmin = 1e11
    shang_Mmin = 1e10
    shang_I0 = 46

    # jiang
    jiang_gamma_1 = 0.13
    jiang_alpha_1 = -0.83
    jiang_gamma_2 = 1.33
    jiang_alpha_2 = -0.02
    jiang_beta_2 = 5.67
    jiang_zeta = 1.19
end


function jiang_shmf(m, M_halo, model)
    dndm = (
        (
            (model.jiang_gamma_1 * ((m / M_halo)^model.jiang_alpha_1)) +
            (model.jiang_gamma_2 * ((m / M_halo)^model.jiang_alpha_2))
        ) * (exp(-(model.jiang_beta_2) * ((m / M_halo)^model.jiang_zeta)))
    )
    return dndm
end

"""
Build a linear interpolation function which maps log(M_h) to N_sat.
"""
function build_shang_interpolator(
    min_log_M::T,
    max_log_M::T,
    model::AbstractCIBModel;
    n_bin = 1000,
) where {T}

    x_m = LinRange(min_log_M, max_log_M, 1000)
    N_sat_i = zero(x_m)

    function integrand_m(lm, lM_halo)
        # lm is natural log of m
        m = exp(lm)
        M_halo = exp(lM_halo)
        dns_dm = jiang_shmf(m, M_halo, model)
        return dns_dm
    end

    for i = 1:size(x_m, 1)
        N_sat_i[i], err = quadgk(
            t -> integrand_m(t, x_m[i]),
            log(model.shang_Msmin),
            x_m[i],
            rtol = 1e-8,
        )
        N_sat_i[i] = convert(T, max(0.0, N_sat_i[i]))
    end

    return LinearInterpolation(x_m, N_sat_i)
end



function sigma_cen(m::T, model::AbstractCIBModel) where {T}
    return (
        exp(
            -(log10(m) - log10(model.shang_Mpeak))^2 /
            (T(2) * model.shang_sigmaM),
        ) * m
    ) / sqrt(T(2π) * model.shang_sigmaM)
end

function nu2theta(nu::T, z::T, model::AbstractCIBModel) where {T}
    phys_h = T(6.62606957e-27)   # erg.s
    phys_k = T(1.3806488e-16)    # erg/K
    Td = model.shang_Td * (one(T) + z)^model.shang_alpha
    xnu = phys_h * nu / phys_k / Td
    return xnu^(T(4) + model.shang_beta) / expm1(xnu) / nu / model.shang_I0
end

"""
<L_sat> interpolation values
"""
function integrand_L(lm, lM_halo, model::AbstractCIBModel)
    m = exp(lm)
    return sigma_cen(m, model) * jiang_shmf(m, exp(lM_halo), model)
end


"""
Build a linear interpolator that takes in ln(M_halo) and returns sigma.
"""
function build_sigma_sat_ln_interpolator(
    max_ln_M::T,
    model::AbstractCIBModel;
    n_bin = 1000,
) where {T}
    x = LinRange(log(model.shang_Mmin), max_ln_M, n_bin)
    L_mean = zero(x)

    for i = 1:n_bin
        L_mean[i], err = quadgk(
            t -> integrand_L(t, x[i], model),
            log(model.shang_Mmin),
            x[i],
            rtol = 1e-6,
        )
    end
    return LinearInterpolation(T.(x), T.(L_mean), extrapolation_bc = zero(T))
end

function build_muofn_interpolator(
    model;
    min_mu::T = -6.0f0,
    max_mu::T = 0.0f0,
    nbin = 1000,
) where {T}
    mu = T(10) .^ LinRange(min_mu, max_mu, nbin)
    n = zero(mu)
    for i = 1:nbin
        n[i], err = quadgk(
            lmu -> jiang_shmf(exp(lmu), one(T), model),
            log(mu[i]),
            0.0f0,
            rtol = 1.0f-6,
        )
    end
    return LinearInterpolation(T.(reverse(n)), T.(reverse(mu)))
end

"""
Compute redshift evolution factor for LF.
"""
function shang_z_evo(z::T, model::AbstractCIBModel) where {T}
    return (one(T) + min(z, model.shang_zplat))^model.shang_eta
end

"""
Construct the necessary interpolator set.
"""
function get_interpolators(
    model::AbstractCIBModel,
    cosmo::Cosmology.FlatLCDM{T},
    min_halo_mass::T,
    max_halo_mass::T,
) where {T}
    return (
        r2z = build_r2z_interpolator(
            model.min_redshift,
            model.max_redshift,
            cosmo,
        ),
        hod_shang = build_shang_interpolator(
            log(min_halo_mass),
            log(max_halo_mass),
            model,
        ),
        c_lnm2r = build_c_lnm2r_interpolator(),
        sigma_sat = build_sigma_sat_ln_interpolator(log(max_halo_mass), model),
        muofn = build_muofn_interpolator(model),
    )
end



# Fill up arrays with information related to CIB central sources.
function process_centrals!(
    model::AbstractCIBModel{T},
    cosmo::Cosmology.FlatLCDM{T},
    Healpix_res::Resolution;
    interp,
    hp_ind_cen,
    dist_cen,
    redshift_cen,
    lum_cen,
    n_sat_bar,
    n_sat_bar_result,
    halo_pos,
    halo_mass,
) where {T}

    N_halos = size(halo_mass, 1)

    Threads.@threads for i = 1:N_halos
        # location information for centrals
        hp_ind_cen[i] = Healpix.vec2pixRing(
            Healpix_res,
            halo_pos[1, i],
            halo_pos[2, i],
            halo_pos[3, i],
        )
        dist_cen[i] =
            sqrt(halo_pos[1, i]^2 + halo_pos[2, i]^2 + halo_pos[3, i]^2)
        redshift_cen[i] = interp.r2z(dist_cen[i])

        # compute HOD
        n_sat_bar[i] = interp.hod_shang(log(halo_mass[i]))
        n_sat_bar_result[i] =
            rand(Distributions.Poisson(Float64.(n_sat_bar[i])))

        # get central luminosity
        lum_cen[i] =
            sigma_cen(halo_mass[i], model) * shang_z_evo(redshift_cen[i], model)
    end
end


# Fill up arrays with information related to CIB satellites.
function process_sats!(
    model::AbstractCIBModel{T},
    cosmo::Cosmology.FlatLCDM{T},
    Healpix_res::Resolution;
    interp,
    hp_ind_sat,
    dist_sat,
    redshift_sat,
    lum_sat,
    cumsat,
    halo_mass,
    halo_pos,
    redshift_cen,
    n_sat_bar,
    n_sat_bar_result,
) where {T}

    N_halos = size(halo_mass, 1)
    Threads.@threads for i_halo = 1:N_halos
        r_cen = m2r(halo_mass[i_halo], cosmo)
        c_cen = mz2c(halo_mass[i_halo], redshift_cen[i_halo], cosmo)
        for j = 1:n_sat_bar_result[i_halo]
            # i is central halo index, j is index of satellite within each halo
            i_sat = cumsat[i_halo] + j # index of satellite in satellite arrays

            log_msat_inner = max(log(rand(T)), T(-7.0))
            r_sat =
                r_cen *
                interp.c_lnm2r(c_cen, log_msat_inner) *
                T(200.0^(-1.0 / 3.0))
            m_sat =
                (interp.muofn(rand(T) * n_sat_bar[i_halo]) * halo_mass[i_halo])

            phi = random_phi(T)
            theta = random_theta(T)
            x_sat = halo_pos[1, i_halo] + r_sat * sin(theta) * cos(phi)
            y_sat = halo_pos[2, i_halo] + r_sat * sin(theta) * sin(phi)
            z_sat = halo_pos[3, i_halo] + r_sat * cos(theta)
            dist_sat[i_sat] = sqrt(x_sat^2 + y_sat^2 + z_sat^2)
            redshift_sat[i_sat] = interp.r2z(dist_sat[i_sat])

            # lum_sat[i_sat] = interp.sigma_sat(log(m_sat)) * shang_z_evo(
            #     redshift_sat[i], model)
            lum_sat[i_sat] =
                sigma_cen(m_sat, model) *
                shang_z_evo(redshift_sat[i_sat], model)
            hp_ind_sat[i_sat] =
                Healpix.vec2pixRing(Healpix_res, x_sat, y_sat, z_sat)
        end
    end
end


"""
    generate_sources(model, cosmo, halo_pos_inp, halo_mass_inp; verbose=true)

Produce a source catalog from a model and halo catalog. This converts the
halo arrays into the type specified by `model`.

# Arguments:
- `model::AbstractCIBModel{T}`: source model parameters
- `cosmo::Cosmology.FlatLCDM{T}`: background cosmology
- `Healpix_res::Resolution`: Healpix map resolution
- `halo_pos_inp::AbstractArray{TH,2}`: halo positions with dims (3, nhalos)
- `halo_mass_inp::AbstractArray{TH,1}`: halo masses

# Keywords
- `verbose::Bool=true`: print out progress details
"""
function generate_sources(
    model::AbstractCIBModel{T},
    cosmo::Cosmology.FlatLCDM{T},
    halo_pos_inp::AbstractArray{TH,2},
    halo_mass_inp::AbstractArray{TH,1};
    verbose = true,
) where {T,TH}

    # make sure halo inputs are the CIB type
    halo_pos = convert(Array{T,2}, halo_pos_inp)
    halo_mass = convert(Array{T,1}, halo_mass_inp)

    # set up basics
    N_halos = size(halo_mass, 1)
    interp =
        get_interpolators(model, cosmo, minimum(halo_mass), maximum(halo_mass))
    res = Resolution(model.nside)

    verbose && println("Allocating for $(N_halos) centrals.")
    hp_ind_cen = Array{Int64}(undef, N_halos)  # healpix index of halo
    lum_cen = Array{T}(undef, N_halos)  # Lum of central w/o ν-dependence
    redshift_cen = Array{T}(undef, N_halos)
    dist_cen = Array{T}(undef, N_halos)
    n_sat_bar = Array{T}(undef, N_halos)
    n_sat_bar_result = Array{Int32}(undef, N_halos)

    # STEP 1: compute central properties -----------------------------------
    verbose && println("Processing centrals on $(Threads.nthreads()) threads.")
    process_centrals!(
        model,
        cosmo,
        res,
        interp = interp,
        hp_ind_cen = hp_ind_cen,
        dist_cen = dist_cen,
        redshift_cen = redshift_cen,
        lum_cen = lum_cen,
        n_sat_bar = n_sat_bar,
        n_sat_bar_result = n_sat_bar_result,
        halo_pos = halo_pos,
        halo_mass = halo_mass,
    )

    # STEP 2: Generate satellite arrays -----------------------------
    cumsat = generate_subhalo_offsets(n_sat_bar_result)
    total_n_sat = cumsat[end]
    hp_ind_sat = Array{Int64}(undef, total_n_sat)  # healpix index of halo
    lum_sat = Array{T}(undef, total_n_sat)  # Lum of central w/o ν-dependence
    redshift_sat = Array{T}(undef, total_n_sat)
    dist_sat = Array{T}(undef, total_n_sat)

    # STEP 3: compute satellite properties -----------------------------------
    verbose && println("Processing $(total_n_sat) satellites.")
    process_sats!(
        model,
        cosmo,
        res,
        interp = interp,
        hp_ind_sat = hp_ind_sat,
        dist_sat = dist_sat,
        redshift_sat = redshift_sat,
        lum_sat = lum_sat,
        cumsat = cumsat,
        halo_mass = halo_mass,
        halo_pos = halo_pos,
        redshift_cen = redshift_cen,
        n_sat_bar = n_sat_bar,
        n_sat_bar_result = n_sat_bar_result,
    )

    return (
        hp_ind_cen = hp_ind_cen,
        lum_cen = lum_cen,
        redshift_cen = redshift_cen,
        dist_cen = dist_cen,
        hp_ind_sat = hp_ind_sat,
        lum_sat = lum_sat,
        redshift_sat = redshift_sat,
        dist_sat = dist_sat,
        N_cen = N_halos,
        N_sat = total_n_sat,
    )
end


"""
    paint!(result_map, nu_obs, model, sources, fluxes_cen, fluxes_sat)

Paint a source catalog onto a map, recording the fluxes in
`fluxes_cen` and `fluxes_sat`.

# Arguments:
- `result_map::Map{T_map, RingOrder}`: Healpix map to paint
- `nu_obs`: frequency in Hz
- `model::AbstractCIBModel{T}`: source model parameters
- `sources`: NamedTuple containing source information from generate_sources
- `fluxes_cen::AbstractArray`: buffer for writing fluxes of centrals
- `fluxes_sat::AbstractArray`: buffer for writing fluxes of satellites
"""
function paint!(
    result_map::Map{T_map,RingOrder},
    nu_obs,
    model::AbstractCIBModel{T},
    sources,
    fluxes_cen::AbstractArray,
    fluxes_sat::AbstractArray,
) where {T_map,T}

    pixel_array = result_map.pixels
    fill!(pixel_array, zero(T))  # prepare the frequency map

    # process centrals for this frequency
    Threads.@threads for i = 1:sources.N_cen
        nu = (one(T) + sources.redshift_cen[i]) * nu_obs
        fluxes_cen[i] = l2f(
            sources.lum_cen[i] * nu2theta(nu, sources.redshift_cen[i], model),
            sources.dist_cen[i],
            sources.redshift_cen[i],
        )
        pixel_array[sources.hp_ind_cen[i]] += fluxes_cen[i]
    end

    # process satellites for this frequency
    Threads.@threads for i = 1:sources.N_sat
        nu = (one(T) + sources.redshift_sat[i]) * nu_obs
        fluxes_sat[i] = l2f(
            sources.lum_sat[i] * nu2theta(nu, sources.redshift_sat[i], model),
            sources.dist_sat[i],
            sources.redshift_sat[i],
        )
        pixel_array[sources.hp_ind_sat[i]] += fluxes_sat[i]
    end

    # divide by healpix pixel size
    per_pixel_steradian = 1 / nside2pixarea(result_map.resolution.nside)
    pixel_array .*= per_pixel_steradian
end

"""
Paint a source catalog onto a map.

This function creates the arrays for you.
"""
function paint!(
    result_map::Map{T,RingOrder},
    nu_obs::T,
    model::AbstractCIBModel,
    sources,
) where {T}

    fluxes_cen = Array{T,1}(undef, sources.N_cen)
    fluxes_sat = Array{T,1}(undef, sources.N_sat)

    paint!(result_map, nu_obs, model, sources, fluxes_cen, fluxes_sat)

    return fluxes_cen, fluxes_sat
end




"""
    estimate_fluxes!(  nu_obs, model, sources, fluxes_cen, fluxes_sat)

Paint a source catalog onto a map, recording the fluxes in
`fluxes_cen` and `fluxes_sat`.

# Arguments:
- `nu_obs`: frequency in Hz
- `model::AbstractCIBModel{T}`: source model parameters
- `sources`: NamedTuple containing source information from generate_sources
- `fluxes_cen::AbstractArray`: buffer for writing fluxes of centrals
- `fluxes_sat::AbstractArray`: buffer for writing fluxes of satellites
"""
function estimate_fluxes!(
    nu_obs,
    model::AbstractCIBModel{T},
    sources,
    fluxes_cen::AbstractArray,
    fluxes_sat::AbstractArray,
) where {T}

    # process centrals for this frequency
    Threads.@threads for i = 1:sources.N_cen
        nu = (one(T) + sources.redshift_cen[i]) * nu_obs
        fluxes_cen[i] = l2f(
            sources.lum_cen[i] * nu2theta(nu, sources.redshift_cen[i], model),
            sources.dist_cen[i],
            sources.redshift_cen[i],
        )
    end

    # process satellites for this frequency
    Threads.@threads for i = 1:sources.N_sat
        nu = (one(T) + sources.redshift_sat[i]) * nu_obs
        fluxes_sat[i] = l2f(
            sources.lum_sat[i] * nu2theta(nu, sources.redshift_sat[i], model),
            sources.dist_sat[i],
            sources.redshift_sat[i],
        )
    end
end



"""
    paint_with_mask!(result_map, nu_obs, model, sources, fluxes_cen, fluxes_sat,mask_cen, mask_sat)

Paint a source catalog onto a map, recording the fluxes in
`fluxes_cen` and `fluxes_sat`.

# Arguments:
- `result_map::Map{T_map, RingOrder}`: Healpix map to paint
- `nu_obs`: frequency in Hz
- `model::AbstractCIBModel{T}`: source model parameters
- `sources`: NamedTuple containing source information from generate_sources
- `fluxes_cen::AbstractArray`: buffer for writing fluxes of centrals
- `fluxes_sat::AbstractArray`: buffer for writing fluxes of satellites
- `mask_cen::AbstractArray`: buffer encoding  the index of the mask which satisfies the condition for  centrals
- `mask_sat::AbstractArray`: buffer encoding  the index of the mask which satisfies the condition for satellites
"""

function paint_with_mask!(
    result_map::Map{T_map,RingOrder},
    nu_obs,
    model::AbstractCIBModel{T},
    sources,
    fluxes_cen::AbstractArray,
    fluxes_sat::AbstractArray,
    mask_cen::AbstractArray,
    mask_sat::AbstractArray,
) where {T_map,T}

    pixel_array = result_map.pixels
    fill!(pixel_array, zero(T))  # prepare the frequency map

    # process centrals for this frequency
    Threads.@threads for i in mask_cen
        nu = (one(T) + sources.redshift_cen[i]) * nu_obs
        fluxes_cen[i] = l2f(
            sources.lum_cen[i] * nu2theta(nu, sources.redshift_cen[i], model),
            sources.dist_cen[i],
            sources.redshift_cen[i],
        )
        pixel_array[sources.hp_ind_cen[i]] += fluxes_cen[i]

    end

    # process satellites for this frequency
    Threads.@threads for i in mask_sat
        nu = (one(T) + sources.redshift_sat[i]) * nu_obs
        fluxes_sat[i] = l2f(
            sources.lum_sat[i] * nu2theta(nu, sources.redshift_sat[i], model),
            sources.dist_sat[i],
            sources.redshift_sat[i],
        )
        pixel_array[sources.hp_ind_sat[i]] += fluxes_sat[i]

    end

    # divide by healpix pixel size
    per_pixel_steradian = 1 / nside2pixarea(result_map.resolution.nside)
    pixel_array .*= per_pixel_steradian
end


"""
    estimate_normalization(
    result_map,
    ν_ref,
    unit_out)

estimate  normalization factor   so that both the fluxes and the map have consistent monopole term with
Fixen et al. 1998 model.

# Arguments:
- `result_map::Map{T_map, RingOrder}`: Healpix CIB map
- `ν_ref`: frequency in GHz (Default 857 GHz )
- `unit_out::Unitful.FreeUnits`: set the output units of the Fixsen monopole


"""

function estimate_normalization(
    result_map::Map{T,RingOrder},
    ν_ref::Unitful.Frequency,
    unit_out::Unitful.FreeUnits,
) where {T}

    fixsen_monopole = fixsen_model(ν_ref; unit_out = unit_out)
    maskpatch = findall(x -> (x != 0), result_map)

    fsky = size(maskpatch)[1] / size(result_map)[1]
    Smean = mean(result_map) / fsky * unit_out
    L0 = fixsen_monopole / Smean
    return L0
end



"""
    make_multifrequency_dataset(
    ν_obs,
    model,
    sources;
    ν_ref= 857u"GHz",
    output_μm= false,
    normalize = false,
    outputpath  = "./cib_data.hdf5",
    unit_out  = u"MJy/sr",
)

This function involves many subroutines aimed at generating at multiple frequency
a dataset of frequency maps at `nu_obs` frequencies, given the CIB model and the
luminosities painted from a halo catalog. As the painting procedure is _normalization free_,
the user can set  the fluxes and the maps  to match the Fixsen monopole at a reference frequency.

At a reference freq. `ν_ref`, we scale   both maps and fluxes generated with XGPaint so that :


`<I(ν_ref)> =  F( ν_ref, T)`


with `F( ν_ref, T) ` being the model of Fixsen. We then apply the normalization factor

`L0 = F( ν_ref, T) /<I(ν_ref)>`

to the other frequencies.


# Arguments:

- `ν_obs::AbstractArray`: array of frequencies to be simulated (in units of GHz)
- `ν_ref::Unitful.Frequency`: reference frequency in GHz (Default 857 GHz )
- `model::AbstractCIBModel{T}`: source model parameters
- `sources::NamedTuple`: containing source information from generate_sources
- `output_μm::Bool`: save the output catalog in units of μm (default outputs in GHz)
- `normalize::Bool`: apply the normalization if `true` (default `false`)
- `outputpath::String`: generic string encoding the path to the folder to save the
    data. It   looks for the string sequence `data` and replaces it with the frequency
    (or wavelength ) of the specific simulated dataset.
- `unit_out::Unitful.FreeUnits`: set the output units of the map and of the fluxes
    (default MJy/sr)
"""

function make_multifrequency_dataset(
    ν_obs::AbstractArray,
    model,
    sources;
    ν_ref::Unitful.Frequency = 857u"GHz",
    output_μm::Bool = false,
    normalize::Bool = false,
    outputpath::String = "./cib_data.hdf5",
    unit_out::Unitful.FreeUnits = u"MJy/sr",
)

    m = Map{Float64,RingOrder}(model.nside)  # create a Healpix map
    fluxes_cen = Array{Float32,1}(undef, sources.N_cen)
    fluxes_sat = Array{Float32,1}(undef, sources.N_sat)
    if normalize
        @time paint!(
            m,
            convert(Float32, ustrip(uconvert(u"Hz", ν_ref))),
            model,
            sources,
            fluxes_cen,
            fluxes_sat,
        )
        @time L0 = estimate_normalization(m, ν_ref, unit_out)
        print("estimated the normalization constant :$(ustrip(L0)).  \n")
    end

    c = uconvert(u"μm* GHz", float(Unitful.c0))
    lambdas = c ./ ν_obs

    for i in eachindex(ν_obs)
        ν = ustrip(uconvert(u"Hz", ν_obs[i]))

        if output_μm == true
            λ = round(typeof(1u"μm"), lambdas[i])
            fname = replace(outputpath, "data" => "$(ustrip(λ))$(unit(λ))")
            print("Saving map at $(λ)  in $fname \n")
        else
            fname = replace(
                outputpath,
                "data" => "$(ustrip(ν_obs[i]))$(unit(ν_obs[i]))",
            )
            print("Saving map at $(ν_obs[i])  in $fname \n")

        end
        @time paint!(
            m,
            convert(Float32, ν),
            model,
            sources,
            fluxes_cen,
            fluxes_sat,
        )
        if normalize
            fluxes_cen *= L0
            fluxes_sat *= L0
            m.pixels *= L0
        end
        write_to_hdf5(fname, fluxes_cen, fluxes_sat, sources, m.pixels)

    end

end
