text:
# MultipleBranches.jl
#  Extending the Feynman theory to multiple phonon branches

"""
    frohlichPartial((f, ϵ_mode); ϵ_o, ϵ_s, meff)

Calculate a (partial) dielectric electron-phonon coupling element.

# Arguments
- `f`: frequency of mode in THz.
- `ϵ_mode`: this mode's contribution to dielectric.
- `ϵ_o`: optical dielectric.
-  `ϵ_s`: total static dielectric contribution.
"""
function frohlichPartial((f, ϵ_mode); ϵ_o, ϵ_s, meff)
    ω = f * 1e12 * 2π
    α = 1 / (4π * ϵ_0) * ϵ_mode / (ϵ_o * ϵ_s) * (q^2 / ħ) * sqrt(meff * me / (2 * ω * ħ))
    return α
end

# deprecated signature wrapped via multiple dispatch
frohlichPartial(ϵ_o, ϵ_s, ϵ_mode, f, meff) = frohlichPartial((f, ϵ_mode), ϵ_o = ϵ_o, ϵ_s = ϵ_s, meff = meff)

rows(M::Matrix) = map(x -> reshape(getindex(M, x, :), :, size(M)[2]), 1:size(M)[1])

"""
    IRtoDielectric(IRmodes, volume)

From absolute value of IR activities of phonon modes, generate a per-mode
contribution to the low-frequency dielectric constant.

IRmodes are tuples f, S with Frequency in THz; InfraRed activity in e^2 amu^-1.
"""
function IRtoDielectric(IRmodes,volume)
    ϵ = 0.0 #* q^2 * THz^-2 * amu^-1 * m^-3
    for r in rows(IRmodes)
        f, S = r # frequency in THz; activity in e^2 amu^-1
        f = f * 1E12 #* THz
        ω = 2π * f
        S = S * q^2 / amu
        ϵ_mode = S / ω^2 / volume
        ϵ_mode /= 3 # take isotropic average = divide by 3
        ϵ += ϵ_mode
        println("Mode f= $f S= $S ϵ_mode = $(ϵ_mode / ϵ_0)")
    end
    println("Raw ionic dielectric contribution: $ϵ absolute $(ϵ / ϵ_0) relative")
    return ϵ / ϵ_0
end

"""
    IRtoalpha(IR, volume)

Calculates contribution to dielectric constant for each polar phonon mode, and thereby the Frohlich alpha contribution for this mode.
"""
function IRtoalpha(IR; volume, ϵ_o, ϵ_s, meff)
    ϵ = 0.0 #* q^2 * THz^-2 * amu^-1 * m^-3
    α_sum = 0.0
    for r in rows(IR)
        f, S = r # frequency in THz; activity in e^2 amu^-1
        ω = 2π * f * 1e12
        S = S * q^2 / amu
        ϵ_mode = S / ω^2 / volume
        ϵ_mode /= 3 # take isotropic average = divide by 3
        ϵ_mode /= ϵ_0 # reduced dielectric constant
        ϵ += ϵ_mode
        #println("Mode f= $f S= $S ϵ_mode = $(upreferred(ϵ_mode/u"ϵ0"))")
        α = frohlichPartial(ϵ_o, ϵ_s, ϵ_mode, f, meff)
        α_sum += α
        if (α > 0.1)
            println("Notable Mode f = $f, α_partial = $α.")
        end
    end
    println("Sum alpha: $(α_sum)")
    return α_sum
end

"""
    DieletricFromIRmode(IRmode)

Calculate dielectric from an individual mode.

IRmode is a tuple f, S with Frequency in THz; InfraRed activity in e^2 amu^-1.
"""
function DielectricFromIRmode(IRmode; volume)
    f, S = IRmode # Assumes Frequency in THz; InfraRed activity in e^2 amu^-1
    ω = 2π * f * 1e12
    S = S * q^2 / amu
    ϵ_mode = S / ω^2 / volume
    ϵ_mode /= 3 # take isotropic average = divide by 3
    ϵ_mode /= ϵ_0 # reduced dielectric
    return ϵ_mode
end

"""
    Hellwarth1999mobilityRHS((α, (v, w) ,f), effectivemass, T)

Calculates the DC mobility using Hellwarth et al. 1999 Eqn. (2).

See Hellwarth et a. 1999: https://doi.org/10.1103/PhysRevB.60.299.
"""
function Hellwarth1999mobilityRHS((α, (v, w) ,f), effectivemass, T)
    mb = effectivemass * MassElectron
    ω = f * 1e12 * 2π
    βred = ħ * ω / (kB * T)

    R = (v^2 - w^2) / (w^2 * v) # inline, page 300 just after Eqn (2)
    b = R * βred / sinh(βred * v / 2) # Feynman1962 version; page 1010, Eqn (47b)
    a = sqrt((βred / 2)^2 + R * βred * coth(βred * v / 2))
    k(u, a, b, v) = (u^2 + a^2 - b * cos(v * u))^(-3/2) * cos(u) # integrand in (2)
    K = quadgk(u -> k(u, a, b, v), 0, Inf)[1] # numerical quadrature integration of (2)

    # Right-hand-side of Eqn 1 in Hellwarth 1999 // Eqn (4) in Baggio1997
    RHS = α / (3 * sqrt(π)) * βred^(5/2) / sinh(βred / 2) * (v^3 / w^3) * K
    μ = RHS^(-1) * q / (ω * mb)

    return 1 / μ
end

"""
----------------------------------------------------------------------
Multiple Branch Frohlich Alpha
----------------------------------------------------------------------

Partial dielectric electron-phonon coupling parameter, decomposed into ionic dielectric contributions from each phonon mode of the matieral.
"""

"""
    ϵ_ionic_mode(phonon_mode_freq, ir_activity, volume)

Calculate the ionic contribution to the dielectric function for a given phonon mode.

# Arguments
- `phonon_mode_freq::Float64`: is the frequency of the mode in THz.
- `ir_activity::Float64`: is the infra-red activity of the mode in e²amu⁻¹.
- `volume::Float64`: is the volume of the unit cell of the material in m³.
"""
function ϵ_ionic_mode(phonon_mode_freq, ir_activity, volume) # single ionic mode

    # Angular phonon frequency for the phonon mode (rad Hz)
    ω_j = 2π * phonon_mode_freq * 1e12

    # Dielectric contribution from a single ionic phonon mode
    ϵ_mode = eV^2 * ir_activity / (3 * volume * ω_j^2 * amu)

    # Normalise ionic dielectric contribution with 1 / (4π ϵ_0) (NB: the 4π has been pre-cancelled)
    return ϵ_mode / ϵ_0
end

"""
    ϵ_total(freqs_and_ir_activity, volume)

Calculate the total ionic contribution to the dielectric function from all phonon modes.

# Arguments
- `freqs_and_ir_activity::Matrix{Float64}`: is a matrix containeing the phonon mode frequencies (in THz) in the first column and the infra-red activities (in e²amu⁻¹) in the second column.
- `volume::Float64`: is the volume of the unit cell of the material in m^3.
"""
function ϵ_total(freqs_and_ir_activity, volume) # total ionic contribution to dielectric

    # Extract phonon frequencies (THz)
    phonon_freqs = freqs_and_ir_activity[:, 1]

    # Extra infra-red activities (e^2 amu^-1)
    ir_activity = freqs_and_ir_activity[:, 2]

    # Sum over all ionic contribution from each phonon mode
    total_ionic = 0.0

    for t in 1:length(phonon_freqs)
        total_ionic += ϵ_ionic_mode(phonon_freqs[t], ir_activity[t], volume)
    end

    return total_ionic
end

"""
    effective_freqs(freqs_and_ir_activity, num_var_params)

Generates a matrix of effective phonon modes with frequencies and infra-red activities derived from a larger matrix using the Principal Component Analysis (PCA) method.

# Arguments
- `freqs_and_ir_activity::Matrix{Float64}`: is a matrix containing the phonon mode frequencies (in THz) in the first column and the infra-red activities (in e²amu⁻¹) in the second column.
- `num_var_params::Integer`: is the number of effective modes required (which needs to be less than the number of modes in `freqs_and_ir_activity``).

*** POSSIBLY REDUNDANT ***
"""
function effective_freqs(freqs_and_ir_activity, num_var_params) # PCA Algorithm

    # Check that the number of effective modes is less than the number of actual phonon modes.
    if num_var_params >= size(freqs_and_ir_activity)[1]

        println("The number of effective phonon modes has to be less than the total number of phonon modes.")

    else

        # Centralise data by subtracting the columnwise mean
        standardized_matrix = freqs_and_ir_activity' .- mean(freqs_and_ir_activity', dims = 2)

        # Calculate the covariance matrix S' * S. Matrix size is (n - 1) x (n - 1) for number of params (here n = 2)
        covariance_matrix = standardized_matrix' * standardized_matrix

        # Extract eigenvectors of the covariance matrix
        eigenvectors = eigvecs(covariance_matrix)

        # Project the original data along the covariance matrix eigenvectors and undo the centralisation
        reduced_matrix = standardized_matrix[:, 1:num_var_params] * eigenvectors[1:num_var_params, 1:num_var_params] *
        eigenvectors[1:num_var_params, 1:num_var_params]' .+ mean(freqs_and_ir_activity', dims = 2)

        # Resultant matrix is positive definite and transposed.
        return abs.(reduced_matrix')
    end
end

"""
    multi_frohlichalpha(ϵ_optic, ϵ_ionic, ϵ_total, phonon_mode_freq, m_eff)

Calculates the partial dielectric electron-phonon coupling parameter for a given longitudinal optical phonon mode.

This decomposes the original Frohlich alpha coupling parameter (defined for a single phonon branch) into contributions from multiple phonon
branches.

# Arguments
- `ϵ_optic::Float64`: is the optical dielectric constant of the material.
- `ϵ_ionic::Float64`: is the ionic dielectric contribution from the phonon mode.
- `ϵ_total::Float64`: is the total ionic dielectric contribution from all phonon modes of the material.
- `phonon_mode_freq::Float64`: is the frequency of the phonon mode (THz).
- `m_eff::Float64` is the band mass of the electron (in units of electron mass m_e).
"""
function multi_frohlichalpha(ϵ_optic, ϵ_ionic, ϵ_total, phonon_mode_freq, m_eff)

    # The Rydberg energy unit
    Ry = eV^4 * me / (2 * ħ^2)

    # Angular phonon frequency for the phonon mode (rad Hz).
    ω = 2π * 1e12 * phonon_mode_freq

    # The static dielectric constant. Calculated here instead of inputted so that ionic modes are properly normalised.
    ϵ_static = ϵ_total + ϵ_optic

    # The contribution to the electron-phonon parameter from the currrent phonon mode. 1 / (4π ϵ_0) is the dielectric normalisation.
    α_j  = (m_eff * Ry / (ħ * ω))^(1 / 2) * ϵ_ionic / (4π * ϵ_0) / (ϵ_optic * ϵ_static)

    return α_j
end

"""
----------------------------------------------------------------------
Multiple Branch Polaron Free Energy
----------------------------------------------------------------------

Calculate the polaron free energy, generalised from Osaka's expression to the case where multiple phonon modes are present in the material.
"""

"""
    κ_i(i, v, w)

Calculates the spring-constant coupling the electron to the 'ith' fictitious mass that approximates the exact electron-phonon interaction with a harmonic coupling to a massive fictitious particle.

Required for calculating the polaron free energy.

Note: Not to be confused with the number of physical phonon branches; many phonon branches could be approximated with one or two etc. fictitious masses for example. The number of fictitious mass does not necessarily need to match the number of phonon branches.

# Arguments
- `i::Integer`: enumerates the current fictitious mass.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
"""
function κ_i(i, v, w)
    κ = v[i]^2 - w[i]^2
    if length(v) > 1
        for j in 1:length(v)
            if j != i
                κ *= (v[j]^2 - w[i]^2) / (w[j]^2 - w[i]^2)
            end
        end
    end
    return κ
end

"""
    h_i(i, v, w)

Calculates the normal-mode (the eigenmodes) frequency of the coupling between the electron and the `ith' fictitious mass that approximates the exact electron-phonon interaction with a harmonic coupling to a massive fictitious particle.

Required for calculating the polaron free energy.

Note: Not to be confused with the number of physical phonon branches; many phonon branches could be approximated with one or two etc. fictitious masses for example. The number of fictitious mass does not necessarily need to match the number of phonon branches.

# Arguments
- `i::Integer`: enumerates the current fictitious mass.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
"""
function h_i(i, v, w)
    h = v[i]^2 - w[i]^2
    if length(v) > 1
        for j in 1:length(v)
            if j != i
                h *= (w[j]^2 - v[i]^2) / (v[j]^2 - v[i]^2)
            end
        end
    end
    return h
end

"""
    C_ij(i, j, v, w)

Calculates the element to the coupling matrix C_ij (a generalisation of Feynman's `C` coupling variational parameter in Feynman 1955) between the electron and the `ith' and `jth' fictitious masses that approximates the exact electron-phonon interaction with a harmonic coupling to a massive fictitious particle.

Required for calculating the polaron free energy.

Note: Not to be confused with the number of physical phonon branches; many phonon branches could be approximated with one or two etc. fictitious masses for example. The number of fictitious mass does not necessarily need to match the number of phonon branches.

# Arguments
- `i::Integer, j::Integer`: enumerate the current fictitious masses under focus (also the index of the element in the coupling matrix C)
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.

See Feynman 1955: http://dx.doi.org/10.1103/PhysRev.97.660.
"""
function C_ij(i, j, v, w)
    C = w[i] * κ_i(i, v, w) * h_i(j, v, w) / (4 * (v[j]^2 - w[i]^2))
    return C
end

"""
    D_j(τ, β, v, w)

Calculates the recoil function (a generalisation of D(u) in Eqn. (35c) in FHIP 1962) that approximates the exact influence (recoil effects) of the phonon bath on the electron with the influence of the fictitious masses attached by springs to the electron. It appears in the exponent of the intermediate scattering function.

# Arguments
- `τ::Float64`: is the imaginary time variable.
- `β::Float64`: is the reduced thermodynamic temperature ħωⱼ/(kT) associated with the 'jth' phonon mode.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.

See FHIP 1962: https://doi.org/10.1103/PhysRev.127.1004.
"""
function D_j(τ, β, v, w)
    D = τ * (1 - τ / β)
    for i in 1:length(v)
        D += (h_i(i, v, w) / v[i]^2) * ((1 + exp(-v[i] * β) - exp(-v[i] * τ) - exp(v[i] * (τ - β))) / (v[i] * (1 - exp(-v[i] * β))) - τ * (1 - τ / β))
    end
    return D
end

"""
    D_j(τ, v, w)

Calculates the recoil function at zero-temperature.

# Arguments
- `τ::Float64`: is the imaginary time variable.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.

See also ['D_j'](@ref).
"""
function D_j(τ, v, w)
    D = τ
    for i in 1:length(v)
        D += (h_i(i, v, w) / v[i]^2) * ((1 - exp(-v[i] * τ)) / v[i] - τ)
    end
    return D
end

"""
    B_j(α, β, v, w; rtol = 1e-3)

Generalisation of the B function from Eqn. (62c) in Hellwarth et al. 1999. This is the expected value of the exact action <S_j> taken w.r.t trial action, given for the 'jth' phonon mode.

Required for calculating the polaron free energy.

# Arguments
- `α::Float64`: is the partial dielectric electron-phonon coupling parameter for the 'jth' phonon mode.
- `β::Float64`: is the reduced thermodynamic temperature ħωⱼ/(kT) associated with the 'jth' phonon mode.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `rtol`: relative tolerance for the accuracy of any quadrature integrations.

See Hellwarth, R. W., Biaggio, I. (1999): https://doi.org/10.1103/PhysRevB.60.299.

See also ['B'](@ref).
"""
function B_j(α, β, v, w; rtol = 1e-3)
    B_integrand(τ) = (exp(-τ) + exp(τ - β)) / sqrt(abs(D_j(τ, β, v, w)))
    B = α / (√π * (1 - exp(-β))) * quadgk(τ -> B_integrand(τ), 0.0, β / 2, rtol = rtol)[1]
    return B
end

"""
    B_j(α, v, w; rtol = 1e-3)

Calculates `B_j(α, β, v, w)` but at zero-temperature, `β = Inf`.

# Arguments
- `α::Float64`: is the partial dielectric electron-phonon coupling parameter for the 'jth' phonon mode.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `rtol`: relative tolerance for the accuracy of any quadrature integrations.

See also [`B_j`](@ref).
"""
function B_j(α, v, w; rtol = 1e-3)
    B_integrand(τ) = exp(-abs(τ)) / sqrt(abs(D_j(abs(τ), v, w)))
    B = α / √π * quadgk(τ -> B_integrand(τ), 0.0, Inf64, rtol = rtol)[1]
    return B
end

"""
    C_j(β, v, w, n)

Generalisation of the C function from Eqn. (62e) in Hellwarth et al. 1999. This is the expected value of the trial action <S_0> taken w.r.t trial action.

Required for calculating the polaron free energy.

# Arguments
- `β::Float64`: is the reduced thermodynamic temperature ħωⱼ/(kT) associated with the 'jth' phonon mode.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `n::Integer`: is the number of phonon modes.

See Hellwarth, R. W., Biaggio, I. (1999): https://doi.org/10.1103/PhysRevB.60.299.

See also ['C'](@ref).
"""
function C_j(β, v, w, n)
    s = 0.0
    # Sum over the contributions from each fictitious mass.
    for i in 1:length(v)
        for j in 1:length(v)
            s += C_ij(i, j, v, w) / (v[j] * w[i]) * (coth(β * v[j] / 2)  - 2 / (β * v[j]))
        end
    end
    # Divide by the number of phonon modes to give an average contribution per phonon mode.
    return 3 * s / n
end

"""
    C_j(v, w, n)

Calculates `C_j(β, v, w, n)` but at zero-temperature, `β = Inf`.

# Arguments
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `n::Integer`: is the number of phonon modes.

See also [`C_j`](@ref).
"""
function C_j(v, w, n)
    s = 0.0
    # Sum over the contributions from each fictitious mass.
    for i in 1:length(v)
        for j in 1:length(v)
            s += C_ij(i, j, v, w) / (v[j] * w[i])
        end
    end
    # Divide by the number of phonon modes to give an average contribution per phonon mode.
    return 3 * s / n
end

"""
    A_j(β, v, w, n)

Generalisation of the A function from Eqn. (62b) in Hellwarth et al. 1999. This is the Helmholtz free energy of the trial model.

Required for calculating the polaron free energy.

# Arguments
- `β::Float64`: is the reduced thermodynamic temperature ħωⱼ/(kT) associated with the 'jth' phonon mode.
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `n::Integer`: is the number of phonon modes.

See Hellwarth, R. W., Biaggio, I. (1999): https://doi.org/10.1103/PhysRevB.60.299.

See also ['A'](@ref).
"""
function A_j(β, v, w, n)
    s = -log(2π * β) / 2
    # Sum over the contributions from each fictitious mass.
    for i in 1:length(v)
        if v[i] != w[i]
            s += v[i] == w[i] ? 0.0 : log(v[i] / w[i]) - log(sinh(v[i] * β / 2) / sinh(w[i] * β / 2))
        end
    end
    # Divide by the number of phonon modes to give an average contribution per phonon mode.
    3 / β * s / n
end

"""
    A_j(v, w, n)

Calculates `A_j(β, v, w, n)` but at zero-temperature, `β = Inf`.

# Arguments
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `n::Integer`: is the number of phonon modes.

See also [`A_j`](@ref).
"""
function A_j(v, w, n)
    s = sum(v .- w)
    return -3 * s / (2 * n)
end

"""
    multi_F(v, w, α, β; ω = 1.0, rtol = 1e-3, T = nothing, verbose = false)

Calculates the Helmholtz free energy of the polaron for a material with multiple phonon branches.

This generalises the Osaka 1959 (below Eqn. (22)) and Hellwarth. et al 1999 (Eqn. (62a)) free energy expressions.

# Arguments
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `α::Float64`: is the partial dielectric electron-phonon coupling parameter for the 'jth' phonon mode.
- `β::Float64`: is the reduced thermodynamic temperature ħωⱼ/(kT) associated with the 'jth' phonon mode.
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `rtol`: relative tolerance for the accuracy of any quadrature integrations.
- `T`: is a token used by `make_polaron()` to keep track of the temperature for printing during a calculation. Do not alter.
- `verbose`: is used by `make_polaron()` to specify whether or not to print. Ignore.

See Osaka, Y. (1959): https://doi.org/10.1143/ptp.22.437 and Hellwarth, R. W., Biaggio, I. (1999): https://doi.org/10.1103/PhysRevB.60.299.

See also [`F`](@ref).
"""
function multi_F(v, w, α, β; ω = 1.0, rtol = 1e-3, T = nothing, verbose = false)

    num_modes = length(ω)

    # Add contribution to the total free energy from the phonon mode.
    F = sum(-(B_j(α[j], β[j], v, w; rtol = rtol) + C_j(β[j], v, w, num_modes) + A_j(β[j], v, w, num_modes)) * ω[j] for j in 1:num_modes)

    # Print the free energy.
    if verbose
        println("\e[2K", "Process: $(count) / $processes ($(round.(count / processes * 100, digits = 1)) %) | T = $(round.(T, digits = 3)) | F = $(round.(F, digits = 3))")
        print("\033[F")

        global count += 1
    end

    # Free energy in units of meV
    return F
end

"""
    multi_F(v, w, α; ω = 1.0, rtol = 1e-3, verbose = false)

Calculates the zero-temperature ground-state energy of the polaron for a material with multiple phonon branches. Similar to `multi_F(v, w, α, β)` but with `β = Inf`. Generalises Eqn. (33) in Feynman 1955.

# Arguments
- `v::Vector{Float64}`: is a vector of the v variational parameters.
- `w::Vector{Float64}`: is a vector of the w variational parameters.
- `α::Float64`: is the partial dielectric electron-phonon coupling parameter for the 'jth' phonon mode.
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `verbose`: is used by `make_polaron()` to specify whether or not to print. Ignore.

See Feynman 1955: http://dx.doi.org/10.1103/PhysRev.97.660.

See also [`multi_F`](@ref).
"""
function multi_F(v, w, α; ω = 1.0, rtol = 1e-3, verbose = false)

    num_modes = length(ω)

    # Add contribution to the total free energy from the phonon mode.
	F = sum(-(B_j(α[j], v, w; rtol = rtol) + C_j(v, w, num_modes) + A_j(v, w, num_modes)) * ω[j] for j in 1:num_modes)

    # Print the free energy.
    if verbose
        println("\e[2K", "Process: $(count) / $processes ($(round.(count / processes * 100, digits = 1)) %) | T = 0.0 | F = $(round.(F, digits = 3))")
        print("\033[F")
        global count += 1
    end

    # Free energy in units of meV
    return F
end

"""
    var_params(α, β; v = 0.0, w = 0.0, ω = 1.0, N = 1, rtol = 1e-3, show_trace = false, T = nothing, verbose = false)

Minimises the multiple phonon mode free energy function for a set of vₚ and wₚ variational parameters. The variational parameters follow the inequality: v₁ > w₁ > v₂ > w₂ > ... > vₙ > wₙ. Generalises `feynmanvw` to multiple variational parameters.

# Arguments
- `α::Vector{Float64}`: is the partial dielectric electron-phonon coupling parameter for the 'jth' phonon mode.
- `β::Vector{Float64}`: is the reduced thermodynamic temperature ħωⱼ/(kT) associated with the 'jth' phonon mode.
- `v::Float64, w::Float64`: determines if the function should start with a random initial set of variational parameters (v, w = 0.0) or a given set of variational parameter values.
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `N::Integer`: specifies the number of variational parameter pairs, v_p and w_p, to use in minimising the free energy.
- `rtol`: relative tolerance for the accuracy of any quadrature integrations.
- `show_trace::Bool`: shows the optimsation trace from `Optim.jl`.
- `T`: is a token used by `make_polaron()` to keep track of the temperature for printing during a calculation. Do not alter.
- `verbose`: is used by `make_polaron()` to specify whether or not to print. Ignore.

See also [`multi_F`](@ref), [`feynmanvw`](@ref).
"""
function var_params(α, β; v = 0.0, w = 0.0, ω = 1.0, N = 1, rtol = 1e-3, show_trace = false, T = nothing, verbose = false) # N number of v and w params

    if N != length(v) != length(w)
        return error("The number of variational parameters v & w must be equal to N.")
    end

    # Use a random set of N initial v and w values.
    if v == 0.0 || w == 0.0
		# Intial guess for v and w parameters.
    	initial = sort(rand(2 * N), rev=true) .* 4.0 .+ 1.0 # initial guess around 4 and ≥ 1.
	else
        Δv = v .- w
        initial = vcat(Δv .+ rtol, w)
    end

    # Limits of the optimisation.
    lower = fill(0.0, 2 * N)
    upper = fill(1000.0, 2 * N)

	# The multiple phonon mode free energy function to minimise.
	f(x) = multi_F([x[2 * n - 1] for n in 1:N] .+ [x[2 * n] for n in 1:N], [x[2 * n] for n in 1:N], α, β; ω = ω, rtol = rtol)

	# Use Optim to optimise the free energy function w.r.t the set of v and w parameters.
	solution = Optim.optimize(
		Optim.OnceDifferentiable(f, initial; autodiff = :forward),
		lower,
		upper,
		initial,
		Fminbox(BFGS()),
		Optim.Options(show_trace = show_trace), # Set time limit for asymptotic convergence if needed.
	)

	# Extract the v and w parameters that minimised the free energy.
	var_params = Optim.minimizer(solution)

    if Optim.converged(solution) == false
        @warn "Failed to converge feynmanvw solution."
    end

	# Separate the v and w parameters into one-dimensional arrays (vectors).
	Δv = [var_params[2 * n - 1] for n in 1:N]
	w = [var_params[2 * n] for n in 1:N]

    # Print the variational parameters that minimised the free energy.
	if verbose
        println("\e[2K", "Process: $(count) / $processes ($(round.(count / processes * 100, digits = 1)) %) | T = $(round.(T, digits = 3)) | v = $(round.(Δv .+ w, digits = 3)) | w = $(round.(w, digits = 3))")
        print("\033[F")

        global count += 1
    end

    # Return the variational parameters that minimised the free energy.
    return hcat(Δv .+ w, w)
end

"""
    var_params(α; v = 0.0, w = 0.0, ω = 1.0, N = 1, rtol = 1e-3, show_trace = false, verbose = false)

Minimises the multiple phonon mode free energy function for a set of vₚ and wₚ variational parameters at zero-temperature. Similar to `var_params(α, β)` but with `β = Inf`.

# Arguments
- `α::Vector{Float64}`: is the partial dielectric electron-phonon coupling parameter for the 'jth' phonon mode.
- `v::Float64, w::Float64`: determines if the function should start with a random initial set of variational parameters (v, w = 0.0) or a given set of variational parameter values.
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `N::Integer`: specifies the number of variational parameter pairs, v_p and w_p, to use in minimising the free energy.
- `rtol`: relative tolerance for the accuracy of any quadrature integrations.
- `show_trace::Bool`: shows the optimsation trace from `Optim.jl`.
- `verbose`: is used by `make_polaron()` to specify whether or not to print. Ignore.

See also [`multi_F`](@ref), [`feynmanvw`](@ref), [`var_param`](@ref).
"""
function var_params(α; v = 0.0, w = 0.0, ω = 1.0, N = 1, rtol = 1e-3, show_trace = false, verbose = false) # N number of v and w params

    if N != length(v) != length(w)
        return error("The number of variational parameters v & w must be equal to N.")
    end

    # Use a random set of N initial v and w values.
    if v == 0.0 || w == 0.0
		# Intial guess for v and w parameters.
    	initial = sort(rand(2 * N), rev=true) .* 4.0 .+ 1.0 # initial guess around 4 and ≥ 1.
	else
        Δv = v .- w
        initial = vcat(Δv .+ rtol, w)
    end

    # Limits of the optimisation.
    lower = fill(0.0, 2 * N)
    upper = fill(1000, 2 * N)

	# The multiple phonon mode free energy function to minimise.
	f(x) = multi_F([x[2 * n - 1] for n in 1:N] .+ [x[2 * n] for n in 1:N], [x[2 * n] for n in 1:N], α; ω = ω, rtol = rtol)

	# Use Optim to optimise the free energy function w.r.t the set of v and w parameters.
	solution = Optim.optimize(
		Optim.OnceDifferentiable(f, initial; autodiff = :forward),
		lower,
		upper,
		initial,
		Fminbox(BFGS()),
		Optim.Options(show_trace = show_trace), # Set time limit for asymptotic convergence if needed.
	)

	# Extract the v and w parameters that minimised the free energy.
	var_params = Optim.minimizer(solution)

	# Separate the v and w parameters into one-dimensional arrays (vectors).
	Δv = [var_params[2 * n - 1] for n in 1:N]
	w = [var_params[2 * n] for n in 1:N]

	# Print the variational parameters that minimised the free energy.
	if verbose
        println("\e[2K", "Process: $(count) / $processes ($(round.(count / processes * 100, digits = 1)) %) | T = 0.0 | v = $(round.(Δv .+ w, digits = 3)) | w = $(round.(w, digits = 3))")
        print("\033[F")

        global count += 1
    end

    # Return the variational parameters that minimised the free energy.
    return hcat(Δv .+ w, w)
end


text:
module ChainRulesTestUtils

@nospecialize

using ChainRulesCore
using ChainRulesCore: frule, rrule
using Compat: only
using FiniteDifferences
using FiniteDifferences: to_vec
using LinearAlgebra
using Random
using Test

export TestIterator
export test_approx, test_scalar, test_frule, test_rrule, generate_well_conditioned_matrix
export ⊢, rand_tangent
export @maybe_inferred

__init__() = init_test_inferred_setting!()

include("global_config.jl")

include("rand_tangent.jl")
include("generate_tangent.jl")
include("data_generation.jl")
include("iterator.jl")

include("output_control.jl")
include("check_result.jl")

include("rule_config.jl")
include("finite_difference_calls.jl")
include("testers.jl")
end # module

ext: jl


text:
const AbstractTuple = Tuple #Union{Tuple, Flux.Tracker.TrackedTuple}

"abstract type AbstractModel <: ControlSystems.LTISystem end"
abstract type AbstractModel <: ControlSystems.LTISystem end
Base.Broadcast.broadcastable(p::AbstractModel) = Ref(p)

"""
    struct AR <: AbstractModel

Represents an all-pole transfer function, i.e., and AR model

# Arguments:
- `a`: denvec
- `ac`: denvec cont. time
- `p`: discrete time poles
- `pc`: continuous time poles
- `b`: Numerator scalar
"""
struct AR{T,Rt <: DiscreteRoots,Ct <: ContinuousRoots} <: AbstractModel
    a::T
    ac::T
    p::Rt
    pc::Ct
    b::Float64
    # function AR(xo::AbstractTuple,λ=1e-2)
    #     a = ls(getARregressor(xo[1], xo[2]),λ) |> polyvec
    #     r = DiscreteRoots(hproots(reverse(nograd(a))))
    #     rc = ContinuousRoots(r)
    #     ac = roots2poly(rc)
    #     new{typeof(a), typeof(r), typeof(rc)}(a, ac, r, rc)
    # end
end
function AR(a::AbstractVector, σ²=nothing)
    a = real.(a) # SVector{length(a)}(a)# This is a zygote workaround
    r = DiscreteRoots(hproots(rev(a)))
    rc = ContinuousRoots(r)
    ac = roots2poly(rc)
    ac = real.(ac)
    b = σ² === nothing ? one(eltype(a)) : scalefactor(Continuous(), ac, σ²)
    AR{typeof(a), typeof(r), typeof(rc)}(a, ac, r, rc, b)
end

function AR(::Continuous, ac::AbstractVector, σ²=nothing)
    ac = real.(ac) # This is a zygote workaround
    rc = ContinuousRoots(hproots(rev(ac)))
    r = DiscreteRoots(rc)
    a = roots2poly(r)
    a = real.(a)
    b = σ² === nothing ? one(eltype(a)) : scalefactor(Continuous(), ac, σ²)
    AR{typeof(a), typeof(r), typeof(rc)}(a, ac, r, rc, b)
end
function AR(rc::ContinuousRoots, σ²=nothing)
    r = DiscreteRoots(rc)
    a = roots2poly(r)
    ac = roots2poly(rc)
    b = σ² === nothing ? one(eltype(a)) : scalefactor(Continuous(), ac, σ²)
    AR{typeof(a), typeof(r), typeof(rc)}(a, ac, r, rc, b)
end
function AR(r::DiscreteRoots, σ²=nothing)
    rc = ContinuousRoots(r)
    a = roots2poly(r)
    ac = roots2poly(rc)
    b = σ² === nothing ? one(eltype(a)) : scalefactor(Continuous(), ac, σ²)
    AR{typeof(a), typeof(r), typeof(rc)}(a, ac, r, rc, b)
end

function Base.show(io::IO, ::MIME"text/plain", m::AR{T,Rt,Ct}) where {T,Rt,Ct}
    println(io,"AR{coeff type: $(eltype(T)), root type: $(eltype(Ct))}(")
    println(io,"b: ", m.b)
    printstyled(io,"Cont. poles: ", bold=true)
    println(io, repr("text/plain", m.pc))
    println(io, "Abs disc: ", repr(real.(abs.(m.p))))
    printstyled(io,"Cont. coeffs: ", bold=true)
    println(io, repr("text/plain", m.ac))
    print(io,"Spectral energy: ", spectralenergy(Continuous(), m))
end

function Base.show(io::IO, m::AR{T,Rt,Ct}) where {T,Rt,Ct}
    print(io,"AR{coeff: $(eltype(T)), root: $(eltype(Ct))} Abs disc ∈: ", repr(round.(extrema(real.(abs.(m.p))), sigdigits=5)))
end

"`checkroots(r::DiscreteRoots)` prints a warning if there are roots on the negative real axis."
checkroots(r::DiscreteRoots) = any(imag(r) == 0 && real(r) < 0 for r in r) && @warn("Roots on negative real axis, no corresponding continuous time representation exists. Consider prefiltering the signal or decreasing the regularization factor.", maxlog=5)

"""
    struct ARMA{T} <: AbstractModel

Represents an ARMA model, i.e., transfer function

# Arguments:
- `b`: numvec
- `bc`: numvec cont. time
- `a`: denvec
- `ac`: denvec cont. time
- `z`: zeros
- `p`: poles
"""
struct ARMA{T1,T2,T3,T4,Rt <: DiscreteRoots,Crt <: ContinuousRoots} <: AbstractModel
    b::T1
    bc::T2
    a::T3
    ac::T4
    z::Rt
    zc::Crt
    p::Rt
    pc::Crt

    function ARMA(b::AbstractVector,a::AbstractVector)
        r = DiscreteRoots(hproots(rev(a)))
        rc = ContinuousRoots(r)
        ac = roots2poly(rc)
        ac = real.(ac) # This is a zygote workaround
        z  = DiscreteRoots(hproots(rev(b)))
        zc = ContinuousRoots(z)
        b   = roots2poly(z)
        bc  = roots2poly(zc)
        new{typeof(b), typeof(bc), typeof(a), typeof(ac), typeof(r), typeof(rc)}(b, bc, a, ac, z, zc, r, rc)
    end

    function ARMA(::Continuous, bc, ac::AbstractVector)
        rc = ContinuousRoots(hproots(rev(ac)))
        zc = ContinuousRoots(hproots(rev(bc)))
        r = DiscreteRoots(rc)
        z = DiscreteRoots(zc)
        a = roots2poly(r) #|> Vector
        b = roots2poly(z) #|> Vector
        new{typeof(b), typeof(bc), typeof(a), typeof(ac), typeof(r), typeof(rc)}(b, bc, a, ac, z, zc, r, rc)
    end
    function ARMA(zc::ContinuousRoots, rc::ContinuousRoots)
        r = DiscreteRoots(rc)
        z = DiscreteRoots(zc)
        a = roots2poly(r)
        ac = roots2poly(rc)
        b = roots2poly(z)
        bc = roots2poly(zc)
        new{typeof(b), typeof(bc), typeof(a), typeof(ac), typeof(r), typeof(rc)}(b, bc, a, ac, z, zc, r, rc)
    end
    function ARMA(z::DiscreteRoots, r::DiscreteRoots)
        rc = ContinuousRoots(r)
        zc = ContinuousRoots(z)
        a = roots2poly(r)
        ac = roots2poly(rc)
        b = roots2poly(z)
        bc = roots2poly(zc)
        new{typeof(b), typeof(bc), typeof(a), typeof(ac), typeof(r), typeof(rc)}(b, bc, a, ac, z, zc, r, rc)
    end

end

hproots(a::AbstractVector{T}) where T = eigsort(Complex{float(T)}.(roots(Double64.(a))))
hproots(a::AbstractVector{<:Double64}) = eigsort(roots(a))
"""
    ControlSystems.tf(m::AR, ts)

Convert model to a transfer function compatible with ControlSystems.jl
"""
ControlSystems.tf(m::AR, ts) = tf(1, Vector(m.a), ts)
"""
    ControlSystems.tf(m::AR)

Convert model to a transfer function compatible with ControlSystems.jl
"""
ControlSystems.tf(m::AR) = tf(1, Vector(m.ac))
ControlSystems.tf(m::ARMA, ts) = tf(Vector(m.c), Vector(m.a), ts)
ControlSystems.tf(m::ARMA) = tf(Vector(m.cc), Vector(m.ac))

Base.convert(::Type{T}, m::AbstractModel) where T <: TransferFunction{Continuous, <:ControlSystems.SisoRational} = tf(m)

Base.convert(::Type{T}, m::AbstractModel) where T <: TransferFunction{Continuous, <:ControlSystems.SisoZpk} = zpk(tf(m))

function ARMA(g::ControlSystems.TransferFunction)
    ControlSystems.issiso(g) || error("Can only convert SISO systems to ARMA")
    b,a = numvec(g)[], denvec(g)[]
    ControlSystems.iscontinuous(g) && return ARMA(Continuous(), b, a)
    ARMA(b, a)
end

Base.convert(::Type{ControlSystems.TransferFunction}, m::AbstractModel) = tf(m)
Base.promote_rule(::Type{<:ControlSystems.TransferFunction}, ::Type{<:AbstractModel}) = ControlSystems.TransferFunction

function Base.isapprox(m1::AR, m2::AR, args...; kwargs...)
    all(isapprox(getfield(m1,field), getfield(m2,field), args...; kwargs...) for field in fieldnames(typeof(m1)))
end

function change_precision(F, m::AR)
    CF = Complex{F}
    a  = F.(m.a)
    r  = change_precision(F,m.p)
    rc = change_precision(F,m.pc)
    AR{typeof(a), typeof(r), typeof(rc)}(a, F.(m.ac), r, rc, m.b)
end

"""
    roots(m::AbstractModel)

Returns the roots of a model
"""
roots

PolynomialRoots.roots(::Discrete, m::AR) = m.p
PolynomialRoots.roots(::Continuous, m::AR) = m.pc
PolynomialRoots.roots(::Discrete, m::ARMA) = m.p
PolynomialRoots.roots(::Continuous, m::ARMA) = m.pc
PolynomialRoots.roots(::Continuous, r::ContinuousRoots) = r
ControlSystems.pole(d::TimeEvolution, m::AbstractModel) = roots(d,m)
ControlSystems.pole(m::AbstractModel) = roots(Continuous(), m)
ControlSystems.tzero(m::AbstractModel) = []
ControlSystems.tzero(::Discrete, m::ARMA) = m.z
ControlSystems.tzero(::Continuous, m::ARMA) = m.zc
"""
    ControlSystems.denvec(::TimeEvolution, m::AbstractModel)

Get the denominator polynomial vector
"""
ControlSystems.denvec(::Discrete, m::AbstractModel) = m.a
ControlSystems.numvec(::Discrete, m::AR) = error("Not yet implemented")#[m.b]
ControlSystems.numvec(::Continuous, m::AR) = [m.b]
ControlSystems.numvec(::Discrete, m::ARMA) = m.b
ControlSystems.denvec(::Continuous, m::AbstractModel) = m.ac
ControlSystems.numvec(::Continuous, m::ARMA) = m.bc

for T in (:AR, :ARMA)
    @eval begin
        ControlSystems.bode(m::$T, w::AbstractVector{<:Real}, args...; kwargs...) = bode(tf(m), w, args...; kwargs...)
        ControlSystems.nyquist(m::$T, w::AbstractVector{<:Real}, args...; kwargs...) = nyquist(tf(m), w, args...; kwargs...)
        ControlSystems.freqresp(m::$T, w::AbstractVector{<:Real}, args...; kwargs...) = freqresp(tf(m), w, args...; kwargs...)
        ControlSystems.step(m::$T, Tf::Real, args...; kwargs...) = step(tf(m), Tf, args...; kwargs...)
    end
end

function Base.getproperty(m::AbstractModel, p::Symbol)
    p === :Ts && return 0.0
    p === :nx && return length(m.ac)
    p === :nu && return 1
    p === :ny && return 1
    p === :timeevol && return Continuous()
    getfield(m,p)
end

"""
    coefficients(::Domain, m::AbstractModel)

Return all fitted coefficients
"""
coefficients(::Discrete, m::AR) = m.a[2:end]
coefficients(::Discrete, m::ARMA) = [m.a[2:end]; m.b]
coefficients(::Continuous, m::AR) = m.ac[2:end]
coefficients(::Continuous, m::ARMA) = [m.ac[2:end]; m.bc]

Base.getindex(m::AbstractModel, i) = i < length(m.ac) ? m.ac[1+i] : m.b[i-length(m.ac)+1]
Base.length(m::AbstractModel) = length(m.a)+length(m.b)-1
Base.size(m::AbstractModel) = (1,1)

function domain_transform(d::Continuous, m::AR)
    p = domain_transform(d, roots(m))
    roots2poly(p)
end

"Abstract type that represents a way to fit a model to data"
abstract type FitMethod end
Base.Broadcast.broadcastable(p::FitMethod) = Ref(p)
(fm::FitMethod)(X) = fitmodel(fm, X)

fitmodel(fm,X::AbstractModel) = X

"""
    LS <: FitMethod

This fitmethod is a good default option.

# Arguments:
- `na::Int`: number of roots (order of the system). The number of peaks in the spectrum will be `na÷2`.
- `λ::Float64 = 0.01`: reg factor

# Computational performance improvements
Estimating a lot (1000s) of models might take a while. The bulk of the time is spent performing a matrix factorization, something that can be significantly sped up by
- Using `Flaot32` instead of `Float64`
- Use [MKL.jl](https://github.com/JuliaComputing/MKL.jl) instead of the default OpenBLAS (can yield about 2x performance improvement).

"""
LS
@kwdef struct LS <: FitMethod
    na::Int
    λ::Float64 = 1e-2
end

"""
    fitmodel(fm::LS, X::AbstractArray)

## Computational performance improvements
Estimating a lot (1000s) of models might take a while. The bulk of the time is spent performing a matrix factorization, something that can be significantly sped up by
- Using `Flaot32` instead of `Float64`
- Use [MKL.jl](https://github.com/JuliaComputing/MKL.jl) instead of the default OpenBLAS (can yield about 2x performance improvement).
"""
function fitmodel(fm::LS,X::AbstractArray)
    y,A = getARregressor(X, fm.na)
    a = ls(A, y, fm.λ) |> polyvec
    AR(a, var(X))
end

function fitmodel!(AS, fm::LS, X::AbstractArray)
    y,A = getARregressor!(AS, X, fm.na)
    a = ls(A, y, fm.λ, true) |> polyvec
    AR(a, var(X))
end

# diffpol(n) = [(-1)^k*binomial(n,k) for k in 1:n]
diffpol(n) = diagm(0=>ones(n), 1=>-ones(n))[:,1:n]

"""
    ls(A, y, λ=0)

Regularized Least-squares
"""
function ls(A, y, λ=0, regspace=false)
    isderiving() && (return _ls(A,y,λ))
    n = size(A,2)
    if λ > 0
        if regspace # There's room for the regularization
            A[end-n+1:end, :] .= (sqrt(λ)*I)(n)
            y[end-n+1:end]    .= 0
        else
           A = [A; sqrt(λ)*I]
           y = [y;zeros(n)]
        end
    end
    # @show size(A2), size(y)
    # svd(A2)\y
    # (A'A - γ*I)\(A'y)
    A\y
end

function _ls(A,y,λ)
    # if λ > 0
        # A2 = [A; sqrt(λ)*I]
    # else
        # A2 = A
    # end
    # @show size(A2), size(y), size(A)
   # (A2'A2)\(A'y)
   (A'A + λ*I)\(A'y)
end

"""
    TLS <: FitMethod

Total least squares. This fit method is good if the spectrum has sharp peaks, in particular if the number of peaks is known in advance.

# Arguments:
- `na::Int`: number of roots (order of the system). The number of peaks in the spectrum will be `na÷2`.
- `λ::Float64 = 0`: reg factor
"""
TLS
@kwdef struct TLS <: FitMethod
    na::Int
    λ::Float64 = 0
end

"""
    fitmodel(fm::TLS, X::AbstractArray)
"""
function fitmodel(fm::TLS,X::AbstractArray)
    isderiving() && return fitmodel(fm,X,true)
    Ay = getARregressor_(X, fm.na)
    if fm.λ > 0
        Ay = [Ay;sqrt(fm.λ)*I]
        Ay[end,end] = 0
    end
    a = tls!(Ay, size(Ay,2)-1) |> vec |> rev |> polyvec
    AR(a, var(X))
end

function fitmodel(fm::TLS,X::AbstractArray,diff::Bool)
    y,A = getARregressor(X, fm.na)
    if fm.λ > 0
        A = [A;sqrt(fm.λ)*I]
        y = [y;zeros(size(A,2))]
    end
    a = tls(A,y) |> vec |> rev |> polyvec
    AR(a, var(X))
end

"""
    IRLS <: FitMethod

Iteratively reqeighted least squares. This fitmethod is currently not recommended, it does not appear to produce nice spectra. (feel free to try it out though).

# Arguments:
- `na::Int`: number of roots (order of the system)
"""
IRLS
@kwdef struct IRLS <: FitMethod
    na::Int
end

"""
    fitmodel(fm::IRLS, X::AbstractArray)
"""
function fitmodel(fm::IRLS,X::AbstractArray)
    y,A = getARregressor(X, fm.na)
    a = irls(A,y) |> vec |> rev |> polyvec
    AR(a, var(X))
end

"""
    PLR <: FitMethod

Pseudo linear regression. Estimates the noise components by performing an initial fit of higher order. Tihs fitmethod produces an [`ARMA`](@ref) model. Support for ARMA models is not as strong as for [`AR`](@ref) models.

# Arguments:
- `nc::Int`: order of numerator
- `na::Int`: order of denomenator
- `initial::T = TLS(na=80)`: fitmethod for the initial fit. Can be, e.g., [`LS`](@ref), [`TLS`](@ref) or any function that returns a coefficient vector
- `λ::Float64 = 0.0001`: reg factor
"""
PLR
@kwdef struct PLR{T} <: FitMethod
    nc::Int
    na::Int
    initial::T = TLS(na=80)
    λ::Float64 = 1e-4
end
"""
    fitmodel(fm::PLR, X::AbstractArray)
"""
function fitmodel(fm::PLR,X::AbstractArray)
    plr(X,fm.na,fm.nc, fm.initial, λ=fm.λ)
end

"""
    plr(y, na, nc, initial; λ=0.01)

Performs pseudo-linear regression to estimate an ARMA model.

# Arguments:
- `y`: signal
- `na`: denomenator order
- `nc`: numerator order
- `initial`: fitmethod for the initial fit. Can be, e.g., [`LS`](@ref), [`TLS`](@ref) or any function that returns a coefficient vector
- `λ`: reg
"""
function plr(y,na,nc,initial; λ = 1e-2)
    na >= 1 || throw(ArgumentError("na must be positive"))
    y_train,A = getARregressor(y,initial.na)
    w1 = initial isa TLS ? tls(A,y_train) : ls(A,y_train,initial.λ)
    yhat = A*vec(w1)
    ehat = y_train - yhat
    ΔN = length(y)-length(ehat)
    y_train,A = getARXregressor(y[ΔN:end-1],ehat,na,nc)
    w = tls(A,y_train)
    a,c = params2poly(w,na,nc)
    c[1] = 1
    # TODO: scalefactor for PLR
    ARMA(c,a)
end

function params2poly(w,na,nb)
    a = [1; -w[1:na]]
    w = w[na+1:end]
    b = map(nb) do nb
        b = w[1:nb]
        w = w[nb+1:end]
        b
    end
    a,b
end

"""
    toeplitz(c, r)

Returns a toepliz matrix with first column and row specified (c[1] == r[1]).
"""
function toeplitz(c,r, A = similar(c, length(c), length(r)))
    @assert c[1] == r[1]
    nc = length(c)
    nr = length(r)
    A[1:nc,1] = c
    A[1,:] = r
    for i in 2:nr
        A[2:nc,i] = @view A[1:nc-1,i-1]
    end
    A
end

function getARregressor(y, na)
    m    = na+1 # Start of yr
    n    = length(y) - m + 1 # Final length of yr
    A    = @views toeplitz(y[m:m+n-1],y[m:-1:m-na])
    @assert size(A,2) == na+1
    y    = A[:,1] # extract yr
    A    = A[:,2:end]
    return y,A
end

function getARregressor!(AS, y, na)
    m    = na+1 # Start of yr
    @assert size(AS,2) == m "Expected size(A,2) == m got size(A,2) == $((size(AS,2), m))"
    n    = length(y) - m + 1 # Final length of yr
    A    = @views toeplitz(y[m:m+n-1],y[m:-1:m-na], AS)
    y    = @view A[:,1] # extract yr
    A    = @view A[:,2:end]
    return y,A
end

function getARregressor_(y, na)
    y    = rev(y)
    m    = na+1 # Start of yr
    n    = length(y) - m + 1 # Final length of yr
    Ay   = toeplitz(y[m:m+n-1],y[m:-1:m-na])
    return Ay
end

# getARregressor(a::TrackedArray, b) = Flux.Tracker.track(getARregressor, a, b)

# @grad reshape(xs, dims) = reshape(data(xs), dims), Δ -> (reshape(Δ, size(xs)),nothing)
function getARXregressor(y::AbstractVector,u::AbstractVector, na, nb)
    m    = max(na,nb)+1 # Start of yr
    @assert m >= 1
    n    = length(y) - m + 1 # Final length of yr
    @assert n <= length(y)
    A    = toeplitz(y[m:m+n-1],y[m:-1:m-na])
    @assert size(A,2) == na+1
    y = A[:,1] # extract yr
    A = A[:,2:end]
    s = m-1
    A = [A toeplitz(u[s:s+n-1],u[s:-1:s-nb+1])]
    return y,A
end

poly(w) = [-rev(w); 1]
polyvec(w) = [1; -w]
polyroots(w) = roots(poly(w))

function polyconv(a,b)
    na,nb = length(a),length(b)
    if na > nb
        (a,b) = (b,a)
    end
    c = zeros(promote_type(eltype(a), eltype(b)), na+nb-1)
    for i = 1:length(c)
        for j = 1:min(i,na,nb)
            av = j ∈ eachindex(a) ? a[j] : zero(eltype(a))
            bv = i-j+1 ∈ eachindex(b) ? b[i-j+1] : zero(eltype(b))
            c[i] += av*bv
        end
    end
    c
end

"""
    roots2poly(roots)

Accepts a vector of complex roots and returns the polynomial with those roots
"""
function roots2poly(roots::AbstractVector{CT})::Vector{real(CT)} where CT
    isderiving() && return roots2poly_zygote(roots)
    p = @MVector [1.]
    for r in 1:length(roots)
        p = _roots2poly_kernel(p, roots[r])
    end
    Vector{real(CT)}(real(p))
end

function _roots2poly_kernel(a::StaticVector{N,T},b) where {N,T}
    T2 = promote_type(T,typeof(b))
    vT = T2 <: Real ? Complex{T} : T2
    c = MVector{N+1,vT}(ntuple(_->0, N+1))
    c[1] = 1
    n = length(a)
    for i in 2:n
        c[i] = -b*a[i-1] + a[i]
    end
    c[end] = -b*a[end]
    c
end

function roots2poly_zygote(roots)
    p = [complex(1.)]
    for r in roots
        p = polyconv(p, [1, -r])
    end
    real(p)
end

@inline function polyderivative(a)
    n = length(a)-1
    powers = n:-1:1
    ap     = powers.*a[1:end-1]
end

# @inline function polyval(ap,r)
#     n = length(ap)
#     sum(j->ap[j]*r^(n-j), 1:n)
# end # Below is workaround for Zygote bug

@inline function polyval(ap,r)
    n = length(ap)
    s = zero(promote_type(Float64, eltype(ap), typeof(r)))
    @inbounds for j = 1:n
        s += ap[j]*r^(n-j)
    end
    s
end

@inline polyval(b::Number,_) = b

function rev(x)
    if isderiving()
        x[end:-1:1] # Workaround for Zygote not being able to reverse vectors
    else
        reverse(x)
    end
end

"""
    residues(a::AbstractVector, b, r=roots(reverse(a)))

Returns a vector of residues for the system represented by denominator polynomial `a`

Ref: slide 21 https://stanford.edu/~boyd/ee102/rational.pdf
Tihs methid is numerically sensitive. Note that if two poles are on top of each other, the residue is infinite.
"""
function residues(a::AbstractVector, b, r::AbstractVector{CT} = hproots(rev(a))) where CT
    # a[1] ≈ 1 || println("Warning: a[1] is ", a[1]) # This must not be @warn because Zygote can't differentiate that
    n = length(a)-1
    ap = polyderivative(a)
    map(r) do r
        CT(polyval(b, r)/polyval(ap, r))
    end
end

function abs2residues!(rw, a::AbstractVector, b, r::AbstractVector{CT} = hproots(rev(a))) where CT
    # a[1] ≈ 1 || println("Warning: a[1] is ", a[1]) # This must not be @warn because Zygote can't differentiate that
    n = length(a)-1
    ap = polyderivative(a)
    map!(rw, r) do r
        abs2(polyval(b, r)/polyval(ap, r))
    end
end

residues(d::TimeEvolution, m::AbstractModel) = residues(denvec(d,m), numvec(d,m), roots(d,m))

abs2residues!(rw, d::TimeEvolution, m::AbstractModel) = abs2residues!(rw, denvec(d,m), numvec(d,m), roots(d,m))

"""
    residues(r::AbstractRoots)

Returns a vector of residues for the system represented by roots `r`

"""
function residues(r::AbstractRoots)
    a = isderiving() ? complex.(roots2poly(r)) : roots2poly(r)
    residues(a, 1, r)
end

function residueweight(m::AbstractModel)
    res = residues(Continuous(), m)
    e = roots(Continuous(), m)
    rw = abs.(π*abs2.(res)./ real.(e))
    isderiving() ? complex.(rw) : rw
end

function residueweight(rw, m::AbstractModel)
    abs2residues!(rw, Continuous(), m)
    e = roots(Continuous(), m)
    @. rw = abs(π*rw/ real(e))
    isderiving() ? complex.(rw) : rw
end

function unitweight(m::AbstractModel)
    r = roots(Continuous(), m)
    RT = float(real(eltype(r)))
    N = length(r)
    Fill(RT(1/N), N)
end

# Base.delete_method.(methods(residues))
"""
    poles2model(r::AbstractVector{<:Real}, i::AbstractVector{<:Real})

Returns a transer function with the desired poles. There will be twice as many poles as the length of `i` and `r`.
"""
function poles2model(r::AbstractVector{<:Real}, i::AbstractVector{<:Real})
    roots = [complex.(r, i); complex.(r, -i)]
    AR(Continuous(),roots2poly(roots))
end

function Base.rand(::Type{AR}, dr, di, n=2)
    @assert n % 2 == 0
    n = n ÷ 2
    poles2model(rand(dr,n), rand(di,n))
end

"""
    spectralenergy(G::LTISystem)

Calculates the energy in the spectrum associated with `G`
Ref: Robust and optimal control, Kemin Zhou, John C. Doyle, Keith Glover Lemma 4.6
"""
function spectralenergy(G::LTISystem)
    sys = ss(G)
    A = sys.A
    eltype(A) <: Complex && @warn "This function is known to be incorrect for complex systems"
    e = eigen(A)
    e.values .= reflect(ContinuousRoots(e.values))
    sys.A .= real((e.vectors) * Diagonal(e.values) * inv(e.vectors))
    sys,X = balreal(sys)
    2π*tr(sys.C*X*sys.C')
end

# spectralenergy(a::AbstractArray) = spectralenergy(determine_domain(roots(reverse(a))), a)
function spectralenergy(d::TimeEvolution, m::AbstractModel)
    spectralenergy(d, denvec(d, m), numvec(d,m))
end

"""
    spectralenergy(d::TimeEvolution, a::AbstractVector, b)

Calculates the energy in the spectrum associated with transfer function `b/a`
"""
function spectralenergy(d::TimeEvolution, ai::AbstractVector{T}, b)::T where T
    a = Double64.(ai)
    ac = a .* (-1).^(length(a)-1:-1:0)
    a2 = polyconv(ac,a)
    r2 = roots(rev(a2))
    filterfun = d isa Continuous ? r -> real(r) < 0 : r -> abs(r) < 1
    # r2 = filter(filterfun, r2)
    r2 = r2[filterfun.(r2)] # Zygote can't diff the filter above
    res = residues(a2, b, r2)
    e = 2π*sum(res)
    ae = real(e)
    if ae < 1e-3  && !(T <: BigFloat) && !isderiving() # In this case, accuracy is probably compromised and we should do the calculation with higher precision.
        return spectralenergy(d, big.(a), b)

    end
    abs(imag(e))/abs(real(e)) > 1e-3 && println("Got a large imaginary part in the spectral energy $(abs(imag(e))/abs(real(e)))")
    @assert ae >= 0 "Computed energy was negative: $ae"
    ae
end

"""
    determine_domain(r)

Tries to automatically figure out which time domain the roots represent. Throws an error if ambiguous.
"""
function determine_domain(r)
    if all(<(0), real.(r)) # Seems like Cont
        if all(<(1), abs.(r)) # Seems like Disc as well
            error("I can't determine if domain is discrete or continuous. Wrap in the correct wrapper (ContinuousRoots / DiscreteRoots)")
        end
        return Continuous()
    end
    all(<(1), abs.(r)) || error("I can't determine if domain is discrete or continuous. Wrap in the correct wrapper (ContinuousRoots / DiscreteRoots)")
    Discrete()
end

function AutoRoots(d::TimeEvolution, r)
    d isa Discrete ? DiscreteRoots(r) : ContinuousRoots(r)
end
function AutoRoots(r)
    determine_domain(r) isa Discrete ? DiscreteRoots(r) : ContinuousRoots(r)
end
normalization_factor(r) = normalization_factor(AutoRoots(r))

"""
    normalize_energy(r)

Returns the factor that, when used to multiply the poles, results in a system with unit spectral energy.
"""
function normalization_factor(r::AbstractRoots)
    e = spectralenergy(domain(r), roots2poly(r), 1)
    n = length(r)
    s = e^(1/(2n-1))
end

"""
    normalize_energy(r)

Returns poles scaled to achieve unit spectral energy.
"""
function normalize_energy(r)
    normalization_factor(r)*r
end

"""
    scalefactor(::TimeEvolution, a, [b,] σ²)

Returns `k` such that the system with numerator `kb` and denomenator `a` has spectral energy `σ²`. `σ²` should typically be the variance of the corresponding time signal.
"""
function scalefactor(d::TimeEvolution, a, b, σ²)
    e = spectralenergy(d, a, b)
    (σ²/(e))
end

function scalefactor(d::TimeEvolution, a, σ²)
    e = spectralenergy(d, a, 1)
    (σ²/(e))
end

"""
    examplemodels(n = 10)

Return `n` random models with 6 complex poles each.
"""
function examplemodels(n=10)
    ζ = [0.1, 0.3, 0.7]
    models = map(1:n) do _
        pol = [1]
        for i = eachindex(ζ)
            pol = SpectralDistances.polyconv(pol, [1,2ζ[i] + 0.1randn(),1+0.1randn()])
        end
        AR(Continuous(),pol)
    end
end


text:
## Generic ##

_istracked(x) = false
_istracked(x::TrackedArray) = false
_istracked(x::AbstractArray{<:TrackedReal}) = true
function mapvcat(f, args...)
    out = map(f, args...)
    if _istracked(out)
        init = vcat(out[1])
        return reshape(reduce(vcat, drop(out, 1); init = init), size(out))
    else
        return out
    end
end
@adjoint function mapvcat(f, args...)
    g(f, args...) = map(f, args...)
    return pullback(g, f, args...)
end

function Base.fill(
    value::TrackedReal,
    dims::Vararg{Union{Integer, AbstractUnitRange}},
)
    return track(fill, value, dims...)
end
@grad function Base.fill(value::Real, dims...)
    return fill(data(value), dims...), function(Δ)
        size(Δ) ≢  dims && error("Dimension mismatch")
        return (sum(Δ), map(_->nothing, dims)...)
    end
end

## StatsFuns ##

logsumexp(x::TrackedArray) = track(logsumexp, x)
@grad function logsumexp(x::TrackedArray)
    lse = logsumexp(data(x))
    return lse, Δ -> (Δ .* exp.(x .- lse),)
end

## Linear algebra ##

# Work around https://github.com/FluxML/Tracker.jl/pull/9#issuecomment-480051767

upper(A::AbstractMatrix) = UpperTriangular(A)
lower(A::AbstractMatrix) = LowerTriangular(A)
function upper(C::Cholesky)
    if C.uplo == 'U'
        return upper(C.factors)
    else
        return copy(lower(C.factors)')
    end
end
function lower(C::Cholesky)
    if C.uplo == 'U'
        return copy(upper(C.factors)')
    else
        return lower(C.factors)
    end
end

LinearAlgebra.LowerTriangular(A::TrackedMatrix) = lower(A)
lower(A::TrackedMatrix) = track(lower, A)
@grad lower(A) = lower(Tracker.data(A)), ∇ -> (lower(∇),)

LinearAlgebra.UpperTriangular(A::TrackedMatrix) = upper(A)
upper(A::TrackedMatrix) = track(upper, A)
@grad upper(A) = upper(Tracker.data(A)), ∇ -> (upper(∇),)

function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return track(copy, A)
end
@grad function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return copy(data(A)), ∇ -> (copy(∇),)
end

function LinearAlgebra.cholesky(A::TrackedMatrix; check=true)
    factors_info = turing_chol(A, check)
    factors = factors_info[1]
    info = data(factors_info[2])
    return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
end
function turing_chol(A::AbstractMatrix, check)
    chol = cholesky(A, check=check)
    (chol.factors, chol.info)
end
turing_chol(A::TrackedMatrix, check) = track(turing_chol, A, check)
@grad function turing_chol(A::AbstractMatrix, check)
    C, back = pullback(_turing_chol, data(A), data(check))
    return (C.factors, C.info), Δ->back((factors=data(Δ[1]),))
end
_turing_chol(x, check) = cholesky(x, check=check)

# Specialised logdet for cholesky to target the triangle directly.
logdet_chol_tri(U::AbstractMatrix) = 2 * sum(log, U[diagind(U)])
logdet_chol_tri(U::TrackedMatrix) = track(logdet_chol_tri, U)
@grad function logdet_chol_tri(U::AbstractMatrix)
    U_data = data(U)
    return logdet_chol_tri(U_data), Δ->(Matrix(Diagonal(2 .* Δ ./ diag(U_data))),)
end

function LinearAlgebra.logdet(C::Cholesky{<:TrackedReal, <:TrackedMatrix})
    return logdet_chol_tri(C.U)
end

# Tracker's implementation of ldiv isn't good. We'll use Zygote's instead.

zygote_ldiv(A::AbstractMatrix, B::AbstractVecOrMat) = A \ B
function zygote_ldiv(A::TrackedMatrix, B::TrackedVecOrMat)
    return track(zygote_ldiv, A, B)
end
function zygote_ldiv(A::TrackedMatrix, B::AbstractVecOrMat)
    return track(zygote_ldiv, A, B)
end
zygote_ldiv(A::AbstractMatrix, B::TrackedVecOrMat) =  track(zygote_ldiv, A, B)
@grad function zygote_ldiv(A, B)
    Y, back = pullback(\, data(A), data(B))
    return Y, Δ->back(data(Δ))
end

function Base.:\(a::Cholesky{<:TrackedReal, <:TrackedArray}, b::AbstractVecOrMat)
    return (a.U \ (a.U' \ b))
end

# SpecialFunctions

SpecialFunctions.logabsgamma(x::TrackedReal) = track(logabsgamma, x)
@grad function SpecialFunctions.logabsgamma(x::Real)
    return logabsgamma(data(x)), Δ -> (digamma(data(x)) * Δ[1],)
end
@adjoint function SpecialFunctions.logabsgamma(x::Real)
    return logabsgamma(x), Δ -> (digamma(x) * Δ[1],)
end

# Some Tracker fixes

for i = 0:2, c = Tracker.combinations([:AbstractArray, :TrackedArray, :TrackedReal, :Number], i), f = [:hcat, :vcat]
    if :TrackedReal in c
        cnames = map(_ -> gensym(), c)
        @eval Base.$f($([:($x::$c) for (x, c) in zip(cnames, c)]...), x::Union{TrackedArray,TrackedReal}, xs::Union{AbstractArray,Number}...) =
            track($f, $(cnames...), x, xs...)
    end
end
@grad function vcat(x::Real)
    vcat(data(x)), (Δ) -> (Δ[1],)
end
@grad function vcat(x1::Real, x2::Real)
    vcat(data(x1), data(x2)), (Δ) -> (Δ[1], Δ[2])
end
@grad function vcat(x1::AbstractVector, x2::Real)
    vcat(data(x1), data(x2)), (Δ) -> (Δ[1:length(x1)], Δ[length(x1)+1])
end

# Zygote fill has issues with non-numbers

@adjoint function fill(x::T, dims...) where {T}
    function zfill(x, dims...,)
        return reshape([x for i in 1:prod(dims)], dims)
    end
    pullback(zfill, x, dims...)
end


text:
#Program to find whether the number is even or odd:

function eoro(num)
	if num % 2 == 0
		return ("$num is Even")
	else
		return ("$num is odd")
	end
end

print("Enter the number to find Even or Odd")
num = parse(Int64,readline())
println(eoro(num))

max_stars_repo_name: cashewcodes/Basic-Julia-Projects

max_issues_repo_name: cashewcodes/Basic-Julia-Projects

max_forks_repo_name: cashewcodes/Basic-Julia-Projects


text:
using Calculus
using NLsolve
using Roots
using Interpolations
using LaTeXStrings
include("plot_utils.jl")

RTOL = 1e-5
myapprox(x,y) = abs(x-y)<RTOL  # This is defined to identify small intervals

function give_interval(d2f_zeros, a, b)
	"""
	this function divides interval [a,b] into subintervals
		[(a, def_zeros[1]), (d2f_zeros[1], d2f_zeros[2]), ... , (d2f_zeros[end], b) ]
	if the interval length is less that RTOL, it will be ignored.
	"""
	if isempty(d2f_zeros)
		return [(a,b)]
	else
		intervals = []
		if !myapprox(d2f_zeros[1], a)
			d2f_zeros = vcat(a, d2f_zeros)
		end
		if !myapprox(d2f_zeros[end], b)
			d2f_zeros = vcat(d2f_zeros, b)
		end
		for i in 1:length(d2f_zeros)-1
			if !myapprox(d2f_zeros[i], d2f_zeros[i+1])
				push!(intervals, (d2f_zeros[i], d2f_zeros[i+1]))
			end
		end
		return intervals
	end
end

# Two auxillary functions that will be used later.
aux_func_1(f, df, alpha, beta) = (beta*df(beta)-alpha*df(alpha) -(f(beta)-f(alpha)))/(df(beta)-df(alpha))
line_equation(f, df, p, x) = df(p)*(x-p)+f(p)

# generate y_i candidates for concave upper bound / convex lower bound
function get_yi_candidates(xx, f, df, N)
	y_i_m = zeros(N+2)
	y_i_p = zeros(N+2)
	# first get all points "looking forwards" (eval right endpoint of the line)
	############################################################################
	y_i_m[1] = f(xx[1])
	y_i_m[2] = line_equation(f,df, xx[1],xx[2]) # args are: tangent point, eval point
	for i = 3:N+1
		y_i_m[i] = line_equation(f, df, (xx[i-1]+xx[i])/2 , xx[i])
	end
	y_i_m[N+2] = line_equation(f, df, xx[N+2], xx[N+2]) # should just be f(xx[N+2])
	# then get all points "looking backwards" (eval left endpoint of the line)
	############################################################################
	y_i_p[1] = line_equation(f, df, xx[1], xx[1])
	for i = 3:N+1
		y_i_p[i-1] = line_equation(f, df, (xx[i-1] + xx[i])/2, xx[i-1])
	end
	y_i_p[N+1] = line_equation(f, df, xx[N+2], xx[N+1])
	y_i_p[N+2] = f(xx[N+2])
	return y_i_m, y_i_p
end

function check_concave_upper_optimality(y_i_m, y_i_p; ζ=1e-6)
	"""
	Check if continuity conditions have been successfully enforced.
	How: check if two y_i generated by two candidate tangent bounds are within 1e-5 of each other.
	"""
	## DEBUG:
	# if maximum(abs.(y_i_m - y_i_p)) > ζ
	# 	@warn "Max abs continuity violation of $(maximum(abs.(y_i_m - y_i_p))) is greater than continuity violation threshold of $ζ ,"
	# end
	return all(abs.(y_i_m - y_i_p) .<= ζ)
end

function get_safe_yi(y_i_m, y_i_p, xx, f, this_interval_convex)
	yy = zeros(length(xx))
	yy[1] = f(xx[1])
	yy[end] = f(xx[end])
	if this_interval_convex # means we are computing lower bound for convex function
		yy[2:end-1] = min.(y_i_m[2:end-1], y_i_p[2:end-1])
	else # means we are computing upper bound for concave function
		yy[2:end-1] = max.(y_i_m[2:end-1], y_i_p[2:end-1])
	end
	return yy
end

function bound_jensen(f, df, N, aa, bb)
	"""
	This function generates a function that computes the system of equations
		to be solved for a convex case
	"""
	function obj_jensen!(F, z)
		if N==1
			F[1] = f(bb) - f(aa) - df(z[1])*(bb-aa)
		else
			F[1] = f(z[2]) - f(aa) - df(z[1])*(z[2]-aa)
			for i in 2:N-1
				F[i] = f(z[i+1]) - f(z[i-1]) - df(z[i])*(z[i+1]-z[i-1])
			end
			F[N] = f(bb) - f(z[N-1]) - df(z[N])*(bb-z[N-1])
		end
	end
	return obj_jensen!
end

function bound_tangent(f, df, N, aa, bb, method)
	"""
	This function generates a function that computes the system of equations
		to be solved for a concave case.

	method can be "continuous" or "optimal", depending on what algorithm is preferred.
	"optimal" method produces the tightest overapproximator, but it might be discontinuous
	at inflection points of f. Method "continuous" takes care of the continuity, but
	is sub-optimal.
	"""

	function obj_tangent_continuous!(F, z)
		if N == 1
			F[1] = aux_func_1(f, df, aa, bb) - z[1]
		else
			for i=1:N
				if i == 1
					xm1 = aa
				else
					xm1 = (z[i]+z[i-1])/2
				end
				if i == N
					xm2 = bb
				else
					xm2 = (z[i]+z[i+1])/2
				end
				F[i] = aux_func_1(f, df, xm1, xm2) - z[i]
			end
		end

	end

	function obj_tangent_optimal!(F, z)
		if N == 1
			F[1] = (aa+bb)/2 - z[1]
		else
			for i in 1:N
				if i == 1
					xm1 = (z[i]+aa)/2
					dz1 = z[i] - aa
				else
					xm1 = (z[i]+z[i-1])/2
					dz1 = z[i] - z[i-1]
				end
				if i == N
					xm2 = (z[i]+bb)/2
					dz2 = z[i] - bb
				else
					xm2 = (z[i]+z[i+1])/2
					dz2 = z[i] - z[i+1]
				end
				F[i] = 1/2*df(xm1)*dz1 + f(xm1) - 1/2*df(xm2)*dz2 - f(xm2)
			end
		end
	end

	if method == "continuous"
		return obj_tangent_continuous!
	elseif method == "optimal"
		return obj_tangent_optimal!
	else
		error(" the method is not defined")
	end
end

"""
Interface to the bound function that accepts symbolic functions, not executable ones, so
that symbolic differentiation can be used.
"""

function bound(f, a, b, N; conc_method="continuous", lowerbound=false, df=nothing,
	d2f=nothing, d2f_zeros=nothing, convex=nothing, plot=false)
	"""
	This function over(under)-approximate function f(x).

	f:         the function to be bounded
	[a,b]:     domain of f
	N:         N-1 is the number of linear segments within each sub-interval.
	           sub-intervals are regions where second derivative does not change sign.
	conc_method: determines the concave algorithm. can take continuous or optimal.
	lowerbound:  true if you want a lower bound
	df:        derivative of f, if available
	d2f:       second derivative of f, if available
	d2f_zeros: zeros of the second derivative, if available
	convex:    true for convex, false for concave, nothing for mixed.
	plot:      true to plot
	"""

	"""
	Example:
		bound(cos, 0, π, 3)
		bound(x->exp(x^2), 0, 2, 3, convex=true)
		bound(sin, 0, π, 3, lowerbound=true)
		bound(sin, -π/2, π, 3, d2f_zeros=[0])
		bound(x-> x^3-sin(x), 0, 2, 3)
	"""
	@debug "Number of points in bound, N= $N"
	if N == -1 # optimally choose N
		return bound_optimal(f, a, b; conc_method=conc_method,
			lowerbound=lowerbound, df=df, d2f=d2f, d2f_zeros=d2f_zeros, convex=convex,
			plot=plot)
	end

	try
		f(a)
		f(b)
	catch
		error("$a or $b is not in the domanin of $f")
	end

	if isnothing(df)
    	df = Calculus.derivative(f)
	end
	if isnothing(d2f)
		d2f = Calculus.second_derivative(f)
	end
    if ! isnothing(convex) # if convexity is specified, no sub-intervals necessary.
		intervals = [(a,b)]
		@debug "convexity specified, no sub intervals necessary." convex
    else #if concavity is not uniform, calculate intervals of uniform concavity
		if isnothing(d2f_zeros) # calculate zeros of second derivative, if not given.
			println("WARNING: d2f_zeros have not been specified for function $f. Convex and concave regions will be identified using a numerical procedure. Soundness not guaranteed. ")
			d2f_zeros = fzeros(d2f, a, b)
        end
        intervals = give_interval(d2f_zeros, a, b)  # find subintervals.
    end

	# println("OVERT applied on $f within ranges [$a, $b]")
    xp = []
    yp = []
    for interval in intervals
		aa, bb = interval

		# initial guess for nonlinear solver. we divide [aa,bb] by N equally spaced points.
        zGuess = range(aa, stop=bb, length=N+2)
        zGuess = reshape(zGuess, N+2,1)[2:end-1]

		if isnothing(convex)
			this_interval_convex = d2f((aa+bb)/2) >= 0
		else
			@assert convex == (d2f((aa+bb)/2) >= 0)
			this_interval_convex = convex
		end
        if xor(this_interval_convex, lowerbound)  # upper bound for convex or lower bound for concave
			obj = bound_jensen(f, df, N, aa, bb)
		else
			obj = bound_tangent(f, df, N, aa, bb, conc_method)
		end

		# solving nonlinear system.
		# z = nlsolve(obj, zGuess, autodiff = :forward)

		z = 0
		itr_nlsolve = 0
		try # This may overshoot to outside of domain for log
			z = nlsolve(obj, zGuess, autodiff = :forward)

			# check zeros are within the interval
			@assert z.zero[1] > aa
			@assert z.zero[end] < bb
			for ii=1:length(z.zero)-1
				@assert z.zero[ii] < z.zero[ii+1]
			end
    	catch # if overshoots, change the zGuess until it works. 10 iterations allowed
			zGuess_orig = copy(zGuess)
			# shifting zGuess towards left
			while itr_nlsolve <= 10
				zGuess[1:end-1] = zGuess[2:end]
				zGuess[end] = 0.5*(bb + zGuess[end])
				try
					z = nlsolve(obj, zGuess, autodiff = :forward)

					# check zeros are within the interval, unique, and in ascending order
					@assert z.zero[1] > aa
					@assert z.zero[end] < bb
					for ii=1:length(z.zero)-1
						@assert z.zero[ii] < z.zero[ii+1]
					end
					break
				catch
					itr_nlsolve += 1
				end
				#z = mcpsolve(obj, sol_lb, sol_ub , zGuess_exp, autodiff = :forward)
			end

			# reset zGuess
			itr_nlsolve == 11 ? zGuess = zGuess_orig : nothing
			# shifting zGuess towards right
			while (itr_nlsolve <= 20) && (itr_nlsolve >= 11)
				zGuess[2:end] = zGuess[1:end-1]
				zGuess[1] = 0.5*(aa + zGuess[1])
				try
					z = nlsolve(obj, zGuess, autodiff = :forward)
					# check zeros are within the interval
					@assert z.zero[1] > aa
					@assert z.zero[end] < bb
					for ii=1:length(z.zero)-1
						@assert z.zero[ii] < z.zero[ii+1]
					end
					break
				catch
					itr_nlsolve += 1
				end
				#z = mcpsolve(obj, sol_lb, sol_ub , zGuess_exp, autodiff = :forward)
			end
			if itr_nlsolve == 21
				error("nlsolve could not converge")
			end
		end

        xx = vcat(aa, z.zero, bb)  # add endpoints to the intermediate points
		if xor(this_interval_convex, lowerbound) # upper bound for convex or lower bound for concave
            yy = [f(x) for x in xx]  # for the convex case, points lie on the function
        else # for the concave case, y-values should be calculated.
			yy = zeros(N+2)
			if conc_method == "continuous"
				y_i_m, y_i_p = get_yi_candidates(xx, f, df, N)
				@assert check_concave_upper_optimality(y_i_m, y_i_p)
				yy = get_safe_yi(y_i_m, y_i_p, xx, f, this_interval_convex)
			elseif conc_method == "optimal"
				for i in 1:N+1
					xm2 = (xx[i]+xx[i+1])/2
					yy[i] = line_equation(f, df, xm2, xx[i])
				end
				yy[N+2] = line_equation(f, df, (xx[N+1]+xx[N+2])/2, xx[N+2])
			end
        end
        push!(xp, xx)
        push!(yp, yy)
    end
    if plot
		#plot_bound(f, a, b, xp, yp; existing_plot=existing_plot)
		return xp, yp
	else
		#println("no plotting 4 u")
		return xp, yp
	end
end

# function overapprox(f,a,b,N; conc_method="continuous", df=nothing, d2f=nothing,
# 	d2f_zeros=nothing, convex=nothing, plot=false)

# 	LB = bound(f, a, b, N; conc_method=conc_method, lowerbound=true, df=df, d2f=d2f,
# 	d2f_zeros=d2f_zeros, convex=convex, plot=plot, reuse=false)

# 	UB = bound(f, a, b, N; conc_method=conc_method, lowerbound=false, df=df, d2f=d2f,
# 	d2f_zeros=d2f_zeros, convex=convex, plot=plot, reuse=true)
# 	return LB, UB
# end

function bound_optimal(f, a, b; rel_error_tol=0.02, Nmax = 20, conc_method="continuous",
	lowerbound=false, df=nothing, d2f=nothing, d2f_zeros=nothing, convex=nothing,
	plot=false)

	try
		f(a)
		f(b)
	catch
		error("$a or $b is not in the domanin of $f")
	end

	if isnothing(df)
		df = Calculus.derivative(f)
	end
	if isnothing(d2f)
		d2f = Calculus.second_derivative(f)
	end
	if ! isnothing(convex) # if convexity is specified, no sub-intervals necessary.
		intervals = [(a,b)]
		@debug "convexity specified, no sub intervals necessary."
	else
		if isnothing(d2f_zeros) # calculate zeros of second derivative, if not given.
			println("WARNING: d2f_zeros have not been specified for function $f. Convex and concave regions will be identified using a numerical procedure. Soundness not guaranteed. ")
			d2f_zeros = fzeros(d2f, a, b)
		end
		intervals = give_interval(d2f_zeros, a, b)  # find subintervals.
	end
	xp = []
	yp = []
	for interval in intervals
		aa, bb = interval
		# specify convexity
		if isnothing(convex)
			this_interval_convex = (d2f((aa+bb)/2) >= 0)
		else
			@assert convex == (d2f((aa+bb)/2) >= 0)
			this_interval_convex = convex
		end
		for N = 1:Nmax
			xp_candidate, yp_candidate = bound(f, aa, bb, N; conc_method="continuous", lowerbound=lowerbound, df=df,
			d2f=d2f, d2f_zeros=d2f_zeros, convex=this_interval_convex, plot=plot)

			# interpolate can sometimes give an error because it thinks the
			# points are outside the domain.
			# see this issue https://github.com/JuliaMath/Interpolations.jl/issues/158
			# hence shorten the range of xtest by epsilon
			ϵ_itp = 1E-8

			itp = Interpolations.interpolate((xp_candidate[1],), yp_candidate[1], Gridded(Linear()))
			xtest = range(aa + ϵ_itp, stop=bb - ϵ_itp, length=100)
			ytest = [itp(xt) for xt in xtest]
			ftest = [f(xt)   for xt in xtest]
			sc = maximum(ftest) - minimum(ftest)
			error = [abs(itp(xt) - f(xt))/sc for xt in xtest]
			if (maximum(error) < rel_error_tol) || (N == Nmax)
				xp = vcat(xp, xp_candidate)
				yp = vcat(yp, yp_candidate)
				break
			end
		end
	end

	xp = unique(xp)
	yp = unique(yp)

	if plot
		#plot_bound(f, a, b, xp, yp; existing_plot=existing_plot, saveflag=true);
		return xp, yp
	else
		return xp, yp
	end
end


text:
"""
    fetch(family)

Returns a family as a vector of `FASTA.Record`.
"""
function fetch(family::String)
    path = joinpath(cache, "fasta_files", "$family.fa")
    if !isfile(path)
        mkpath(dirname(path))
        url = "$rfam_baseurl/fasta_files/$family.fa.gz"
        download(url, path * ".gz"; timeout = Inf)
        gunzip(path) # decompress
    end

    records = FASTA.Record[]
    open(FASTA.Reader, path) do reader
        for record in reader
            push!(records, record)
        end
    end
    return records
end


text:
using IRTools
using IRTools: var, xcall, insert!, insertafter!, func, recurse!, @dynamo
include("loggingprofiler.jl")
LoggingProfiler.resize!(LoggingProfiler.to, 10000)

function timable(ex::Expr)
    ex.head != :call && return(false)
    length(ex.args) < 2 && return(false)
    ex.args[1] isa Core.GlobalRef && return(true)
    ex.args[1] isa Symbol && return(true)
    return(false)
end
timable(ex) = false

recursable(gr::GlobalRef) = gr.name ∉ [:profile_fun, :record_start, :record_end]
recursable(ex::Expr) = ex.head == :call && recursable(ex.args[1])
recursable(ex) = false

exportname(ex::GlobalRef) = QuoteNode(ex.name)
exportname(ex::Symbol) = QuoteNode(ex)
exportname(ex::Expr) = exportname(ex.args[1])
exportname(i::Int) = QuoteNode(Symbol("Int(",i,")"))

profile_fun(f::Core.IntrinsicFunction, args...) = f(args...)
profile_fun(f::Core.Builtin, args...) = f(args...)

@dynamo function profile_fun(f, args...)
    ir = IRTools.Inner.IR(f, args...)
    for (v, ex) in ir
        if timable(ex.expr)
            fname = exportname(ex.expr)
            insert!(ir, v, xcall(LoggingProfiler, :record_start, fname))
            insertafter!(ir, v, xcall(LoggingProfiler, :record_end, fname))
        end
    end
    for (x, st) in ir
        recursable(st.expr) || continue
        ir[x] = xcall(profile_fun, st.expr.args...)
    end
    # recurse!(ir)
    return ir
end

macro record(ex)
    esc(Expr(:call, :profile_fun, ex.args...))
end

function foo(x, y)
  z =  x * y
  z + sin(y)
end

LoggingProfiler.reset!()
@record foo(1.0, 1.0)
LoggingProfiler.to


text:
#=
    Utility function to
    - Create a temporary file,
    - Invoke a given function that presumably writes to the same file,
    - Clean up the temporary file, and
    - Return the file's full content.

    @param f
        Function f(fh::IOStream) that fills the file represented by fh
    @return String
=#
function readback(f::Function)
    (fname, fh) = mktemp()
    try
        f(fh)
        seek(fh, 0)
        return join(readlines(fh), "\n")
    finally
        close(fh)
        rm(fname)
    end
end


text:
include("./assim_useTBsave.jl")

#function do_assim(obs_mask, norm_factor, outdir, vodA, vodB, vodC)

#trueTB = get_TB_2(sim_res1[2][:,1], tsoil, tcan, canpot_true, laiM, vodA, vodB, vodC);
#obsH = trueTB[1] + noise_std*randn(Float64, length(trueTB[1]));
#obsV = trueTB[2] + noise_std*randn(Float64, length(trueTB[1]));

datalen = 24*365;
mask1 = zeros(datalen);
mask1[2:(24*3):end] .= 1;
mask1[(2+12):(24*3):end] .= 1;
mask1 = mask1 .== 1;

mask1o = zeros(datalen);
mask1o[2:(24*3):end] .= 1;
mask1o = mask1o .== 1;

mask6 = zeros(datalen);
mask6[7:(24*3):end] .= 1;
mask6[(7+12):(24*3):end] .= 1;
mask6 = mask6 .== 1;

maskAll = ones(datalen) .== 1;

vA, vB, vC = (0.067, 0.82, 0.051);

function run_scenario(scenario, sout, s_in)
  #sout = string(scenario,"/")
  if scenario == "o1AMPM"
    do_assim(mask1, 1, sout, vA, vB, vC, s_in)

  elseif scenario == "o6AMPM"
    do_assim(mask6, 1, sout, vA, vB, vC, s_in)

  elseif scenario == "o1AM"
    do_assim(mask1o, 2, sout, vA, vB, vC, s_in)

  elseif scenario == "oAll"
    do_assim(maskAll, 1.0/36, sout, vA, vB, vC, s_in)
  end
end

run_scenario(ARGS[1],ARGS[2],ARGS[3])


text:
@testset "threading" begin
    @testset "$fancy_map" for fancy_map in (tmap, tmap_with_warmup)
        @testset "basesize=$basesize" for basesize in (1, 2, 3, 100)
            @test fancy_map(x->2x, [1, 2, 3, 4]; basesize=basesize) == [2, 4, 6, 8]

            @test isequal(
                fancy_map(+, [1, 2, 3, 4, 5], [10, 20, 30, 40, 50]; basesize=basesize),
                [11, 22, 33, 44, 55]
            )
        end

        # Make sure that do blocks like in the pipeline work
        # got to make sure brackets are in right place
        @test  [11, 22, 33] == fancy_map([1, 2, 3], [10, 20, 30]) do x, y
            x+y
        end

        # Test that incorrectly brackets do blocks do not
        # Note: for tmap it is a TaskFailedException wrapping a MethodError
        # for tmap_with_warmup it is a MethodError
        # We are actually fine with either from either so we just test the Union
        ExpectedException = Union{TaskFailedException, MethodError}
        @test_throws ExpectedException fancy_map([1, 2, 3], [10, 20, 30]) do (x, y)
            x+y
        end
    end
end


text:
Base.one(::Type{P}) where {C, T, P <: Polynomial{C, T}} = Polynomial(one(T))
Base.one(p::Polynomial) = one(typeof(p))

Base.zero(::Type{Polynomial{C, T, A}}) where {C, T, A} = Polynomial(A())
Base.zero(t::PolynomialLike) = zero(typeof(t))

combine(t1::Term, t2::Term) = combine(promote(t1, t2)...)
function MA.promote_operation(::typeof(combine), ::Type{Term{S, M1}}, ::Type{Term{T, M2}}) where {S, T, M1, M2}
    return Term{MA.promote_operation(+, S, T), promote_type(M1, M2)}
end
combine(t1::T, t2::T) where {T <: Term} = Term(t1.coefficient + t2.coefficient, t1.monomial)
compare(t1::Term, t2::Term) = monomial(t1) > monomial(t2)

# Graded Lexicographic order
# First compare total degree, then lexicographic order
function isless(m1::Monomial{V}, m2::Monomial{V}) where {V}
    d1 = degree(m1)
    d2 = degree(m2)
    if d1 < d2
        return true
    elseif d1 > d2
        return false
    else
        return exponents(m1) < exponents(m2)
    end
end
function grlex(m1::Monomial{V}, m2::Monomial{V}) where {V}
    d1 = degree(m1)
    d2 = degree(m2)
    if d1 != d2
        return d1 - d2
    else
        if exponents(m1) == exponents(m2)
            return 0
        elseif exponents(m1) < exponents(m2)
            return -1
        else
            return 1
        end
    end
end
grlex(m1::Monomial, m2::Monomial) = grlex(promote(m1, m2)...)

jointerms(terms1::AbstractArray{<:Term}, terms2::AbstractArray{<:Term}) = mergesorted(terms1, terms2, compare, combine)
function jointerms!(output::AbstractArray{<:Term}, terms1::AbstractArray{<:Term}, terms2::AbstractArray{<:Term})
    resize!(output, length(terms1) + length(terms2))
    Sequences.mergesorted!(output, terms1, terms2, compare, combine)
end

(+)(p1::Polynomial, p2::Polynomial) = Polynomial(jointerms(terms(p1), terms(p2)))
(+)(p1::Polynomial, p2::Polynomial{<:UniformScaling}) = p1 + MP.mapcoefficientsnz(J -> J.λ, p2)
(+)(p1::Polynomial{<:UniformScaling}, p2::Polynomial) = MP.mapcoefficientsnz(J -> J.λ, p1) + p2
(+)(p1::Polynomial{<:UniformScaling}, p2::Polynomial{<:UniformScaling}) = MP.mapcoefficientsnz(J -> J.λ, p1) + p2
function MA.mutable_operate_to!(result::Polynomial, ::typeof(+), p1::Polynomial, p2::Polynomial)
    if result === p1 || result === p2
        error("Cannot call `mutable_operate_to!(output, +, p, q)` with `output` equal to `p` or `q`, call `mutable_operate!` instead.")
    end
    jointerms!(result.terms, terms(p1), terms(p2))
    result
end
function MA.mutable_operate_to!(result::Polynomial, ::typeof(*), p::Polynomial, t::AbstractTermLike)
    if iszero(t)
        MA.mutable_operate!(zero, result)
    else
        resize!(result.terms, nterms(p))
        for i in eachindex(p.terms)
            # TODO could use MA.mul_to! for indices that were presents in `result` before the `resize!`.
            result.terms[i] = p.terms[i] * t
        end
        return result
    end
end
function MA.mutable_operate_to!(result::Polynomial, ::typeof(*), t::AbstractTermLike, p::Polynomial)
    if iszero(t)
        MA.mutable_operate!(zero, result)
    else
        resize!(result.terms, nterms(p))
        for i in eachindex(p.terms)
            # TODO could use MA.mul_to! for indices that were presents in `result` before the `resize!`.
            result.terms[i] = t * p.terms[i]
        end
        return result
    end
end
(-)(p1::Polynomial, p2::Polynomial) = Polynomial(jointerms(terms(p1), (-).(terms(p2))))
(-)(p1::Polynomial, p2::Polynomial{<:UniformScaling}) = p1 - MP.mapcoefficientsnz(J -> J.λ, p2)
(-)(p1::Polynomial{<:UniformScaling}, p2::Polynomial) = MP.mapcoefficientsnz(J -> J.λ, p1) - p2
(-)(p1::Polynomial{<:UniformScaling}, p2::Polynomial{<:UniformScaling}) = MP.mapcoefficientsnz(J -> J.λ, p1) - p2

function MA.mutable_operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{T}, q::Polynomial) where T
    get1(i) = p.terms[i]
    function get2(i)
        t = q.terms[i]
        Term(MA.scaling_convert(T, MA.operate(op, t.coefficient)), t.monomial)
    end
    set(i, t::Term) = p.terms[i] = t
    push(t::Term) = push!(p.terms, t)
    compare_monomials(t::Term, j::Int) = grlex(q.terms[j].monomial, t.monomial)
    compare_monomials(i::Int, j::Int) = compare_monomials(get1(i), j)
    combine(i::Int, j::Int) = p.terms[i] = Term(MA.operate!(op, p.terms[i].coefficient, q.terms[j].coefficient), p.terms[i].monomial)
    combine(t::Term, j::Int) = Term(MA.operate!(op, t.coefficient, q.terms[j].coefficient), t.monomial)
    resize(n) = resize!(p.terms, n)
    # We can modify the coefficient since it's the result of `combine`.
    keep(t::Term) = !MA.iszero!(t.coefficient)
    keep(i::Int) = !MA.iszero!(p.terms[i].coefficient)
    MP.polynomial_merge!(
        nterms(p), nterms(q), get1, get2, set, push,
        compare_monomials, combine, keep, resize
    )
    return p
end

(==)(::Variable{N}, ::Variable{N}) where {N} = true
(==)(::Variable, ::Variable) = false
(==)(m1::Monomial{V}, m2::Monomial{V}) where {V} = exponents(m1) == exponents(m2)

# Multiplication is handled as a special case so that we can write these
# definitions without resorting to promotion:
MP.multconstant(α, v::Monomial) = Term(α, v)
MP.multconstant(α, v::Variable) = Term(α, Monomial(v))

(*)(v1::V, v2::V) where {V <: Variable} = Monomial{(V(),), 1}((2,))
(*)(v1::Variable, v2::Variable) = (*)(promote(v1, v2)...)

function MP.divides(m1::Monomial{V, N}, m2::Monomial{V, N}) where {V, N}
    reduce((d, exp) -> d && (exp[1] <= exp[2]), zip(m1.exponents, m2.exponents), init=true)
end
MP.divides(m1::Monomial, m2::Monomial) = divides(promote(m1, m2)...)
function MP.mapexponents(op, m1::M, m2::M) where M<:Monomial
    M(map(op, m1.exponents, m2.exponents))
end
MP.mapexponents(op, m1::Monomial, m2::Monomial) = mapexponents(op, promote(m1, m2)...)
#function MP.mapexponents_to!(output::M, op, m1::M, m2::M) where M<:Monomial
#    map!(op, output.exponents, m1.exponents, m2.exponents)
#    return output
#end
#function MP.mapexponents!(op, m1::M, m2::M) where M<:Monomial
#    map!(op, m1.exponents, m1.exponents, m2.exponents)
#    return m1
#end

function MA.mutable_operate_to!(output::Polynomial, ::typeof(*), p::Polynomial, q::Polynomial)
    empty!(output.terms)
    MP.mul_to_terms!(output.terms, p, q)
    sort!(output.terms, lt=(>))
    MP.uniqterms!(output.terms)
    return output
end
function MA.mutable_operate!(::typeof(*), p::Polynomial, q::Polynomial)
    return MA.mutable_operate_to!(p, *, MA.mutable_copy(p), q)
end

function MA.mutable_operate!(::typeof(zero), p::Polynomial)
    empty!(p.terms)
    return p
end
function MA.mutable_operate!(::typeof(one), p::Polynomial{T}) where T
    if isempty(p.terms)
        push!(p.terms, constantterm(one(T), p))
    else
        t = p.terms[1]
        p.terms[1] = Term(MA.one!(coefficient(t)), constantmonomial(t))
        resize!(p.terms, 1)
    end
    return p
end

# The exponents are stored in a tuple, this is not mutable.
# We could remove these methods since it is the default.
MA.mutability(::Type{<:Monomial}) = MA.NotMutable()
MA.mutability(::Type{<:Term}) = MA.NotMutable()
# The polynomials can be mutated.
MA.mutability(::Type{<:Polynomial}) = MA.IsMutable()

^(v::V, x::Integer) where {V <: Variable} = Monomial{(V(),), 1}((x,))

# dot(v1::AbstractVector{<:TermLike}, v2::AbstractVector) = dot(v1, v2)
# dot(v1::AbstractVector, v2::AbstractVector{<:TermLike}) = dot(v1, v2)
# dot(v1::AbstractVector{<:TermLike}, v2::AbstractVector{<:TermLike}) = dot(v1, v2)

# All of these types are immutable, so there's no need to copy anything to get
# a shallow copy.
Base.copy(x::TermLike) = x

# Terms are not mutable under the MutableArithmetics API
MA.mutable_copy(p::Polynomial) = Polynomial([Term(MA.copy_if_mutable(t.coefficient), t.monomial) for t in terms(p)])
Base.copy(p::Polynomial) = MA.mutable_copy(p)

adjoint(v::Variable) = v
adjoint(m::Monomial) = m
adjoint(t::Term) = Term(adjoint(coefficient(t)), monomial(t))
adjoint(x::Polynomial) = Polynomial(adjoint.(terms(x)))

function MP.mapcoefficientsnz_to!(output::Polynomial, f::Function, p::Polynomial)
    resize!(output.terms, nterms(p))
    for i in eachindex(p.terms)
        t = p.terms[i]
        output.terms[i] = Term(f(coefficient(t)), monomial(t))
    end
    return output
end
function MP.mapcoefficientsnz_to!(output::Polynomial, f::Function, p::AbstractPolynomialLike)
    return MP.mapcoefficientsnz_to!(output, f, polynomial(p))
end


text:
"""
## mocouf

This subroutine calculates the value of the yield function
for a Mohr-Coulomb material (phi in degrees).

### Function
```julia
mocouf(phi, c, sigm, dsbar, theta)
```

### Arguments
```julia
* psi::Float64              : Local force and momemts (Updated)
* c::Float64                : Globale forces and moments
* sigm::Float64             : Element orientation angle (3D)
* dsbar::Float64            : Globale forces and moments
* theta::Float64            : Element orientation angle (3D)
```

### Return value
```julia
* ::Float64                 : Value of yield function
```
"""
function mocouf(phi, c, sigm, dsbar, theta)
  snph = sind(phi)
  csph = cosd(phi)
  csth = cos(theta)
  snth = sin(theta)
  snph*sigm + dsbar*(csth/sqrt(3) - snth*snph/3) - c*csph
end


text:
# Michelle Medina
# July 8, 2020
# 2D code for Nodal Basis

export maxlevel, getindex, calcNodal2D, basis, x, evaluate, H_2_Nodal, Nodal_2_H_sparse

#********** INDEX **********
"""
Function gives the highest level (it's log2 because length must be a multiple of 2 + 1)
"""
function maxlevel(N::NodalBasis2D{T}, i::Int)::Int where {T}
	return log2(size(N.values, i) - 1)
end

function maxlevel(H::HierarchicalBasis2D{T}, i::Int)::Int where {T}
	return size(H.levels, i) - 1
end

function Base.getindex(N::NodalBasis2D{T}, i::Int, j::Int)::T where {T}
	@assert 0 <= i <= (2^maxlevel(N, 1))
	@assert 0 <= j <= (2^maxlevel(N, 2))
	return N.values[i+1, j+1]
end

function Base.getindex(H::HierarchicalBasis2D{T}, l::Int, m::Int, i::Int, j::Int)::T where {T}
	@assert 0 <= l <= maxlevel(H, 1)
	@assert 0 <= m <= maxlevel(H, 2)
	@assert 0 <= i <= (2^l)
	@assert 0 <= j <= (2^m)
	return H.levels[l+1, m+1].coefficients[i+1, j+1]
end

function Base.setindex!(N::NodalBasis2D{T}, X::T, i::Int, j::Int)::T where {T}
	@assert 0 <= i <= (2^maxlevel(N, 1))
	@assert 0 <= j <= (2^maxlevel(N, 2))
	N.values[i+1, j+1] = X
end

function Base.setindex!(H::HierarchicalBasis2D{T}, X::T, l::Int, m::Int, i::Int, j::Int)::T where {T}
	@assert 0 <= l <= maxlevel(H, 1)
	@assert 0 <= m <= maxlevel(H, 2)
	@assert 0 <= i <= (2^l)
	@assert 0 <= j <= (2^m)
	H.levels[l+1, m+1].coefficients[i+1, j+1] = X
end

function Base.getindex(N::NodalBasis2D{T}, l::Int, m::Int, i::Int, j::Int)::T where {T}
	@assert 0 <= l <= maxlevel(N, 1)
	@assert 0 <= m <= maxlevel(N, 2)
	@assert 0 <= i <= (2^maxlevel(N, 1))
	@assert 0 <= j <= (2^maxlevel(N, 2))
	I = 2^(maxlevel(N, 1)-l)*i
	J = 2^(maxlevel(N, 2)-m)*j
	return N.values[I+1, J+1]
end

function calcNodal2D(u::Function, x::Array{T,1}, y::Array{T,1})::NodalBasis2D{T} where {T}
	coefficients = zeros(length(x), length(y))
 	for index in CartesianIndices(coefficients)
		i, j = index.I
 		coefficients[index] = u(x[i], y[j])
	end
 	return NodalBasis2D(coefficients)
end

function basis(l::Int, m::Int, i::Int, j::Int, x::T, y::T)::T where {T}
	return basis(l, i, x)*basis(m, j, y)
end

function x(l::Int, m::Int, i::Int, j::Int)::NTuple{2,Float64}
	return (x(l,i), x(m,j))
end

function evaluate(H::HierarchicalBasis2D{T}, x::T, y::T)::T where {T}
	value = T(0)
	for l in 0:maxlevel(H,1), m in 0:maxlevel(H,2)
		for i in 0:2^l, j in 0:2^m
			value = value + H[l, m, i, j]*basis(l, m, i, j, x, y)
		end
	end
	return value
end

function Base.zeros(H::HierarchicalBasis2D{T})::NodalBasis2D{T} where {T}
	return NodalBasis2D(zeros(T, (2^maxlevel(H,1) + 1, 2^maxlevel(H,2) + 1)))
end

function Base.zeros(N::NodalBasis2D{T})::HierarchicalBasis2D{T} where {T}
	H = Array{Any, 2}(missing, maxlevel(N,1) + 1, maxlevel(N,2) + 1)
	for l in 0:maxlevel(N,1), m in 0:maxlevel(N,2)
		H[l+1,m+1] = zeros(T, (2^l + 1, 2^m + 1))
	end
	return HierarchicalBasis2D{T}(Level2D.(H))
end

function H_2_Nodal(H::HierarchicalBasis2D{T})::NodalBasis2D{T} where {T}
  	N = zeros(H)
 	for i in 0:2^maxlevel(H, 1)
		for j in 0:2^maxlevel(H, 2)
 			N[i,j] = evaluate(H, x(maxlevel(H, 1), maxlevel(H,2), i, j)...)
		end
 	end
 	return N
end

function Nodal_2_H(N::NodalBasis2D{T})::HierarchicalBasis2D{T} where {T}
 	H = zeros(N)
 	for l in 0:maxlevel(N, 1), m in 0:maxlevel(N,2)
		for i in 0:2^l, j in 0:2^m
  			H[l, m, i, j] = N[l, m, i, j] - evaluate(H, x(l,m,i,j)...)
		end
  	end
  	return H
end

function Nodal_2_H_sparse(N::NodalBasis2D{T}, s::Int)::HierarchicalBasis2D{T} where {T}
 	H = zeros(N)
 	for l in 0:maxlevel(N, 1), m in 0:maxlevel(N,2)
		for i in 0:2^l, j in 0:2^m
			if l+m <= s
  				H[l, m, i, j] = N[l, m, i, j] - evaluate(H, x(l,m,i,j)...)
			end
		end
  	end
  	return H
end

 # function Nodal_2_H_new(N::NodalBasis2D{T})::HierarchicalBasis2D{T} where {T}
 # 	 H = zeros(N)
 # 	 H[0,0,0] = N[0,0,0]
 # 	 H[0,0,1] = N[0,0,1]
 # 	 H[0,1,0] = N[0,1,0]
 # 	 H[0,1,1] = N[0,1,1]
 # 	 for l in 1:maxlevel(N, 1), i in 1:2^l-1, j in 1:2^l-1
 # 		 p1 = (x(l, m, i-1, j)..., N[l, m, i-1, j])
 # 		 p2 = (x(l, i-1, j-1)...,N[l, i-1, j-1])
 # 		 p3 = (x(l, i, j+1)...,N[l, i, j+1])
 # 		 p4 = (x(l, i+1, j+1)...,N[l, i-+, j+1])
 # 		 p5 = x(l,i,j)...
 # 		 H[l,i,j] = N[l,i,j] - interpolate(p1, p2, p3, p4, p5)
 # 	 end
 # 	 return H
 # end


text:
function logq(real::State{A1, A2, A3}, mps::StateMP) where {A1, A2, A3}
  lq = 0.0

  for key in fieldnames(State)
    if !(key in (:y_m, :y_log_s))
      mp = getfield(mps, key)
      r = getfield(real, key)

      lq += ADVI.log_q(mp, r)
    end
  end

  return lq
end

max_stars_repo_name: luiarthur/CytofResearch

max_issues_repo_name: luiarthur/CytofResearch