text:
module model

export write_grid, write_model, write_lambda, write_dust, Grid, size_au
export AbstractParameters, ParametersStandard, ParametersTruncated, ParametersCavity, ParametersVertical, convert_vector, convert_dict
export lnprior

# The double dot is because we are now inside the model module, and we want to import the
# constants module, which is part of the enclosing DiskJockey package.
using ..constants
using Dierckx
using Optim
using ODE

# Write the wavelength sampling file. Only run on setup
function write_lambda(lams::AbstractArray, basedir::AbstractString)
    fcam = open(basedir * "camera_wavelength_micron.inp", "w")
    nlam = length(lams)
    @printf(fcam, "%d\n", nlam)
    for lam in lams
        @printf(fcam, "%.9e\n", lam) # [microns]
    end

    close(fcam)
end

const eqmirror = true # mirror the grid about the z=0 midplane ?
# if we decide to mirror, then ncells = 1/2 of the true value

# Define a grid object which stores all of these variables
# This will not change for the duration of the run
immutable Grid
    nr::Int
    ntheta::Int
    nphi::Int
    ncells::Int
    # cell edges
    Rs::Vector{Float64} # [cm]
    Thetas::Vector{Float64}
    Phis::Vector{Float64}
    # cell centers
    rs::Vector{Float64} # [cm]
    thetas::Vector{Float64}
    phis::Vector{Float64}
end

function Grid(nr::Int, ntheta::Int, r_in::Real, r_out::Real, eqmirror::Bool=true)
    # Specify a 2D axisymmetric *separable* grid in spherical coordinates:
    # {r, theta, phi}, where theta is angle from zenith, phi is azimuth

    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # [cm] Inner extent of disk
    r_out = convert(Float64, r_out) * AU # [cm] Outer extent of disk

    #Define the cell *walls*
    Rs = logspace(log10(r_in), log10(r_out), nr+1) # [cm] logarithmically spaced

    if eqmirror
        ped = 0.1
        #Thetas = linspace(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        #Spaced closer near the z=0
    else
        Thetas = linspace(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end]) # [cm]
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)

end

# Create a grid object using a logarithmic then linear then logarithmic radial spacing
function Grid(r_in::Real, r_linstart::Real, r_linend::Real, r_out::Real, n_in::Int, n_mid::Int, n_out::Int, ntheta::Int, eqmirror::Bool=true)
    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    nr = n_in + n_mid + n_out
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # [cm] Inner extent of disk
    r_linstart = convert(Float64, r_linstart) * AU # [cm] Start of linear section
    r_linend = convert(Float64, r_linend) * AU # [cm] End of linear section
    r_out = convert(Float64, r_out) * AU # [cm] Outer extent of disk

    #Define the cell *walls*
    # logarithmically spaced inner grid
    Rs_in = logspace(log10(r_in), log10(r_linstart), n_in + 1) # [cm]

    # linearly spaced middle grid
    Rs_mid = linspace(r_linstart, r_linend, n_mid + 1) # [cm]

    # logarithmically spaced outer grid
    Rs_out = logspace(log10(r_linend), log10(r_out), n_out + 1) # [cm]

    Rs = cat(1, Rs_in[1:end-1], Rs_mid, Rs_out[2:end])

    if eqmirror
        ped = 0.1
        #Thetas = linspace(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        #Spaced closer near the z=0
    else
        Thetas = linspace(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end]) # [cm]
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)
end

# Read from a dictionary, then choose how to make the grid based upon the arguments
function Grid(d::Dict)
    if "r_linstart" in keys(d)
        # We're going for a log-linear-log grid
        names = ["r_in", "r_linstart", "r_linend", "r_out", "n_in", "n_mid", "n_out", "ntheta"]
    else
        # We're just going for a log spaced grid
        names = ["nr", "ntheta", "r_in", "r_out"]
    end
    vec = [d[name] for name in names]
    # Unroll these into an actual parameter
    return Grid(vec...)
end

#This function only needs to be run once, upon setup.
function write_grid(basedir::AbstractString, grid::Grid)
    #amr_grid.inp
    f = open(basedir * "amr_grid.inp", "w")

    #Write the header
    @printf(f, "%d\n", 1) #iformat
    @printf(f, "%d\n", 0) #regular grid (no AMR or Oct-tree)
    @printf(f, "%d\n", 100) #spherical coordiantes
    @printf(f, "%d\n", 0) #gridinfo (none needed for now)
    #incl_r incl_phi incl_z #use this axis?
    @printf(f, "%d %d %d \n", 1, 1, 0) # 2D axisymmetric
    #n_r    n_phi   n_z #number of cells in this dimension
    @printf(f, "%d %d %d \n", grid.nr, grid.ntheta, grid.nphi)

    for R in grid.Rs
        @printf(f, "%.9e\n", R)
    end

    for Theta in grid.Thetas
        @printf(f, "%.9e\n", Theta)
    end

    for Phi in grid.Phis
        @printf(f, "%.9e\n", Phi)
    end

    close(f)
end

# Define an abstract Parameters type, then subset of parameters that can be used to dispatch specifics

abstract AbstractParameters

type ParametersStandard <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface temperature gradient exponent
    Sigma_c::Float64 # [g/cm^2] surface density at characteristic radius
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

type ParametersTruncated <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] Characteristic radius
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface temperature gradient exponent
    gamma_e::Float64 # exponent for outer exponential surface density taper
    Sigma_c::Float64 # [g/cm^2] surface density at characteristic radius
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

type ParametersCavity <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    r_cav::Float64 # [AU] inner radius of the disk, where an exponentially depleted cavity starts
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface density gradient exponent
    gamma_cav::Float64 # surface density gradient exponent for inner cavity
    Sigma_c::Float64 # [g/cm^2] surface density at characteristic radius
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

type ParametersVertical <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10m::Float64 # [K] temperature at 10 AU, midplane
    q_m::Float64 # midplane temperature gradient exponent
    T_10a::Float64 # [K] temperature at 10 AU, atmosphere
    q_a::Float64 # atmosphere temperature gradient exponent
    T_freeze::Float64 # [K] temperature below which to reduce CO abundance
    X_freeze::Float64 # [ratio] amount to reduce CO abundance
    sigma_s::Float64 # Photodissociation boundary in units of A_V.
    gamma::Float64 # surface temperature gradient exponent
    h::Float64 # Number of scale heights that z_q is at, typically fixed to 4
    delta::Float64 # Shape exponent, currently fixed to 2
    Sigma_c::Float64 # [g/cm^2] surface density at characteristic radius
    ksi::Float64 # [cm s^{-1}] micsroturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

"""A dictionary of parameter lists for conversion."""
registered_params = Dict([("standard", ["M_star", "r_c", "T_10", "q", "gamma", "Sigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("truncated", ["M_star", "r_c", "T_10", "q", "gamma", "gamma_e", "Sigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("cavity", ["M_star", "r_c", "r_cav", "T_10", "q", "gamma", "gamma_cav", "Sigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("vertical", ["M_star", "r_c", "T_10m", "q_m", "T_10a", "q_a", "T_freeze", "X_freeze", "sigma_s", "gamma", "h", "delta", "Sigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"])])

registered_types = Dict([("standard", ParametersStandard), ("truncated", ParametersTruncated), ("cavity", ParametersCavity), ("vertical", ParametersVertical)])

"""Unroll a vector of parameter values into a parameter type."""
function convert_vector(p::Vector{Float64}, model::AbstractString, fix_params::Vector; args...)
    args = Dict{Symbol}{Float64}(args)

    # The goal is to assemble a length-nparam vector that can be unrolled into a parameter type
    # e.g., ParametersStandard(M_star, r_c, T_10, q, gamma, Sigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)

    # Select the registerd parameters corresponding to this model
    reg_params = registered_params[model]

    # fit_parameters are the ones in reg_params that are not in (∉) fix_params
    fit_params = filter(x->∉(x,fix_params), reg_params)

    nparams = length(reg_params)

    # Make an empty vector of this same length
    par_vec = Array(Float64, nparams)

    # This requires assigning p to fit_params
    # Find the indexes that correspond to fit_params
    par_indexes = findin(reg_params, fit_params)
    # Stuff p directly into these
    par_vec[par_indexes] = p

    # Then reading fix_params from args
    # First, create an array of fixed values analogous to p
    p_fixed = Float64[args[convert(Symbol, par)] for par in fix_params]
    par_indexes = findin(reg_params, fix_params)
    par_vec[par_indexes] = p_fixed

    # Find the location of logSigma_c and make it Sigma_c
    indSigma_c = findin(reg_params, ["Sigma_c"])
    @assert length(indSigma_c) == 1 "Could not find Sigma_c in order to convert from logSigma_c."
    par_vec[indSigma_c] = 10.^par_vec[indSigma_c]

    # Then assembling these in the same orignial order as registered_params, into the parameter
    # type corresponding to the model.
    return registered_types[model](par_vec...)

end
#
# """Used to unroll a vector of parameter values into a parameter type."""
# function convert_vector(p::Vector{Float64}, model::AbstractString, fix_d::Bool; args...)
#     args = Dict{Symbol}{Float64}(args)
#     if model == "standard"
#         # Take gamma from the args
#         gamma = args[:gamma]
#         if fix_d
#             # distance is fixed and provided by the args
#             dpc = args[:dpc]
#             M_star, r_c, T_10, q, logSigma_c, ksi, incl, PA, vel, mu_RA, mu_DEC = p
#         else
#             # we are fitting with distance as a parameter
#             M_star, r_c, T_10, q, logSigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC = p
#         end
#
#         Sigma_c = 10^logSigma_c
#         return ParametersStandard(M_star, r_c, T_10, q, gamma, Sigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)
#
#     elseif model == "truncated"
#         # Take gamma from the args
#         gamma = args[:gamma]
#         if fix_d
#             dpc = args[:dpc]
#             M_star, r_c, T_10, q, gamma_e, logSigma_c, ksi, incl, PA, vel, mu_RA, mu_DEC = p
#         else
#             M_star, r_c, T_10, q, gamma_e, logSigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC = p
#         end
#
#         Sigma_c = 10^logSigma_c
#         return ParametersTruncated(M_star, r_c, T_10, q, gamma, gamma_e, Sigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)
#
#     elseif model == "cavity"
#         gamma = args[:gamma]
#         if fix_d
#             dpc = args[:dpc]
#             M_star, r_c, r_cav, T_10, q, gamma_cav, logSigma_c, ksi, incl, PA, vel, mu_RA, mu_DEC = p
#         else
#             M_star, r_c, r_cav, T_10, q, gamma_cav, logSigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC = p
#         end
#
#         Sigma_c = 10^logSigma_c
#         return ParametersCavity(M_star, r_c, r_cav, T_10, q, gamma, gamma_cav, logSigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)
#
#     elseif model == "vertical"
#         gamma = args[:gamma]
#         T_freeze = args[:T_freeze]
#         X_freeze = args[:X_freeze]
#         sigma_s = args[:sigma_s]
#         h = args[:h]
#         delta = args[:delta]
#         if fix_d
#             dpc = args[:dpc]
#             M_star, r_c, T_10m, q_m, T_10a, q_a, logSigma_c, ksi, incl, PA, vel, mu_RA, mu_DEC = p
#         else
#             M_star, r_c, T_10m, q_m, T_10a, q_a, logSigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC = p
#         end
#
#         Sigma_c = 10^logSigma_c
#         return ParametersVertical(M_star, r_c, T_10m, q_m, T_10a, q_a, T_freeze, X_freeze, sigma_s, gamma, h, delta, Sigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)
#     else
#         # Raise an error that we don't know this model.
#         throw(ErrorException("Model type $model not yet implemented in model.jl"))
#     end
# end

"""Used to turn a dictionary of parameter values (from config.yaml) directly into a parameter type. Generally used for synthesis and plotting command line scripts."""
function convert_dict(p::Dict, model::AbstractString)
    # Select the registerd parameters corresponding to this model
    reg_params = registered_params[model]
    nparams = length(reg_params)

    # add a new field, which is the conversion of logSigma_c to Sigma_c
    p["Sigma_c"] = 10^p["logSigma_c"]

    # Using this order of parameters, unpack the dictionary p into a vector
    # reg_params reads Sigma_c, not logSigma_c
    par_vec = Float64[p[par_name] for par_name in reg_params]

    return registered_types[model](par_vec...)

end

# """Used to turn a dictionary of parameter values (from config.yaml) into a parameter type. Generally used for synthesis and plotting command line scripts."""
# function convert_dict(p::Dict, model::AbstractString)
#     if model == "standard"
#         names = ["M_star", "r_c", "T_10", "q", "gamma", "logSigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]
#
#         vec = Float64[p[name] for name in names]
#         vec[6] = 10^vec[6]# 10^ gas
#
#         # Unroll these into an actual parameter
#         return ParametersStandard(vec...)
#
#     elseif model == "truncated"
#         names = ["M_star", "r_c", "T_10", "q", "gamma", "gamma_e", "logSigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]
#
#         vec = Float64[p[name] for name in names]
#         vec[7] = 10^vec[7]# 10^ gas
#
#         # Unroll these into an actual parameter
#         return ParametersTruncated(vec...)
#     elseif model == "cavity"
#         names = ["M_star", "r_c", "r_cav", "T_10", "q", "gamma", "gamma_cav", "logSigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]
#
#         vec = Float64[p[name] for name in names]
#         vec[8] = 10^vec[8]# 10^ gas
#
#         # Unroll these into an actual parameter
#         return ParametersCavity(vec...)
#     elseif model == "vertical"
#         names = ["M_star", "r_c", "T_10m", "q_m", "T_10a", "q_a", "T_freeze", "X_freeze", "sigma_s", "gamma", "h", "delta", "logSigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]
#
#         vec = Float64[p[name] for name in names]
#         vec[13] = 10^vec[13]# 10^ gas
#
#         # Unroll these into an actual parameter
#         return ParametersVertical(vec...)
#
#     else
#         # Raise an error that we don't know this model.
#         throw(ErrorException("Model type $model not yet implemented in model.jl"))
#     end
# end

"""The common sense priors that apply to all parameter values"""
function lnprior_base(pars::AbstractParameters, dpc_mu::Float64, dpc_sig::Float64)
    # Create a giant short-circuit or loop to test for sensical parameter values.
    if pars.M_star <= 0.0 || pars.ksi <= 0. || pars.T_10 <= 0. || pars.r_c <= 0.0  || pars.T_10 > 1500. || pars.q < 0. || pars.q > 1.0 || pars.incl < 0. || pars.incl > 180. || pars.PA < -180. || pars.PA > 520.
        # println("M_star ", pars.M_star)
        # println("r_c ", pars.r_c)
        # println("T_10 ", pars.T_10)
        # println("q ", pars.q)
        # println("incl ", pars.incl)
        # println("PA ", pars.PA)
        throw(ModelException("Parameters outside of prior range."))
    end

    # Impose distance prior
    dlow = dpc_mu - 3. * dpc_sig
    dhigh = dpc_mu + 3. * dpc_sig

    # hard +/- 3 sigma cutoff
    if (pars.dpc < dlow) || (pars.dpc > dhigh)
        throw(ModelException("Distance outside of 3 sigma prior range."))
    end

    # If we've passed all the hard-cut offs by this point, return the sum of the distance prior and the geometrical inclination prior.
    return -0.5 * (pars.dpc - dpc_mu)^2 / dpc_sig^2 + log(0.5 * sind(pars.incl))

end

function lnprior(pars::ParametersStandard, dpc_mu::Float64, dpc_sig::Float64, grid::Grid)

    lnp = lnprior_base(pars, dpc_mu, dpc_sig)

    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.
    if (3 * pars.r_c) > r_out
        throw(ModelException("Model radius too large for grid size."))
    else
        return lnp
    end

end

function lnprior(pars::ParametersTruncated, dpc_mu::Float64, dpc_sig::Float64, grid::Grid)
    lnp = lnprior_base(pars, dpc_mu, dpc_sig)

    r_out = grid.Rs[end]/AU # [AU]

    if (3 * pars.r_c) > r_out || pars.gamma_e > 2.0
        throw(ModelException("Disk radii outside model grid radius, or outer power law too flat."))
    else
        return lnp
    end

end

function lnprior(pars::ParametersCavity, dpc_mu::Float64, dpc_sig::Float64, grid::Grid)
    lnp = lnprior_base(pars, dpc_mu, dpc_sig)

    r_in = grid.Rs[1]/AU # [AU]
    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.

    # Also check to make sure that r_cav is less than r_c but larger than r_in.
    if (3 * pars.r_c) > r_out || pars.r_cav < r_in || pars.r_cav > pars.r_c || pars.gamma_cav < 0.0
        throw(ModelException("Model radius too large, grid size or cavity too large, or gamma_cav less than zero."))
    else
        return lnp
    end

end

function lnprior(pars::ParametersVertical, dpc_mu::Float64, dpc_sig::Float64, grid::Grid)
    # Create a giant short-circuit or loop to test for sensical parameter values.
    if pars.M_star <= 0.0 || pars.ksi <= 0. || pars.T_10a <= 0. || pars.T_10m <= 0. || pars.r_c <= 0.0  || pars.T_10a > 1500. || pars.q_m < 0. || pars.q_a < 0. || pars.q_m > 1.0 || pars.q_a > 1.0 || pars.incl < 0. || pars.incl > 180. || pars.PA < -180. || pars.PA > 520. || pars.X_freeze > 1.0 || pars.sigma_s < 0.0
        throw(ModelException("Parameters outside of prior range."))
    end

    # Check to see that the temperatures make sense
    if pars.T_10m > pars.T_10a
        throw(ModelException("Atmosphere is cooler than midplane."))
    end

    # Impose distance prior
    dlow = dpc_mu - 3. * dpc_sig
    dhigh = dpc_mu + 3. * dpc_sig

    # hard +/- 3 sigma cutoff
    if (pars.dpc < dlow) || (pars.dpc > dhigh)
        throw(ModelException("Distance outside of 3 sigma prior range."))
    end

    # If we've passed all the hard-cut offs by this point, return the sum of the distance prior and the geometrical inclination prior.
    lnp = -0.5 * (pars.dpc - dpc_mu)^2 / dpc_sig^2 + log(0.5 * sind(pars.incl))

    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.
    if (3 * pars.r_c) > r_out
        throw(ModelException("Model radius too large for grid size."))
    else
        return lnp
    end

end

# Determine the physical size of the image. Size_arcsec is the full width of the image,
# and so sizeau is the full size of the image as well (RADMC3D conventions).
# Because there seems to be a small but constant shift in the size of the image
function size_au(size_arcsec::Real, dpc::Real, grid::Grid)
    outer_model = (1.1 * grid.rs[end] / AU) # [AU]
    sizeau_desired = size_arcsec * dpc # [AU]

    if sizeau_desired < outer_model
        # We don't want this error to be caught by the lnprob routine, since this should
        # always halt execution of the script.
        throw(ErrorException("Image size ($size_au AU) must be larger than 110% the model grid size ($outer_model AU)."))
    end

    # because there seems to be a slight offset between what sizeau is specified and the actual
    # size of the image, we will also specifiy sizeau_command, which is the value to give to RADMC
    sizeau_command = sizeau_desired/(1. + RADMC_SIZEAU_SHIFT)

    return (sizeau_desired, sizeau_command) # [AU]

end

# Assume all inputs to these functions are in CGS units and in *cylindrical* coordinates.
# Parametric type T allows passing individual Float64 or Vectors.
# # Alternate functions accept pars passed around, where pars is in M_star, AU, etc...
function velocity{T}(r::T, M_star::Float64)
    sqrt(G * M_star ./ r)
end
velocity{T}(r::T, pars::AbstractParameters) = velocity(r, pars.M_star * M_sun)

# For the vertical temperature gradient
function velocity(r::Float64, z::Float64, M_star::Float64)
    sqrt(G * M_star / (r^2 + z^2)^(3./2)) * r
end
velocity(r::Float64, z::Float64, pars::ParametersVertical) = velocity(r, z, pars.M_star * M_sun)

function temperature{T}(r::T, T_10::Float64, q::Float64)
    T_10 * (r ./ (10. * AU)).^(-q)
end
temperature{T}(r::T, pars::AbstractParameters) = temperature(r, pars.T_10, pars.q)

function Hp{T}(r::T, M_star::Float64, T_10::Float64, q::Float64)
    temp = temperature(r, T_10, q)
    sqrt(kB * temp .* r.^3./(mu_gas * m_H * G * M_star))
end
Hp{T}(r::T,  pars::AbstractParameters) = Hp(r, pars.M_star * M_sun, pars.T_10, pars.q)
# Scale height computed from midplane temperature for vertical temperature gradient model
Hp{T}(r::T,  pars::ParametersVertical) = Hp(r, pars.M_star * M_sun, pars.T_10m, pars.q_m)

# For the vertical temperature gradient model
T_mid{T}(r::T, pars::ParametersVertical) = temperature(r, pars.T_10m, pars.q_m)
T_atm{T}(r::T, pars::ParametersVertical) = temperature(r, pars.T_10a, pars.q_a)

# Atmosphere height computed as multiple of scale height (computed at midplane)
function z_q{T}(r::T, pars::ParametersVertical)
    return pars.h * Hp(r, pars)
end

function temperature(r::Float64, z::Float64, pars::ParametersVertical)
    zq = z_q(r, pars)
    Ta = T_atm(r, pars)
    if z >= zq
        return Ta
    else
        Tm = T_mid(r, pars)
        # return Ta + (Tm - Ta) * cos(pi * z/ (2 * zq))^pars.delta # Dartois
        return Tm + (Ta - Tm) * sin(pi * z/ (2 * zq))^(2 * pars.delta) # Williams and Best 14
    end
end

# Calculate the gas surface density
function Sigma(r::Float64, pars::Union{ParametersStandard, ParametersVertical})
    r_c = pars.r_c * AU

    gamma = pars.gamma
    Sigma_c = pars.Sigma_c

    S = Sigma_c * (r/r_c)^(-gamma) * exp(-(r/r_c)^(2 - gamma))

    return S
end

function Sigma(r::Float64, pars::ParametersTruncated)
    r_c = pars.r_c * AU # [cm]
    gamma = pars.gamma
    gamma_e = pars.gamma_e
    Sigma_c = pars.Sigma_c

    S = Sigma_c * (r/r_c)^(-gamma) * exp(-(r/r_c)^(2 - gamma_e))
    return S
end

function Sigma(r::Float64, pars::ParametersCavity)
    r_c = pars.r_c * AU
    r_cav = pars.r_cav * AU

    gamma = pars.gamma
    gamma_cav = pars.gamma_cav
    Sigma_c = pars.Sigma_c

    inner_taper = exp(-(r_cav/r)^gamma_cav)
    outer_taper = exp(-(r/r_c)^(2 - gamma))
    power_law = (r/r_c)^(-gamma)

    S = Sigma_c * inner_taper * power_law * outer_taper

    return S
end

# Delivers a gas density in g/cm^3
function rho_gas(r::Float64, z::Float64, pars::AbstractParameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    # Calculate the density
    rho = S/(sqrt(2. * pi) * H) * exp(-0.5 * (z/H)^2)

    # If the density is less than our cutoff, just return 0.0
    if rho < rho_gas_critical
        return constants.rho_gas_zero
    # Otherwise, return the density at this height
    else
        return rho
    end
end

# Determines the midplane density assuming vertically isothermal
function rho_gas_mid(r::Float64, pars::AbstractParameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    # Calculate the density at the midplane
    rho = S/(sqrt(2. * pi) * H)

    # If the density is less than our cutoff, just return 0.0
    if rho < rho_gas_critical
        return constants.rho_gas_zero
    # Otherwise, return the density at this height
    else
        return rho
    end
end

# The temperature change, dT/dz
function dT(r::Float64, z::Float64, pars::ParametersVertical)
    zq = z_q(r, pars)
    Ta = T_atm(r, pars)

    if z >= zq
        return 0.
    else
        Tm = T_mid(r, pars)
        # return - pi * pars.delta / (2 * zq) * (Tm - Ta) * (cos(pi * z/(2 * zq)))^(pars.delta - 1) * sin(pi * z / (2 * zq))
        return pars.delta * pi/zq * (Ta - Tm) * cos(pi * z/(2 * zq)) * (sin(pi * z/(2 * zq)))^(2 * pars.delta - 1)
    end
end

# The integrand
function dlnrho(r::Float64, z::Float64, pars::ParametersVertical)
    T = temperature(r, z, pars)
    dTemp = dT(r, z, pars)
    return -(dTemp/T + (mu_gas * m_H)/(kB * T) * (G * pars.M_star * M_sun * z)/(r^2 + z^2)^1.5)
end

# The integrand for a vertically isothermal testcase
function dlnrho(r::Float64, z::Float64, pars::ParametersStandard)
    return -(mu_gas * m_H)/(kB * T) * (G * pars.M_star * M_sun * z)/r^3
end

function z_top(r::Float64, pars::ParametersVertical)
    zq = z_q(r, pars) # The height of the disk "atmosphere"

    return zq * 5
end

function rho_gas(r::Float64, z::Float64, pars::ParametersVertical)

    # Calculate the "top" of the atmosphere,
    # a height where we are sure we can enforce that density = 0
    ztop = z_top(r, pars)

    # If we are querying a height above this, just return 0 now
    if z > ztop
        return constants.rho_gas_zero
    end

    # Calculate what the column density should be at the midplane
    sigma = Sigma(r, pars) / 2

    # Calculate the photodissociation height
    # Threshold column density
    thresh_H2 = pars.sigma_s * Av_sigmaH # [n_H2/cm^2]
    thresh = thresh_H2 * (mu_gas * amu / X_H2) # [g/cm^2] of gas

    # If there is not even enough column density at the midplane in order to exceed the threshold
    # to protect against photodissociation of CO, just return 0 now
    if sigma < thresh
        return constants.rho_gas_zero
    end

    # Create an array of heights that goes from the midplane (0, 0.001, ..., ztop)
    zs = cat(1, [0], logspace(log10(0.001 * AU), log10(ztop), 64 - 1))

    # Define the ODE at this radius r
    function f(z, y)
        # calculate y'
        dlnrho(r, z, pars) * y
    end

    # Choose a starting guess corresponding to the midplane density for an isothermal model.
    # This helps us get in the ballpark and avoid a lot of the numerical errors that result
    start = rho_gas_mid(r, pars)

    # Now solve the ODE on the grid of postulated values
    # For the inner disk, we require a larger value of abstol
    # We need to experiment with what is best
    z_out, y_out = ode45(f, start, zs; abstol=1e-19)

    # Integrate the output from the ODE solver to find the total (un-normalized) column density
    spl = Spline1D(z_out, y_out)
    tot = integrate(spl, z_out[1], z_out[end])

    # Apply a correction factor to the un-normalized densities
    cor = sigma / tot
    rhos = y_out .* cor

    # This is the function that will be minimized
    function g(zvar::Float64)
        return abs(thresh - cor * integrate(spl, zvar, z_out[end]))
    end

    z_phot = optimize(g, z_out[1], z_out[end]).minimum

    # If we are above this height, return the mimimum gas density
    if z > z_phot
        return constants.rho_gas_zero
    # If we are below it, interpolate rhos to this point
    else
        rho = cor * evaluate(spl, z)

        return rho
    end
end

# Now, replace these functions to simply multiply rho_gas by X_12CO/m_12CO, or X_13CO/m_13CO, etc.
n_12CO(r::Float64, z::Float64, pars::AbstractParameters) = number_densities["12CO"] * rho_gas(r, z, pars)

n_13CO(r::Float64, z::Float64, pars::AbstractParameters) = number_densities["13CO"] * rho_gas(r, z, pars)

n_C18O(r::Float64, z::Float64, pars::AbstractParameters) = number_densities["C18O"] * rho_gas(r, z, pars)

# It is realistic to include freezout of CO onto dust grains.
# This is the amount by which the number density of the CO is reduced (X_freeze) relative to
# the nominal value.
function X_freeze(temp::Float64, pars::AbstractParameters)
    # If it's cooler than the freezout temperature, reduce the number density by the given factor
    if temp <= pars.T_freeze
        return pars.X_freeze
    # Otherwise, just keep it as is.
    else
        return 1.0
    end
end

function rho_dust(r::Float64, z::Float64, pars::AbstractParameters)

    # Use the rho_gas function to get the total density in [g/cm^3]
    mGas = rho_gas(r, z, pars)

    # Convert from mGas to mDust using Gas/Dust ratio of 100
    mDust = mGas * 0.01 # [g]

    return mDust
end

# Ksi is microturbulent broadining width in units of km/s. Output of this function
# is in cm/s for RADMC (RADMC manual, eqn 7.12)
function microturbulence(ksi::Float64)
    return ksi * 1.e5 # convert from km/s to cm/s
end

microturbulence(pars::AbstractParameters) = microturbulence(pars.ksi)

function write_model(pars::AbstractParameters, basedir::AbstractString, grid::Grid, species::AbstractString)

    funcs = Dict([("12CO", n_12CO), ("13CO", n_13CO), ("C18O", n_C18O)])
    n_CO = funcs[species]

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", n_CO(r_cyl, z, pars))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, pars))
                @printf(ftemp, "%.9e\n", temperature(r_cyl, pars))
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function write_model(pars::ParametersVertical, basedir::AbstractString, grid::Grid, species::AbstractString)

    function n_CO(r_cyl, z)
        return number_densities[species] * rho_gas(r_cyl, z, pars)
    end

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                temp = temperature(r_cyl, z, pars)

                @printf(fdens, "%.9e\n", X_freeze(temp, pars) * n_CO(r_cyl, z))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, z, pars))
                @printf(ftemp, "%.9e\n", temp)
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function write_dust(pars::AbstractParameters, basedir::AbstractString, grid::Grid)
    fdens = open(basedir * "dust_density.inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)
    @printf(fdens, "%d\n", 1) # number of dust species

    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", rho_dust(r_cyl, z, pars))
            end
        end
    end

    close(fdens)
end

end

max_stars_repo_name: elnjensen/DiskJockey

max_issues_repo_name: elnjensen/DiskJockey

max_forks_repo_name: elnjensen/DiskJockey


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    EstimationSolution

A solution to a spatial estimation problem.
"""
struct EstimationSolution{𝒟,ℳ,𝒱}
  domain::𝒟
  mean::ℳ
  variance::𝒱
end

# -------------
# VARIABLE API
# -------------

Base.getindex(solution::EstimationSolution, var::Symbol) =
  (mean=solution.mean[var], variance=solution.variance[var])

# ------------
# IO methods
# ------------
function Base.show(io::IO, solution::EstimationSolution)
  N = ncoords(solution.domain)
  print(io, "$(N)D EstimationSolution")
end

function Base.show(io::IO, ::MIME"text/plain", solution::EstimationSolution)
  println(io, solution)
  println(io, "  domain: ", solution.domain)
  print(  io, "  variables: ", join(keys(solution.mean), ", ", " and "))
end

ext: jl


text:
const NaNValue    = "NaN"
const NaNName     = "NaN65"

const InfValue    = "∞"
const PosInfValue = "+∞"
const NegInfValue = "-∞"
const InfName     = "Inf65"

const TinyName  = "TiNY"
const HugeName  = "HuGE"

const Exactly   = "∣"
const Inexactly = "⍿"

function show(io::IO, x::Float65)

    state = getstate(x)
    fp = convert(Float64, x)

    if isfinite(x)
        if Tiny(Float64) < fp < Huge(Float64)
           s = prepshow(fp, state)
        elseif fp == 0.0
           s = prepshow0(fp, state)
        elseif fp <= Tiny(Float64)
          s = prepshowTiny(fp, state)
        else
          s = prepshowHuge(fp, state)
        end
    elseif isinf(x.fp)  # ±Inf or NaN (restore from clearing the state)
       s = prepshowInf(x.fp,state)
    else
       s = prepshowNaN(x.fp,state)
    end

    print(io, s)
end

@inline prepshowState(state::Bool) = state ? Inexactly : Exactly
@inline function prepnumsign(fp::Float64)
    isneg = signbit(fp)
    if isneg
        "-"
    else
        "+"
    end
end

function prepshow(fp::Float64, state::Bool)
   string(fp, prepshowState(state))
end
function prepshow0(fp::Float64, state::Bool)
   string(fp, prepshowState(state))
end
function prepshowTiny(fp::Float64, state::Bool)
   string( prepnumsign(fp), TinyName )
end
function prepshowHuge(fp::Float64, state::Bool)
   string( prepnumsign(fp), HugeName )
end
function prepshowInf(fp::Float64, state::Bool)
   string( prepnumsign(fp), InfValue )
end
function prepshowNaN(fp::Float64, state::Bool)
   NaNValue
end

# read(IO, x), write(IO, x)

function read(io::IO, ::Type{Float65})
    fp = read(io, T)
    convert(Float65, v)
end
function write(io::IO, x::Float65)
    write(io, x.fp)
end


text:
@testset "Element2MILP: $(fct_type), dimension $(dim), $(T)" for fct_type in ["vector of variables", "vector affine function"], dim in [2, 3], T in [Int, Float64]
    mock = MOIU.MockOptimizer(MILPModel{T}())
    model = COIB.Element2MILP{T}(mock)

    @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.Integer)
    @test MOI.supports_constraint(
        model,
        MOI.ScalarAffineFunction{T},
        MOI.EqualTo{T},
    )
    @test MOIB.supports_bridging_constraint(
        model,
        MOI.VectorAffineFunction{T},
        CP.Element{T},
    )

    x_index, _ = MOI.add_constrained_variable(model, MOI.Integer())
    if T == Int
        x_value, _ = MOI.add_constrained_variable(model, MOI.Integer())
    elseif T == Float64
        x_value = MOI.add_variable(model)
    end
    values = T.(collect(T(4) .+ (1:dim)))

    fct = if fct_type == "vector of variables"
        MOI.VectorOfVariables([x_value, x_index])
    elseif fct_type == "vector affine function"
        MOIU.vectorize([x_value, x_index])
    else
        @assert false
    end
    c = MOI.add_constraint(model, fct, CP.Element(values))

    @test MOI.is_valid(model, x_value)
    @test MOI.is_valid(model, x_index)
    @test MOI.is_valid(model, c)

    bridge = MOIBC.bridges(model)[MOI.ConstraintIndex{MOI.VectorOfVariables, CP.Element}(-1)]

    @testset "Bridge properties" begin
        @test MOIBC.concrete_bridge_type(typeof(bridge), MOI.VectorOfVariables, CP.Element) == typeof(bridge)
        @test MOIB.added_constrained_variable_types(typeof(bridge)) == [(MOI.ZeroOne,)]
        @test MOIB.added_constraint_types(typeof(bridge)) == [
            (MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}),
        ]

        @test MOI.get(bridge, MOI.NumberOfVariables()) == dim
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.VariableIndex, MOI.ZeroOne}()) == dim
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}()) == 3

        @test MOI.get(bridge, MOI.ListOfVariableIndices()) == bridge.vars
        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.VariableIndex, MOI.ZeroOne}()) == bridge.vars_bin
        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}()) == [bridge.con_unary, bridge.con_choose_one, bridge.con_value]
    end

    @testset "Set of variables" begin
        @test length(bridge.vars) == dim
        for i in 1:dim
            @test MOI.is_valid(model, bridge.vars[i])
            @test MOI.is_valid(model, bridge.vars_bin[i])
            @test MOI.get(model, MOI.ConstraintFunction(), bridge.vars_bin[i]) == bridge.vars[i]
            @test MOI.get(model, MOI.ConstraintSet(), bridge.vars_bin[i]) == MOI.ZeroOne()
        end
    end

    @testset "Unary decomposition of the index" begin
        @test MOI.is_valid(model, bridge.con_unary)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con_unary)
        @test length(f.terms) == dim + 1
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con_unary) == MOI.EqualTo(zero(T))

        t = f.terms[1]
        @test t.coefficient === -one(T)
        @test t.variable == x_index

        for i in 1:dim
            t = f.terms[i + 1]
            @test t.coefficient === T(i)
            @test t.variable == bridge.vars[i]
        end
    end

    @testset "Unicity of the unary decomposition" begin
        @test MOI.is_valid(model, bridge.con_choose_one)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con_choose_one)
        @test length(f.terms) == dim
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con_choose_one) == MOI.EqualTo(one(T))

        for i in 1:dim
            t = f.terms[i]
            @test t.coefficient === one(T)
            @test t.variable == bridge.vars[i]
        end
    end

    @testset "Value at the index" begin
        @test MOI.is_valid(model, bridge.con_value)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con_value)
        @test length(f.terms) == dim + 1
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con_value) == MOI.EqualTo(zero(T))

        t = f.terms[1]
        @test t.coefficient === -one(T)
        @test t.variable == x_value

        for i in 1:dim
            t = f.terms[i + 1]
            @test t.coefficient === values[i]
            @test t.variable == bridge.vars[i]
        end
    end
end


text:
@testset "named tuple difference" begin
    modvals = (x1=1, x2=[0.0, 1.0])
    addvals = NamedTuple()
    remvals = (z=3,)
    a = NamedTupleDifference(modvals, addvals, remvals)
    b = NamedTupleDifference(modvals, addvals, remvals)
    c = NamedTupleDifference(modvals, addvals, remvals)

    @testset "constructors" begin
        @test isa(a, NamedTupleDifference)
        @test a.modvals == modvals && a.addvals == addvals && a.remvals == remvals
    end

    @testset "equality operator" begin
        @test a == a
        @test a == b && b == a
        @test a == b && b == c && a == c
    end

    @testset "hash" begin
        @test hash(a) == hash(a)
        @test a == b && hash(a) == hash(b)
    end

    @testset "accessors" begin
        @test modified(a) == modvals
        @test added(a) == addvals
        @test removed(a) == remvals
    end

    @testset "difference" begin
        a = (x=1,)
        b = (x=(s=1, t=2),)
        c = (x=[1, 2],)
        d = (x=2, y=1)
        e = []
        ei = Int[]
        ent = NamedTuple()
        @test diff(a, a) == NamedTupleDifference((x=0,), ent, ent)
        @test diff(b, b) == NamedTupleDifference(
            (x=NamedTupleDifference((s=0, t=0), ent, ent),), ent, ent
        )
        @test diff(c, c) == NamedTupleDifference(
            (x=VectorDifference([1, 2], ei, ei, [0, 0], e, e),), ent, ent
        )
        @test diff(a, d) == NamedTupleDifference((x=-1,), (y=1,), ent)
        @test diff(d, a) == NamedTupleDifference((x=1,), ent, (y=1,))
    end
end


text:
module Hankel
import SpecialFunctions: besselj, gamma
using GSL: GSL
using Roots: Roots
import LinearAlgebra: mul!, ldiv!, dot
import Base: *, \
using ChainRulesCore

export QDHT, integrateK, integrateR, onaxis, symmetric, Rsymmetric

const J₀₀ = besselj(0, 0)

include("utils.jl")
include("qdht.jl")
include("chainrules.jl")

end

max_stars_repo_name: chrisbrahms/Hankel

max_issues_repo_name: chrisbrahms/Hankel

max_forks_repo_name: chrisbrahms/Hankel


text:
# This file was generated, do not modify it. # hide
using MLJ
using UrlDownload
using PyPlot
ioff() # hide
import DataFrames: DataFrame, describe, names, select!
using Statistics

@load LinearRegressor pkg=MLJLinearModels


text:
# Built-in Exceptions

# sqrt(-1)

# Throwing Exceptions

a = 1; b = 0
# b != 0 ? a / b : throw(DivideError())

# Errors

# b != 0 ? a / b : error("Divide by zero not allowed!")

# Try-Catch-Finally

try
  sqrt(-1)
catch e
  println("Square root of negative number not allowed!")
end

try
  sqrt(-1)
catch e
  println("Square root of negative number not allowed!")
finally
  println("This is always executed!")
end

max_stars_repo_name: cadamsmith/julia-programming

max_issues_repo_name: cadamsmith/julia-programming

max_forks_repo_name: cadamsmith/julia-programming


text:
using Test: @test_nowarn, @test
using SHA
using Trixi
using Trixi2Vtk

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
const EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples")

function run_trixi(elixir; parameters...)
  @test_nowarn trixi_include(joinpath(EXAMPLES_DIR, elixir); parameters...)
end

function sha1file(filename)
  open(filename) do f
    bytes2hex(sha1(f))
  end
end

function test_trixi2vtk(filenames, outdir; hashes=nothing, kwargs...)
  @test_nowarn trixi2vtk(joinpath(outdir, filenames); output_directory=outdir, kwargs...)

  if !isnothing(hashes)
    for (filename, hash_expected) in hashes
      hash_measured = sha1file(joinpath(outdir, filename))
      @test hash_expected == hash_measured
    end
  end
end


text:
using Test
using CoarseGraining3DPrisms
using LinearAlgebra

const K = 2
const x = K+1
const y = K/2

a = 0.0

@test delta(1.0,1.0,1.0) == 0

#@show PrA

PrA = zeros(x,x,x,x,x,x,x,x,x,x,x,x)
for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jd1 in 0.:0.5:y, je1 in 0.:0.5:y, jf1 in 0.:0.5:y, ja2 in 0.:0.5:y, jb2 in 0.:0.5:y, jd2 in 0.:0.5:y, je2 in 0.:0.5:y, jf2 in 0.:0.5:y, jc in 0.:0.5:y, jg in 0.:0.5:y
    ans = numchop( prismA(ja1,jb1,jd1,je1,jf1,ja2,jb2,jd2,je2,jf2,jc,jg))
    PrA[Int(2*ja1+1),Int(2*jb1+1),Int(2*jd1+1),Int(2*je1+1),Int(2*jf1+1),Int(2*ja2+1),Int(2*jb2+1),Int(2*jd2+1),Int(2*je2+1),Int(2*jf2+1),Int(2*jc+1),Int(2*jg+1)] = ans
end

#PrA = zeros(x,x,x,x,x,x,x,x,x,x,x,x)
ampsA = Float64[]
indxA = Array{Int64,1}[]
for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jd1 in 0.:0.5:y, je1 in 0.:0.5:y, jf1 in 0.:0.5:y, ja2 in 0.:0.5:y, jb2 in 0.:0.5:y, jd2 in 0.:0.5:y, je2 in 0.:0.5:y, jf2 in 0.:0.5:y, jc in 0.:0.5:y, jg in 0.:0.5:y
    ans = numchop( prismA(ja1,jb1,jd1,je1,jf1,ja2,jb2,jd2,je2,jf2,jc,jg))
    if ans != 0
        push!(ampsA,ans)
        push!(indxA, [Int(2*ja1+1),Int(2*jb1+1),Int(2*jd1+1),Int(2*je1+1),Int(2*jf1+1),Int(2*ja2+1),Int(2*jb2+1),Int(2*jd2+1),Int(2*je2+1),Int(2*jf2+1),Int(2*jc+1),Int(2*jg+1)])
        #push!(lindx,)
    end
end

#PrB = zeros(x,x,x,x,x,x,x,x,x,x,x,x)
#for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jd1 in 0.:0.5:y, je1 in 0.:0.5:y, jf1 in 0.:0.5:y, ja2 in 0.:0.5:y, jb2 in 0.:0.5:y, jd2 in 0.:0.5:y, je2 in 0.:0.5:y, jf2 in 0.:0.5:y, jc in 0.:0.5:y, jg in 0.:0.5:y
#    ans = numchop( prismB(ja1,jb1,jd1,je1,jf1,ja2,jb2,jd2,je2,jf2,jc,jg))
   # PrB[Int(2*ja1+1),Int(2*jb1+1),Int(2*jd1+1),Int(2*je1+1),Int(2*jf1+1),Int(2*ja2+1),Int(2*jb2+1),Int(2*jd2+1),Int(2*je2+1),Int(2*jf2+1),Int(2*jc+1),Int(2*jg+1)] = ans
#end

#PrB = zeros(x,x,x,x,x,x,x,x,x,x,x,x)
ampsB = Float64[]
indxB = Array{Int64,1}[]
for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jd1 in 0.:0.5:y, je1 in 0.:0.5:y, jf1 in 0.:0.5:y, ja2 in 0.:0.5:y, jb2 in 0.:0.5:y, jd2 in 0.:0.5:y, je2 in 0.:0.5:y, jf2 in 0.:0.5:y, jc in 0.:0.5:y, jg in 0.:0.5:y
    ans = numchop( prismB(ja1,jb1,jd1,je1,jf1,ja2,jb2,jd2,je2,jf2,jc,jg))
    if ans != 0
        push!(ampsB,ans)
        push!(indxB,[Int(2*ja1+1),Int(2*jb1+1),Int(2*jd1+1),Int(2*je1+1),Int(2*jf1+1),Int(2*ja2+1),Int(2*jb2+1),Int(2*jd2+1),Int(2*je2+1),Int(2*jf2+1),Int(2*jc+1),Int(2*jg+1)])
    end
end

#pos = findall(x-> x!=0, PrA )
#PrANZ = PrA[pos]
#posB = findall(x-> x!=0, PrB )
#PrANZ = PrB[posB]

#@show pos

function loppo(amps::Array{Float64,1},indx::Array{Array{Int64,1},1},jd1::Float64,jd2::Float64,je1::Float64)
    #posa = findall(x-> x!=0, prA )
    #PrANZ = PrA[pos]
    #a = Float64[]
    #b = Array{Int64,1}[]
    c = []
    for i in 1:length(indx)# position of d1 = 3, d2= 8 and e1 = 4
        if indx[i][3] == Int(2*jd1+1) && indx[i][8] == Int(2*jd2+1) && indx[i][4] == Int(2*je1+1)
            #push!(a,amps[i])
            #push!(b,indx[i])
            push!(c,(indx[i],amps[i]))
        end
    end
    return c
end

function Nsmat(amps::Array{Float64,1},indx::Array{Array{Int64,1},1},jd1::Float64,jd2::Float64,je1::Float64)
    #lopV = loppo(amps,indx,jd1,jd2,je1)[1]
    lopC = loppo(amps,indx,jd1,jd2,je1)
    a = Array{Int64,1}[]
    #for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
    b = Array{Int64,1}[]
    for i in 1:length(lopC) # position of a1 = 1, b1 = 2 g = 12
        unique!(push!(a,[lopC[i][1][1],lopC[i][1][2],lopC[i][1][12]]))
        unique!(push!(b,[lopC[i][1][5],lopC[i][1][6],lopC[i][1][7],lopC[i][1][9],lopC[i][1][10],lopC[i][1][11]]))
    end
    dd = zeros(length(a),length(b))
    for i in 1:length(lopC)
        qa = findall(x-> x == [lopC[i][1][1], lopC[i][1][2], lopC[i][1][12]],a)
        qb = findall(x-> x == [lopC[i][1][5], lopC[i][1][6], lopC[i][1][7], lopC[i][1][9], lopC[i][1][10], lopC[i][1][11] ], b)
        dd[qa,qb] = [lopC[i][2]]
    end
    return a,b,dd
end

function blockReduceA(amps::Array{Float64,1},indx::Array{Array{Int64,1},1},jd1::Float64,jd2::Float64,je1::Float64)
    #a = Float64[]
    #b = Array{Int64,1}[]
    ampInfo = [] # will contain indices and amplitudes
    for i in 1:length(indx)# position of d1 = 3, d2= 8 and e1 = 4
        if indx[i][3] == Int(2*jd1+1) && indx[i][8] == Int(2*jd2+1) && indx[i][4] == Int(2*je1+1)
            push!(ampInfo,(indx[i],amps[i]))
        end
    end
    nzcol = Array{Int64,1}[]
    #for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
    nzrow = Array{Int64,1}[]
    for i in 1:length(ampInfo) # position of a1 = 1, b1 = 2 g = 12
        unique!(push!(nzcol,[ampInfo[i][1][1],ampInfo[i][1][2],ampInfo[i][1][12]]))
        unique!(push!(nzrow,[ampInfo[i][1][5],ampInfo[i][1][6],ampInfo[i][1][7],ampInfo[i][1][9],ampInfo[i][1][10],ampInfo[i][1][11]]))
    end
    dd = zeros(length(nzcol),length(nzrow))
    if length(dd) != 0
        for i in 1:length(ampInfo)
            qa = findall(x-> x == [ampInfo[i][1][1], ampInfo[i][1][2], ampInfo[i][1][12]],nzcol)
            qb = findall(x-> x == [ampInfo[i][1][5], ampInfo[i][1][6], ampInfo[i][1][7], ampInfo[i][1][9], ampInfo[i][1][10], ampInfo[i][1][11] ], nzrow)
            dd[qa,qb] = [ampInfo[i][2]]
        end
        return nzcol,nzrow, dd
    else
        return 0,0,0
    end
end

#get the index for each of the shared triangular face.. return index values and amplitudes
# gives 30 elements for spins(0,0,0)
function loppo(amps::Array{Float64,1},indx::Array{Array{Int64,1},1},jd1::Float64,jd2::Float64,je1::Float64)
    #posa = findall(x-> x!=0, prA )
    #PrANZ = PrA[pos]
    a = Float64[]
    b = Array{Int64,1}[]
    c = []
    for i in 1:length(indx)# position of d1 = 3, d2= 8 and e1 = 4
        if indx[i][3] == Int(2*jd1+1) && indx[i][8] == Int(2*jd2+1) && indx[i][4] == Int(2*je1+1)
            #push!(a,amps[i])
            #push!(b,indx[i])
            push!(c,(indx[i],amps[i]))
        end
    end
    return c
end

function Nsmat(amps::Array{Float64,1},indx::Array{Array{Int64,1},1},jd1::Float64,jd2::Float64,je1::Float64)
    #lopV = loppo(amps,indx,jd1,jd2,je1)[1]
    lopC = loppo(amps,indx,jd1,jd2,je1)
    a = Array{Int64,1}[]
    #for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
    b = Array{Int64,1}[]
    for i in 1:length(lopC) # position of a1 = 1, b1 = 2 g = 12
        unique!(push!(a,[lopC[i][1][1],lopC[i][1][2],lopC[i][1][12]]))
        unique!(push!(b,[lopC[i][1][5],lopC[i][1][6],lopC[i][1][7],lopC[i][1][9],lopC[i][1][10],lopC[i][1][11]]))
    end
    dd = zeros(length(a),length(b))
    for i in 1:length(lopC)
        qa = findall(x-> x == [lopC[i][1][1], lopC[i][1][2], lopC[i][1][12]],a)
        qb = findall(x-> x == [lopC[i][1][5], lopC[i][1][6], lopC[i][1][7], lopC[i][1][9], lopC[i][1][10], lopC[i][1][11] ], b)
        dd[qa,qb] = [lopC[i][2]]
    end
    return a,b,dd
end

function blockReduceA(amps::Array{Float64,1},indx::Array{CartesianIndex{12},1},jd1::Float64,jd2::Float64,je1::Float64)
    #a = Float64[]
    #b = Array{Int64,1}[]
    ampInf = [] # will contain indices and amplitudes
    for i in 1:length(indx)# position of d1 = 3, d2= 8 and e1 = 4
        if indx[i][3] == Int(2*jd1+1) && indx[i][8] == Int(2*jd2+1) && indx[i][4] == Int(2*je1+1)
            push!(ampInf,(indx[i],amps[i]))
        end
    end
    # a and b are index and amplitudes
    c = []
    for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
        d = []
        for i in 1:length(b) # position of a1 = 1, b1 = 2 g = 12
            if b[i][1] == Int(2*ja1+1) && b[i][2] == Int(2*jb1+1) && b[i][12] == Int(2*jg+1)
                push!(d,a[i])
            end
        end
        if d != []
            push!(c,d)
        end
    end
    if length(c) !=0
        dd = zeros(length(c),length(c[1]))
        for i in 1:length(c)
            c[i] = convert(Array{Float64},c[i])
            dd[i,:] = c[i]
        end
        return dd
    else
        return 0
    end
end

function Nsmate(amps::Array{Float64,1},indx::Array{Array{Int64,1},1},jd1::Float64,jd2::Float64,je1::Float64)
    lopV = loppo(amps,indx,jd1,jd2,je1)[1]
    lopC = loppo(amps,indx,jd1,jd2,je1)[2]
    a = []
    #for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
    b = []
    for i in 1:length(lopC) # position of a1 = 1, b1 = 2 g = 12
        unique!(push!(a,(lopC[i][1],lopC[i][2],lopC[i][12])))
        unique!(push!(b,(lopC[i][5],lopC[i][6],lopC[i][7],lopC[i][9],lopC[i][10],lopC[i][11])))
    end
    dd = zeros(length(a),length(b))
    if length(dd) != 0
        if length(a)*length(b) == length(lopV)
            for i in 1:length(a), j in 1:length(b)
               dd[i,j] = lopV[(i-1)*length(b)+1]
            end
            return dd
        elseif length(a) == length(b) == length(lopV)
            for i in 1:length(a)
                dd[i,i] = lopV[i]
            end
            return dd
        end
    else
        return 0
    end
    #return dd
end

#get reduced matrix
function smat(amps::Array{Float64,1},indx::Array{CartesianIndex{12},1},jd1::Float64,jd2::Float64,je1::Float64)
    lopV = loppo(amps,indx,jd1,jd2,je1)[1]
    lopC = loppo(amps,indx,jd1,jd2,je1)[2]
    b = []
    #for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
    #a = []
    for i in 1:length(lopC) # position of a1 = 1, b1 = 2 g = 12
        if lopC[i][1] == Int(2*ja1+1) && lopC[i][2] == Int(2*jb1+1) && lopC[i][12] == Int(2*jg+1)
            push!(a,lopV[i])
        end
    end
    if a != []
        push!(b,a)
    end
    #end
    if length(b) !=0
        dd = zeros(length(b),length(b[1]))
        for i in 1:length(b)
            b[i] = convert(Array{Float64},b[i])
            dd[i,:] = b[i]
        end
        return dd
    else
        return 0
    end
end

#Combine into one function
function blockReduceA(amps::Array{Float64,1},indx::Array{CartesianIndex{12},1},jd1::Float64,jd2::Float64,je1::Float64)
    #posa = findall(x-> x!=0, prA )
    #PrANZ = PrA[pos]
    a = []
    b = []
    for i in 1:length(indx)# position of d1 = 3, d2= 8 and e1 = 4
        if indx[i][3] == Int(2*jd1+1) && indx[i][8] == Int(2*jd2+1) && indx[i][4] == Int(2*je1+1)
            push!(a,amps[i])
            push!(b,indx[i])
        end
    end
    # a and b are index and amplitudes
    c = []
    for ja1 in 0.:0.5:y, jb1 in 0.:0.5:y, jg in 0.:0.5:y
        d = []
        for i in 1:length(b) # position of a1 = 1, b1 = 2 g = 12
            if b[i][1] == Int(2*ja1+1) && b[i][2] == Int(2*jb1+1) && b[i][12] == Int(2*jg+1)
                push!(d,a[i])
            end
        end
        if d != []
            push!(c,d)
        end
    end
    if length(c) !=0
        dd = zeros(length(c),length(c[1]))
        for i in 1:length(c)
            c[i] = convert(Array{Float64},c[i])
            dd[i,:] = c[i]
        end
        return dd
    else
        return 0
    end
end

function svdR(prA::Array{Float64,12},ja::Float64,jb::Float64,jc::Float64)# parameters d1,d2,e1
    mat = smat(prA,ja,jb,jc)
    U, s, V = svd(mat)
    #V = ctranspose(V)
    return (U,s,V)
end

function svdRA(amps::Array{Float64,1},indx::Array{CartesianIndex{12},1},ja::Float64,jb::Float64,jc::Float64)# parameters d1,d2,e1
    mat = blockReduceA(amps,indx,ja,jb,jc)
    U, s, V = svd(mat)
    #V = ctranspose(V)
    return (U,s,V)
end

@time for ja in 0.:0.5:y, jb in 0.:0.5:y, jc in 0.:0.5:y
    sing = svdRA(ampsA,indxA,ja,jb,jc)[2]
    if sing[1] != 0
        println(sing)
    end
end

#@time println(svdA(PrA,0.0,0.5,0.5)[2][1:6])
@time for ja in 0.:0.5:y, jb in 0.:0.5:y, jc in 0.:0.5:y
    sing = svdA(PrA,ja,jb,jc)[2]
    if sing[1] != 0
        println(sing[1:3])
    end
end

#@time fullUtensorA(PrA)
#@time fullUtensorB(PrB)
#@time fullVtensorA(PrA)
#@time prsmUAVB(PrA,PrB)
#@time println(svdR(PrA,0.,0.5,0.5)[2])
#@time println(svdRA(PrA,0.,0.5,0.5)[2])

#@time for ja in 0.:0.5:y, jb in 0.:0.5:y, jc in 0.:0.5:y
#    sing = svdRA(ampsB,indxB,ja,jb,jc)[2]
#    if sing[1] != 0
#        println(sing)
#    end
#end

#@time blockReduceA(ampsB,indxB,0.,0.5,1.)
#@time println(loppoF(PrA,0.,0.5,0.5))

CoarseGraining3DPrisms.greet()


text:
using Documenter, Highlights
using InteractiveUtils

# Generate demo pages.

let dir = joinpath(dirname(@__FILE__), "..", "test", "samples")
    destination = joinpath(dirname(@__FILE__), "src", "demo", "lexers")
    ispath(destination) && rm(destination, recursive = true)
    mkpath(destination)
    for file in readdir(dir)
        source = read(joinpath(dir, file), String)
        lexer = Highlights.lexer(file)
        def = Highlights.Compiler.metadata(lexer)
        open(joinpath(destination, file * ".md"), "w") do buf
            println(buf, "# ", def.name, "\n")
            println(buf, "`$lexer` -- *", def.description, "*", "\n")
            println(buf, "```@raw html")
            stylesheet(buf, MIME("text/html"))
            highlight(buf, MIME("text/html"), source, lexer)
            println(buf, "```")
        end
    end
    source = read(joinpath(dir, "julia"), String)
    destination = joinpath(dirname(@__FILE__), "src", "demo", "themes")
    ispath(destination) && rm(destination, recursive = true)
    mkpath(destination)
    for theme in subtypes(Highlights.AbstractTheme)
        def = Highlights.Themes.metadata(theme)
        open(joinpath(destination, replace(lowercase(def[:name]), " " => "-") * ".md"), "w") do buf
            println(buf, "# ", def[:name], "\n")
            println(buf, "`$theme` -- *", def[:description], "*", "\n")
            println(buf, "```@raw html")
            stylesheet(buf, MIME("text/html"), theme)
            highlight(buf, MIME("text/html"), source, Lexers.JuliaLexer, theme)
            println(buf, "```")
        end
    end
end

const LEXERS = readdir(joinpath(dirname(@__FILE__), "src", "demo", "lexers"))
const THEMES = readdir(joinpath(dirname(@__FILE__), "src", "demo", "themes"))

# Build docs.

makedocs(
    format = Documenter.HTML(),
    modules = [Highlights],
    sitename = "Highlights.jl",
    authors = "Michael Hatherly",
    pages = [
        "Home" => "index.md",
        "Manual" => [
            "man/guide.md",
            "man/theme.md",
            "man/lexer.md",
            "man/formatting.md",
        ],
        "Demos" => [
            hide("demo/lexers.md", ["demo/lexers/$f" for f in LEXERS])
            hide("demo/themes.md", ["demo/themes/$f" for f in THEMES])
        ],
        "Library" => [
            "lib/public.md",
            "lib/internals.md",
        ],
    ],
)

deploydocs(
    repo = "github.com/JuliaDocs/Highlights.jl.git",
    target = "build",
    deps = nothing,
    make = nothing,
)


text:
### A Pluto.jl notebook ###
# v0.19.5

using Markdown
using InteractiveUtils

# ╔═╡ ed6834a3-a7fc-44fe-9bec-f803135b2ee4
using AbstractPlutoDingetjes, HypertextLiteral, Parameters, PlutoDevMacros, Deno_jll

# ╔═╡ 86a8ff20-f99b-41b5-8598-ccbe24735474
@only_in_nb using PlutoUI

# ╔═╡ ca6c2091-0e06-4f75-b762-27d3d65580af
@only_in_nb PlutoUI.TableOfContents()

# ╔═╡ 6ae7ce1d-fd4c-4884-a49d-b2318bfa9881
md"# Ingredients"

# ╔═╡ ca2c5d3d-d77f-462f-949b-ef85c6db5910
@plutoinclude "./context.jl" "all"

# ╔═╡ 07535f14-08ac-4779-8c08-2bf295c28f9f
md"# Axis"

# ╔═╡ d4fdafe7-97c8-4e5d-969f-cfa20a831955
md"## Structure"

# ╔═╡ ee072da1-9c81-4310-a98a-638e782bb3a0
@enum Direction Right Left Top Bottom

# ╔═╡ 0c147c8a-df2f-4e0f-9760-83c7bab2e7e9
begin
  struct Axis <: D3Component
	direction    :: Direction
	attributes   :: D3Attr
	id           :: String
  end

  Axis(direction::Direction, id::String;
	attributes=D3Attr()
  ) = Axis(
	  direction,
	  attributes,
	  id
  )
end;

# ╔═╡ 7dc05859-dc1f-428d-83b4-d6275db2e247
function bounds(axis::Axis)
	xs = (v -> v.x).(data)
	ys = reduce(vcat, (v -> [v.y0, v.y1]).(data))
	return (minx=0, miny=0, maxx=1, maxy=1)
end

# ╔═╡ c13a4acc-dd77-4c91-b408-c68cea240684
md"## Javascript Snippet"

# ╔═╡ 88d1881f-1865-4c46-ac50-5e26eacb46d6
Base.show(io::IO, m::MIME"text/javascript", axis::Axis) =
    show(io, m, @js """
	axis(
		ctx,
		$(PublishToJS(string(axis.direction))),
		x_scale,
		y_scale,
		$(PublishToJS(to_named_tuple(axis.attributes))),
		$(axis.id)
	);
	""")

# ╔═╡ 25e92ec2-eaf2-46ba-8ec7-83fe19d2b9f1
Base.show(io::IO, m::MIME"text/html", axis::Axis) =	show(io, m, @htl("""
<b style="color: red">You should try rendering axis with other things</b>
"""))

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
AbstractPlutoDingetjes = "6e696c72-6542-2067-7265-42206c756150"
Deno_jll = "04572ae6-984a-583e-9378-9577a1c2574d"
HypertextLiteral = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
PlutoDevMacros = "a0499f29-c39b-4c5c-807c-88074221b949"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
AbstractPlutoDingetjes = "~1.1.4"
Deno_jll = "~1.20.4"
HypertextLiteral = "~0.9.4"
Parameters = "~0.12.3"
PlutoDevMacros = "~0.4.5"
PlutoUI = "~0.7.39"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.3"
manifest_format = "2.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "0f4e115f6f34bbe43c19751c90a38b2f380637b9"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.3"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.Deno_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "970da1e64a94f13b51c81691c376a1d5a83a0b3c"
uuid = "04572ae6-984a-583e-9378-9577a1c2574d"
version = "1.20.4+0"

[[deps.Downloads]]
deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.FileWatching]]
uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
deps = ["Tricks"]
git-tree-sha1 = "c47c5fa4c5308f27ccaac35504858d8914e102f9"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.4"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "abc9885a7ca2052a736a600f7fa66209f96506e1"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.4.1"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "3c837543ddb02250ef42f4738347454f95079d4e"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.3"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "3d3e902b31198a27340d0bf00d6ac452866021cf"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.9"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.Parameters]]
deps = ["OrderedCollections", "UnPack"]
git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
version = "0.12.3"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "1285416549ccfcdf0c50d4997a94331e88d68413"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.3.1"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlutoDevMacros]]
deps = ["MacroTools", "Requires"]
git-tree-sha1 = "994167def8f46d3be21783a76705228430e29632"
uuid = "a0499f29-c39b-4c5c-807c-88074221b949"
version = "0.4.5"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "ColorTypes", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "8d1f54886b9037091edf146b517989fc4a09efec"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.39"

[[deps.Preferences]]
deps = ["TOML"]
git-tree-sha1 = "47e5f437cc0e7ef2ce8406ce1e7e24d44915f88d"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.3.0"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.3.0"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.Tricks]]
git-tree-sha1 = "6bac775f2d42a611cdfcd1fb217ee719630c4175"
uuid = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
version = "0.1.6"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.UnPack]]
git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
version = "1.0.2"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─ca6c2091-0e06-4f75-b762-27d3d65580af
# ╟─86a8ff20-f99b-41b5-8598-ccbe24735474
# ╠═6ae7ce1d-fd4c-4884-a49d-b2318bfa9881
# ╠═ed6834a3-a7fc-44fe-9bec-f803135b2ee4
# ╠═ca2c5d3d-d77f-462f-949b-ef85c6db5910
# ╟─07535f14-08ac-4779-8c08-2bf295c28f9f
# ╟─d4fdafe7-97c8-4e5d-969f-cfa20a831955
# ╠═ee072da1-9c81-4310-a98a-638e782bb3a0
# ╠═0c147c8a-df2f-4e0f-9760-83c7bab2e7e9
# ╠═7dc05859-dc1f-428d-83b4-d6275db2e247
# ╟─c13a4acc-dd77-4c91-b408-c68cea240684
# ╠═88d1881f-1865-4c46-ac50-5e26eacb46d6
# ╠═25e92ec2-eaf2-46ba-8ec7-83fe19d2b9f1
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
# Types for Julia implementation of pseudolikelihood estimation (Stage2)
# Claudia August 2014
# in julia: include("types.jl")
#
# Types: Edge, Node, HybridNetwork
################################################################

# procedure to create hybrid network:
# 1) create edges defined as hybrid or not
# 2) create nodes with such edges
# 3) setNode! to add nodes into edges
# 4) create hybrid network
# 5) updateGammaz!

# -------------- EDGE -------------------------#
"""
 `ANode`

Abstract node. An object of type `Edge` has a `node` attribute,
which is an vector of (2) ANode objects.
The object of type `Node` is an `ANode`, and has an `edge` attribute,
which is vector of `Edge` objects.
"""
abstract type ANode end

# warning: inCycle and containRoot are updated after edge is part of a network
mutable struct Edge
    number::Int
    length::Float64 #default 1.0
    hybrid::Bool
    y::Float64 # exp(-t), cannot set in constructor for congruence
    z::Float64 # 1-y , cannot set in constructor for congruence
    gamma::Float64 # set to 1.0 for tree edges, hybrid?gamma:1.0
    node::Array{ANode,1} # we can also leave blank: node (see issues.jl)
    isChild1::Bool # used for hybrid edges to set the direction (default true)
    isMajor::Bool  # major edge treated as tree edge for network traversal
                   # true if gamma>.5, or if it is the original tree edge
    inCycle::Int # = Hybrid node number if this edge is part of a cycle created by such hybrid node
    # -1 if not part of cycle. used to add new hybrid edge. updated after edge is part of a network
    containRoot::Bool # true if this edge can contain a root given the direction of hybrid edges
                      # used to add new hybrid edge. updated after edge is part of a network
    istIdentifiable::Bool # true if the parameter t (length) for this edge is identifiable as part of a network
                          # updated after part of a network
    fromBadDiamondI::Bool # true if the edge came from deleting a bad diamond I hybridization
    # inner constructors: ensure congruence among (length, y, z) and (gamma, hybrid, isMajor), and size(node)=2
    Edge(number::Int) = new(number,1.0,false,exp(-1.0),1-exp(-1.0),1.,[],true,true,-1,true,true,false)
    Edge(number::Int, length::Float64) = new(number,length,false,exp(-length),1-exp(-length),1.,[],true,true,-1,true,true,false)
    Edge(number::Int, length::Float64,hybrid::Bool,gamma::Float64)= new(number,length,hybrid,exp(-length),1-exp(-length),hybrid?gamma:1.,[],true,(!hybrid || gamma>0.5)?true:false,-1,!hybrid,true,false)
    Edge(number::Int, length::Float64,hybrid::Bool,gamma::Float64,isMajor::Bool)= new(number,length,hybrid,exp(-length),1-exp(-length),hybrid?gamma:1.,[],true,isMajor,-1,!hybrid,true,false)
    function Edge(number::Int, length::Float64,hybrid::Bool,gamma::Float64,node::Array{ANode,1})
        size(node,1) != 2 ?
        error("vector of nodes must have exactly 2 values") :
        new(number,length,hybrid,exp(-length),1-exp(-length),hybrid?gamma:1.,node,true,(!hybrid || gamma>0.5)?true:false,-1,!hybrid,true,false)
    end
    function Edge(number::Int, length::Float64,hybrid::Bool,gamma::Float64,node::Array{ANode,1},isChild1::Bool, inCycle::Int, containRoot::Bool, istIdentifiable::Bool)
        size(node,1) != 2 ?
        error("vector of nodes must have exactly 2 values") :
        new(number,length,hybrid,exp(-length),1-exp(-length),hybrid?gamma:1.,node,isChild1,(!hybrid || gamma>0.5)?true:false,inCycle,containRoot,istIdentifiable,false)
    end
end

# warning: gammaz, inCycle, isBadTriangle/Diamond updated until the node is part of a network
mutable struct Node <: ANode
    number::Int
    leaf::Bool
    hybrid::Bool
    gammaz::Float64  # notes file for explanation. gammaz if tree node, gamma2z if hybrid node.
                     # updated after node is part of network with updateGammaz!
    edge::Array{Edge,1}
    hasHybEdge::Bool #is there a hybrid edge in edge? only needed when hybrid=false (tree node)
    isBadDiamondI::Bool # for hybrid node, is it bad diamond case I, update in updateGammaz!
    isBadDiamondII::Bool # for hybrid node, is it bad diamond case II, update in updateGammaz!
    isExtBadTriangle::Bool # for hybrid node, is it extremely bad triangle, udpate in updateGammaz!
    isVeryBadTriangle::Bool # for hybrid node, is it very bad triangle, udpate in updateGammaz!
    isBadTriangle::Bool # for hybrid node, is it very bad triangle, udpate in updateGammaz!
    inCycle::Int # = hybrid node if this node is part of a cycle created by such hybrid node, -1 if not part of cycle
    prev # previous node in cycle, used in updateInCycle. defined as "Any", set as "nothing" to begin with
    k::Int # num nodes in cycle, only stored in hybrid node, updated after node becomes part of network
           # default -1
    typeHyb::Int8 # type of hybridization (1,2,3,4, or 5), needed for quartet network only. default -1
    name::AbstractString
    # inner constructor: set hasHybEdge depending on edge
    Node() = new(-1.,false,false,-1.,[],false,false,false,false,false,false,-1.,nothing,-1,-1,"")
    Node(number::Int, leaf::Bool) = new(number,leaf,false,-1.,[],false,false,false,false,false,false,-1.,nothing,-1,-1,"")
    Node(number::Int, leaf::Bool, hybrid::Bool) = new(number,leaf,hybrid,-1.,[],hybrid,false,false,false,false,false,-1.,nothing,-1,-1,"")
    Node(number::Int, leaf::Bool, hybrid::Bool, edge::Array{Edge,1})=new(number,leaf,hybrid,-1.,edge,!all((e->!e.hybrid),edge),false,false,false,false,false,-1.,nothing,-1,-1,"")
    Node(number::Int, leaf::Bool, hybrid::Bool,gammaz::Float64, edge::Array{Edge,1}) = new(number,leaf,hybrid,gammaz,edge,!all((e->!e.hybrid), edge),false,false,false,false,false,-1.,nothing,-1,-1,"")
end

# partition type
mutable struct Partition
    cycle::Vector{Int} #hybrid node number for cycle (or cycles)
    edges::Vector{Edge} #edges in partition
end

abstract type Network end

# warning: no attempt to make sure the direction of edges matches with the root
# warning: no check if it is network or tree, node array can have no hybrids
# warning: nodes and edges need to be defined and linked before adding to a network
"""
`HybridNetwork type`
Explicit network or tree with the following attributes:

- numTaxa
- numNodes (total number of nodes)
- numEdges
- numHybrids (number of hybrid nodes)
- edge (array of Edges)
- node (array of Nodes)
- root (index of root in vector 'node'. May be artificial, for printing and traversal purposes only.)
- hybrid (array of Nodes: those are are hybrid nodes)
- leaf (array of Nodes: those that are leaves)
- loglik (negative log pseudolik after estimation)
- isRooted (true or false)
"""
mutable struct HybridNetwork <: Network
    numTaxa::Int  # cannot set in constructor for congruence
    numNodes::Int
    numEdges::Int
    node::Array{Node,1}
    edge::Array{Edge,1}
    root::Int # node[root] is the root node, default 1
    names::Array{String,1} # translate table for taxon names --but also includes hybrid names...
    hybrid::Array{Node,1} # array of hybrid nodes in network
    numHybrids::Int # number of hybrid nodes
    cladewiseorder_nodeIndex::Vector{Int} # index in 'node' for "cladewise" preorder in main tree
    visited::Array{Bool,1} # reusable array of booleans
    edges_changed::Array{Edge,1} # reusable array of edges
    nodes_changed::Array{Node,1} # reusable array of nodes. used for preorder traversal
    leaf::Array{Node,1} # array of leaves
    ht::Vector{Float64} # vector of parameters to optimize
    numht::Vector{Int} # vector of number of the hybrid nodes and edges in ht e.g. [3,6,8,...], 2 hybrid nodes 3,6, and edge 8 is the 1st identifiable
    numBad::Int # number of bad diamond I hybrid nodes, set as 0
    hasVeryBadTriangle::Bool # true if the network has extremely/very bad triangles that should be ignored
    index::Vector{Int} #index in net.edge, net.node of elements in net.ht to make updating easy
    loglik::Float64 # value of the min -loglik after optBL
    blacklist::Vector{Int} # reusable array of integers, used in afterOptBL
    partition::Vector{Partition} # to choose edges from a partition only to avoid intersecting cycles
    cleaned::Bool # attribute to know if the network has been cleaned after readm default false
    isRooted::Bool # to know if network is rooted, e.g. after directEdges! (which updates isChild1 of each edge)
    # inner constructor
    function HybridNetwork(node::Array{Node,1},edge::Array{Edge,1})
        hybrid=Node[];
        leaf=Node[];
        [n.hybrid?push!(hybrid,n):nothing for n in node];
        [n.leaf?push!(leaf,n):nothing for n in node];
        new(size(leaf,1),size(node,1),size(edge,1),node,edge,1,[],hybrid,size(hybrid,1), #numTaxa,...,numHybrids
            [],[],[],[],leaf,[],[], #cladewiseorder,...,numht
            0,false,[],0,[],[],false,false) #numBad...
    end
    function HybridNetwork(node::Array{Node,1},edge::Array{Edge,1},root::Int)
        hybrid=Node[];
        leaf=Node[];
        [n.hybrid?push!(hybrid,n):nothing for n in node];
        [n.leaf?push!(leaf,n):nothing for n in node];
        new(size(leaf,1),size(node,1),size(edge,1),node,edge,root,[],hybrid,size(hybrid,1), #numTaxa,...,numHybrids
            [],[],[],[],leaf,[],[], #cladewiseorder,...,numht
            0,false,[],0,[],[],false,false) #numBad...
    end
    HybridNetwork() = new(0,0,0,[],[],0,[],[],0, # numTaxa ... numHybrid
                          [],[],[],[],[],[],[], # cladewiseorder...
                          0,false,[],0,[],[],false,false); # numBad ...
end

# type created from a HybridNetwork only to extract a given quartet
mutable struct QuartetNetwork <: Network
    numTaxa::Int
    numNodes::Int
    numEdges::Int
    node::Array{Node,1}
    edge::Array{Edge,1}
    hybrid::Array{Node,1} # array of hybrid nodes in network
    leaf::Array{Node,1} # array of leaves
    numHybrids::Int # number of hybrid nodes
    hasEdge::Array{Bool,1} # array of boolean with all the original identifiable edges of HybridNetwork and gammas (net.ht)
    quartetTaxon::Array{String,1} # the quartet taxa in the order it represents. Points to same array as its Quartet.taxon
    which::Int8 # 0 it tree quartet, 1 is equivalent to tree quartet and 2 if two minor CF different, default -1
    typeHyb::Array{Int8,1} #array with the type of hybridization of each hybrid node in the quartet
    t1::Float64 # length of internal edge, used when qnet.which=1, default = -1
    names::Array{String,1} # taxon and node names, same order as in network.node
    split::Array{Int8,1} # split that denotes to which side each leaf is from the split, i.e. [1,2,2,1] means that leaf1 and 4 are on the same side of the split, default -1,-1,-1,-1
    formula::Array{Int8,1} # array for qnet.which=1 that indicates if the expCf is major (2) or minor (1) at qnet.expCF[i] depending on qnet.formula[i], default -1,-1,-1
    expCF::Array{Float64,1} # three expected CF in order 12|34, 13|24, 14|23 (matching obsCF from qnet.quartet), default [0,0,0]
    indexht::Vector{Int} # index in net.ht for each edge in qnet.ht
    changed::Bool # true if the expCF would be changed with the current parameters in the optimization, to recalculate, default true
    index::Vector{Int} # index in qnet.edge (qnet.node for gammaz) of the members in qnet.indexht to know how to find quickly in qnet
    # inner constructor
    function QuartetNetwork(net::HybridNetwork)
        net2 = deepcopy(net); #fixit: maybe we dont need deepcopy of all, maybe only arrays
        new(net2.numTaxa,net2.numNodes,net2.numEdges,net2.node,net2.edge,net2.hybrid,net2.leaf,net2.numHybrids, [true for e in net2.edge],[],-1,[], -1.,net2.names,Int8[-1,-1,-1,-1],Int8[-1,-1,-1],[0,0,0],[],true,[])
        #new(sum([n.leaf?1:0 for n in net.node]),size(net.node,1),size(net.edge,1),copy(net.node),copy(net.edge),copy(net.hybrid),size(net.hybrid,1), [true for e in net2.edge],[],-1,[],-1.,net2.names,[-1,-1,-1,-1],[-1,-1,-1],[],true,[])
    end
    function QuartetNetwork(net::HybridNetwork,quartet::Array{String,1})
        net2 = deepcopy(net);
        new(net2.numTaxa,net2.numNodes,net2.numEdges,net2.node,net2.edge,net2.hybrid,net2.leaf,net2.numHybrids, [true for e in net2.edge],quartet,-1,[],-1.,net2.names,Int8[-1,-1,-1,-1],Int8[-1,-1,-1],[0,0,0],[],true,[])
    end
    QuartetNetwork() = new(0,0,0,[],[],[],[],0,[],[],-1,[],-1.0,[],[],[],[],[],true,[])
end

"""
`Quartet type`

type that saves the information on a given 4-taxon subset. It contains the following attributes:

- number
- taxon (vector of taxon names)
- obsCF (vector of observed CF)
- logPseudoLik
- ngenes (number of gene trees used to compute the observed CF: -1 if unknown)
- qnet (internal topological structure that saves the expCF after snaq estimation to emphasize that the expCF depend on a specific network, not the data)
"""
mutable struct Quartet
    number::Int
    taxon::Array{String,1} # taxa 1234. qnet.quartetTaxon points to the same array.
    obsCF::Array{Float64,1} # three observed CF in order 12|34, 13|24, 14|23
    qnet::QuartetNetwork # quartet network for the current network (want to keep as if private attribute)
    logPseudoLik::Float64 # log pseudolik value for the quartet
    ngenes::Number # number of gene trees used to compute the obsCV, default -1; changed to Number in case ngenes is average, so float
    # inner constructor: to guarantee obsCF are only three and add up to 1
    function Quartet(number::Integer,t1::AbstractString,t2::AbstractString,t3::AbstractString,t4::AbstractString,obsCF::Array{Float64,1})
        size(obsCF,1) != 3 ? error("observed CF vector should have size 3, not $(size(obsCF,1))") : nothing
        0.99 < sum(obsCF) < 1.02 || warn("observed CF should add up to 1, not $(sum(obsCF))")
        new(number,[t1,t2,t3,t4],obsCF,QuartetNetwork(),0,-1);
    end
    function Quartet(number::Integer,t1::Array{String,1},obsCF::Array{Float64,1})
        size(obsCF,1) != 3 ? error("observed CF vector should have size 3, not $(size(obsCF,1))") : nothing
        0.99< sum(obsCF) < 1.02 || warn("observed CF should add up to 1, not $(sum(obsCF))")
        size(t1,1) != 4 ? error("array of taxa should have size 4, not $(size(t1,1))") : nothing
        0.0 <= obsCF[1] <= 1.0 || error("obsCF must be between (0,1), but it is $(obsCF[1]) for $(t1)")
        0.0 <= obsCF[2] <= 1.0 || error("obsCF must be between (0,1), but it is $(obsCF[2]) for $(t1)")
        0.0 <= obsCF[3] <= 1.0 || error("obsCF must be between (0,1), but it is $(obsCF[3]) for $(t1)")
        new(number,t1,obsCF,QuartetNetwork(),0,-1);
    end
    Quartet() = new(0,[],[],QuartetNetwork(),0,-1)
end

# Data -------

"""
`DataCF type`

type that contains the following attributes:

- quartet (vector of Quartets)
- numQuartets
- tree (vector of trees: empty if a table of CF was input instead of list of trees)
- numTrees (-1 if a table CF was input instead of list of trees)
- repSpecies (taxon names that were repeated in table of CF or input gene trees: used inside snaq for multiple alleles case)

The list of Quartet may be accessed with the attribute .quartet.
If the input was a list of trees, the HybridNetwork's can be accessed with the attribute .tree.
For example, if the DataCF object is named d, d.quartet[1] will show the first quartet
and d.tree[1] will print the first input tree.
"""
mutable struct DataCF # fixit
    quartet::Array{Quartet,1} # array of quartets read from CF output table or list of quartets in file
    numQuartets::Integer # number of quartets
    tree::Vector{HybridNetwork} #array of input gene trees
    numTrees::Integer # number of gene trees
    repSpecies::Vector{String} #repeated species in the case of multiple alleles
    DataCF(quartet::Array{Quartet,1}) = new(quartet,length(quartet),[],-1,[])
    DataCF(quartet::Array{Quartet,1},trees::Vector{HybridNetwork}) = new(quartet,length(quartet),trees,length(trees),[])
    DataCF() = new([],0,[],-1,[])
end

# aux type for the updateBL function
mutable struct EdgeParts
    edgenum::Int
    part1::Vector{Node}
    part2::Vector{Node}
    part3::Vector{Node}
    part4::Vector{Node}
end

struct RootMismatch <: Exception
    msg::String
end
RootMismatch() = RootMismatch("")
Base.showerror(io::IO, e::RootMismatch) = print(io, "RootMismatch: ", e.msg);


text:
module RocketSkipNextOperatorTest

using Test
using Rocket

include("../test_helpers.jl")

@testset "operator: skip_next()" begin

    println("Testing: operator skip_next()")

    run_testset([
        (
            source = from(1:5) |> skip_next(),
            values = @ts(c)
        ),
        (
            source = completed() |> skip_next(),
            values = @ts(c)
        ),
        (
            source = faulted(1) |> skip_next(),
            values = @ts(e(1))
        ),
        (
            source = never() |> skip_next(),
            values = @ts()
        )
    ])

end

end


text:
using ParSpMatVec
using Test
using SparseArrays
using Printf
using LinearAlgebra

n = 50000
numProcs =4;
nvec = 5
A = sprand(n,n, 2.e-6);

x = rand(n,nvec);  x = x*10 .- 5;
y = rand(n,nvec);  y = y*10 .- 5;

println("Real")

alpha = 123.56
beta = 543.21
y2 = copy(y)
println("y = beta*y + alpha * A'*x")
BaseTime = @elapsed y2 = beta*y2 + alpha * A'*x;

for k=0:numProcs
	y3 = copy(y)
	println("ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3,",k,")")
	PSMVtime = @elapsed ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, k );
	@printf "Base=%1.4f\t ParSpMatVec=%1.4f\t speedup=%1.4f\n" BaseTime PSMVtime BaseTime/PSMVtime
	@test norm(y3-y2) / norm(y) < 1.e-12
end
println()

# test error handling for nthreads
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, -1);
catch E
	@test isa(E,ArgumentError)
end

# test error handling for sizes
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x[1:10,:], beta, y3, 1);
catch E
	@test isa(E,DimensionMismatch)
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3[:,2], 1);
catch E
	@test isa(E,DimensionMismatch)
end

println("Complex Scalars, Real matrix")

alpha = 123.56 .+ 1im*randn()
beta = 543.21 .+ 1im*randn()
y = rand(n,nvec) + 1im* rand(n,nvec)
x  = rand(n,nvec) + 1im* rand(n,nvec)
println("y = beta*y + alpha * A'*x")

BaseTime = @elapsed y2 = beta*y + alpha * A'*x;

for k=0:numProcs
	y3 = copy(y)
	println("ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3,",k,")")
	PSMVtime = @elapsed ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, k );
	@printf "Base=%1.4f\t ParSpMatVec=%1.4f\t speedup=%1.4f\n" BaseTime PSMVtime BaseTime/PSMVtime
	@test norm(y3-y2) / norm(y) < 1.e-12
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, -1);
catch E
	@test isa(E,ArgumentError)
end
# test error handling for sizes
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x[1:10,:], beta, y3, 1);
catch E
	@test isa(E,DimensionMismatch)
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3[:,2], 1);
catch E
	@test isa(E,DimensionMismatch)
end
println()
#-----------------------------------
# Complex

y2=0;
ii,jj,vv = findnz(A)

ai = ones(length(ii));
vv = vv + im*ai
A = sparse(ii,jj, vv, n,n)
ii=0; jj=0; vv=0; ai=0;

xi = rand(n,nvec);  xi = xi*10 .- 5;
x = x + im*xi
xi=0
yi = rand(n,nvec);  yi = yi*10 .- 5;
y = y + im*yi
yi=0

println("Complex")

alpha = complex(123.56, 333.444)
beta  = complex(543.21, 111.222)

y2 = copy(y)
println("y = beta*y + alpha * A'*x")

BaseTime = @elapsed y2 = beta*y2 + alpha * A'*x;

for k=0:numProcs
	y3 = copy(y)
	println("ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3,",k,")")
	PSMVtime = @elapsed ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, k );
	@printf "Base=%1.4f\t ParSpMatVec=%1.4f\t speedup=%1.4f\n" BaseTime PSMVtime BaseTime/PSMVtime
	@test norm(y3-y2) / norm(y) < 1.e-12
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, -1);
catch E
	@test isa(E,ArgumentError)
end
# test error handling for sizes
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x[1:10,:], beta, y3, 1);
catch E
	@test isa(E,DimensionMismatch)
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3[:,2], 1);
catch E
	@test isa(E,DimensionMismatch)
end
println()

println("Complex short")
alpha = convert(ComplexF32, alpha)
beta  = convert(ComplexF32,beta);
A = convert(SparseMatrixCSC{ComplexF32,Int64},A);
x = convert(Array{ComplexF32},x);
y = convert(Array{ComplexF32},y);

y2 = copy(y)
println("y = beta*y + alpha * A'*x")

BaseTime = @elapsed y2 = beta*y2 + alpha * A'*x;

for k=0:numProcs
	y3 = copy(y)
	println("ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3,",k,")")
	PSMVtime = @elapsed ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, k );
	@printf "Base=%1.4f\t ParSpMatVec=%1.4f\t speedup=%1.4f\n" BaseTime PSMVtime BaseTime/PSMVtime
	@test norm(y3-y2) / norm(y) < 1.e-4
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, -1);
catch E
	@test isa(E,ArgumentError)
end
# test error handling for sizes
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x[1:10,:], beta, y3, 1);
catch E
	@test isa(E,DimensionMismatch)
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3[:,2], 1);
catch E
	@test isa(E,DimensionMismatch)
end
println()

println("Complex short with a real matrix")
alpha = convert(ComplexF32, alpha)
beta  = convert(ComplexF32,beta);
A = convert(SparseMatrixCSC{Float32,Int64},real(A));
x = convert(Array{ComplexF32},x);
y = convert(Array{ComplexF32},y);

y2 = copy(y)
println("y = beta*y + alpha * A'*x")

BaseTime = @elapsed y2 = beta*y2 + alpha * A'*x;

for k=0:numProcs
	y3 = copy(y)
	println("ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3,",k,")")
	PSMVtime = @elapsed ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, k );
	@printf "Base=%1.4f\t ParSpMatVec=%1.4f\t speedup=%1.4f\n" BaseTime PSMVtime BaseTime/PSMVtime
	@test norm(y3-y2) / norm(y) < 1.e-5
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, -1);
catch E
	@test isa(E,ArgumentError)
end
# test error handling for sizes
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x[1:10,:], beta, y3, 1);
catch E
	@test isa(E,DimensionMismatch)
end
try
	y3 = copy(y)
	ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3[:,2], 1);
catch E
	@test isa(E,DimensionMismatch)
end

println()
println("Complex single with a complex matrix but double target and source")
alpha = convert(ComplexF64, alpha)
beta  = convert(ComplexF64,beta);
A = convert(SparseMatrixCSC{ComplexF32,Int64},real(A) + 1im*A);
x = convert(Array{ComplexF64},x);
y = convert(Array{ComplexF64},y);

y2 = copy(y)
println("y = beta*y + alpha * A'*x")

BaseTime = @elapsed y2 = beta*y2 + alpha * A'*x;

for k=0:numProcs
	y3 = copy(y)
	println("ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3,",k,")")
	PSMVtime = @elapsed ParSpMatVec.Ac_mul_B!( alpha, A, x, beta, y3, k );
	@printf "Base=%1.4f\t ParSpMatVec=%1.4f\t speedup=%1.4f\n" BaseTime PSMVtime BaseTime/PSMVtime
	@test norm(y3-y2) / norm(y) < 1.e-5
	 norm(y3-y2) / norm(y)
end
println()


text:
# Tests for Gradient Bandits

@testset "Gradient Bandits" begin

noOfArms = 5

@testset "Gradient Bandit" begin
    srand(1729)
    agent = GradientBandit( noOfArms, 0.1 )

    @test agent.noOfSteps == 0
    @test agent.lastPlayedArm == 0
    @test agent.avgValue == 0
    @test agent.pDist.p == 1/noOfArms * ones(noOfArms)

    @test info_str( agent ) == "GradientBandit (α=0.100)"
    @test info_str( agent, false ) == "GradientBandit (α=0.100)"
    @test info_str( agent, true ) == "GradientBandit (\$\\alpha=0.100\$)"

    for t = 1:10
        get_arm_index( agent ); update_reward!( agent, rand() )
    end
    @test agent.noOfSteps == 10

    reset!( agent )
    @test agent.noOfSteps == 0 &&
            agent.lastPlayedArm == 0 &&
            agent.avgValue == 0 &&
            agent.pDist.p == 1/noOfArms * ones(noOfArms) &&
            agent.α == 0.100 &&
            agent.preference == zeros(noOfArms)
end

end