text:
@doc """
    Lattice

# Fields
- `dim :: Int`
    - dimension of lattice
- `size :: Vector{Int}`
    - length of lattice in each dimension
- `nsitetypes :: Int`
    - the number of `sitetype`s
- `nbondtypes :: Int`
    - the number of `bondtype`s
- `sites :: Vector{Vector{Int}}`
    - list of site indecies for each `sitetype`
- `bonds :: Vector{Vector{Int}}`
    - list of bond indecies for each `sitetype`
- `nsites :: Int`
    - the number of sites
- `nbonds :: Int`
    - the number of bonds
- `sitetypes :: Vector{Int}`
    - `sitetype` of each site
- `bondtypes :: Vector{Int}`
    - `bondtype` of each bond
- `transvector :: Matrix{Float64}`
    - lattice vector represented in Cartesian system
    - `@assert size(transvector) == (dim, dim)`
- `site_coords :: Matrix{Float64}`
    - coordinate of each site represented in lattice system
    - `@assert size(site_coords) == (dim, nsites)`
- `bond_dirs :: Matrix{Float64}`
    - displacement of each bond represented in lattice system
    - `@assert size(bond_dirs) == (dim, nbonds)`
- `neighborsites :: Vector{Vector{Int}}`
    - list of indecies of neighbor sites of each site
- `neighborbonds :: Vector{Vector{Int}}`
    - list of indecies of neighbor bonds of each site
- `source :: Vector{Int}`
    - the index of an end site of each bond index
- `target :: Vector{Int}`
    - the index of the other end site of each bond index
"""
mutable struct Lattice
    dim :: Int
    size :: Vector{Int}
    nsitetypes :: Int
    nbondtypes :: Int
    sites :: Vector{Vector{Int}}
    bonds :: Vector{Vector{Int}}
    nsites :: Int
    nbonds :: Int
    sitetypes :: Vector{Int}
    bondtypes :: Vector{Int}
    transvector :: Matrix{Float64}
    site_coords :: Matrix{Float64}
    bond_dirs :: Matrix{Float64}
    neighborsites :: Vector{Vector{Int}}
    neighborbonds :: Vector{Vector{Int}}
    source :: Vector{Int}
    target :: Vector{Int}
end

import Base.size
import Distributions.dim

@doc """
    dim(lat::Lattice)
    dim(model::Model)

Returns the dimension of lattice.
"""
dim(lat::Lattice) = lat.dim
dim(model::Model) = dim(model.lat)

@doc """
    size(lat::Lattice, [dim::Integer])
    size(model::Model, [dim::Integer])

Returns the size of lattice.
"""
size(lat::Lattice) = lat.size
size(lat::Lattice, dim::Integer) = lat.size[dim]
size(model::Model) = size(model.lat)
size(model::Model, dim::Integer) = size(model.lat,dim)

@doc """
    sites(lat::Lattice, sitetype::Integer)
    sites(model::Model, sitetype::Integer)

Returns sites with `sitetype`
"""
sites(lat::Lattice, sitetype::Integer) = lat.sites[sitetype]
sites(model::Model, sitetype::Integer) = sites(model.lat, sitetype)

@doc """
    numsites(lat::Lattice)
    numsites(model::Model)

Returns the number of all sites.
"""
numsites(lat::Lattice) = lat.nsites
numsites(model::Model) = numsites(model.lat)

@doc """
    numsites(lat::Lattice, sitetype::Integer)
    numsites(model::Model, sitetype::Integer)

Returns the number of `sitetype` sites.
"""
numsites(lat::Lattice, sitetype::Integer) = length(lat.sites[sitetype])
numsites(model::Model, sitetype::Integer) = numsites(model.lat, sitetype)

@doc """
    bonds(lat::Lattice, bondtype::Integer)
    bonds(model::Model, bondtype::Integer)

Returns bonds with `bondtype`
"""
bonds(lat::Lattice, bondtype::Integer) = lat.bonds[bondtype]
bonds(model::Model, bondtype::Integer) = bonds(model.lat, bondtype)

@doc """
    numbonds(lat::Lattice)
    numbonds(model::Model)

Returns the number of all bonds.
"""
numbonds(lat::Lattice) = lat.nbonds
numbonds(model::Model) = numbonds(model.lat)

@doc """
    numbonds(lat::Lattice, bondtype::Integer)
    numbonds(model::Model, bondtype::Integer)

Returns the number of `bondtype` bonds.
"""
numbonds(lat::Lattice, bondtype::Integer) = length(lat.bonds[bondtype])
numbonds(model::Model, bondtype::Integer) = numbonds(model.lat, bondtype)

@doc """
    numsitetypes(lat::Lattice)
    numsitetypes(model::Model)

Returns the number of sitetypes.
"""
numsitetypes(lat::Lattice) = lat.nsitetypes
numsitetypes(model::Model) = numsitetypes(model.lat)

@doc """
    numbondtypes(lat::Lattice)
    numbondtypes(model::Model)

Returns the number of bondtypes.
"""
numbondtypes(lat::Lattice) = lat.nbondtypes
numbondtypes(model::Model) = numbondtypes(model.lat)

@doc """
    neighborsites(lat::Lattice, site::Integer)
    neighborsites(model::Model, site::Integer)

Returns the neighbor sites of `site`.
"""
neighborsites(lat::Lattice, site::Integer) = lat.neighborsites[site]
neighborsites(model::Model, site::Integer) = neighborsites(model.lat, site)

@doc """
    neighborbonds(lat::Lattice, site::Integer)
    neighborbonds(model::Model, site::Integer)

Returns the neighbor bonds of `site`.
"""
neighborbonds(lat::Lattice, site::Integer) = lat.neighborbonds[site]
neighborbonds(model::Model, site::Integer) = neighborbonds(model.lat, site)

@doc """
    neighbors(lat::Lattice, site::Integer)
    neighbors(model::Model, site::Integer)

Returns the neighbor sites and bonds of `site`.
"""
neighbors(lat::Lattice, site::Integer) = zip(neighborsites(lat,site), neighborbonds(lat,site))
neighbors(model::Model, site::Integer) = neighbors(model.lat, site)

@doc """
    source(lat::Lattice, bond::Integer)
    source(model::Model, bond::Integer)

Returns the source site of `bond`.
"""
source(lat::Lattice, bond::Integer) = lat.source[bond]
source(model::Model, bond::Integer) = source(model.lat, bond)

@doc """
    target(lat::Lattice, bond::Integer)
    target(model::Model, bond::Integer)

Returns the target site of `bond`.
"""
target(lat::Lattice, bond::Integer) = lat.target[bond]
target(model::Model, bond::Integer) = target(model.lat, bond)

@doc """
    sitecoordinate(lat::Lattice, site::Integer)
    sitecoordinate(model::Model, site::Integer)

Returns the coordinate of the `site` in the Cartesian system
"""
sitecoordinate(lat::Lattice, site::Integer) = lat.transvector * lat.site_coords[:,site]
sitecoordinate(model::Model, site::Integer) = sitecoordinate(model.lat, site)

@doc """
    lattice_sitecoordinate(lat::Lattice, site::Integer)
    lattice_sitecoordinate(model::Model, site::Integer)

Returns the coordinate of the `site` in the lattice system
"""
lattice_sitecoordinate(lat::Lattice, site::Integer) = lat.site_coords[:, site]
lattice_sitecoordinate(model::Model, site::Integer) = lattice_sitecoordinate(model.lat, site)

@doc """
    sitetype(lat::Lattice, site::Integer)
    sitetype(model::Model, site::Integer)

Returns the type of `site`
"""
sitetype(lat::Lattice, site::Integer) = lat.sitetypes[site]
sitetype(model::Model, site::Integer) = sitetype(model.lat, site)

@doc """
    bondtype(lat::Lattice, bond::Integer)
    bondtype(model::Model, bond::Integer)

Returns the type of `bond`
"""
bondtype(lat::Lattice, bond::Integer) = lat.bondtypes[bond]
bondtype(model::Model, bond::Integer) = bondtype(model.lat, bond)

@doc """
    bonddirection(lat::Lattice, bond::Integer)
    bonddirection(model::Model, bond::Integer)

Returns the direction of the `bond` as vector in the Cartesian system
"""
bonddirection(lat::Lattice, bond::Integer) = lat.transvector * lat.bond_dirs[:, bond]
bonddirection(model::Model, bond::Integer) = bonddirection(model.lat, bond)

@doc """
    lattice_bonddirection(lat::Lattice, bond::Integer)
    lattice_bonddirection(model::Model, bond::Integer)

Returns the direction of the `bond` as vector in the lattice system
"""
lattice_bonddirection(lat::Lattice, bond::Integer) = lat.bond_dirs[:, bond]
lattice_bonddirection(model::Model, bond::Integer) = lattice_bonddirection(model.lat, bond)


text:
module DistributionsLH

using DocStringExtensions, Lazy, Parameters, Random, StatsBase, Statistics
using Distributions, LinearAlgebra
using ModelObjectsLH, ModelParams

export AbstractDistributionSwitches, AbstractDistributionLH

# Generic functions
export init_distribution, draw, check_draws, isbounded, min_value, max_value, quantiles, scale_normal_draws
# Uniform
export AbstractUniform, AbstractUniformSwitches, Uniform, UniformSwitches, UniformCentered, UniformCenteredSwitches, UniformFixedBounds
export init_uniform
# Beta
export AbstractBeta, AbstractBetaSwitches, Beta, BetaSwitches
export init_beta, alpha_value, beta_value
# Normal
export NormalSwitches, Normal, LogNormalSwitches, LogNormal
export init_normal, init_log_normal, init_normal_switches, init_log_normal_switches

abstract type AbstractDistributionSwitches{T1} <: ModelSwitches end
abstract type AbstractDistributionLH{T1} <: ModelObject end

include("uniform.jl");
include("beta.jl");
include("normal.jl");
include("log_normal.jl");
include("mv_normal.jl")

## --------  Generic

"""
	init_distribution(objId, switches)

Initialize a distribution from its switches.
"""
function init_distribution end

"""
    draw(distrib, nDims, rng)
    draw(distrib, rng)

Draw from the distribution.
"""
function draw end

"""
    check_draws(distrib, drawM)

Check draws against bounds (if any) and `Inf` values.
"""
function check_draws end

"""
	quantiles(distrib, pctV)

Quantiles.
"""
function quantiles end

"""
	$(SIGNATURES)

Converts N(0,1) draws into the specified distribution.
"""
function scale_normal_draws(b :: AbstractDistributionLH{F},  inM :: T)  where {F, T}
    cdfM = cdf.(Distributions.Normal(), inM);
    drawM = quantiles(b,  cdfM);
    return drawM
end

"""
    isbounded(distrib)

Is this a bounded distribution?
"""
function isbounded end

"""
    min_value(distrib)

Minimum value (if any).
"""
function min_value end

"""
    max_value(distrib)

Maximum value (if any).
"""
function max_value end

function pretty_print(x)
    show(stdout, MIME("text/plain"), x);
    println("");
end

end # module

ext: jl


text:
using Test
using Dates, DataFrames, CSV, YAML
using Streamfall

DATA_PATH = joinpath(@__DIR__, "data/hymod")

# Load and generate stream network
network = YAML.load_file(joinpath(DATA_PATH, "hymod_network.yml"))
sn = create_network("HyMod Network", network)

# Load climate data
date_format = "YYYY-mm-dd"
obs_data = CSV.File(joinpath(DATA_PATH, "leaf_river_data.csv"),
                    comment="#",
                    dateformat=date_format) |> DataFrame

hist_streamflow = obs_data[:, "leaf_river_outflow"]
climate_data = obs_data[:, ["Date", "leaf_river_P", "leaf_river_ET"]]
climate = Climate(climate_data, "_P", "_ET")

@testset "Single node calibration" begin
    @test calibrate!(sn, climate, hist_streamflow; MaxTime=30) isa Any
end


text:
# Autogenerated wrapper script for boost_jll for x86_64-w64-mingw32-cxx03
export libboost_atomic, libboost_chrono, libboost_container, libboost_context, libboost_contract, libboost_coroutine, libboost_date_time, libboost_filesystem, libboost_graph, libboost_iostreams, libboost_log, libboost_log_setup, libboost_prg_exec_monitor, libboost_program_options, libboost_random, libboost_regex, libboost_serialization, libboost_stacktrace_basic, libboost_stacktrace_noop, libboost_system, libboost_thread, libboost_timer, libboost_type_erasure, libboost_unit_test_framework, libboost_wave, libboost_wserialization

JLLWrappers.@generate_wrapper_header("boost")
JLLWrappers.@declare_library_product(libboost_atomic, "libboost_atomic.dll")
JLLWrappers.@declare_library_product(libboost_chrono, "libboost_chrono.dll")
JLLWrappers.@declare_library_product(libboost_container, "libboost_container.dll")
JLLWrappers.@declare_library_product(libboost_context, "libboost_context.dll")
JLLWrappers.@declare_library_product(libboost_contract, "libboost_contract.dll")
JLLWrappers.@declare_library_product(libboost_coroutine, "libboost_coroutine.dll")
JLLWrappers.@declare_library_product(libboost_date_time, "libboost_date_time.dll")
JLLWrappers.@declare_library_product(libboost_filesystem, "libboost_filesystem.dll")
JLLWrappers.@declare_library_product(libboost_graph, "libboost_graph.dll")
JLLWrappers.@declare_library_product(libboost_iostreams, "libboost_iostreams.dll")
JLLWrappers.@declare_library_product(libboost_log, "libboost_log.dll")
JLLWrappers.@declare_library_product(libboost_log_setup, "libboost_log_setup.dll")
JLLWrappers.@declare_library_product(libboost_prg_exec_monitor, "libboost_prg_exec_monitor.dll")
JLLWrappers.@declare_library_product(libboost_program_options, "libboost_program_options.dll")
JLLWrappers.@declare_library_product(libboost_random, "libboost_random.dll")
JLLWrappers.@declare_library_product(libboost_regex, "libboost_regex.dll")
JLLWrappers.@declare_library_product(libboost_serialization, "libboost_serialization.dll")
JLLWrappers.@declare_library_product(libboost_stacktrace_basic, "libboost_stacktrace_basic.dll")
JLLWrappers.@declare_library_product(libboost_stacktrace_noop, "libboost_stacktrace_noop.dll")
JLLWrappers.@declare_library_product(libboost_system, "libboost_system.dll")
JLLWrappers.@declare_library_product(libboost_thread, "libboost_thread.dll")
JLLWrappers.@declare_library_product(libboost_timer, "libboost_timer.dll")
JLLWrappers.@declare_library_product(libboost_type_erasure, "libboost_type_erasure.dll")
JLLWrappers.@declare_library_product(libboost_unit_test_framework, "libboost_unit_test_framework.dll")
JLLWrappers.@declare_library_product(libboost_wave, "libboost_wave.dll")
JLLWrappers.@declare_library_product(libboost_wserialization, "libboost_wserialization.dll")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libboost_atomic,
        "bin\\libboost_atomic.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_chrono,
        "bin\\libboost_chrono.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_container,
        "bin\\libboost_container.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_context,
        "bin\\libboost_context.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_contract,
        "bin\\libboost_contract.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_coroutine,
        "bin\\libboost_coroutine.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_date_time,
        "bin\\libboost_date_time.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_filesystem,
        "bin\\libboost_filesystem.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_graph,
        "bin\\libboost_graph.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_iostreams,
        "bin\\libboost_iostreams.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_log,
        "bin\\libboost_log.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_log_setup,
        "bin\\libboost_log_setup.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_prg_exec_monitor,
        "bin\\libboost_prg_exec_monitor.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_program_options,
        "bin\\libboost_program_options.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_random,
        "bin\\libboost_random.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_regex,
        "bin\\libboost_regex.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_serialization,
        "bin\\libboost_serialization.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_stacktrace_basic,
        "bin\\libboost_stacktrace_basic.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_stacktrace_noop,
        "bin\\libboost_stacktrace_noop.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_system,
        "bin\\libboost_system.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_thread,
        "bin\\libboost_thread.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_timer,
        "bin\\libboost_timer.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_type_erasure,
        "bin\\libboost_type_erasure.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_unit_test_framework,
        "bin\\libboost_unit_test_framework.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_wave,
        "bin\\libboost_wave.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_wserialization,
        "bin\\libboost_wserialization.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
module GroupsCore

import Random
import Markdown
import AbstractAlgebra
import AbstractAlgebra: elem_type, gens, ngens, order, parent_type
import AbstractAlgebra: inv!, mul!
const Group = AbstractAlgebra.Group
const GroupElement = AbstractAlgebra.GroupElem

# abstract type Group end
# abstract type GroupElement end

export Group, GroupElement
export commutator, gens, hasgens, isfiniteorder, ngens, order
# export one!, inv!, mul!, conj!, commutator!, div_left!, div_right!

include("exceptions.jl")
include("groups.jl")
include("group_elements.jl")

include("extensions.jl")

include("constructions/constructions.jl")
using .Constructions

end


text:
#environment.jl - gets the esizesize and fsizesize variables.
################################################################################

doc"""
`Unums.environment` outputs the current environment as the appropriate Unum type
"""
function environment()
  Unum{options[:env_ESS], options[:env_FSS]}
end

doc"""
  `Unums.setenvironment` sets the current environment as the appropriate Unum type
"""
function setenvironment(ESS, FSS)
  options[:env_ESS] = ESS
  options[:env_FSS] = FSS
  Unum{ESS, FSS}
end

max_stars_repo_name: REX-Computing/unumjl

max_issues_repo_name: REX-Computing/unumjl

max_forks_repo_name: REX-Computing/unumjl


text:
# calculation of T-matrix using complex numbers.
# the only drawback of this is that it is not compatible with automatic differentiation

# TODO:
# 1- Use StaticArrays when possible.

export calculate_Tmatrix_for_spheroid

include("utils.jl")
include("geometry.jl")

# import
using VectorSphericalWaves
using StaticArrays
using EllipsisNotation
using LinearAlgebra
using Trapz
using HDF5

"""
    Calculate M_mn_wave as an array
This is the same as VectorSphericalHarmonics.M_mn_wave, but accept kr_array, θ_array, ϕ_array
Parameters
==========
kr_array, θ_array, ϕ_array : arrays of arbitrary shape

return
======
M_mn_wave_array_ : M_mn_wave with shape same as any of kr_array, θ_array, ϕ_array, with an added dimension to represent the three components
"""
function M_mn_wave_array(
        m::Int, n::Int, kr_array::AbstractVecOrMat{<:Complex{<:Real}}, θ_array::AbstractVecOrMat{R},
        ϕ_array::AbstractVecOrMat{R}; kind="regular"
    ) where {R <: Real}
    # Alok way is faster indeed!
    # TODO: @Alok, I think if we use boradcast it would be faster. I think avoiding preallocation makes the code cleaner and faster
    M_mn_wave_array_ = (_ -> zero(SVector{3,Complex})).(kr_array)
    for idx in eachindex(kr_array)
        M_mn_wave_array_[idx] = M_mn_wave(m, n, kr_array[idx], θ_array[idx], ϕ_array[idx], kind=kind)
    end
    return M_mn_wave_array_
end

"""
    Calculate N_mn_wave as an array
This is the same as VectorSphericalHarmonics.N_mn_wave, but accept kr_array, θ_array, ϕ_array
Parameters
==========
kr_array, θ_array, ϕ_array : arrays of arbitrary shape

return
======
N_mn_wave_array_ : N_mn_wave with shape same as any of kr_array, θ_array, ϕ_array, with an added dimension to represent the three components
"""
function N_mn_wave_array(
        m::Int, n::Int, kr_array::AbstractVecOrMat{<:Complex{<:Real}}, θ_array::AbstractVecOrMat{R},
        ϕ_array::AbstractVecOrMat{R}; kind="regular"
    ) where {R <: Real}
    # Alok way
    N_mn_wave_array_ = (_ -> zero(SVector{3,Complex})).(kr_array)
    for idx in eachindex(kr_array)
        N_mn_wave_array_[idx] = N_mn_wave(m, n, kr_array[idx], θ_array[idx], ϕ_array[idx], kind=kind)
    end
    return N_mn_wave_array_
end

#############################################################
# calculate J and Rg J, from equations 5.184 and 5.190
"""
    Calculate the integrand of J (including dS=r²sin(θ)dθdϕ) from equations 5.184 and 5.190.
It calculate all the integrand, and multiply it by r²sin(θ)

Inputs
======
m, n, m_, n_ : int, order and rank of VSWF inside and outside the particle
k1r_array, k2r_array : complex 2D array, product of wavevector and r, evaluated at points on the surface of the particle
θ_array, ϕ_array : complex 2D array, spherical coordinates of points on the particle surface
n̂_array : 3D array, unit vector normal to the surface of the particle
kind : string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
J_superscript : superscript at the top of J, it can be any of [11,12,21,22]

Outputs
=======
J
"""
function J_mn_m_n__integrand(
        m::Int, n::Int, m_::Int, n_::Int,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", J_superscript=11
    ) where {R <: Real, C <: Complex{R}}

    # determining the type of the first the second VSWF
    if J_superscript == 11 # TODO: this if-statement can be done more nicely. We separate J_superscript into two pieces, the number 1 represents M_mn_wave_array, while number 2 represents N_mn_wave_array
        first_function = M_mn_wave_array
        second_function = M_mn_wave_array
    elseif J_superscript == 12
        first_function = M_mn_wave_array
        second_function = N_mn_wave_array
    elseif J_superscript == 21
        first_function = N_mn_wave_array
        second_function = M_mn_wave_array
    elseif J_superscript == 22
        first_function = N_mn_wave_array
        second_function = N_mn_wave_array
    else
        throw(DomainError("J_superscript must be any of [11,12,21,22]"))
    end

    # determining the type of the first and second VSWF
    kind_first_function = "regular"
    if kind == "irregular"
        kind_second_function = "irregular"
    elseif kind == "regular"
        kind_second_function = "regular"
    else
        throw(DomainError("""kind must be any of ["regular", "irregular"]"""))
    end

    # the cross product
    cross_product_MN = cross.(
        first_function(m_, n_, k2r_array, θ_array, ϕ_array, kind=kind_first_function), # I can directly call M,N waves, with dots.
        second_function(-m, n, k1r_array, θ_array, ϕ_array, kind=kind_second_function)
    )

    # dot multiplying the cross product by the unit vector, and multiplying by (-1)^m
    cross_product_MN_dot_n̂ = (-1).^m .* vector_dot_product.(cross_product_MN, n̂_array)

    # multiplying by dS=r²sin(θ)
    J_integrand = surface_integrand(cross_product_MN_dot_n̂, r_array, θ_array)

    return J_integrand
end

"""
    Same as J_mn_m_n__integrand, but uses M_mn_wave_SVector and N_mn_wave_SVector
"""
function J_mn_m_n__integrand_SVector(
    m::Int, n::Int, m_::Int, n_::Int,
    k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
    r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
    kind="regular", J_superscript=11
) where {R <: Real, C <: Complex{R}}

# determining the type of the first the second VSWF
if J_superscript == 11 # TODO: this if-statement can be done more nicely. We separate J_superscript into two pieces, the number 1 represents M_mn_wave_SVector, while number 2 represents N_mn_wave_SVector
    first_function = M_mn_wave_SVector
    second_function = M_mn_wave_SVector
elseif J_superscript == 12
    first_function = M_mn_wave_SVector
    second_function = N_mn_wave_SVector
elseif J_superscript == 21
    first_function = N_mn_wave_SVector
    second_function = M_mn_wave_SVector
elseif J_superscript == 22
    first_function = N_mn_wave_SVector
    second_function = N_mn_wave_SVector
else
    throw(DomainError("J_superscript must be any of [11,12,21,22]"))
end

# determining the type of the first and second VSWF
kind_first_function = "regular"
if kind == "irregular"
    kind_second_function = "irregular"
elseif kind == "regular"
    kind_second_function = "regular"
else
    throw(DomainError("""kind must be any of ["regular", "irregular"]"""))
end

# the cross product
cross_product_MN = cross.(
    first_function.(m_, n_, k2r_array, θ_array, ϕ_array, kind=kind_first_function), # I can directly call M,N waves, with dots.
    second_function.(-m, n, k1r_array, θ_array, ϕ_array, kind=kind_second_function)
)

# dot multiplying the cross product by the unit vector, and multiplying by (-1)^m
cross_product_MN_dot_n̂ = (-1).^m .* vector_dot_product.(cross_product_MN, n̂_array)

# multiplying by dS=r²sin(θ)
J_integrand = surface_integrand(cross_product_MN_dot_n̂, r_array, θ_array)

return J_integrand
end

function J_mn_m_n_(
        m::Int, n::Int, m_::Int, n_::Int,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", J_superscript=11, rotationally_symmetric=false,
    ) where {R <: Real, C <: Complex{R}}
    if rotationally_symmetric
        # make sure that θ_array is 1D
        if length(size(θ_array)) != 1
            throw(DomainError("Since you have indicated << rotationally_symmetric = true >>, θ_array has to be 1D. Now it is $(length(size(θ_array)))D"))
        end
        ϕ_array = convert(typeof(θ_array), zeros(size(θ_array)))
    end

    # getting the integrand
    if rotationally_symmetric && (m != m_)
        # the integral over ϕ is 2π * δ_m_m_, so it is zero if m != m_
        J_integrand_dS = zeros(size(θ_array))

    else
        J_integrand_dS = J_mn_m_n__integrand_SVector(
            m, n,m_,n_,
            k1r_array,k2r_array,
            r_array,θ_array,ϕ_array,n̂_array;
            kind=kind,J_superscript=J_superscript
        )
    end

    # calculate the surface integral
    if rotationally_symmetric
        # integrate over θ only
        J = 2π * trapz((θ_array), J_integrand_dS)
    else
        # integrate over θ and ϕ
        # TODO: replace this integral with surface mesh quadrature, like this one: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5095443/ , https://github.com/jareeger/Smooth_Closed_Surface_Quadrature_RBF-julia
        # assuming that θ_array, ϕ_array were created with meshgrid function
        J = trapz((θ_array[:,1], ϕ_array[1,:]), J_integrand_dS)
end

    return J
end

function Q_mn_m_n_(
        m::Int, n::Int, m_::Int, n_::Int,
        k1::C, k2::C,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", Q_superscript=11, rotationally_symmetric=false,
    ) where {R <: Real, C <: Complex{R}}
    if Q_superscript == 11; J_superscript_1 = 21 ; J_superscript_2 = 12
    elseif Q_superscript == 12; J_superscript_1 = 11 ; J_superscript_2 = 22
    elseif Q_superscript == 21; J_superscript_1 = 22 ; J_superscript_2 = 11
    elseif Q_superscript == 22; J_superscript_1 = 12 ; J_superscript_2 = 21
    end

    Q = (
        -im .* k1 .* k2 .* J_mn_m_n_(m, n, m_, n_, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array; kind=kind, J_superscript=J_superscript_1, rotationally_symmetric=rotationally_symmetric)
        - im .* k1.^2    .* J_mn_m_n_(m, n, m_, n_, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array; kind=kind, J_superscript=J_superscript_2, rotationally_symmetric=rotationally_symmetric)
    )

    return Q
end

function Q_matrix(
        n_max::Int,
        k1::C, k2::C,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        kind="regular", rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false,
        verbose=false,
    ) where {R <: Real, C <: Complex{R}}
    idx_max = get_max_single_index_from_n_max(n_max)
    Q_mn_m_n_11 = zeros(typeof(k1), idx_max, idx_max) # TODO: @Alok, should I replace arrays with SMatrix? I am afraid it may get slower, as I have seen that StaticArray may get slower for arrays larger than 100 elements
    Q_mn_m_n_12 = zeros(typeof(k1), idx_max, idx_max)
    Q_mn_m_n_21 = zeros(typeof(k1), idx_max, idx_max)
    Q_mn_m_n_22 = zeros(typeof(k1), idx_max, idx_max)

    idx = 0;
    for n = 1:n_max
        for m = -n:n
            idx += 1
            idx_ = 0;
            for n_ = 1:n_max
                for m_ = -n_:n_
                    idx_ += 1
                    if verbose; println("n,m,idx = $n,$m,$idx  n_,m_,idx_ = $n_,$m_,$idx_"); end

                    if rotationally_symmetric
                        if m == m_
                            if symmetric_about_plane_perpendicular_z
                                # apply equations 5.208 and 5.209
                                if iseven(n + n_)
                                    Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                                    Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                                else
                                    Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                                    Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                                end
                            else
                                Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                            end
                        end
                    else
                        if symmetric_about_plane_perpendicular_z && (m == m_)
                            # apply equations 5.208 and 5.209
                            if iseven(n + n_)
                                Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                            else
                                Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                                Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                            end
                        else
                            Q_mn_m_n_11[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=11, rotationally_symmetric=rotationally_symmetric)
                            Q_mn_m_n_12[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=12, rotationally_symmetric=rotationally_symmetric)
                            Q_mn_m_n_21[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=21, rotationally_symmetric=rotationally_symmetric)
                            Q_mn_m_n_22[idx, idx_] = Q_mn_m_n_(m, n, m_, n_, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind=kind,Q_superscript=22, rotationally_symmetric=rotationally_symmetric)
                        end
            end
        end
    end
        end
    end
    Q = vcat(
        (hcat(Q_mn_m_n_11, Q_mn_m_n_12)),
        (hcat(Q_mn_m_n_21, Q_mn_m_n_22))
    )
    return Q
end

function T_matrix(
        n_max::Int,
        k1::C, k2::C,
        k1r_array::AbstractVecOrMat{C}, k2r_array::AbstractVecOrMat{C},
        r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
        rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false, HDF5_filename=nothing,
        verbose=false, create_new_arrays=false, BigFloat_precision = nothing
    ) where {R <: Real, C <: Complex{R}}

    if BigFloat_precision != nothing
        return setprecision(BigFloat_precision) do
            return T_matrix(
                n_max, big(k1), big(k2), big.(k1r_array), big.(k2r_array), big.(r_array), big.(θ_array), big.(ϕ_array), [big.(n) for n in n̂_array];
                HDF5_filename=HDF5_filename, rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z,
            )
        end
    else
        if create_new_arrays
            RgQ =       Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="regular"  , rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose)
            Q_inv = inv(Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="irregular", rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose))
            T = -1 .* RgQ * Q_inv

        else
            T = (
                -     Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="regular"  , rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose)
                * inv(Q_matrix(n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;kind="irregular", rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z, verbose=verbose))
            )
        end

        if HDF5_filename != nothing
            save_Tmatrix_to_HDF5_file(T, HDF5_filename)
        end

        return T
    end
end

"""
    Calculate T-matrix for a spheroid, with an extra kwarg `BigFloat_precision`

BigFloat_precision : set to `nothing` by default. When BigFloat_precision == nothing, then the function will use input types as they are.
"""
function calculate_Tmatrix_for_spheroid(
        rx::R, rz::R, n_max::Int,
        k1::Complex{R}, k2::Complex{R};
        n_θ_points=10, n_ϕ_points=20, HDF5_filename=nothing,
        rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false,
        BigFloat_precision = nothing
    ) where {R <: Real}

    # create a grid of θ_ϕ
    θ_array, ϕ_array = meshgrid_θ_ϕ(n_θ_points, n_ϕ_points; min_θ=1e-16, min_ϕ=1e-16, rotationally_symmetric=rotationally_symmetric)

    # calculate r and n̂ for the geometry
    r_array, n̂_array = ellipsoid(rx, rz, θ_array);

    # calculate T-matrix
    k1r_array = k1 .* r_array;
    k2r_array = k2 .* r_array;
    T = T_matrix(
        n_max, k1, k2, k1r_array, k2r_array, r_array, θ_array, ϕ_array, n̂_array;
        HDF5_filename=HDF5_filename, rotationally_symmetric=rotationally_symmetric, symmetric_about_plane_perpendicular_z=symmetric_about_plane_perpendicular_z,
        BigFloat_precision = BigFloat_precision
    )
    return T
end

"""
    General and nice wrapper for T-matric
Inputs
surrounding material (e.g., dielectric constant, refractive index, wavevector)
particle material (e.g., dielectric constant, refractive index, wavevector)
angular_resolution
n_max
particle geometry: it can be:
- function ``r(θ)`` (i.e., axi-symmetric particle)
- a function ``r(θ,ϕ)`` (i.e., arbitrary 3D particle)
- a meshfile
- a string defining geometry class (e.g., "cylinder", "spheroid", "ellipsoid", etc.)

# output
# - T_matrix,
"""
function Tmatrix_nice()
    # TODO
end

### precompiling
"""
R =
precompile(T_matrix, (
    Int, Complex{R}, k2::Complex{R},
    k1r_array::AbstractVecOrMat{<:Complex{<:Real}}, k2r_array::AbstractVecOrMat{<:Complex{<:Real}},
    r_array::AbstractVecOrMat{R}, θ_array::AbstractVecOrMat{R}, ϕ_array::AbstractVecOrMat{R}, n̂_array::Any; # TODO: I don't know why I get an error when I use n̂_array::AbstractVecOrMat{Vector{Float64}}
    rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false, HDF5_filename=nothing,
    verbose=false, create_new_arrays=false,)
"""


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# * https://discourse.julialang.org/t/for-monte-carlo-simulation-with-same-code-same-algorithm-how-fast-is-julia-compared-with-fortran/67021/41
# * https://discourse.julialang.org/t/for-monte-carlo-simulation-with-same-code-same-algorithm-how-fast-is-julia-compared-with-fortran/67021/43
# * https://github.com/genkuroki/public/blob/main/0018/%40defkwargs.ipynb
# * https://github.com/genkuroki/public/blob/main/0018/%40defunpack.ipynb
# * https://github.com/genkuroki/public/blob/main/0018/How%20to%20define%20unpacking%20macros.ipynb
# * https://github.com/genkuroki/public/blob/main/0018/How%20to%20define%20unpacking%20macros%20Part%202.ipynb

# %%
using ConcreteStructs

@concrete struct Foo a; b; c end

macro defkwargs(typename::Symbol)
    T = Core.eval(__module__, typename)
    names = fieldnames(T)
    kwargs = Expr(:parameters, names...)
    :($typename($kwargs) = $typename($(names...))) |> esc
end

macro defkwargs(typename::Symbol, kwargs_expr)
    T = Core.eval(__module__, typename)
    names = fieldnames(T)
    D = Core.eval(__module__, kwargs_expr)
    kwargs = Expr(:parameters)
    for name in names
        val = get(D, name, nothing)
        if isnothing(val)
            push!(kwargs.args, name)
        else
            push!(kwargs.args, Expr(:kw, name, val))
        end
    end
    :($typename($kwargs) = $typename($(names...))) |> esc
end

# %%
@macroexpand @defkwargs Foo

# %%
@defkwargs Foo

# %%
methods(Foo)

# %%
Foo(a = 1, b = 2.0, c = "three")

# %%
@macroexpand @defkwargs Foo (a = 1, c = "three")

# %%
@defkwargs Foo (a = 1, c = "three")

# %%
methods(Foo)

# %%
Foo(b = 2.0)

# %%
default = (a = 1, b = 2.0, c = "three")
@macroexpand @defkwargs Foo default

# %%
@defkwargs Foo default

# %%
methods(Foo)

# %%
Foo()

# %%

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
# ------------------------------------------------------------------------------
# Quick version of functions (for testing) (drawn from some of SJV's old code)
q_loglik(lr, w) = sum(-log.(1 .+ exp.(-lr.y .* (lr.X * w))))
q_gradloglik(lr, w) =
    sum(lr.y[i] * (1 / (1 + exp.(lr.y[i] * dot(lr.X[i,:], w)))) * X[i,:]
            for i ∈ 1:lr.n )

function q_gradloglik_cv(lr, wstar)
    gll_star = q_gradloglik(lr, wstar)
    function gll_cv(w; i=rand(1:lr.n))
        gll_star + lr.n * (
          lr.y[i] .* 1/(1 + exp.(lr.y[i] .* (dot(lr.X[i,:], w)))) .* lr.X[i, :] - lr.y[i] .* 1/(1 + exp.(lr.y[i] .* (dot(lr.X[i,:], wstar)))) .* lr.X[i, :])
    end
end
# -----------------------------------------------------------------------------
Random.seed!(123)

xvals = 10000 * randn(100)

safemaximum(x, s) = (length(s) == 1) ? max(x, s) : max(x, s[1])

@test maximum(logistic(x) - 1.0/(1.0 + exp(-x)) for x in xvals) <= 1e-10
vals = [loglogistic(x) - (-log(1.0 + exp(-x))) for x in xvals]
@test maximum(vals[broadcast(!, isinf.(vals))]) <= 1e-10
@test maximum(norm(gradloglogistic(x) - 1.0 / (1.0 + exp(x)))
        for x in xvals) <= 1e-10

# Generate data for Likelihood testing etc.
n = 1000               # number of observations
p = 10                 # number of dimensions (covariates)
X = randn(n, p) .+ 0.1  # feature matrix
w = 10*rand(p)         # true vector of parameters
# observations according to a logistic thresholded to {-1,1}
y = (logistic.(X * w) .> rand(n)) .* 2.0 .- 1.0
# proxy for N*L upper bound
b  = sum(mapslices(r->norm(r)^2, X, dims=1))/4
lr = LogReg(X, y, b)

xvals = [randn(p) for i in 1:5]

@test maximum(loglik(lr, x) - q_loglik(lr, x)
        for x in xvals) <= 1e-10

@test maximum(norm(gradloglik(lr, x) - q_gradloglik(lr, x))
        for x in xvals) <= 1e-10

xstar = randn(p)
i = rand(1:n)

gll1(x) = gradloglik_cv(lr, xstar)(x, i)
gll2(x) = q_gradloglik_cv(lr, xstar)(x, i=i)
# some aggregation of numerical noise
@test maximum(norm(gll1(x) - gll2(x))
        for x in xvals) <= 1e-10


text:
"""
`AbstractPoint{N,T} <: AbstractVector{T}`

Supertype for fixedsize `N`-dimensional points with elements of type `T`.
"""
abstract type AbstractPoint{N,T} <: AbstractVector{T} end

Base.haslength(::AbstractPoint) = true
Base.length(::AbstractPoint{N}) where N = N
Base.size(::AbstractPoint{N}) where N = (N,)
Base.eltype(::AbstractPoint{N,T}) where {N,T} = T
Base.IndexStyle(::Type{<:AbstractPoint}) = Base.IndexLinear()

"""
`Point{N,T} <: AbstractPoint{N,T}`

Fixedsize `N`-dimensional point with elements of type `T`.
"""
struct Point{N,T} <: AbstractPoint{N,T}
  data :: NTuple{N,T}
end

function Point(x, y)
  values = promote(x, y)
  Point{2,eltype(values)}(values)
end

function Point(x, y, z)
  values = promote(x, y, z)
  Point{3,eltype(values)}(values)
end

using Base: unsafe_load, unsafe_convert

Point{N,T}(vec::Vector{T}) where {N,T} =
  unsafe_load(unsafe_convert(Ptr{Point{N,T}}, vec))

Base.getindex(point::Point, i) = point.data[i]

struct MetaPoint{N,T,A,M} <: AbstractPoint{N,T}
  data :: NTuple{N,T}
  meta :: NamedTuple{A,M}
end

Base.getindex(point::MetaPoint, i) = point.data[i]
Base.getindex(point::MetaPoint, sym::Symbol) = point.meta[sym]
meta(point::MetaPoint) = point.meta


text:
# MIT License
#
# Copyright (c) 2020 Joshua E Gentry

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# =================================================================================================
# Test writing a PDF image.
@testset "pix_write_pdf" begin
    local pix = pix_with()

    # ---------------------------------------------------------------------------------------------
    # To a file.
    local filename = safe_tmp_file()
    @test pix_write_pdf(filename, pix) == true
    @test filesize(filename) > 0
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To an IOBuffer.
    local buffer = IOBuffer()
    @test pix_write_pdf(buffer, pix) == true
    @test length(take!(buffer)) > 0

    # ---------------------------------------------------------------------------------------------
    # To an IOStream
    local filename = safe_tmp_file()
    local file     = open(filename, read=true, write=true)
    @test pix_write_pdf(file, pix) == true
    @test position(file) != 0
    close(file)
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To a byte array.
    local bytes = pix_write_pdf(pix)
    @test length(bytes) > 0
end

# =================================================================================================
# Test writing a PDF image.
@testset "pix_write_pdf | ppi = $ppi, title = $title" for
        ppi in [ 300, 600, 1200 ], title in [ "", "My Title"]
    local pix = pix_with()

    # ---------------------------------------------------------------------------------------------
    # To a file.
    local filename = safe_tmp_file()
    @test pix_write_pdf(filename, pix; ppi = ppi, title = title) == true
    @test filesize(filename) > 0
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To an IOBuffer.
    local buffer = IOBuffer()
    @test pix_write_pdf(buffer, pix; ppi = ppi, title = title) == true
    @test length(take!(buffer)) > 0

    # ---------------------------------------------------------------------------------------------
    # To a byte array.
    local bytes = pix_write_pdf(pix; ppi = ppi, title = title)
    @test length(bytes) > 0
end

# =================================================================================================
# Test writing a PDF image.
@testset "pix_write_pdf /w invalid values" begin
    local pix = pix_with()
    local err = "PPI needs to be greater than 0."

    # ---------------------------------------------------------------------------------------------
    # To a file.
    local filename = safe_tmp_file()
    @test (@test_logs (:error, err) pix_write_pdf(filename, pix; ppi = 0)) == false
    @test filesize(filename) == 0
    @test (@test_logs (:error, err) pix_write_pdf(filename, pix; ppi = -300)) == false
    @test filesize(filename) == 0
    @test (@test_logs (:error, err) pix_write_pdf(filename, pix; ppi = 0, title = "Testing the test")) == false
    @test filesize(filename) == 0
    @test (@test_logs (:error, err) pix_write_pdf(filename, pix; ppi = -300, title = "Testing the test")) == false
    @test filesize(filename) == 0
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To an IOBuffer.
    local buffer = IOBuffer()
    @test (@test_logs (:error, err) pix_write_pdf(buffer, pix; ppi = 0)) == false
    @test length(take!(buffer)) == 0
    @test (@test_logs (:error, err) pix_write_pdf(buffer, pix; ppi = -300)) == false
    @test length(take!(buffer)) == 0
    @test (@test_logs (:error, err) pix_write_pdf(buffer, pix; ppi = 0, title = "Testing the test")) == false
    @test length(take!(buffer)) == 0
    @test (@test_logs (:error, err) pix_write_pdf(buffer, pix; ppi = -300, title = "Testing the test")) == false
    @test length(take!(buffer)) == 0

    # ---------------------------------------------------------------------------------------------
    # To a byte array.
    local bytes = Vector{UInt8}()
    @test (@test_logs (:error, err) pix_write_pdf(pix; ppi = 0)) == nothing
    @test (@test_logs (:error, err) pix_write_pdf(pix; ppi = -300)) == nothing
    @test (@test_logs (:error, err) pix_write_pdf(pix; ppi = 0, title = "Testing the test")) == nothing
    @test (@test_logs (:error, err) pix_write_pdf(pix; ppi = -300, title = "Testing the test")) == nothing
    @test length(bytes) == 0
end

# =================================================================================================
# Test writing a Pix that has been freed.
@testset "pix_write_pdf /w freed Pix" begin
    local err = "Pix has been freed."
    local pix = pix_with()
    pix_delete!(pix)

    # ---------------------------------------------------------------------------------------------
    # To a file.
    local filename = safe_tmp_file()
    @test (@test_logs (:error, err) pix_write_pdf(filename, pix)) == false
    rm(filename)

    # ---------------------------------------------------------------------------------------------
    # To an IOBuffer.
    local buffer = IOBuffer()
    @test (@test_logs (:error, err) pix_write_pdf(buffer, pix)) == false

    # ---------------------------------------------------------------------------------------------
    # To a byte array.
    @test (@test_logs (:error, err) pix_write_pdf(pix)) == nothing
end


text:
include("twoplayergames.jl")

function threatpoint(x::Array{Float64, 2}})

    threat_point = minimum(x, 2)[:]

    return threat_point
end

function Q_operator(tpg::TwoPlayerGame, a1::Int, a2::Int, W::Matrix{Float64}, u::Vector{Float64})

    δ = tpg.δ
    # Flow from action profile and gain from deviation
    _g = g(tpg, a1, a2)
    _h = h(tpg, a1, a2)

    # Find the lines to intersect set with
    w1 = ((1-δ)/δ)*_h[1] + u[1]
    w2 = ((1-δ)/δ)*_h[2] + u[2]

    Q_aWu = intersect_AS(W, w1, w2)

    return Q_aWu
end

function C_operator(tpg::TwoPlayerGame, a1::Int, a2::Int, W::Matrix{Float64}, u::Vector{Float64})

    δ = tpg.δ

    # Flow from action profile and gain from deviation
    _g = g(tpg, a1, a2)
    _h = h(tpg, a1, a2)

    if all(δ*(_g - u) .>= (1-δ)*_h)
        # Only one relevant point if the above holds
        out_ep = [ga]
        out_nep = 1
    else
        # Allocate space for points and give it a size hint
        out_ep = []
        out_nep = 0
        sizehint!(out_ep, 3)

        # Compute elements in Q(a, W, u)
        Q_aWu = Q_operator(tpg, a1, a2, W, u)

        # Check each extreme point of Q(a, W, u) for binding IC
        for ep in Q_aWu
            # Checking binding IC for both agents
            bIC_1 = isapprox((1-δ)*_h[1] - δ*(g[1] - u[1]), 0.0) ? true : false
            bIC_2 = isapprox((1-δ)*_h[2] - δ*(g[2] - u[2]), 0.0) ? true : false

            # Keep points for which either IC is binding
            if bIC_1 || bIC_2
                push!(out_ep, ep)
                out_nep += 1
            end
        end
    end

    return out_ep, out_nep
end

function R_operator(tpg::TwoPlayerGame, W::Matrix{Float64}, u::Vector{Float64})
    # Create space for the extreme points of R(W, u) -- Abreu Sannikov
    # prove that it will be less than 3|A| so we can give a size hint.
    R_Wu = []
    sizehint!(R_Wu, 3*na1*na2)

    for a1=1:na1
        for a2=1:na2
            # Compute the flow payoff
            flow_today = (1-δ)*g(tpg, a1, a2)

            # Get set of possible continuation values
            C_aWu, nC_aWu = C_operator(tpg, a1, a2, W, u)
            for cind=1:nC_aWu
                val_today = flow_today + δ*C_aWu[cind]
                push!(R_Wu, val_today)
            end
        end
    end

    # Take convex hull
    R_Wu = convexhull(R_Wu)

    return R_Wu
end


text:
@testset "Unimplemented Error" begin

    # TODO: This is a pain to debug if it fails. Try to fix it.

    struct InterfaceFunction
        f::Function
        args::Tuple
    end

    function IF(f::Function, args...)
        return InterfaceFunction(f, args)
    end

    struct Unimplemented <: PH.AbstractSubproblem
    end

    to_test = [
        IF(PH.add_ph_objective_terms, Vector{PH.VariableID}(), 2.25),
        IF(PH.objective_value),
        IF(PH.report_values, Vector{PH.VariableID}()),
        IF(PH.report_variable_info, build_scen_tree()),
        IF(PH.solve_subproblem),
        IF(PH.update_ph_terms, Dict{PH.VariableID,Float64}(), Dict{PH.VariableID,Float64}()),
        IF(PH.warm_start),
        IF(PH.report_penalty_info, PH.ProportionalPenaltyParameter),
        IF(PH.add_lagrange_terms, Vector{VariableID}()),
        IF(PH.update_lagrange_terms, Dict{VariableID,Float64}()),
    ]

    subprob = Unimplemented()
    for interface_function in to_test
        @test_throws PH.UnimplementedError begin
            interface_function.f(subprob, interface_function.args...)
        end
    end

    @test_throws PH.UnimplementedError begin
        PH.ef_copy_model(JuMP.Model(),
                         subprob,
                         PH.scid(0),
                         build_scen_tree(),
                         Dict{PH.NodeID,Any}()
                         )
    end

    @test_throws PH.UnimplementedError begin
        PH.ef_node_dict_constructor(typeof(subprob))
    end
end

@testset "Scenario Subproblem Functions" begin
    st = build_scen_tree()

    js = create_model(PH.scid(0))
    org_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)

    vid_name_map = PH.report_variable_info(js, st)
    @test length(vid_name_map) == length(JuMP.all_variables(js.model))
    for (vid, vinfo) in pairs(vid_name_map)
        if vid.stage == PH.stid(1)
            @test occursin("x", vinfo.name)
        elseif vid.stage == PH.stid(2)
            @test vinfo.name == "y"
        else
            @test occursin("z", vinfo.name)
        end
    end

    ref_map = invert_map(js.vars)
    observed_vids = Set{PH.VariableID}()
    for var in JuMP.all_variables(js.model)
        vid = ref_map[var]
        @test !(vid in observed_vids)
        push!(observed_vids, vid)
    end
    @test length(observed_vids) == length(JuMP.all_variables(js.model))

    (br_vids, lf_vids) = PH._split_variables(st, collect(keys(vid_name_map)))

    sts = PH.solve_subproblem(js)
    @test sts == MOI.LOCALLY_SOLVED
    @test sts == JuMP.termination_status(js.model)
    @test PH.objective_value(js) == JuMP.objective_value(js.model)
    org_obj_val = PH.objective_value(js)

    vals = PH.report_values(js, collect(keys(vid_name_map)))
    @test length(vals) == length(JuMP.all_variables(js.model))
    for var in JuMP.all_variables(js.model)
        @test JuMP.value(var) == vals[ref_map[var]]
    end

    r = PH.ScalarPenaltyParameter(10.0)
    rhalf = 0.5 * 10.0
    PH.add_ph_objective_terms(js, br_vids, r.value)
    ph_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)

    diff = ph_obj_func - org_obj_func
    JuMP.drop_zeros!(diff)

    @test length(JuMP.linear_terms(diff)) == 0
    @test length(JuMP.quad_terms(diff)) == (4 * length(br_vids))
    for qt in JuMP.quad_terms(diff)
        coef = qt[1]
        @test (isapprox(coef, rhalf) || isapprox(coef, -r.value) || isapprox(coef, 1.0))
    end

    w_vals = Dict{PH.VariableID,Float64}()
    for (k,w) in enumerate(keys(js.w_vars))
        w_vals[w] = k * rand()
    end
    xhat_vals = Dict{PH.VariableID,Float64}()
    for (k,xhat) in enumerate(keys(js.xhat_vars))
        xhat_vals[xhat] = k * rand()
    end
    PH.update_ph_terms(js, w_vals, xhat_vals)

    PH.warm_start(js)
    sts = PH.solve_subproblem(js)
    @test sts == MOI.LOCALLY_SOLVED
    @test org_obj_val != PH.objective_value(js)

    for (wid, wref) in pairs(js.w_vars)
        @test JuMP.value(wref) == w_vals[wid]
    end
    for (xhid, xhref) in pairs(js.xhat_vars)
        @test JuMP.value(xhref) == xhat_vals[xhid]
    end
end

# @testset "Extensive Form" begin
#     #TODO: Make tests for this case
# end

@testset "Penalties" begin
    st = build_scen_tree()

    # Scalar Penalty
    js = create_model(PH.scid(0))
    vid_name_map = PH.report_variable_info(js, st)
    (br_vids, lf_vids) = PH._split_variables(st, collect(keys(vid_name_map)))
    obj = JuMP.objective_function(js.model)

    PH.add_ph_objective_terms(js, br_vids, 10.0)

    ph_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)
    diff = ph_obj_func - obj
    JuMP.drop_zeros!(diff)

    @test length(JuMP.linear_terms(diff)) == 0
    @test length(JuMP.quad_terms(diff)) == (4 * length(br_vids))
    for vid in br_vids
        x_ref = js.vars[vid]
        w_ref = js.w_vars[vid]
        xhat_ref = js.xhat_vars[vid]
        for (c, v1, v2) in JuMP.quad_terms(diff)
            if (x_ref == v1 && x_ref == v2) || (xhat_ref == v1 && xhat_ref == v2)
                @test isapprox(c, 5.0)
            elseif (x_ref == v1 && xhat_ref == v2) || (xhat_ref == v1 && x_ref == v2)
                @test isapprox(c, -10.0)
            elseif (x_ref == v1 && w_ref == v2) || (w_ref == v1 && x_ref == v2)
                @test isapprox(c, 1.0)
            end
        end
    end

    # Proportional Penalty
    js = create_model(PH.scid(0))
    vid_name_map = PH.report_variable_info(js, st)
    (br_vids, lf_vids) = PH._split_variables(st, collect(keys(vid_name_map)))
    coef_dict = PH.report_penalty_info(js, br_vids, PH.ProportionalPenaltyParameter)
    obj = JuMP.objective_function(js.model)
    for (vid, coef) in pairs(coef_dict)
        ref = js.vars[vid]
        for lt in JuMP.linear_terms(obj)
            if lt[2] == ref
                @test lt[1] == coef
                break
            end
        end
        coef_dict[vid] = 10.0 * coef
    end
    PH.add_ph_objective_terms(js, br_vids, coef_dict)

    ph_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)
    diff = ph_obj_func - obj
    JuMP.drop_zeros!(diff)

    @test length(JuMP.linear_terms(diff)) == 0
    @test length(JuMP.quad_terms(diff)) == (4 * length(br_vids))

    for vid in br_vids
        x_ref = js.vars[vid]
        w_ref = js.w_vars[vid]
        xhat_ref = js.xhat_vars[vid]
        for (c, v1, v2) in JuMP.quad_terms(diff)
            if (x_ref == v1 && x_ref == v2) || (xhat_ref == v1 && xhat_ref == v2)
                @test isapprox(c, 0.5 * coef_dict[vid])
            elseif (x_ref == v1 && xhat_ref == v2) || (xhat_ref == v1 && x_ref == v2)
                @test isapprox(c, -1.0 * coef_dict[vid])
            elseif (x_ref == v1 && w_ref == v2) || (w_ref == v1 && x_ref == v2)
                @test isapprox(c, 1.0)
            end
        end
    end

    # SEP Penalty
    js = create_model(PH.scid(0))
    vid_name_map = PH.report_variable_info(js, st)
    (br_vids, lf_vids) = PH._split_variables(st, collect(keys(vid_name_map)))
    coef_dict = PH.report_penalty_info(js, br_vids, PH.SEPPenaltyParameter)
    obj = JuMP.objective_function(js.model)
    for (vid, coef) in pairs(coef_dict)
        ref = js.vars[vid]
        for lt in JuMP.linear_terms(obj)
            if lt[2] == ref
                @test lt[1] == coef
                break
            end
        end
        coef_dict[vid] = rand() * coef
    end
    PH.add_ph_objective_terms(js, br_vids, coef_dict)

    ph_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)
    diff = ph_obj_func - obj
    JuMP.drop_zeros!(diff)

    @test length(JuMP.linear_terms(diff)) == 0
    @test length(JuMP.quad_terms(diff)) == (4 * length(br_vids))

    for vid in br_vids
        x_ref = js.vars[vid]
        w_ref = js.w_vars[vid]
        xhat_ref = js.xhat_vars[vid]
        for (c, v1, v2) in JuMP.quad_terms(diff)
            if (x_ref == v1 && x_ref == v2) || (xhat_ref == v1 && xhat_ref == v2)
                @test isapprox(c, 0.5 * coef_dict[vid])
            elseif (x_ref == v1 && xhat_ref == v2) || (xhat_ref == v1 && x_ref == v2)
                @test isapprox(c, -1.0 * coef_dict[vid])
            elseif (x_ref == v1 && w_ref == v2) || (w_ref == v1 && x_ref == v2)
                @test isapprox(c, 1.0)
            end
        end
    end

end

@testset "Fix Variables" begin

    js = create_model(PH.scid(0))
    PH.report_variable_info(js, build_scen_tree())

    is_fixed = Dict{PH.VariableID,Float64}()
    is_free = Vector{PH.VariableID}()
    for vid in keys(js.vars)
        if PH.value(PH.index(vid)) % 2 == 0 # r > 0.5
            is_fixed[vid] = rand()
        else
            push!(is_free, vid)
        end
    end

    PH.fix_variables(js, is_fixed)
    sts = PH.solve_subproblem(js)
    @test sts == MOI.LOCALLY_SOLVED

    for (vid,var) in pairs(js.vars)
        if haskey(is_fixed, vid)
            @test JuMP.is_fixed(var)
            @test JuMP.value(var) == is_fixed[vid]
        elseif vid in is_free
            @test !JuMP.is_fixed(var)
        else
            error("Vid $vid is neither fixed nor free")
        end
    end

end

@testset "Lagrange Terms" begin
    st = build_scen_tree()

    js = create_model(PH.scid(0))
    org_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)

    vid_name_map = PH.report_variable_info(js, st)
    (br_vids, lf_vids) = PH._split_variables(st, collect(keys(vid_name_map)))

    PH.add_lagrange_terms(js, br_vids)
    ph_obj_func = JuMP.objective_function(js.model, JuMP.QuadExpr)

    diff = ph_obj_func - org_obj_func
    JuMP.drop_zeros!(diff)

    @test length(JuMP.linear_terms(diff)) == 0
    @test length(JuMP.quad_terms(diff)) == length(br_vids)
    for qt in JuMP.quad_terms(diff)
        coef = qt[1]
        @test isapprox(coef, 1.0)
    end

    w_vals = Dict{PH.VariableID,Float64}()
    for (k,w) in enumerate(keys(js.w_vars))
        w_vals[w] = k * rand()
    end
    PH.update_lagrange_terms(js, w_vals)

    sts = PH.solve_subproblem(js)
    @test sts == MOI.LOCALLY_SOLVED

    for (wid, wref) in pairs(js.w_vars)
        @test JuMP.value(wref) == w_vals[wid]
    end
end


text:
module ForwardGame

import JuMP
import Ipopt
import ..DynamicsModelInterface
import ..JuMPUtils
import ..ForwardOptimalControl

using JuMP: @variable, @constraint, @objective
using UnPack: @unpack

export IBRGameSolver, KKTGameSolver, solve_game

#================================ Iterated Best Open-Loop Response =================================#

struct IBRGameSolver end

# TODO: We could allow to pass an inner solver
function solve_game(
    ::IBRGameSolver,
    control_system,
    player_cost_models,
    x0,
    T;
    init = (; x = zeros(control_system.n_states, T), u = zeros(control_system.n_controls, T)),
    max_ibr_rounds = 10,
    ibr_convergence_tolerance = 1e-3,
    verbose = false,
    inner_solver_kwargs...,
)
    @unpack n_states, n_controls = control_system
    n_players = length(player_cost_models)

    last_ibr_solution = init
    last_player_solution = last_ibr_solution
    player_opt_models = resize!(JuMP.Model[], n_players)
    converged = false

    for i_ibr in 1:max_ibr_rounds
        for (player_idx, player_cost_model) in enumerate(player_cost_models)
            last_player_converged, last_player_solution, player_opt_models[player_idx] =
                ForwardOptimalControl.solve_optimal_control(
                    control_system,
                    player_cost_model,
                    x0,
                    T;
                    fixed_inputs = filter(i -> i ∉ player_cost_model.player_inputs, 1:n_controls),
                    init = last_player_solution,
                    inner_solver_kwargs...,
                    verbose,
                )
            @assert last_player_converged
        end

        converged =
            sum(Δu -> Δu^2, last_player_solution.u - last_ibr_solution.u) <=
            ibr_convergence_tolerance
        last_ibr_solution = last_player_solution

        if converged
            verbose && @info "Converged at ibr iterate: $i_ibr"
            break
        end
    end

    converged || @warn "IBR terminated pre-maturely."

    converged, last_ibr_solution, player_opt_models
end

#================================= Open-Loop KKT Nash Constraints ==================================#

struct KKTGameSolver end

function solve_game(
    ::KKTGameSolver,
    control_system,
    player_cost_models,
    x0,
    T;
    solver = Ipopt.Optimizer,
    solver_attributes = (; print_level = 3),
    init = (),
    match_equilibrium = nothing,
    verbose = false,
)

    n_players = length(player_cost_models)
    @unpack n_states, n_controls = control_system

    opt_model = JuMP.Model(solver)
    JuMPUtils.set_solver_attributes!(opt_model; solver_attributes...)

    # Decision Variables
    x = @variable(opt_model, [1:n_states, 1:T])
    u = @variable(opt_model, [1:n_controls, 1:T])
    λ = @variable(opt_model, [1:n_states, 1:(T - 1), 1:n_players])

    # Initialization
    JuMPUtils.init_if_hasproperty!(λ, init, :λ)
    JuMPUtils.init_if_hasproperty!(x, init, :x)
    JuMPUtils.init_if_hasproperty!(u, init, :u)

    # constraints
    @constraint(opt_model, x[:, 1] .== x0)
    DynamicsModelInterface.add_dynamics_constraints!(control_system, opt_model, x, u)
    df = DynamicsModelInterface.add_dynamics_jacobians!(control_system, opt_model, x, u)

    for (player_idx, cost_model) in enumerate(player_cost_models)
        @unpack player_inputs, weights = cost_model
        dJ = cost_model.add_objective_gradients!(opt_model, x, u; weights)

        # KKT Nash constraints
        @constraint(
            opt_model,
            [t = 2:(T - 1)],
            dJ.dx[:, t] + λ[:, t - 1, player_idx] - (λ[:, t, player_idx]' * df.dx[:, :, t])' .== 0
        )
        @constraint(opt_model, dJ.dx[:, T] + λ[:, T - 1, player_idx] .== 0)

        @constraint(
            opt_model,
            [t = 1:(T - 1)],
            dJ.du[player_inputs, t] - (λ[:, t, player_idx]' * df.du[:, player_inputs, t])' .== 0
        )
        @constraint(opt_model, dJ.du[player_inputs, T] .== 0)
    end

    if !isnothing(match_equilibrium)
        @objective(
            opt_model,
            Min,
            sum(el -> el^2, x - match_equilibrium.x)
        )
    end

    time = @elapsed JuMP.optimize!(opt_model)
    verbose && @info time

    JuMPUtils.isconverged(opt_model), JuMPUtils.get_values(; x, u, λ), opt_model
end

end


text:
"""
Execute a regular expression without the overhead of `Base.Regex`
"""
exec(re::Regex, bytes, offset::Int=1) =
    Base.PCRE.exec(re.regex, bytes, offset-1, re.match_options, re.match_data)

"""
`SubString` containing the bytes following the matched regular expression.
"""
nextbytes(re::Regex, bytes) = SubString(bytes, re.ovec[2]+1)

"""
`SubString` containing a regular expression match group.
"""
group(i, re::Regex, bytes) = SubString(bytes, re.ovec[2i+1]+1,
                                              prevind(bytes, re.ovec[2i+2]+1))

group(i, re::Regex, bytes, default) =
    re.ovec[2i+1] == Base.PCRE.UNSET ?
    default :
    SubString(bytes, re.ovec[2i+1]+1, prevind(bytes, re.ovec[2i+2]+1))


text:
module Buttons

using Genie, Stipple, StippleUI, StippleUI.API
import Genie.Renderer.Html: HTMLString, normal_element, register_normal_element

export btn, btngroup, btndropdown, Btn

register_normal_element("q__btn", context = @__MODULE__)
register_normal_element("q__btn__group", context = @__MODULE__)
register_normal_element("q__btn__dropdown", context = @__MODULE__)

function btn( label::String = "",
              args...;
              content::Union{String,Vector,Function} = "",
              kwargs...)
  q__btn([isa(content, Function) ? content() : join(content)],
          args...;
          label = label, kw(kwargs)...
  )
end

function btn( content::Union{Function,Vector},
              label::String = "",
              args...;
              kwargs...)
  btn(label, args...; content = content, kwargs...)
end

function btn( args...;
              label::String = "",
              content::Union{String,Vector,Function} = "",
              kwargs...)
  btn(label, args...; content = content, kwargs...)
end

mutable struct Btn
  label
  args
  wrap
  content
  kwargs

  Btn(label::String = "",
      args...;
      content::Union{String,Vector,Function} = "",
      kwargs...) = new(fieldname, args, mask, kwargs)
end

Base.string(bt::Btn) = btn(bt.label, bt.args...; bt.content, bt.kwargs...)

function btngroup(args...; kwargs...)
  q__btn__group(args...; kw(kwargs)...)
end

function btndropdown(args...; kwargs...)
  q__btn__dropdown(args...; kw(kwargs)...)
end

end


text:
module PkgREPL

import ..CondaPkg
import Pkg
import Markdown

### parsing

function parse_pkg(x::String)
    m = match(r"^\s*(([^:]+)::)?([-_.A-Za-z0-9]+)\s*([<>=!0-9].*)?$", x)
    m === nothing && error("invalid conda package: $x")
    channel = m.captures[2]
    name = m.captures[3]
    version = m.captures[4]
    if version === nothing
        version = ""
    end
    if channel === nothing
        channel = ""
    end
    CondaPkg.PkgSpec(name, version=version, channel=channel)
end

function parse_pip_pkg(x::String)
    m = match(r"^\s*([-_.A-Za-z0-9]+)\s*([~!<>=@].*)?$", x)
    m === nothing && error("invalid pip package: $x")
    name = m.captures[1]
    version = m.captures[2]
    if version === nothing
        version = ""
    end
    CondaPkg.PipPkgSpec(name, version=version)
end

function parse_channel(x::String)
    CondaPkg.ChannelSpec(x)
end

### options

const channel_opt = Pkg.REPLMode.OptionDeclaration([
    :name => "channel",
    :short_name => "c",
    :api => :mode => :channel,
])

const pip_opt = Pkg.REPLMode.OptionDeclaration([
    :name => "pip",
    :api => :mode => :pip,
])

const resolve_opt = Pkg.REPLMode.OptionDeclaration([
    :name => "resolve",
    :short_name => "r",
    :api => :resolve => true,
])

const force_opt = Pkg.REPLMode.OptionDeclaration([
    :name => "force",
    :short_name => "f",
    :api => :force => true,
])

### status

function status()
    CondaPkg.status()
end

const status_help = Markdown.parse("""
```
conda st|status
```

Display information about the Conda environment.
""")

const status_spec = Pkg.REPLMode.CommandSpec(
    name = "status",
    short_name = "st",
    api = status,
    help = status_help,
    description = "information about the Conda environment",
)

### resolve

function resolve(; force=false)
    CondaPkg.resolve(force=force, interactive=true)
end

const resolve_help = Markdown.parse("""
```
conda [-f|--force] resolve
```

Ensure all Conda dependencies are installed into the environment.
""")

const resolve_spec = Pkg.REPLMode.CommandSpec(
    name = "resolve",
    api = resolve,
    help = resolve_help,
    description = "ensure all Conda dependencies are installed",
    option_spec = [force_opt],
)

### add

function add(args; mode=:package, resolve=false)
    if mode == :package
        CondaPkg.add(parse_pkg.(args))
    elseif mode == :channel
        CondaPkg.add(parse_channel.(args))
    elseif mode == :pip
        CondaPkg.add(parse_pip_pkg.(args))
    end
    resolve && CondaPkg.resolve()
end

const add_help = Markdown.parse("""
```
conda add [-c|--channel] [--pip] [-r|--resolve] pkg ...
```

Add packages or channels to the environment.

This adds Conda packages by default. Use `--channel` or `-c` to add channels instead, or
`--pip` to add Pip packages.

The Conda environment is not immediately resolved. Use the `--resolve` or `-r`
flag to force resolve.

**Examples**
```
pkg> conda add python
pkg> conda add python>=3.5,<4
pkg> conda add conda-forge::numpy
pkg> conda add --channel anaconda
pkg> conda add --pip build
pkg> conda add --pip build~=0.7.0
```
""")

const add_spec = Pkg.REPLMode.CommandSpec(
    name = "add",
    api = add,
    should_splat = false,
    help = add_help,
    description = "add Conda packages or channels",
    arg_count = 0 => Inf,
    option_spec = [channel_opt, pip_opt, resolve_opt],
)

### rm

function rm(args; mode=:package, resolve=false)
    if mode == :package
        CondaPkg.rm(parse_pkg.(args))
    elseif mode == :channel
        CondaPkg.rm(parse_channel.(args))
    elseif mode == :pip
        CondaPkg.rm(parse_pip_pkg.(args))
    end
    resolve && CondaPkg.resolve()
end

const rm_help = Markdown.parse("""
```
conda rm|remove [-c|--channel] [--pip] [-r|--resolve] pkg ...
```

Remove packages or channels from the environment.

This removes Conda packages by default. Use `--channel` or `-c` to remove channels instead,
or `--pip` to remove Pip packages.

The Conda environment is not immediately resolved. Use the `--resolve` or `-r`
flag to force resolve.

**Examples**
```
pkg> conda rm python
pkg> conda rm --channel anaconda
pkg> conda rm --pip build
```
""")

const rm_spec = Pkg.REPLMode.CommandSpec(
    name = "remove",
    short_name = "rm",
    api = rm,
    should_splat = false,
    help = rm_help,
    description = "remove Conda packages or channels",
    arg_count = 0 => Inf,
    option_spec = [channel_opt, pip_opt, resolve_opt],
)

### gc

function gc()
    CondaPkg.gc()
end

const gc_help = Markdown.parse("""
```
conda gc
```

Delete any files no longer used by Conda.
""")

const gc_spec = Pkg.REPLMode.CommandSpec(
    name = "gc",
    api = gc,
    help = gc_help,
    description = "delete files no longer used by Conda",
)

## run

function run(args)
    CondaPkg.withenv() do
        Base.run(Cmd(args))
    end
end

const run_help = Markdown.parse("""
```
conda run cmd ...
```

Run the given command in the Conda environment.
""")

const run_spec = Pkg.REPLMode.CommandSpec(
    name = "run",
    api = run,
    should_splat = false,
    help = run_help,
    arg_count = 1 => Inf,
    description = "run a command in the Conda environment",
)

### all specs

const SPECS = Dict(
    "st" => status_spec,
    "status" => status_spec,
    "resolve" => resolve_spec,
    "add" => add_spec,
    "remove" => rm_spec,
    "rm" => rm_spec,
    "gc" => gc_spec,
    "run" => run_spec,
)

function __init__()
    # add the commands to the repl
    Pkg.REPLMode.SPECS["conda"] = SPECS
    # update the help with the new commands
    copy!(Pkg.REPLMode.help.content, Pkg.REPLMode.gen_help().content)
end

end


text:
@testset verbose = true "Helper functions Testing" begin

    @test DecisionTrees.isContiguous(collect(1:23))
    @test DecisionTrees.isContiguous(collect(22:23))
    @test DecisionTrees.isContiguous(collect(23:23))
    @test !DecisionTrees.isContiguous([1,9,3])
    @test !DecisionTrees.isContiguous([1,3,2])

    function test_mean_and_var_merge(nloops, sizeofVectors)
        for loopV = 1:nloops
        # nn=150
            nn = sizeofVectors
            weight = rand(nn)
            denom = rand(nn)
            num = rand(nn)
            trni = StatsBase.sample(1:nn, floor(Int, .4 * nn), replace=false)
            sort!(trni)
            strArr = [randstring(1) for i = 1:nn]
            floatArr = float.(abs.(rand(Int, nn) .% 9))
            pdStr = DecisionTrees.PooledArray(strArr)
            pdFloat = DecisionTrees.PooledArray(floatArr)

            if rand() < .5
                pd = pdStr
            else
                pd = pdFloat
            end

            thispd = view(pd, trni)
            cnt, sumnumerator, sumdenominator, sumweight, moments_per_pdaclass = DecisionTrees.aggregate_data_msePointwise(thispd, num, denom, weight)
        # check if the 'merge!' of all compoments matches the total
            total = OnlineStats.Series(OnlineStats.Mean(), OnlineStats.Variance())
            for i = 1:length(moments_per_pdaclass)
                if moments_per_pdaclass[i].stats[1].n > 0 # this is needed because of https://github.com/joshday/OnlineStats.jl/issues/125
                    merge!(total, moments_per_pdaclass[i])
                end
            end
            ratioOfND = num ./ denom
        # @show ratioOfND[trni]
        # @show total
        # @show Statistics.mean(ratioOfND[trni]),total.stats[1].μ
        # @show StatsBase.var(ratioOfND[trni]),total.stats[2].σ2
            if length(trni) > 0
                @test isapprox(Statistics.mean(view(ratioOfND, trni)), total.stats[1].μ)
                @test isapprox(StatsBase.var(view(ratioOfND, trni), corrected=false), total.stats[2].σ2)
            end
        end
        return nothing
    end

    function test_mean_and_var_UNmerge(nloops, sizeofVectors)
        # sizeofVectors=10
        for jj = 1:nloops

            nn = sizeofVectors
            weight = rand(nn)
            denom = rand(nn)
            num = rand(nn)
            trni = StatsBase.sample(1:nn, floor(Int, .4 * nn), replace=false)
            sort!(trni)
            strArr = [randstring(1) for i = 1:nn]
            floatArr = float.(abs.(rand(Int, nn) .% 9))
            pdStr = DecisionTrees.PooledArray(strArr)
            pdFloat = DecisionTrees.PooledArray(floatArr)

            if rand() < .5
                pd = pdStr
            else
                pd = pdFloat
            end

            thispd = view(pd, trni)
            cnt, sumnumerator, sumdenominator, sumweight, moments_per_pdaclass = DecisionTrees.aggregate_data_msePointwise(thispd, num, denom, weight)
        # check if the 'merge!' of all compoments matches the total
            total = OnlineStats.Series(OnlineStats.Mean(), OnlineStats.Variance())
            for i = 1:length(moments_per_pdaclass)
                if moments_per_pdaclass[i].stats[1].n > 0 # this is needed because of https://github.com/joshday/OnlineStats.jl/issues/125
                    merge!(total, moments_per_pdaclass[i])
                end
            end

            if length(moments_per_pdaclass) > 2
        # remove one or two random elements of moments_per_pdaclass
                ratioOfND = num ./ denom
                pick = rand() > .5
                reducedTotal = deepcopy(total)
                if pick
                    el1 = StatsBase.sample(1:length(moments_per_pdaclass))
                    DecisionTrees.unmerge!(reducedTotal, moments_per_pdaclass[el1])
                    valuesRemaining = ratioOfND[trni][(pd[trni] .!= pd.pool[el1])]
                else
                    (el1, el2) = StatsBase.sample(1:length(moments_per_pdaclass), 2, replace=false)
                    DecisionTrees.unmerge!(reducedTotal, moments_per_pdaclass[el1])
                    DecisionTrees.unmerge!(reducedTotal, moments_per_pdaclass[el2])
                    idx1 = (pd.refs .!= pd.pool[2]) .& (pd.refs .!= pd.pool[4])
                    valuesRemaining = ratioOfND[trni][(pd[trni] .!= pd.pool[el1]) .& (pd[trni] .!= pd.pool[el2])]
                end
            #=
            @show valuesRemaining
            @show reducedTotal
            @show Statistics.mean(valuesRemaining),reducedTotal.stats[1].μ
            @show StatsBase.var(valuesRemaining),reducedTotal.stats[2].σ2 =#
                if length(trni) > 0
                    @test isapprox(Statistics.mean(valuesRemaining), reducedTotal.stats[1].μ)
                    @test isapprox(StatsBase.var(valuesRemaining, corrected=false), reducedTotal.stats[2].σ2, rtol=1e-4)
            #    @show eltype(pd),pick
                end
            end
        end
        return nothing

    end

    Random.seed!(99231)
    test_mean_and_var_merge(1000, 1000)
    test_mean_and_var_UNmerge(1000, 1000)

    #some test for coverage
    DecisionTrees.get_sha1()

end #end testset


text:
p_death(a, b, m) = a + ((1 - a) / (1 + exp(-b * (-m))))

p_phage_decay(decay_factor, time_in_state) = 1 - (1 * exp(-decay_factor * time_in_state))

function p_lysis(phage, model)
    nearby_phages = nearby_t(:phage, phage, model)
    nearby_phages = length(nearby_phages)
    return 1 / (1 + model.properties.α * exp(-nearby_phages + model.properties.κ))
end

function p_adsorption(cell, model)
    nearby_phages = nearby_t(:phage, cell, model)
    p_host = model[cell].species === :a ? 0.9 : 0.1
    return p_host / (1 + exp(-length(nearby_phages)))
end

function p_grow(model)
    cells = by_single_type(:bacterium)(model)
    filter!(id -> isempty(model[id].phages_inside), cells)
    cell_count = length(cells)

    properties = model.properties
    ΔN = properties.growth_rate * cell_count * (1 - cell_count / properties.carrying_capacity)
    return ΔN / cell_count
end

function phage_decay(phage, model)
    if rand(model.rng) < p_phage_decay(model.decay_factor, model[phage].time_in_state)
        kill_agent!(phage, model)
    end
end

function tick_phage(phage, model)
    model[phage].time_in_state += 1
end

function infect(phage, cell, model)
    model[phage].state = :in_host
    model[phage].time_in_state = 0

    push!(model[cell].phages_inside, phage)
end

function burst(phage, cell, model)
    environment = model.properties.environment
    if environment === :well_mixed
        nearby = positions(model, :random)
    elseif environment === :spatially_structured
        r = 1
    elseif environment === :semi_solid
        r = 3
    end

    if environment === :spatially_structured || environment === :semi_solid
        nearby = collect(nearby_positions(model[phage].pos, model, r))
    end

    selected = rand(model.rng, nearby, model.properties.burst_size)
    for pos ∈ selected
        roll = rand(model.rng)
        if roll < 0.45
            kind = :temperate
        elseif roll < 0.85
            kind = :virulent
        else
            kind = :deficient
        end
        add_agent!(pos, Organism, model, roll < 0.5 ? :a : :b, Vector{Int}(), kind, :free, 0, :phage)
    end

    genocide!(model, model[cell].phages_inside)
    kill_agent!(cell, model)
end

function diffuse(id, model)
    agent = model[id]

    r = 3
    nearby = collect(nearby_positions(agent.pos, model, r))

    if agent.type === :phage
        selected = rand(model.rng, nearby)
        move_agent!(agent, selected, model)
    elseif agent.type === :bacterium
        move_bacterium_single!(agent, model, nearby)
    end
end

function grow(cell, p_grow, model)
    function has_no_bacteria(pos)
        ids = ids_in_position(pos, model)
        return isempty(ids) || !any(id -> model[id].type === :bacterium, ids)
    end

    !(rand(model.rng) < p_grow) && return

    environment = model.properties.environment
    if environment === :well_mixed
        selected = positions(model, :random)
    elseif environment === :spatially_structured
        r = 1
    elseif environment === :semi_solid
        r = 3
    end

    if environment === :spatially_structured || environment === :semi_solid
        selected = collect(nearby_positions(model[cell].pos, model, r))
    end

    filter!(has_no_bacteria, selected)
    isempty(selected) && return

    target = rand(model.rng, selected)
    add_agent!(target, Organism, model,
        rand(model.rng) < 0.5 ? :a : :b, Vector{Int}(), :nothing, :nothing, -1, :bacterium)
end

function bacteria_death_inherent(bacteria, model)
    for id ∈ bacteria
        if rand(model.rng) < p_death(model.a, model.b, model.m)
            genocide!(model, model[id].phages_inside)
            kill_agent!(id, model)
        end
    end
end

function bacteria_death_lysis(bacteria, model) # TODO: maybe only call this on infected bacteria (check the other (old) methods too)
    for cell ∈ bacteria
        phages_inside = model[cell].phages_inside
        if !isempty(phages_inside)
            for phage ∈ phages_inside
                tick_phage(phage, model)
                agent = model[phage]
                if (agent.time_in_state ≥ model.latent_period) && (agent.kind === :virulent || agent.kind === :induced_temperate)
                    if rand(model.rng) < model.properties.p_burst
                        burst(phage, cell, model)
                        break
                    end
                end
            end
        end
    end
end


text:
mutable struct Bus
    BUS_I::Vector{Int64}    # "bus number"
    BUS_TYPE::Vector{Int64} # "bus type"
    PD::Vector{Float64}     # "real power demand [MW]"
    QD::Vector{Float64}     # "reactive power demand [MVar]"
    GS::Vector{Float64}     # "shunt conductance [MW at V = 1 pu]"
    BS::Vector{Float64}     # "shunt susceptance [MVar at V = 1 pu]"
    BUS_AREA::Vector{Int64} # "area number"
    VM::Vector{Float64}     # "voltage magnitude [pu]"
    VA::Vector{Float64}     # "voltage angle [deg]"
    BASE_KV::Vector{Float64}    # "base voltage [kV]"
    ZONE::Vector{Int64}     # "loss zone"
    VMAX::Vector{Float64}   # "maximum voltage magnitude [pu]"
    VMIN::Vector{Float64}   # "minimum voltage magnitude [pu]"
    LAM_P::Vector{Float64}  # "Lagrange multiplier on real power mismatch (u/MW)"
    LAM_Q::Vector{Float64}  # "Lagrange multiplier on reactive power mismatch (u/MVAr)"
    MU_VMAX::Vector{Float64}    # "Kuhn-Tucker multiplier on upper voltage limit (u/p.u.)"
    MU_VMIN::Vector{Float64}    # "Kuhn-Tucker multiplier on lower voltage limit (u/p.u.)"
end

"""
Extract bus information from mpc Dict. Return instance of Bus.

IN:

* mpc. Instance of MPC dict (loaded, for instance, by loadcase())

OUT:

* Instance of Bus containing all bus data.
"""
function extract_bus(mpc)
    bus = mpc["bus"]
    nc = size(bus, 2)
    nf = length(fieldnames(Bus))
    blanks = fill([NaN], nf - nc)
    return Bus([bus[:, i] for i in 1:size(bus, 2)]..., blanks...)
end

"""
Shortcut for returning bus info for one of the cases listed by
`casenames()`.
"""
function extract_bus(cname::AbstractString)
    mpc = loadcase(cname; describe=false)
    extract_bus(mpc)
end


text:
module NLboxsolve

using ForwardDiff, LinearAlgebra, SparseArrays

include("structs.jl")
include("givens.jl")
include("arnoldi.jl")
include("gmres.jl")
include("jacvec.jl")
include("boxsolvers.jl")

export SolverState,
       SolverTrace,
       SolverResults

export constrained_newton,
       constrained_newton_ms,
       constrained_levenberg_marquardt,
       constrained_levenberg_marquardt_kyf,
       constrained_levenberg_marquardt_fan,
       constrained_levenberg_marquardt_ar,
       constrained_trust_region,
       constrained_dogleg_solver,
       constrained_dogleg_bmp_solver,
       constrained_newton_krylov,
       constrained_newton_krylov_fs,
       constrained_jacobian_free_newton_krylov,
       constrained_newton_sparse,
       constrained_newton_ms_sparse,
       constrained_levenberg_marquardt_sparse,
       constrained_levenberg_marquardt_kyf_sparse,
       constrained_levenberg_marquardt_fan_sparse,
       constrained_levenberg_marquardt_ar_sparse,
       constrained_trust_region_sparse,
       constrained_dogleg_solver_sparse,
       constrained_dogleg_bmp_solver_sparse,
       constrained_newton_krylov_sparse,
       constrained_newton_krylov_fs_sparse,
       nlboxsolve

end


text:
# ---
# title: 1611. Minimum One Bit Operations to Make Integers Zero
# id: problem1611
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming, Bit Manipulation
# link: <https://leetcode.com/problems/minimum-one-bit-operations-to-make-integers-zero/description/>
# hidden: true
# ---
#
# Given an integer `n`, you must transform it into `0` using the following
# operations any number of times:
#
#   * Change the rightmost (`0th`) bit in the binary representation of `n`.
#   * Change the `ith` bit in the binary representation of `n` if the `(i-1)th` bit is set to `1` and the `(i-2)th` through `0th` bits are set to `0`.
#
# Return _the minimum number of operations to transform_`n` _into_`0` _._
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 0
#     Output: 0
#
#
# **Example 2:**
#
#
#
#     Input: n = 3
#     Output: 2
#     Explanation: The binary representation of 3 is "11".
#     " _1_ 1" -> " _0_ 1" with the 2nd operation since the 0th bit is 1.
#     "0 _1_ " -> "0 _0_ " with the 1st operation.
#
#
# **Example 3:**
#
#
#
#     Input: n = 6
#     Output: 4
#     Explanation: The binary representation of 6 is "110".
#     " _1_ 10" -> " _0_ 10" with the 2nd operation since the 1st bit is 1 and 0th through 0th bits are 0.
#     "01 _0_ " -> "01 _1_ " with the 1st operation.
#     "0 _1_ 1" -> "0 _0_ 1" with the 2nd operation since the 0th bit is 1.
#     "00 _1_ " -> "00 _0_ " with the 1st operation.
#
#
# **Example 4:**
#
#
#
#     Input: n = 9
#     Output: 14
#
#
# **Example 5:**
#
#
#
#     Input: n = 333
#     Output: 393
#
#
#
#
# **Constraints:**
#
#   * `0 <= n <= 109`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using BiomolecularStructures.MatrixPlot
A = [1.0 1.0 1.0; 1.0 1.0 1.0; 1.0 1.0 1.0]

x, y, z = matrix_prep(A)

@test x == [1.0, 1.0, 1.0]
@test y == [1.0, 1.0, 1.0]
@test z == [1.0, 1.0, 1.0]

P = [51.65 -1.90 50.07;
    50.40 -1.23 50.65;
    50.68 -0.04 51.54;
    50.22 -0.02 52.85]

Q = [51.30 -2.99 46.54;
     51.09 -1.88 47.58;
     52.36 -1.20 48.03;
     52.71 -1.18 49.38]

@test matrices_plot(P,Q,false) == true


text:
#shared image alignment computational procedure
#1.  Compute transform to align camera2's flipped padded half-image (step 8) with its flipped full image (step 7) (should be almost aligned already)
#2.  Compute transform to align camera2's flipped full image (step 7) with camera1's padded half-image (step 8)
#3.  Compose transform from 2 with transform from 1 and call recenter() on the composed transform.
#4.  For each camera2 half-image:
    # -flip y
    # -pad vertically to full chip size
    # -apply transform from #3
    # -nansum with camera1's corresponding padded half-image
    # -find minimal bounding box that contains all non-nan pixels and return image cropped to that region

function stitch(img_top_fixed::AbstractArray{T,2}, img_bottom_moving::AbstractArray{T,2}, tfm, out_type::DataType; kwargs...) where {T}
    if T!=out_type
        img_top_fixed = out_type.(img_top_fixed)
        img_bottom_moving = out_type.(img_bottom_moving)
    end
    return stitch(img_top_fixed, img_bottom_moving, tfm; kwargs...)
end

#Stitches img_top_fixed and img_bottom_moving using a transformation tfm that, when applied to img_bottom_moving, aligns it with img_top_fixed.
#Since img_top_fixed and img_bottom_moving are expected to have little-to-no overlap they are padded before applying the transformation.
#The transformed img_bottom_moving is added to img_top_fixed and the interesting (non-NaN) region of the summed image is returned.
#For most accurate stitching you must provide ysz_full, the size in the y dimension of full images used to compute tfm with the stitch_tfm function
function stitch(img_top_fixed::AbstractArray{T,2}, img_bottom_moving::AbstractArray{T,2}, tfm; flip_y_bottom = true, ysz_full=2048) where {T}
    yroi = size(img_top_fixed,2)
    @assert size(img_bottom_moving,2) == yroi
    yextra = ceil(Int, abs(tfm.translation[2]))
    #fillvals should be NaN, but currently interpolations doesn't handle it well (imputes NaNs at gridpoints)
    img_bottom_moving_pre = ypad(img_bottom_moving, ysz_full; pad_side=:both, flip_y=flip_y_bottom, fillval=0.0)
    img_top_fixed = ypad(img_top_fixed, ysz_full; pad_side=:both, flip_y=false, fillval=0.0)
    img_bottom_moving = warp_and_resample(img_bottom_moving_pre, tfm)
    imgsummed = nanplus(img_top_fixed,img_bottom_moving)
    ystart = div(ysz_full,2) - div(yroi,2) + 1
    return view(imgsummed, :, ystart:ystart+2*yroi-1)
end

#returns a transform that stitches a pair of corresponding shared images into a single image when passed to the stitch function
#Assumes that moving_roi contains the bottom half of the image and that it needs to be flipped vertically to align with fixed_roi.
function stitch_tfm(moving_full, fixed_roi, moving_roi; trunc_frac=0.1, kwargs...)
    @assert size(fixed_roi) == size(moving_roi)

    #bright bead clumps can bias calculation, so truncate them at this value:
    trunc_thr = trunc_frac * maximum_finite(fixed_roi)
    moving_roi = trunc_above.(moving_roi, trunc_thr)
    fixed_roi = trunc_above.(fixed_roi, trunc_thr)
    moving_full = trunc_above.(moving_full, trunc_thr)

    #flip and pad
    moving_full = flipy(moving_full)
    fixed_roi_pad = ypad(fixed_roi, size(moving_full,2); flip_y = false, fillval=NaN)
    moving_roi_pad = ypad(moving_roi, size(moving_full,2); flip_y = true, fillval=NaN)

    #################### Moving strip to moving full
    print("Registering the moving strip to the moving full image\n")
    mxshift_moving = (50, size(moving_roi,2)+50)
    m_roi_pwr = sum(abs2.(moving_roi_pad[.!(isnan.(moving_roi_pad))]))
    moving_pre_rigid, pre_rigid_mm = qd_rigid(moving_full, moving_roi_pad, mxshift_moving, [pi/60]; thresh=0.85*m_roi_pwr, minwidth_rot=[pi/10000], kwargs...)
    @show moving_pre_rigid

    #refine it further, allowing full affine
    moving_pre_tfm, pre_mm = qd_affine(moving_full, moving_roi_pad, (50,50); thresh=0.5*m_roi_pwr, initial_tfm=moving_pre_rigid, kwargs...)
    @show moving_pre_tfm

    #################### Moving strip to moving full
    print("Registering the moving full image to the fixed strip\n")
    mxshift_xcam = (50, round(Int, size(moving_roi, 2)+50))
    f_roi_pwr = sum(abs2.(fixed_roi_pad[.!(isnan.(fixed_roi_pad))]))
    moving_post_rigid, post_rigid_mm = qd_rigid(fixed_roi_pad, moving_full, mxshift_xcam, [pi/60]; thresh=0.5*f_roi_pwr, minwidth_rot=[pi/10000], kwargs...)
    @show moving_post_rigid

    #refine it further, allowing full affine
    print("Found rigid transform, now optimizing allowing general affine transforms\n")
    moving_post_tfm, post_mm = qd_affine(fixed_roi_pad, moving_full, (50,50); thresh=0.5*f_roi_pwr, initial_tfm=moving_post_rigid, kwargs...)
    @show moving_post_tfm

    return recenter(moving_post_tfm ∘ moving_pre_tfm, RegisterQD.center(moving_full)), post_mm
end

#green and magenta overlay
mg_overlay(img1, img2) = colorview(RGB, nanz.(img1), nanz.(img2), nanz.(img1))
#green and magenta overlay, one color per half-image
function stitched_mg_overlay(strip1::AbstractArray{T,2}, strip2::AbstractArray{T,2}, tfm, ysz_full; flip_y2=true) where {T}
    strip0 = zeros(T, size(strip1))
    simg1 = stitch(strip1, strip0, tfm; flip_y_bottom = false, ysz_full=ysz_full)
    simg2 = stitch(strip0, strip2, tfm; flip_y_bottom = flip_y2, ysz_full=ysz_full)
    return mg_overlay(simg1,simg2)
end

#By convention, we align images to camera1, which has the top portion of the split image on OCPI2
#This returns an initial guess translation mapping camera 2's image to camera1's image after camera2's image has been flipped (flip required on OCPI2)
#ysz_roi is the pixel size of the image in the vertical direction (the dimension that gets split between cameras)
#Note: currently not using this
function initial_guess_tfm(ysz_chip::Int, ysz_roi::Int)
    @assert iseven(ysz_chip)
    @assert iseven(ysz_roi)
    if ysz_roi > ysz_chip
        error("ROI size cannot be larger than chip size")
    end
    return Translation(0, -ysz_roi)
end

#Makes an initial guess based on stripsz
#function full2full(fixed, moving, stripsz::Int, mxshift, mxrot, rotstep; thresh = 0.1*sum(abs2.(moving[.!(isnan.(moving))])), tfm0 = IdentityTransformation(), kwargs...)
#    @assert size(fixed) == size(moving)
#    @show tfm_guess = initial_guess_tfm(size(fixed,2), stripsz) ∘ tfm0
#    #tfm, mm = register_padmatched(fixed, moving, mxshift, mxrot, rotstep; thresh=thresh, tfm0=tfm_guess, kwargs...)
#    tfm, mm = register_padmatched(fixed, moving, mxshift, mxrot, rotstep; thresh=thresh, kwargs...)
#end
#
#function register_padmatched(fixed::AbstractMatrix, moving::AbstractMatrix, mxshift, mxrot, rotstep; thresh = 0.1*sum(abs2.(moving[.!(isnan.(moving))])), tfm0 = IdentityTransformation(), kwargs...)
#    if size(fixed,1) != size(moving, 1)
#        error("Images must be equally sized in the first dimension")
#    end
#    fixed, moving = match_ypadding(fixed, moving)
#    tfm, mm = qd_rigid(copy(fixed), copy(moving), mxshift, [mxrot...], [rotstep...]; thresh=thresh, tfm0=tfm0, kwargs...)
#end
#
#function match_ypadding(img1, img2; padval=NaN)
#    if size(img1,2) > size(img2, 2)
#        img2 = ypad(img2, size(img1,2); flip_y = false, fillval=padval)
#    end
#    if size(img1,2) < size(img2, 2)
#        img1 = ypad(img1, size(img2,2); flip_y = false, fillval=padval)
#    end
#    return img1, img2
#end


text:
Base.bits(b::Ubound) = string("$(bits(b.lower)) -> $(bits(b.upper))")

function Base.show{ESS,FSS}(io::IO, x::Ubound{ESS,FSS})
  @typenames

  print(io, "$bname($(x.lower), $(x.upper))")
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule GSL_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("GSL")
JLLWrappers.@generate_main_file("GSL", UUID("1b77fbbe-d8ee-58f0-85f9-836ddc23a7a4"))
end  # module GSL_jll


text:
export Triangular

"""
    Triangular{Shape, BC} <: BoundedLattice{2, BC}

Triangular lattices.

Ref: https://en.wikipedia.org/wiki/Hexagonal_lattice
"""
struct Triangular{Shape, BC} <: BoundedLattice{2, BC} end

Triangular(x, y; boundary=Periodic) = Triangular{Tuple{x, y}, boundary}()

nameof(::Triangular) = "Triangular Lattice"


text:
using Plots, LaTeXStrings
using PyCall
PyDict(pyimport("matplotlib")."rcParams")["text.usetex"] = [true]
pyplot()
default(titlefont=18, guidefont=18, tickfont=16, legendfont= 13,
      bg_legend = nothing, fg_legend= :black, framestyle = :semi, grid = true)
# ...............................................................#
x = LinRange(0,2π,100)
plot(x, sin.(x), c = :red,label = L"sin(x)", title = "Sine",
    yticks = ([-1,0,1], [L"-1", L"0", L"1"]), ylab = L"f(x)",
    xlab = L"x", xticks = ([0,π/2, π, 3π/2, 2π],
    [L"0", L"\pi/2", L"\pi", L"3\pi/2", L"2\pi"]))
savefig("figSin3.pdf")

max_stars_repo_name: lazarusA/helpfulJulia

max_issues_repo_name: lazarusA/helpfulJulia

max_forks_repo_name: lazarusA/helpfulJulia