# MIT License
#
# Copyright (c) 2017–2022: Justin Willmert
#
# See https://justinwillmert.com/articles/2022/solving-for-spherical-harmonic-analysis-quadrature-weights/
# and https://github.com/jmert/CMB.jl

using AssociatedLegendrePolynomials: λlm!
using FFTW
using FastGaussQuadrature: gausslegendre

struct RingInfo
    offsets::Tuple{Int,Int}
    stride::Int
    nϕ::Int
    cosθ::Float64
    ϕ₀::Float64
    ΔΩ::Float64
end

struct MapRings
    rings::Vector{RingInfo}
end

"""
    maprings = ecp_maprings(nθ, nϕ)

Constructs a `MapRings` structure describing an equidistant cylindrical projection (ECP) grid of dimensions
`nθ × nϕ` pixels in column-major ordering.
"""
function ecp_maprings(nθ, nϕ)
    nθh = (nθ + 1) ÷ 2  # number of rings in N. hemisphere + equator (if nθ is odd)
    ϕ₀ = 2π / 2nϕ  # constant for all rings
    ΔθΔϕ = (2π / nϕ) * (π / nθ)  # constant portion of ΔΩ = sinθ * ΔθΔϕ
    stride = nθ  # for column-major ordering
    rings = Vector{RingInfo}(undef, nθh)
    for ii in 1:nθh
        o₁ = ii           # northern-ring offset
        o₂ = nθ - ii + 1  # southern-ring offset
        offsets = (o₁, o₁ == o₂ ? 0 : o₂)  # for equator of odd nθ, disable southern ring
        θ_π = (ii - 0.5) / nθ  # θ/π for the ring
        sθ, cθ = sincospi(θ_π)
        ΔΩ = sθ * ΔθΔϕ
        rings[ii] = RingInfo(offsets, stride, nϕ, cθ, ϕ₀, ΔΩ)
    end
    return MapRings(rings)
end

"""
    maprings = gl_maprings(nθ, nϕ)

Constructs a `MapRings` structure describing an Gauss-Legendre (GL) grid of dimensions `nθ × nϕ` pixels in
column-major ordering.
"""
function gl_maprings(nθ, nϕ)
    nθh = (nθ + 1) ÷ 2  # number of rings in N. hemisphere + equator (if nθ is odd)
    ϕ₀, Δϕ = 2π / 2nϕ, 2π / nϕ  # constant for all rings
    cθ, sθΔθ = gausslegendre(nθ) # z = cos(θ) and corresponding weights
    stride = nθ  # for column-major ordering
    rings = Vector{RingInfo}(undef, nθh)
    for ii in 1:nθh
        o₁ = ii           # northern-ring offset
        o₂ = nθ - ii + 1  # southern-ring offset
        offsets = (o₁, o₁ == o₂ ? 0 : o₂)  # for equator of odd nθ, disable southern ring
        # N.B. z ∈ [-1, 1] and θ ∈ [0, π] run in opposite directions, and cos(π-θ) == -z
        rings[ii] = RingInfo(offsets, stride, nϕ, -cθ[ii], ϕ₀, sθΔθ[ii] * Δϕ)
    end
    return MapRings(rings)
end

"""
    (; npix, nϕmax) = _mapring_counts(maprings::MapRings)

Counts total number of pixels (`npix`) described by `maprings` and the longest observed ring (`nϕmax`).
"""
function _mapring_counts(maprings::MapRings)
    nϕmax, npix = 0, 0
    for ring in maprings.rings
        mult = mapreduce(!=(0), +, ring.offsets)  # single or paired rings?
        npix += mult * ring.nϕ                    # add 1 or 2 times Nϕ to total count
        nϕmax = max(nϕmax, ring.nϕ)               # keep track of longest seen ring
    end
    return (; npix, nϕmax)
end

"""
    (i′, isconj, isnyq) alias_index(len::Int, i::Int)

Calculates aliased index `i′` given original index `i` and Fourier axis length 'len'.
Also indicates whether conjugation (`isconj`) or twice-real symmetries (`isnyq`) are needed.
"""
@inline function alias_index(len::Int, i::Int)
    isconj, isnyq = false, false
    nyq = max(1, len ÷ 2)
    i < nyq && return (i, isconj, isnyq)
    i = mod(i, len)
    if i > nyq
        i = len - i
        isconj = true
    elseif i == 0 || (iseven(len) && i == nyq)
        isnyq = true
    end
    return (i, isconj, isnyq)
end

"""
    coeffs′ = alias_coeffs(coeffs::NTuple, isconj::Bool, isnyq::Bool)

Applies conjugation (`isconj`) or twice-real (`isnyq`) transformation, if required, in mode aliasing to the
harmonic coefficients `coeffs`.
"""
@inline function alias_coeffs(coeffs::NTuple, isconj::Bool, isnyq::Bool)
    return ifelse(isnyq, complex.(2 .* real.(coeffs)),
                  ifelse(isconj, conj.(coeffs), coeffs))
end

function synthesize(maprings::MapRings, alms::Matrix{C}) where {C<:Complex}
    R = real(C)
    lmax, mmax = size(alms) .- 1
    (; npix, nϕmax) = _mapring_counts(maprings)
    nϕr = nϕmax ÷ 2 + 1  # (maximum) real-symmetric FFT's Nyquist length (index)

    mapbuf = Vector{R}(undef, npix)
    Λ = Matrix{R}(undef, lmax + 1, mmax + 1)
    λ₁ = zeros(C, nϕr) # phase factors for northern ring
    λ₂ = zeros(C, nϕr) # phase factors for southern ring

    for ring in maprings.rings
        nϕ, ϕ₀, cθ = ring.nϕ, ring.ϕ₀, ring.cosθ

        λlm!(Λ, lmax, mmax, cθ)

        for m in 0:mmax
            acc₁, acc₂ = zero(C), zero(C)
            # Σ_{ℓ = m}^{lmax}
            for ℓ in m:lmax
                term = alms[ℓ+1,m+1] * Λ[ℓ+1,m+1]
                acc₁ += term
                acc₂ += isodd(ℓ + m) ? -term : term
            end
            acc₁, acc₂ = (acc₁, acc₂) .* cis(m * ϕ₀) # ring offset rotation
            i, isconj, isnyq = alias_index(nϕ, m)
            acc₁, acc₂ = alias_coeffs((acc₁, acc₂), isconj, isnyq)
            λ₁[i+1] += acc₁
            λ₂[i+1] += acc₂
        end

        stride = ring.stride
        o₁, o₂ = ring.offsets
        nϕh = nϕ ÷ 2 + 1
        if o₁ != 0
            idx = range(o₁, step = stride, length = nϕ)
            @views mapbuf[idx] = brfft(λ₁[1:nϕh], nϕ)
        end
        if o₂ != 0
            idx = range(o₂, step = stride, length = nϕ)
            @views mapbuf[idx] = brfft(λ₂[1:nϕh], nϕ)
        end
        λ₁ .= zero(C)
        λ₂ .= zero(C)
    end

    return mapbuf
end  # function synthesize

function analyze(maprings::MapRings, mapbuf::Vector{R}, lmax::Integer, mmax::Integer = lmax) where {R<:Real}
    C = complex(R)
    (; npix, nϕmax) = _mapring_counts(maprings)
    nϕr = nϕmax ÷ 2 + 1 # (maximum) real-symmetric FFT's Nyquist length (index)

    alms = fill(zero(C), lmax + 1, mmax + 1)
    Λ = zeros(R, lmax + 1, mmax + 1)
    λ₁ = zeros(C, nϕr) # phase factors for northern ring
    λ₂ = zeros(C, nϕr) # phase factors for southern ring

    for ring in maprings.rings
        nϕ, ϕ₀, cθ, ΔΩ = ring.nϕ, ring.ϕ₀, ring.cosθ, ring.ΔΩ

        λlm!(Λ, lmax, mmax, cθ) # λ_ℓ^m(cos θ) factors

        stride = ring.stride
        o₁, o₂ = ring.offsets
        nϕh = nϕ ÷ 2 + 1
        if o₁ != 0
            idx = range(o₁, step = stride, length = nϕ)
            @views λ₁[1:nϕh] = rfft(mapbuf[idx])
        end
        if o₂ != 0
            idx = range(o₂, step = stride, length = nϕ)
            @views λ₂[1:nϕh] = rfft(mapbuf[idx])
        end

        for m in 0:mmax
            i, isconj, isnyq = alias_index(nϕ, m)
            a₁, a₂ = alias_coeffs((λ₁[i+1], λ₂[i+1]), isconj, isnyq)
            a₁, a₂ = (a₁, a₂) .* (ΔΩ * cis(m * -ϕ₀))
            # Σ_{ℓ = m}^{lmax}
            for ℓ in m:lmax
                c = isodd(ℓ+m) ? a₁ - a₂ : a₁ + a₂
                alms[ℓ+1,m+1] += c * Λ[ℓ+1,m+1]
            end
        end
        λ₁ .= zero(C)
        λ₂ .= zero(C)
    end

    return alms
end  # function analyze

function alm2cl(alms::Matrix{<:Complex{R}}) where R
    lmax, mmax = size(alms) .- 1
    Cl = zeros(R, lmax + 1)
    for ll in 0:lmax
        Cl[ll+1] = abs2(alms[ll+1,1])
        for mm in 1:min(ll, mmax)
            Cl[ll+1] += 2 * abs2(alms[ll+1,mm+1])
        end
        Cl[ll+1] /= 2ll + 1
    end
    return Cl
end