using MuJoCo
using LinearAlgebra

struct MPPIStrategy <: LearningStrategy
   costs::Vector{Float64}
   weights::Vector{Float64}

   Σ::Matrix{Float64}
   λ::Float64
   γ::Float64

   f::FunctionSet
   cfunc::Function
   rollout::Function
   c::Vector{Float64}
   o::Vector{Float64}
   function MPPIStrategy(model, sigma, lambda, gamma, funcset,
                         rollout=mjWrap.rollout2)
      dtype   = Float64
      K       = model.samples.K

      costs   = Vector{dtype}(undef, K)
      weights = Vector{dtype}(undef, K)

      new(costs, weights, sigma, lambda, gamma,
          funcset,
          (x...)->nothing,
          rollout,
          zeros(model.mjsys.nu), zeros(model.mjsys.ns))
   end
end

Base.show(io::IO, s::MPPIStrategy) = print(io, "MPPI")

function apply(f::Function, costs::Vector, states::Array{Float64,3})
   for k=1:size(states,3)
      costs[k] -= f(@view states[:,end,k])
   end
   #costs .-= f(states[:,end,:]) # block calculate
end

function setup!(s::MPPIStrategy, model::TrajOptModel)
   samples = model.samples

   # traj and samples both start from same place; send inital state to TrajOptModel
   if isnan(model.s0[1])
      @info("Setting initial state from experiment.")
      s.f.initstate!(model.mjsys, model.traj)
      model.s0 .= model.traj.state[:,1,1]
   else
      model.traj.state[:,1] .= model.s0
   end
   samples.state[:,1,:] .= model.s0

   push!(model.trace, :maxR => Vector{Float64}())
end

function update!(model::TrajOptModel, s::MPPIStrategy, iter, null)
   samples = model.samples
   T       = samples.T
   K       = samples.K
   mjsys   = model.mjsys

   #randn!(samples.ctrl)
   #rmul!(samples.ctrl, s.Σ[1]) # scale by one value
   #broadcast!(+, samples.ctrl, samples.ctrl, model.theta)
   #samples.ctrl[:,:,1] .= model.theta # have at least one thing to do a full theta rollout

   #samples.ctrl .= 0.0
   randn!( @view(samples.ctrl[:,(1+model.delay):end,:]) )
   rmul!(samples.ctrl, s.Σ[1]) # scale by one value
   broadcast!(+, samples.ctrl, samples.ctrl, model.theta)
   samples.ctrl[:,:,1] = model.theta # have at least one thing to do a full theta rollout

   mjWrap.limitcontrols(mjsys, samples)

   if iter == 1
      s.rollout(mjsys, samples,
                      s.cfunc, s.f.observe!, s.f.reward)
   else
      s.rollout(mjsys, samples,
                      s.cfunc, s.f.observe!, s.f.reward;
                      startT=1+model.delay)
   end
end

function applyweights!(theta::Matrix{Float64},
                       ctrl::Array{Float64,3}, weights::Vector{Float64})
   nu, T, K = size(ctrl)

   @inbounds for t=1:T
      #theta[:,t] = ctrl[:,t,:]*weights
      #@views A_mul_B!(theta[:,t], ctrl[:,t,:], weights)
      theta[:,t] .= 0.0
      for k=1:K
         for i=1:nu
            theta[i,t] += ctrl[i,t,k]*weights[k]
         end
      end
   end
end

function hook(s::MPPIStrategy, model::TrajOptModel, iter)

   LinearAlgebra.BLAS.set_num_threads(Threads.nthreads())
   BLAS.set_num_threads( 0 )

   mjsys   = model.mjsys
   samples = model.samples
   T       = samples.T
   K       = samples.K

   s.costs[:] = -1.0*Common.Rtau(samples.reward, s.γ)
   values = model.terminalvalue( samples.obs[:,end,:] )
   if values != nothing
      s.costs .-= values
   end

   baseline = minimum(s.costs)

   push!(model.trace[:stocR], mean(s.costs))
   push!(model.trace[:meanR], baseline)
   push!(model.trace[:maxR], maximum(s.costs))

   @. s.costs = exp((-(s.costs-baseline)/s.λ))
   η = sum(s.costs)
   @. s.weights = s.costs/η

   applyweights!(model.theta, samples.ctrl, s.weights)

   samples.ctrl[:,:,1] = model.theta # for visualization purposes

   if model.MPC && iter <= model.niter # get first controls, mj.step, store in model.traj
      traj    = model.traj

      # MPC: apply first controls to system; shift the rest around
      s.c[:] = model.theta[:,1]
      model.theta[:,1:end-1] .= @view model.theta[:,2:end]
      model.theta[:,end] .= 0.0

      # TODO
      # TODO
      # think about skip and how it should affect mpc mode
      # TODO
      # TODO
      mjWrap.reset(mjsys, model.s0, s.c, s.o)
      #for i=1:mjsys.skip
         mjWrap.step!(mjsys, s.c, model.s0, s.o) # get next state, next observation
      #end
      s.f.observe!(mjsys, model.s0, s.o) # arbitrary observation vector manipulation

      # advance our main trajectory
      traj.state[:,iter] = model.s0
      traj.ctrl[:,iter]  = s.c
      traj.obs[:,iter]   = s.o

      # update the mpc bundle of trajectories to start 'here'
      #for k=1:K
      #   samples.state[:,1,k] = model.s0
      #end
      samples.state[:,1+model.delay,:] .= model.s0
   end
end

function cleanup!(s::MPPIStrategy, model::TrajOptModel) # just rollout and store in traj
   if model.MPC == false
      #info("Running rollout for main trajectory.")
      mjsys   = model.mjsys
      traj    = model.traj

      model.traj.state .= model.s0
      traj.ctrl[:,:,1] = model.theta
      s.rollout(mjsys, traj, s.cfunc, s.f.observe!, s.f.reward)
   end
end


"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	storage_asymmetric(EP::Model, inputs::Dict, Reserves::Int)

Sets up variables and constraints specific to storage resources with asymmetric charge and discharge capacities. See ```storage()``` in ```storage.jl``` for description of constraints.
"""
function storage_asymmetric(EP::Model, inputs::Dict, Reserves::Int)
	# Set up additional variables, constraints, and expressions associated with storage resources with asymmetric charge & discharge capacity
	# (e.g. most chemical, thermal, and mechanical storage options with distinct charge & discharge components/processes)
	# STOR = 2 corresponds to storage with distinct power and energy capacity decisions and distinct charge and discharge power capacity decisions/ratings

	println("Storage Resources with Asmymetric Charge/Discharge Capacity Module")

	dfGen = inputs["dfGen"]

	T = inputs["T"]     # Number of time steps (hours)

	STOR_ASYMMETRIC = inputs["STOR_ASYMMETRIC"]

	### Constraints ###

	# Storage discharge and charge power (and reserve contribution) related constraints for symmetric storage resources:
	if Reserves == 1
		EP = storage_asymmetric_reserves(EP, inputs)
	else
		# Maximum charging rate must be less than charge power rating
		@constraint(EP, [y in STOR_ASYMMETRIC, t in 1:T], EP[:vCHARGE][y,t] <= EP[:eTotalCapCharge][y])
	end

	return EP
end

@doc raw"""
	storage_asymmetric_reserves(EP::Model, inputs::Dict)

Sets up variables and constraints specific to storage resources with asymmetric charge and discharge capacities when reserves are modeled. See ```storage()``` in ```storage.jl``` for description of constraints.
"""
function storage_asymmetric_reserves(EP::Model, inputs::Dict)

	dfGen = inputs["dfGen"]
	T = inputs["T"]

	STOR_ASYMMETRIC = inputs["STOR_ASYMMETRIC"]

	STOR_ASYM_REG = intersect(STOR_ASYMMETRIC, inputs["REG"]) # Set of asymmetric storage resources with REG reserves
	STOR_ASYM_NO_REG = setdiff(STOR_ASYMMETRIC, STOR_ASYM_REG) # Set of asymmetric storage resources without REG reserves

	if !isempty(STOR_ASYM_REG)
		# Storage units charging can charge faster to provide reserves down and charge slower to provide reserves up
		# Maximum charging rate plus contribution to regulation down must be less than charge power rating
		@constraint(EP, [y in STOR_ASYM_REG, t in 1:T], EP[:vCHARGE][y,t]+EP[:vREG_charge][y,t] <= EP[:eTotalCapCharge][y])
	else
		# Storage units charging can charge faster to provide reserves down and charge slower to provide reserves up
		# Maximum charging rate plus contribution to regulation down must be less than charge power rating
		@constraint(EP, [y in STOR_ASYM_NO_REG, t in 1:T], EP[:vCHARGE][y,t] <= EP[:eTotalCapCharge][y])
	end

	return EP
end


"""
A quick realization of MPS-qubit circuit.

For compatibility please also install in Julia Package mode:
pkg> add Yao
pkg> add https://github.com/QuantumBFS/QuAlgorithmZoo.jl.git

Reference:
    Jin-Guo Liu, Yihong Zhang, Yuan Wan and Lei Wang,
    Variational Quantum Eigensolver with Fewer Qubits(arXiv:1902.02663)
"""

module MPSCircuit
export setMPSpar
export DCbuilder
export MPSC
export MPSbuilder

using Yao, Yao.ConstGate
include("Diff.jl")

# Structure of basic parameters.
struct setMPSpar
    nBitA  # Number of lines(bits) in the Extended circuit.
    vBit   # Number of virtual bits in the MPS circuit.
    rBit   # Number of reusable bits in the MPS circuit.
    nBit   # Number of lines(bits) in the MPS circuit.
    nBlock # Number of  MPS blocks in the MPS circuit.

    function setMPSpar(nBitA::Int64, vBit::Int64, rBit::Int64)
        if (nBitA - vBit -rBit) % rBit != 0
            println("Error: nBlock is not integer!")
            return 0
        end
        nBlock = Int((nBitA - vBit) / rBit)
        nBit = rBit + vBit
        new(nBitA, vBit, rBit, nBit, nBlock)
    end
end

# Structure of elements may needed for a Differentiable circuit.
struct DCbuilder
    block   # The block for 1 depth.
    Cblocks # The whole Blocks for n depths.
    body    # The Differentiable circuit for n depths.
    head    # The "head" of the Differentiable circuit(if directly input bits).
    tail    # The "tail" of the Differentiable circuit(2 layers of rotaional gates).

    function DCbuilder(nBit::Int64, ndepth::Int64)
        block = chain(nBit, vcat([repeat(nBit, Rz(0))], [repeat(nBit, Rx(0))], [repeat(nBit, Rz(0))],
                                    [control(nBit, i, (i-1)=>X) for i=nBit:-1:2], control(nBit, 1, nBit=>X)))
        tail = chain(nBit, repeat(nBit, Rz(0)), repeat(nBit, Rx(0)))
        head = block[2:end]
        Cblocks = chain(nBit, repeat([block],ndepth))
        body = chain(nBit, vcat([head], repeat([block],ndepth), [tail]))
        new(block, Cblocks, body, head, tail)
    end
end

# Structure of elements MPSCircuit.jl provides.
struct MPSC
    circuit  # MPS circuit.
    cBlocks  # Array of all the MPS blocks in the MPS circuit.
    cExtend  # The MPS circuit extended back to where it doesn't reuse any qubit.
    cEBlocks # Array of all the MPS blocks in the Extended circuit.
    dGates   # Differentiable gates of the MPS circuit if applicable.
    nBit     # Number of lines(bits) of the MPS circuit.
    nBlock   # Number of blocks in the MPS ciruict.

    function MPSC(MPSblock, nBitA::Int64, vBit::Int64, rBit::Int64=1)
        par2nd = setMPSpar(nBitA, vBit, rBit)
        nBlock = par2nd.nBlock
        nBit = par2nd.nBit
        #println("M1\n")
        MPS = MPSbuilder(nBitA, vBit, rBit, MPSblock)
        #println("M2\n")
        circuit = MPS[1]
        #println("M3\n")
        cExtend = MPS[2]
        dGates = collect_blocks(AbstractDiff, circuit)
        new(circuit, circuit.blocks, cExtend, cExtend.blocks, dGates, nBit, nBlock)
    end

end

# Methods for creating different types of MPS circuits.
function MPSbuilder(nBitA::Int64, vBit::Int64, rBit::Int64, blockT::Tuple{String, Int64})
    par2nd = setMPSpar(nBitA, vBit, rBit)
    nBlock = par2nd.nBlock
    nBit = par2nd.nBit

    if blockT[1] == "DC" # Differentiable circuit.
        depth = blockT[2]
        # println("s1\n")
        swapV1 = chain(nBit, [put(nBit, (inBit,inBit+1)=>ConstGate.SWAP ) for irBit=rBit:-1:1 for inBit=irBit  :(vBit+irBit-1)])
        swap1V = chain(nBit, [put(nBit, (inBit+1,inBit)=>ConstGate.SWAP ) for irBit=1:   rBit for inBit=(vBit+irBit-1):-1:irBit])
        # println("s2\n")
        cBlockHead = chain(nBit, DCbuilder(nBit, depth).block, swapV1) |> diff
        dispatch!(cBlockHead, :random)
        # println("s3\n")
        cBlock =  chain(nBit, vcat([swap1V], cBlockHead.blocks))
        # println("s4\n")
        cBlocks = []
        cEBlocks = []
        for i=1:nBlock
            push!(cBlocks, cBlock)
            # println("s4_1\n")
            push!(cEBlocks, put(nBitA, Tuple((nBitA-i*rBit-vBit+1):(nBitA-(i-1)*rBit),)=>cBlockHead))
            # println("s4_2\n")
        end
        cBlocks[1] = cBlockHead
        circuit = chain(nBit, cBlocks)
        cExtend = chain(nBitA, cEBlocks)
    end

    (circuit, cExtend)
end

function MPSbuilder(nBitA::Int64, vBit::Int64, rBit::Int64, blockT::String)
    par2nd = setMPSpar(nBitA, vBit, rBit)
    nBlock = par2nd.nBlock
    nBit = par2nd.nBit

    if blockT == "CS" # MPS blocks for 1D cluster state.
        if vBit !=1 || rBit !=1
            println("Error: The nBit=vBit+rBit of cluster state MPS blocks should be 2!\n")
            return 0
        end
        cBlocks = []
        push!(cBlocks, chain(nBit, repeat(nBit,H,(2,1)), control(nBit, 1, nBit=>Z)))
        for i =2:nBlock
            push!(cBlocks, chain(nBit,put(nBit, (2,1)=>ConstGate.SWAP), put(nBit, 1=>H), control(nBit, 1, nBit=>Z)))
        end
        circuit = chain(nBit,cBlocks)
        cExtend = chain(nBitA, vcat([repeat(nBitA,H,Tuple(nBitA:-1:1,))], [control(nBitA, i-1, i=>Z) for i=nBitA:-1:2]))
    end

    (circuit, cExtend)
end

end


# script to confirm all AprilTag sightings are consistent

using YAML, JLD, HDF5
using AprilTags
using Images, ImageView, ImageDraw
using CoordinateTransformations, Rotations, StaticArrays
using MeshCat
using GeometryTypes

# temporarily using Python / OpenCV for solvePnP
using PyCall
@pyimport numpy as np
@pyimport cv2

# load utility function for racecar
include(joinpath(Pkg.dir("Caesar"),"examples","wheeled","racecar","racecarUtils.jl"))
include(joinpath(Pkg.dir("Caesar"),"examples","wheeled","racecar","cameraUtils.jl"))
include(joinpath(Pkg.dir("Caesar"),"examples","wheeled","racecar","visualizationUtils.jl"))

# load camera calibration
cfg = loadConfig()
cw, ch = cfg[:intrinsics][:cx], cfg[:intrinsics][:cy]
fx = fy = cfg[:intrinsics][:fx]
camK = [fx 0 cw; 0 fy ch; 0 0 1]
tagsize = 0.172
# k1,k2 = cfg[:intrinsics][:k1], cfg[:intrinsics][:k2] # makes it worse
k1, k2 = 0.0, 0.0

# tag extrinsic rotation
Rx = RotX(-pi/2)
Rz = RotZ(-pi/2)
bTc= LinearMap(Rz) ∘ LinearMap(Rx)

# image loading directory
datafolder = ENV["HOME"]*"/data/racecar/straightrun3/"
imgfolder = "images"

# Figure export folder
currdirtime = now()
imgdir = joinpath(ENV["HOME"], "Pictures", "testimgs", "$(currdirtime)")
mkdir(imgdir)
mkdir(imgdir*"/tags")

# process images
camlookup = prepCamLookup(175:5:370)
IMGS, TAGS = detectTagsViaCamLookup(camlookup, datafolder*imgfolder, imgdir)
tag_bag = prepTagBag(TAGS)

# save the tag bag file for future use
@save imgdir*"/tag_det_per_pose.jld" tag_bag

# 3D visualization
vis = Visualizer()
open(vis)

setobject!(vis["home"], Triad(1.0))

currtag = ""
bTt = LinearMap(Translation(0.0,0,0) ∘ Quat(1.0,0,0,0))

psid = 5
tags = tag_bag[psid]

# tidx = 1
# tag = tags[tidx]
for (tidx, tag) in tags

currtag = "tag$(tidx)"
Q = Quat(tag[:quat]...)
cTt = (Translation(tag[:pos]...) ∘ LinearMap(Q))
bTt = bTc ∘ cTt
drawTagDetection(vis, currtag, Q, tag[:pos], bTc, tag[:bP2t], posename=Symbol("x$psid"))

end

#


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# 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.

@reexport module QuasiGradient

# Standard library
using LinearAlgebra
using SparseArrays
using Distributed
using Printf

# External libraries
using Parameters
using JuMP
using MathOptInterface
using StochasticPrograms
using StochasticPrograms: AcceptableTermination
using StochasticPrograms: UnspecifiedInstantiation, StageDecompositionStructure, AbstractScenarioProblems, ScenarioProblems, DistributedScenarioProblems, DecisionChannel
using StochasticPrograms: AbstractExecution, Serial, Synchronous, Asynchronous
using StochasticPrograms: AbstractStructuredOptimizer, set_master_optimizer!, set_subproblem_optimizer!
using StochasticPrograms: DecisionMap, get_decisions, all_decisions, set_decision!, set_stage!, SingleDecisionSet, NoSpecifiedConstraint
using StochasticPrograms: all_decisions, all_known_decisions, remove_decision!, update_known_decisions!, KnownValuesChange, NotTaken, Known, KnownDecision
using StochasticPrograms: add_subscript
using StochasticPrograms: AbstractPenaltyTerm, Quadratic, InfNorm, ManhattanNorm, initialize_penaltyterm!, update_penaltyterm!, remove_penalty!, remove_penalty_variables!, remove_penalty_constraints!
using ProgressMeter
using ProgressMeter: AbstractProgress

import Base: show, put!, wait, isready, take!, fetch, zero, +, length, size
import StochasticPrograms: supports_structure, num_iterations, default_structure, check_loadable, load_structure!, restore_structure!, optimize!, optimizer_name, master_optimizer, subproblem_optimizer, num_subproblems, remove_penalty_variables!, remove_penalty_constraints!
using ProgressMeter: Progress

const MOI = MathOptInterface
const MOIU = MOI.Utilities
const CI = MOI.ConstraintIndex

export
    AbstractQuasiGradientAttribute,
    QuasiGradientAlgorithm,
    num_iterations,
    SubProblems,
    Unaltered,
    Smoothed,
    StepSize,
    StepParameter,
    RawStepParameter,
    set_step_attribute,
    set_step_attributes,
    Constant,
    Diminishing,
    Polyak,
    BB,
    Prox,
    ProxParameter,
    RawProxParameter,
    set_prox_attribute,
    set_prox_attributes,
    NoProx,
    Polyhedron,
    AndersonAcceleration,
    DryFriction,
    Nesterov,
    Termination,
    TerminationParameter,
    set_termination_attribute,
    set_termination_attributes,
    AfterMaximumIterations,
    AtObjectiveThreshold,
    AtGradientThreshold

# Include files
include("types/types.jl")
include("step/step.jl")
include("prox/prox.jl")
include("termination/termination.jl")
include("execution/execution.jl")
include("solver.jl")
include("MOI_wrapper.jl")

end # module


struct SettlementPhysical <: SettlementType end
struct SettlementCash <: SettlementType end

mutable struct SwaptionResults
  value::Float64
  additionalResults::Dict{String, Float64}
end

SwaptionResults() = SwaptionResults(0.0, Dict("spreadCorrection" => 0.0, "strike" => 0.0, "atmForward" => 0.0, "annuity" => 0.0, "swapLength" => 0.0, "stdDev" => 0.0,
                                              "vega" => 0.0))

function reset!(results::SwaptionResults)
  results.value = 0.0
  for k in keys(results.additionalResults)
    results.additionalResults[k] = 0.0
  end

  return results
end

mutable struct Swaption{E <: Exercise, S <: SettlementType, P <: PricingEngine, VS <: VanillaSwap} <: Option{E}
  lazyMixin::LazyMixin
  swap::VS
  exercise::E
  delivery::S
  results::SwaptionResults
  pricingEngine::P

  # function call{E, S}(::Type{Swaption}, lz::LazyMixin, swap::VanillaSwap, exercise::E, delivery::S, results::SwaptionResults)
  #   new{E, S, PricingEngine}(lz, swap, exercise, delivery, results)
  # end

  # function call{E, S, P}(::Type{Swaption}, lz::LazyMixin, swap::VanillaSwap, exercise::E, delivery::S, results::SwaptionResults, pricingEngine::P = NullSwaptionEngine())
  #   new{E, S, P}(lz, swap, exercise, delivery, results, pricingEngine)
  # end
  # Swaption{E, S}(lz::LazyMixin, swap::VanillaSwap, exercise::E, delivery::S, results::SwaptionResults, pricingEngine::P = NullSwaptionEngine()) =
  #   new{E, S, P}(lz, swap, exercise, delivery, results, pricingEngine)
end

Swaption(swap::VS, exercise::E) where {VS <: VanillaSwap, E <: Exercise} = Swaption{E, SettlementPhysical, NullSwaptionEngine, VS}(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults(), NullSwaptionEngine())

Swaption(swap::VS, exercise::E, pe::P) where {E <: Exercise, P <: PricingEngine, VS <: VanillaSwap} = Swaption{E, SettlementPhysical, P, VS}(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults(), pe)

mutable struct NonstandardSwaption{E <: Exercise, S <: SettlementType, P <: PricingEngine, NS <: NonstandardSwap} <: Option{E}
  lazyMixin::LazyMixin
  swap::NS
  exercise::E
  delivery::S
  results::SwaptionResults
  pricingEngine::P

  # function call{E, S, P}(::Type{NonstandardSwaption}, lz::LazyMixin, swap::NonstandardSwap, exercise::E, delivery::S, results::SwaptionResults, pricingEngine::P = NullSwaptionEngine())
  #   new{E, S, P}(lz, swap, exercise, delivery, results, pricingEngine)
  # end
end

NonstandardSwaption(swap::NS, exercise::E) where {E <: Exercise, NS <: NonstandardSwap} = NonstandardSwaption{E, SettlementPhysical, NullSwaptionEngine, NS}(LazyMixin(), swap, exercise,
                                                                      SettlementPhysical(), SwaptionResults(), NullSwaptionEngine())
NonstandardSwaption(swap::NS, exercise::E, pe::P) where {E <: Exercise, P <: PricingEngine, NS <: NonstandardSwap} =
    NonstandardSwaption{E, SettlementPhysical, P, NS}(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults(), pe)

function perform_calculations!(swaption::Union{Swaption, NonstandardSwaption})
  reset!(swaption.results)
  _calculate!(swaption.pricingEngine, swaption)

  return swaption
end

function npv(swaption::Union{Swaption, NonstandardSwaption})
  calculate!(swaption)

  return swaption.results.value
end

function calibration_basket(swaption::NonstandardSwaption, swaptionEngine::PricingEngine, swapIndex::SwapIndex, swaptionVol::SwaptionVolatilityStructure,
                            basketType::CalibrationBasketType)
  return calibration_basket(swaptionEngine, swaption, swapIndex, swaptionVol, basketType)
end

function clone(swaption::Swaption, pe::PricingEngine = swaption.pricingEngine)
  lazyMixin, res = pe == swaption.pricingEngine ? (swaption.lazyMixin, swaption.results) : (LazyMixin(), SwaptionResults())

  newSwaption = Swaption(lazyMixin, swaption.swap, swaption.exercise, swaption.delivery, res, pe)

  return newSwaption
end

get_pricing_engine_type(::Swaption{E, S, P, NS}) where {E, S, P, NS} = P


# ============================================================
# General Structs
# ============================================================

"""
$(TYPEDEF)

Computational setup

# Fields

$(FIELDS)
"""
@kwdef struct Setup{S,I<:Integer,E<:AV,F<:Real,G<:Real} <: AbstractSetup
    matter::S = "gas"
    case::S = "dev"
    space::S = "1d0f0v"
    flux::S = "kfvs"
    collision::S = "bgk"
    nSpecies::I = 1
    interpOrder::I = 2
    limiter::S = "vanleer"
    boundary::E = ["fix", "fix"]
    cfl::F = 0.5
    maxTime::G = 0.1
end

function Setup(
    matter,
    case,
    space,
    flux,
    collision,
    ns,
    order,
    limiter,
    bc::T,
    cfl,
    time,
) where {T<:Union{AbstractString,Symbol}}
    boundary = begin
        if parse(Int, space[1]) == 1
            [bc, bc]
        elseif parse(Int, space[1]) == 2
            [bc, bc, bc, bc]
        end
    end

    return Setup{typeof(matter),typeof(ns),typeof(boundary),typeof(cfl),typeof(time)}(
        matter,
        case,
        space,
        flux,
        collision,
        ns,
        order,
        limiter,
        boundary,
        cfl,
        time,
    )
end


function FieldOfViewDetector(arg0::PVCoordinatesProvider, arg1::FieldOfView)
    return FieldOfViewDetector((PVCoordinatesProvider, FieldOfView), arg0, arg1)
end

function FieldOfViewDetector(arg0::PVCoordinatesProvider, arg1::jdouble, arg2::VisibilityTrigger, arg3::FieldOfView)
    return FieldOfViewDetector((PVCoordinatesProvider, jdouble, VisibilityTrigger, FieldOfView), arg0, arg1, arg2, arg3)
end

function g(obj::FieldOfViewDetector, arg0::SpacecraftState)
    return jcall(obj, "g", jdouble, (SpacecraftState,), arg0)
end

function get_fov(obj::FieldOfViewDetector)
    return jcall(obj, "getFOV", FieldOfView, ())
end

function get_pv_target(obj::FieldOfViewDetector)
    return jcall(obj, "getPVTarget", PVCoordinatesProvider, ())
end


module TurbulenceConvection

import ClimaCore
import LinearAlgebra
import DocStringExtensions
import StaticArrays
import StatsBase
import NCDatasets
import JSON
import Dierckx
import RootSolvers
import Statistics
import LambertW
import Thermodynamics
import Distributions
import FastGaussQuadrature
import CLIMAParameters
import OrdinaryDiffEq
import CloudMicrophysics
import SurfaceFluxes
import UnPack

const ODE = OrdinaryDiffEq
const CC = ClimaCore
const CCO = CC.Operators
const NC = NCDatasets
const SA = StaticArrays

const TD = Thermodynamics
const SF = SurfaceFluxes
const UF = SF.UniversalFunctions
const RS = RootSolvers

const CM = CloudMicrophysics
const CM0 = CloudMicrophysics.Microphysics_0M
const CM1 = CloudMicrophysics.Microphysics_1M
const liq_type = CM1.LiquidType()
const ice_type = CM1.IceType()
const rain_type = CM1.RainType()
const snow_type = CM1.SnowType()

const CP = CLIMAParameters
const CPP = CP.Planet
const CPSGS = CP.SubgridScale
const APS = CP.AbstractEarthParameterSet

const CPMP = CP.Atmos.Microphysics
const CPEDMF = CP.Atmos.EDMF

# For dispatching to inherited class
struct BaseCase end

up_sum(vals::AbstractArray) = reshape(sum(vals; dims = 1), size(vals, 2))

function parse_namelist(namelist, keys...; default = nothing, valid_options = nothing)
    param = namelist
    for k in keys
        if haskey(param, k)
            param = param[k]
            if valid_options ≠ nothing && !(param isa Dict)
                @assert param in valid_options
            end
        else
            if default == nothing
                error("No default value given for parameter ($(join(keys, ", "))).")
            else
                @info "Using default value, $default, for parameter ($(join(keys, ", ")))."
                return default
            end
        end
    end
    return param
end

include("ClimaParams.jl")
import .ClimaParams
const ICP = ClimaParams # internal clima parameters

#=
    debug_state(state, code_location::String)

A simple function for debugging the entire state,
specifically for when quantities that should remain
positive-definite become negative.

=#
function debug_state(state, code_location::String)
    prog_gm = center_prog_grid_mean(state)
    aux_gm = center_aux_grid_mean(state)
    prog_gm_f = face_prog_grid_mean(state)
    aux_gm_f = face_aux_grid_mean(state)

    prog_up = center_prog_updrafts(state)
    aux_up = center_aux_updrafts(state)
    prog_up_f = face_prog_updrafts(state)
    aux_up_f = face_aux_updrafts(state)

    prog_en = center_prog_environment(state)
    aux_en = center_aux_environment(state)
    aux_en_f = face_aux_environment(state)

    ######
    ###### Positive-definite variables
    ######

    vars_positive = [
        vec(prog_gm.θ_liq_ice),
        vec(prog_gm_f.w),
        vec(prog_up[1].ρarea),
        vec(prog_up[1].ρaθ_liq_ice),
        vec(prog_up_f[1].ρaw),
        vec(aux_en.area),
        vec(aux_en.θ_liq_ice),
    ]
    vars = vars_positive
    vars_conds = map(v -> any(v .< 0), vars)

    if any(vars_conds)
        @show code_location
        for (i, vc, v) in zip(1:length(vars), vars_conds, vars)
            vc || continue
            @show i, v
        end
        @show vars_conds
        error("Negative state for positive-definite field(s)")
    end

    ######
    ###### All listed variables
    ######
    vars = vars_positive
    vars_conds = map(v -> any(isnan.(v)) || any(isinf.(v)), vars)

    if any(vars_conds)
        @show code_location
        for (i, vc, v) in zip(1:length(vars), vars_conds, vars)
            vc || continue
            @show i, v
        end
        @show vars_conds
        error("Nan/Inf state for field(s)")
    end
end

include("parameters.jl")

include("Grid.jl")
include("dycore_api.jl")
include("NetCDFIO.jl")
include("diagnostics.jl")
include("Fields.jl")
include("ReferenceState.jl")
include("types.jl")
include("name_aliases.jl")
include("Operators.jl")

include("microphysics_coupling.jl")
include("turbulence_functions.jl")
include("utility_functions.jl")
include("TimeStepping.jl")
include("variables.jl")
include("EDMF_Precipitation.jl")
include("EDMF_Environment.jl")
include("EDMF_Updrafts.jl")
include("stochastic_closures.jl")
include("update_aux.jl")
include("Turbulence_PrognosticTKE.jl")
include("Forcing.jl")
include("Radiation.jl")
include("Surface.jl")
include("io.jl")
include("thermodynamics.jl")
include("surface_functions.jl")
include("closures/perturbation_pressure.jl")
include("closures/entr_detr.jl")
include("closures/nondimensional_exchange_functions.jl")
include("closures/mixing_length.jl")
include("closures/buoyancy_gradients.jl")

end


using IcanProj
using JLD2

"""
Load input variables directly from the forcing files.
"""
function load_from_forcings(watershed, resolution, res_vic, opt)

    # Path to model forcings for the given resolution

    if resolution == "1km"
        path_forcing = joinpath(opt["base_folder"])
    else
        path_forcing = joinpath(opt["target_folder"], watershed, resolution)
    end

    # Read one forcing file to get the size

    lat = @sprintf("%0.5f", res_vic.lat[1])
    lon = @sprintf("%0.5f", res_vic.lon[1])

    file_src = joinpath(path_forcing, "forcing/data_$(lat)_$(lon)")

    prec, tmin, tmax, wind = read_vic_forcing(file_src)

    # Allocate arrays

    prec_array = zeros(Float64, length(prec), length(res_vic.lat))
    tmin_array = zeros(Float64, length(prec), length(res_vic.lat))
    tmax_array = zeros(Float64, length(prec), length(res_vic.lat))
    wind_array = zeros(Float64, length(prec), length(res_vic.lat))

    # Read rest of forcing files

    for i in 1:length(res_vic.lat)

        lat = @sprintf("%0.5f", res_vic.lat[i])
        lon = @sprintf("%0.5f", res_vic.lon[i])

        file_src = joinpath(path_forcing, "forcing/data_$(lat)_$(lon)")

        prec, tmin, tmax, wind = read_vic_forcing(file_src)

        prec_array[:, i] = prec
        tmin_array[:, i] = tmin
        tmax_array[:, i] = tmax
        wind_array[:, i] = wind

    end

    return prec_array, tmin_array, tmax_array, wind_array

end

# Get settings

opt = get_options()

resolutions = opt["resolutions"]

watersheds = opt["stat_sel"]

# Process all results

for watershed in watersheds

    for resolution in resolutions

        # Load simulation results

        file = joinpath(opt["eval_folder"], watershed, "res_$(resolution).jld2")

        @load file res_vic

        # Load forcing files

        prec_input, tmin_input, tmax_input, wind_input = load_from_forcings(watershed, resolution, res_vic, opt)
        info("$(watershed) $(resolution)")

        # Compare precipitation

        iprec = find(res_vic.var_names .== "prec")

        prec_output = res_vic.data_all[:, iprec[1], :]

        check = all(round.(prec_input,2) .== round.(prec_output,2))

        info("Precipitation $(check) $(mean(prec_output))")

        # Compare air temperature (not useful to check exact match due to interpolation within mtclim)

        iair = find(res_vic.var_names .== "air_temp")

        air_output = res_vic.data_all[:, iair[1], :]

        air_input = (tmin_input + tmax_input)/2

        info("Air temperature $(round(mean(abs.(air_output-air_input)),1))")

        # Compare wind speed

        iwind = find(res_vic.var_names .== "wind")

        wind_output = res_vic.data_all[:, iwind[1], :]

        check = all(round.(wind_input,2) .== round.(wind_output,2))

        info("Wind $(watershed) $(mean(wind_output))\n\n\n")

    end

end


module GalsimBenchmark
using DataFrames
import FITSIO
import StaticArrays
import WCS

import Celeste: Model, DeterministicVI, ParallelRun, Infer
import Celeste.Model: CatalogEntry
import Celeste.ParallelRun: one_node_single_infer, one_node_joint_infer

const GALSIM_BENCHMARK_DIR = joinpath(Pkg.dir("Celeste"), "benchmark", "galsim")
const FILENAME = "output/galsim_test_images.fits"

function make_psf(psf_sigma_px)
    alphaBar = [1.; 0.]
    xiBar = [0.; 0.]
    tauBar = [psf_sigma_px^2 0.; 0. psf_sigma_px^2]
    [
        Model.PsfComponent(
            alphaBar[k],
            StaticArrays.SVector{2, Float64}(xiBar),
            StaticArrays.SMatrix{2, 2, Float64, 4}(tauBar)
        )
        for k in 1:2
    ]
end

immutable FitsExtension
    pixels::Matrix{Float32}
    header::FITSIO.FITSHeader
end

function read_fits(filename; read_sdss_psf=false)
    println("Reading $filename...")
    @assert isfile(filename)
    fits = FITSIO.FITS(filename)
    println("Found $(length(fits)) extensions.")

    extensions::Vector{FitsExtension} = []
    for extension in fits
        push!(extensions, FitsExtension(read(extension), FITSIO.read_header(extension)))
    end

    # assume WCS same for each extension
    wcs = WCS.from_header(FITSIO.read_header(fits[1], String))[1]

    close(fits)
    extensions, wcs
end

function make_images(band_pixels, psf, wcs, epsilon, iota)
    # assume dimensions equal for all images
    height, width = size(band_pixels[1])
    [
        Model.Image(
            height,
            width,
            band_pixels[band],
            band,
            wcs,
            psf,
            0, # SDSS run
            0, # SDSS camcol
            0, # SDSS field
            fill(epsilon, height, width),
            fill(iota, height),
            Model.RawPSF(Array(Float64, 0, 0), 0, 0, Array(Float64, 0, 0, 0)),
        )
        for band in 1:5
    ]
end

function typical_band_relative_intensities(is_star::Bool)
    source_type_index = is_star ? 1 : 2
    prior_parameters::Model.PriorParams = Model.load_prior()
    # Band relative intensities are a mixture of lognormals. Which mixture component has the most
    # weight?
    dominant_component = indmax(prior_parameters.k[:, source_type_index])
    # What are the most typical log relative intensities for that component?
    inter_band_ratios = exp(
        prior_parameters.c_mean[:, dominant_component, source_type_index]
        - diag(prior_parameters.c_cov[:, :, dominant_component, source_type_index])
    )
    Float64[
        1 / inter_band_ratios[2] / inter_band_ratios[1],
        1 / inter_band_ratios[2],
        1,
        inter_band_ratios[3],
        inter_band_ratios[3] * inter_band_ratios[4],
    ]
end

function typical_reference_brightness(is_star::Bool)
    source_type_index = is_star ? 1 : 2
    prior_parameters::Model.PriorParams = Model.load_prior()
    exp(
        prior_parameters.r_mean[source_type_index]
        - prior_parameters.r_var[source_type_index]
    )
end

# Since we're considering elliptical galaxy shapes, angle is only meaningful up to rotations of 180
# deg. This finds an equivalent angle in [0, 180) deg.
function canonical_angle(params)
    angle_radians = params[Model.ids.e_angle]
    while angle_radians < 0
        angle_radians += pi
    end
    while angle_radians > pi
        angle_radians -= pi
    end
    angle_radians
end

const BENCHMARK_PARAMETER_LABELS = String[
    "X center (world coords)",
    "Y center (world coords)",
    "Minor/major axis ratio",
    "Angle (degrees)",
    "Half-light radius (pixels)",
    "Brightness (nMgy)",
    "Color band 1-2 ratio",
    "Color band 2-3 ratio",
    "Color band 3-4 ratio",
    "Color band 4-5 ratio",
    "Probability of galaxy",
]

function get_field(header::FITSIO.FITSHeader, key::String)
    if haskey(header, key)
        header[key]
    else
        NA
    end
end

function actual_values(ids, star_galaxy_index, params)
    Float64[
        params[ids.u[1]],
        params[ids.u[2]],
        params[ids.e_axis],
        canonical_angle(params) * 180 / pi,
        params[ids.e_scale] * sqrt(params[ids.e_axis]),
        exp(params[ids.r1[star_galaxy_index]]),
        exp(params[ids.c1[1, star_galaxy_index]]),
        exp(params[ids.c1[2, star_galaxy_index]]),
        exp(params[ids.c1[3, star_galaxy_index]]),
        exp(params[ids.c1[4, star_galaxy_index]]),
        params[ids.a[2]],
    ]
end

function get_expected_dataframe(header)
    DataFrame(
        label=fill(header["CL_DESCR"], length(BENCHMARK_PARAMETER_LABELS)),
        field=BENCHMARK_PARAMETER_LABELS,
        expected=Any[
            get_field(header, "CL_X1"),
            get_field(header, "CL_Y1"),
            get_field(header, "CL_RTIO1"),
            get_field(header, "CL_ANGL1"),
            get_field(header, "CL_RADP1"),
            get_field(header, "CL_FLUX1"),
            get_field(header, "CL_C12_1"),
            get_field(header, "CL_C23_1"),
            get_field(header, "CL_C34_1"),
            get_field(header, "CL_C45_1"),
            header["CL_TYPE1"] == "star" ? 0 : 1,
        ])
end

function benchmark_comparison_data(single_infer_params, joint_infer_params, header)
    ids = Model.ids
    star_galaxy_index = header["CL_TYPE1"] == "star" ? 1 : 2
    comparison_dataframe = get_expected_dataframe(header)
    comparison_dataframe[:single_infer_actual] =
        actual_values(ids, star_galaxy_index, single_infer_params)
    comparison_dataframe[:joint_infer_actual] =
        actual_values(ids, star_galaxy_index, joint_infer_params)
    comparison_dataframe
end

function assert_counts_match_expected_flux(band_pixels::Vector{Matrix{Float32}},
                                           header::FITSIO.FITSHeader,
                                           iota::Float64)
    if !header["CL_NOISE"]
        expected_flux_nmgy = prod(size(band_pixels[3])) * header["CL_SKY"]
        for source_index in 1:header["CL_NSRC"]
            expected_flux_nmgy += header[string("CL_FLUX", source_index)]
        end
        expected_flux_counts = expected_flux_nmgy * iota
        @assert abs(sum(band_pixels[3]) - expected_flux_counts) / expected_flux_counts < 1e-3
    end
end

function make_catalog_entry(x_position_world_coords, y_position_world_coords)
    CatalogEntry(
        [x_position_world_coords, y_position_world_coords],
        false, # is_star
        # sample_star_fluxes
        typical_band_relative_intensities(true) .* typical_reference_brightness(true),
        # sample_galaxy_fluxes
        typical_band_relative_intensities(false) .* typical_reference_brightness(false),
        0.1, # gal_frac_dev
        0.7, # gal_ab
        pi / 4, # gal_angle
        4., # gal_scale
        "sample", # objid
        0, # thing_id
    )
end

function make_catalog(header::FITSIO.FITSHeader)
    catalog = CatalogEntry[]
    num_sources = header["CL_NSRC"]
    for source_index in 1:num_sources
        if num_sources == 1
            initial_position = [0.005335 for i in 1:2] # center of image
        else
            initial_position = [
                header[string("CL_X", source_index)],
                header[string("CL_Y", source_index)],
            ]
        end
        push!(catalog, make_catalog_entry(initial_position[1], initial_position[2]))
    end
    catalog
end

# Returns a data frame with one row for each test case and  parameter name, with columns
# * label: test case name
# * field: human-readable parameter name
# * expected: "ground truth" expected value
# * single_infer_actual: inferred MAP value or posterior median from single-source inference
# * joint_infer_actual: ditto, for multi-source joint inference
function main(; test_case_names=String[], print_fn=println,
              infer_source_callback=DeterministicVI.infer_source)
    all_benchmark_data = []
    extensions, wcs = read_fits(joinpath(GALSIM_BENCHMARK_DIR, FILENAME))
    @assert length(extensions) % 5 == 0 # one extension per band for each test case

    for test_case_index in 1:div(length(extensions), 5)
        first_band_index = (test_case_index - 1) * 5 + 1
        header = extensions[first_band_index].header
        this_test_case_name = header["CL_DESCR"]
        if !isempty(test_case_names) && !in(this_test_case_name, test_case_names)
            continue
        end
        println("Running test case '$this_test_case_name'")
        iota = header["CL_IOTA"]
        psf = make_psf(header["CL_SIGMA"])
        n_sources = header["CL_NSRC"]

        band_pixels = [
            extensions[index].pixels for index in first_band_index:(first_band_index + 4)
        ]
        assert_counts_match_expected_flux(band_pixels, header, iota)

        images = make_images(band_pixels, psf, wcs, header["CL_SKY"], iota)
        catalog = make_catalog(header)

        # we're only scoring one object per image
        target_sources = [1,]
        neighbor_map = Infer.find_neighbors(target_sources, catalog, images)

        target_sources_joint = collect(1:n_sources)
        neighbor_map_joint = Infer.find_neighbors(target_sources_joint, catalog, images)

        single_results = one_node_single_infer(
            catalog, target_sources, neighbor_map, images;
            infer_source_callback=infer_source_callback)
        joint_results = one_node_joint_infer(catalog, target_sources_joint,
                                             neighbor_map_joint, images)

        benchmark_data = benchmark_comparison_data(single_results[1].vs,
                                                   joint_results[1].vs,
                                                   header)
        print_fn(repr(benchmark_data))
        push!(all_benchmark_data, benchmark_data)
    end

    full_data = vcat(all_benchmark_data...)
    print_fn(repr(full_data[!isna(full_data[:expected]), :]))
    full_data
end

end # module GalsimBenchmark


using JLD
using HTTP
using ..Config: Configuration
using ..Serde: serialize_payload, deserialize_payload, unpack
using ..Fed: STATS, update_stats!, initialize_stats

"""
    CentralNode(host, port, weights, strategy, evaluate, config)

Orchestrate the entire federated learning process. It handles the clients
registration and requests them to train on their local data.
"""
struct CentralNode
    # networking
    host::String
    port::Int
    client_manager::ClientManager

    # machine learning
    weights::Vector{Float32}
    strategy::Function

    # proxy
    evaluate::Function

    config::Configuration

    CentralNode(
        host::String,
        port::Int,
        weights::Vector{Float32},
        strategy::Function,
        evaluate::Function,
        config::Configuration
    ) = new(host, port, ClientManager(), weights, strategy, evaluate, config)
end

function get_store_url(url)
    protocol, url, port = split(url, ":")
    return protocol * ":" * url * ":" * string(parse(Int, port) + 1010)
end

"""
    fit(central_node::CentralNode)

Fit the global model by orchestrating the federated training.
"""
function fit(central_node::CentralNode)
    global_weights = central_node.weights

    initialize_stats(central_node.config, length(global_weights))

    jldopen("data/request.jld", "w") do req_file
        jldopen("data/response.jld", "w") do res_file
            for round_num in 1:central_node.config.base.num_comm_rounds
                @info "Communication round $round_num"

                # uplink communication
                # 1. chose the clients subset for the round.
                clients = sample_clients(central_node.client_manager, central_node.config.base.fraction_clients)
                # 2. serialize the global model.
                payload = serialize_payload(central_node.config.payload_serde, global_weights)
                # 3. ask clients to train on the global model.
                clients_payload = fit_clients(central_node.config.base.fit_node, clients, payload)

                write(req_file, "$round_num", global_weights)
                client_w = map(p -> p.data, unpack.(clients_payload))
                write(res_file, "$round_num", client_w)

                # downlink communication
                # 4. deserialize the results.
                round_weights = [
                    deserialize_payload(central_node.config.payload_serde, payload, get_store_url(clients[i]))
                    for (i, payload) in enumerate(clients_payload)
                ]

                # update, convergeance tracking, ...
                # 5. update the global model.
                global_weights = central_node.strategy(round_weights)
                # 6. evaluate global model.
                loss, acc = central_node.evaluate(global_weights)
                @info "loss: $loss, acc: $acc"

                # 7. record statistics.
                update_stats!(STATS, round_num, loss, acc)

                save("stats.jld", "stats", STATS)
            end
        end
    end
end

"""
    fit_clients(endpoint::String, clients::Vector{String}, payload::Vector{UInt8})

Ask asynchronously each client from `clients` to train on their local data by
using the model weights contained in the `payload`. Return a `Vector` where each
element is the serialized weights of one client.
"""
function fit_clients(
    endpoint::String,
    clients::Vector{String},
    payload::Vector{UInt8}
)::Vector{Vector{UInt8}}
    # asynchronously ask the clients subset to train, and wait on the results.
    tasks = [@async fit_client(client * endpoint, payload) for client in clients]
    wait.(tasks)

    # extract the results.
    round_weights = map(task -> task.result, tasks)

    return round_weights
end

"""
    fit_client(url::String, payload::Vector{UInt8})

Ask one `client` to train on its local data, by using the model weights
contained in the `payload`. Return the serialized updated `weights`.
"""
function fit_client(
    url::String,
    payload::Vector{UInt8}
)::Vector{UInt8}
    response = HTTP.request("POST", url, [], payload)
    return response.body
end


import Nemo.sub!, Base.gcd
export induce_rational_reconstruction, induce_crt, root, roots,
       is_monic, radical_extension, is_radical_extension,
       is_kummer_extension, cyclotomic_field, wildanger_field,
       compositum

add_verbose_scope(:PolyFactor)
add_assert_scope(:PolyFactor)
add_verbose_scope(:CompactPresentation)
add_assert_scope(:CompactPresentation)

if Int==Int32
  global const p_start = next_prime(2^30)
else
  global const p_start = next_prime(2^60)
end

include("NfAbs/NfAbs.jl")
include("NfAbs/Conjugates.jl")
include("NfAbs/ComplexEmbeddings.jl")
include("NfAbs/CompactRepresentation.jl")
include("NfAbs/Elem.jl")
include("NfAbs/NonSimple.jl")
include("NfAbs/Poly.jl")
include("NfAbs/Simplify.jl")
include("NfAbs/DiscriminantBounds.jl")
include("NfAbs/NormRelation.jl")
include("NfAbs/PolyFact.jl")
include("NfAbs/MPolyGcd.jl")
include("NfAbs/MPolyFactor.jl")
include("NfAbs/MPolyAbsFact.jl")
include("NfAbs/ConjugatesNS.jl")
include("NfAbs/ComplexEmbeddingsNS.jl")


using NeuroIO, MAT, BinDeps, Base.Test

# Unfortunately, these test files come from Black Rock Microsystems and
# are assumed not to be redistributable
testdir = joinpath(dirname(@__FILE__), "NSS Data")
if !isdir(testdir)
    parentdir = dirname(@__FILE__)
    bundlezip = joinpath(parentdir, "NSS Data.zip")
    run(download_cmd("http://cl.ly/1B3y1C0U380G/download/NSS%20Data.zip", bundlezip))
    cd(parentdir) do
        run(`unzip $bundlezip`)
    end
end

if !isfile(joinpath(testdir, "nev.mat"))
    using MATLAB
    npmkdir = joinpath(dirname(@__FILE__), "NPMK")
    if !isdir(npmkdir)
        cd(dirname(@__FILE__)) do
            run(`git clone https://github.com/BlackrockMicrosystems/NPMK.git`)
        end
    end
    @eval mat"""
    addpath(genpath($(joinpath(npmkdir, "NPMK"))))
    cd($testdir);
    NEV = openNEV('./SampleNSSData.nev');
    save('nev.mat', 'NEV');
    NS5 = openNSx('./SampleNSSData.ns5');
    save('ns5.mat', 'NS5');
    """
end

nev = readnev(open(joinpath(testdir, "SampleNSSData.nev")))
nev_mat = matread(joinpath(testdir, "nev.mat"))["NEV"]
@test nev.header.bytes_in_headers == nev_mat["MetaTags"]["HeaderOffset"]
@test nev.header.bytes_in_data_packets == nev_mat["MetaTags"]["PacketBytes"]
@test nev.header.time_resolution_of_time_stamps == nev_mat["MetaTags"]["TimeRes"]
@test nev.header.time_resolution_of_samples == nev_mat["MetaTags"]["SampleRes"]
@test nev.header.application_to_create_file == nev_mat["MetaTags"]["Application"]

info = nev_mat["ElectrodesInfo"]
spikes = nev_mat["Data"]["Spikes"]
for i = 1:length(info["ElectrodeID"])
    id = info["ElectrodeID"][i]
    ch = nev.spike_channels[id]
    @test ch.physical_connector == info["ConnectorBank"][i][1]
    @test ch.connector_pin == info["ConnectorPin"][i]
    # What is this?
    # @test ch.digitization_factor == info["DigitalFactor"][i]
    @test ch.energy_threshold == info["EnergyThreshold"][i]
    @test ch.high_threshold == info["HighThreshold"][i]
    @test ch.low_threshold == info["LowThreshold"][i]
    @test ch.number_of_sorted_units == info["Units"][i]
    @test ch.label == rstrip(string(info["ElectrodeLabel"][i]...), '\0')
    @test ch.high_pass.cutoff == info["HighFreqCorner"][i]
    @test ch.high_pass.order == info["HighFreqOrder"][i]
    @test ch.high_pass.filter_type == info["HighFilterType"][i]
    @test ch.low_pass.cutoff == info["LowFreqCorner"][i]
    @test ch.low_pass.order == info["LowFreqOrder"][i]
    @test ch.low_pass.filter_type == info["LowFilterType"][i]

    spike_index = find(spikes["Electrode"] .== id)
    @test NeuroIO.times(ch) == spikes["TimeStamp"][spike_index]/nev.header.time_resolution_of_time_stamps
    @test ch.unit_numbers == spikes["Unit"][spike_index]
    @test NeuroIO.data(ch) == spikes["Waveform"][:, spike_index]
end

ns5 = readnsx(open(joinpath(testdir, "SampleNSSData.ns5")))
ns5_mat = matread(joinpath(testdir, "ns5.mat"))["NS5"]
@test ns5.header.label == rstrip(string(ns5_mat["MetaTags"]["SamplingLabel"]...), '\0')
@test ns5.header.period == 30000/ns5_mat["MetaTags"]["SamplingFreq"]
@test ns5.header.time_resolution_of_time_stamps == ns5_mat["MetaTags"]["TimeRes"]

info = ns5_mat["ElectrodesInfo"]
for i = 1:length(info["ElectrodeID"])
    id = info["ElectrodeID"][i]
    ch = ns5.continuous_channels[id]
    @test ch.label == rstrip(string(info["Label"][i]...), '\0')
    @test ch.physical_connector == info["ConnectorBank"][i][1]
    @test ch.connector_pin == info["ConnectorPin"][i]
    @test ch.min_digital_value == info["MinDigiValue"][i]
    @test ch.max_digital_value == info["MaxDigiValue"][i]
    @test ch.min_analog_value == info["MinAnalogValue"][i]
    @test ch.max_analog_value == info["MaxAnalogValue"][i]
    @test ch.units == rstrip(string(info["AnalogUnits"][i]...), '\0')
    @test ch.high_pass.cutoff == info["HighFreqCorner"][i]
    @test ch.high_pass.order == info["HighFreqOrder"][i]
    @test ch.high_pass.filter_type == info["HighFilterType"][i]
    @test ch.low_pass.cutoff == info["LowFreqCorner"][i]
    @test ch.low_pass.order == info["LowFreqOrder"][i]
    @test ch.low_pass.filter_type == info["LowFilterType"][i]

    @test NeuroIO.data(ch) == vec(ns5_mat["Data"][i, :])
end


using SinGAN

using FileIO
using Flux

function main()
    # hp = HyperParams()
    hp = load_hyperparams("hyperparams.json")
    # save_hyperparams("hyperparams.json", hp)
    dscrp, genp = setup_models(hp)
    @info dscrp
    @info genp

    img_name = "floral_shoppe.jpg"
    orig_rgb_img = load(img_name)
    orig_img = SinGAN.rgb_to_array(orig_rgb_img) |> gpu
    img_shapes = image_shapes(hp)
    real_img_p = SinGAN.build_image_pyramid(orig_img, img_shapes, img_shapes)
    # load_model_params!(dscrp, genp, Base.length(img_shapes))

    amplifiers, z_rec = train!(dscrp, genp, real_img_p, hp)
    generate_animation(genp, 0.1f0, 0.9f0, first(amplifiers), z_rec)
end

main()


using PCLCommon
using PCLIO
using PCLVisualization
using Libfreenect2
using Cxx

const w = 512
const h = 424
save_pcd = false

genfilename(ext=".pcd") =
    joinpath(dirname(@__FILE__), string(now(), "_", time_ns(), ext))

"""Get point cloud from undistored depth and reigstered color"""
function getPointCloudXYZRGB(registration::Registration, undistorted, registered)
    w = Libfreenect2.width(undistorted)
    h = Libfreenect2.height(undistorted)
    cloud = PointCloud{PointXYZRGB}(w, h)
    icxx"$(cloud.handle)->is_dense = false;"
    pointsptr = icxx"&$(cloud.handle)->points[0];"
    icxx"""
    for (size_t ri = 0; ri < $h; ++ri) {
        for (size_t ci = 0; ci < $w; ++ci) {
            auto p = $(pointsptr) + $w * ri + ci;
            $(registration.handle)->getPointXYZRGB($(undistorted.handle),
                $(registered.handle), ri, ci, p->x, p->y, p->z, p->rgb);
        }
    }
    """
    cloud
end

function ppcallback(event::cxxt"pcl::visualization::PointPickingEvent&",
        args::Ptr{Void})
    if Int(icxx"$(event).getPointIndex();") == -1
        return
    end
    p = PointXYZ()
    icxx"$(event).getPoint($p.x, $p.y, $p.z);"
    @show p
    return nothing::Void
end

f = Freenect2()
device = openDefaultDevice(f, OpenGLPacketPipeline())
listener = SyncMultiFrameListenerPtr()
setIrAndDepthFrameListener(device, listener)
setColorFrameListener(device, listener)

start(device)

# NOTE: must be called after start(device)
registration = Registration(getIrCameraParams(device),
    getColorCameraParams(device))
undistorted = FramePtr(w, h, 4, key=Libfreenect2.FRAME_DEPTH)
registered = FramePtr(w, h, 4, key=Libfreenect2.FRAME_COLOR)

info("Prepare PCL visualizer...")
global viewer = PCLVisualizer("pcl visualizer")
addCoordinateSystem(viewer, 1.2, 0,0,0)

registerPointPickingCallback(viewer, ppcallback)

global should_save = false
if !isdefined(:viewer_cb_defined)
cxx"""
void viewer_cb(const pcl::visualization::KeyboardEvent &event) {
    std::cout << "key event:" << event.getKeyCode() << std::endl;
    if (event.getKeyCode() == 's') {
        $:(global should_save = true; nothing);
    }
}
"""
const viewer_cb_defined = true
end

icxx"$(viewer.handle)->registerKeyboardCallback(viewer_cb);"

try
    while !wasStopped(viewer)
        frames = waitForNewFrame(listener)
        color = frames[FrameType.COLOR]
        ir = frames[FrameType.IR]
        depth = frames[FrameType.DEPTH]

        # Depth and color registration
        t = @elapsed begin
            apply(registration, color, depth, undistorted, registered)
            cloud = getPointCloudXYZRGB(registration, undistorted, registered)
        end
        println("Registration and getPointCloudXYZRGB: $(1/t) Hz")

        if should_save
            info("save pcd file...")
            savePCDFile(genfilename(), cloud; binary_mode=true)
            should_save = false
        end

        color_handler = PointCloudColorHandlerRGBField(cloud)
        if !updatePointCloud(viewer, cloud, color_handler, id="libfreenect2")
            addPointCloud(viewer, cloud, color_handler, id="libfreenect2")
        end
        try
            spinOnce(viewer, 1)
        catch e
            @show e
        end

        Libfreenect2.release(listener, frames)

        rand() > 0.95 && gc(false)
    end
finally
    stop(device)
    close(device)
    close(viewer)
    device=0;viewer=0;gc()
end


include("plot_domain.jl")
include("plot_field.jl")
include("plot_moments.jl")

# Plot the result across angular frequency for a specific position (x)
@recipe function plot(simres::SimulationResult;
        x = simres.x,
        x_indices = [findmin([norm(z - y) for z in simres.x])[2] for y in x],
        ω_indices = Colon(), field_apply = real)

    for x_ind in x_indices

        fs = field_apply.(field(simres)[x_ind, ω_indices])
        xguide = ((typeof(simres) <: FrequencySimulationResult) ? "ω" : "t")

        @series begin
            label --> "$field_apply x=$(simres.x[x_ind])"
            xguide --> xguide
            (getfield(simres, 3)[ω_indices], fs)
        end
    end
end

"Plot just the particles"
@recipe function plot(sim::FrequencySimulation; bounds = :none)

    # println("Plotting a simulation on its own")

    @series begin
        if bounds != :none
            # bounds = bounding_box(sim.particles)
            xlims --> (bottomleft(bounds)[1], topright(bounds)[1])
            ylims --> (bottomleft(bounds)[2], topright(bounds)[2])
        end

        sim.particles
    end

end

"Plot the field for a particular wavenumber"
@recipe function plot(sim::FrequencySimulation{T,3}, ω::T;
        res=10, xres=res, yres=res,
        y = :auto,
        field_apply=real,
        region_shape = :auto,
        bounds = :auto,
        exclude_region = EmptyShape{T,3}(),
        drawparticles=false) where {T}

    # If user wants us to, generate bounding rectangle around particles
    region_shape = (region_shape != :auto) ? region_shape :
        if isempty(sim.particles)
            if bounds == :auto
                @warn "What region to plot? For example, use keyword bounds = Box([[-1.0,-1.0],[1.0,1.0]])"
                Box([[-one(T),-one(T)],[one(T),one(T)]])
            else bounds
            end
        else
            region_shape = bounding_box(sim.particles)
        end

    bounds = bounding_box(region_shape)

    # If user has not set xlims and ylims, set them to the rectangle
    cs = corners(bounds)
    xlims --> (cs[1][1], cs[end][1])
    ylims --> (cs[1][end], cs[end][end])

    if y == :auto
        y = cs[1][2]
    end

    # Incase the user did set the xlims and ylims, generate a new bounding box with them
    # p_xlims = plotattributes[:xlims]
    # p_ylims = plotattributes[:ylims]
    # bounds = Box([[T(p_xlims[1]),T(p_ylims[1])], [T(p_xlims[2]),T(p_ylims[2])]])
    #
    # region_shape = (bounds ⊆ region_shape) ? bounds : region_shape

    field_sim = run(sim, region_shape, [ω]; y=y, xres=xres, zres=yres, exclude_region=exclude_region)
    xy_mat = reshape(field_sim.x, (xres+1, yres+1))
    x_pixels = [x[1] for x in xy_mat[:,1]]
    y_pixels = [x[end] for x in xy_mat[1,:]]

    @series begin

        # Turn the responses (a big long vector) into a matrix, so that the heatmap will understand us
        response_mat = transpose(reshape(field(field_sim), (xres+1, yres+1)))
        seriestype --> :contour
        fill --> true
        grid --> false
        aspect_ratio := 1.0
        seriescolor --> :balance
        title --> "Field at ω=$ω"

        (x_pixels, y_pixels, field_apply.(response_mat))
    end

    if drawparticles
        @series begin
            sim.particles
        end
    end

end

"Plot the field for a particular wavenumber"
@recipe function plot(sim::FrequencySimulation{T,2}, ω::T;
        res=10, xres=res, yres=res,
        field_apply=real,
        region_shape = :auto,
        bounds = :auto,
        exclude_region = EmptyShape{T,2}(),
        drawparticles=false) where {T}

    # If user wants us to, generate bounding rectangle around particles
    region_shape = (region_shape != :auto) ? region_shape :
        if isempty(sim.particles)
            if bounds == :auto
                @warn "What region to plot? For example, use keyword bounds = Box([[-1.0,-1.0],[1.0,1.0]])"
                Box([[-one(T),-one(T)],[one(T),one(T)]])
            else bounds
            end
        else
            region_shape = bounding_box(sim.particles)
        end

    bounds = bounding_box(region_shape)
    # If user has not set xlims and ylims, set them to the rectangle
    xlims --> (bottomleft(bounds)[1], topright(bounds)[1])
    ylims --> (bottomleft(bounds)[2], topright(bounds)[2])

    # Incase the user did set the xlims and ylims, generate a new bounding
    # rectangle with them
    p_xlims = plotattributes[:xlims]
    p_ylims = plotattributes[:ylims]
    bounds = Box([[T(p_xlims[1]),T(p_ylims[1])], [T(p_xlims[2]),T(p_ylims[2])]])

    region_shape = (bounds ⊆ region_shape) ? bounds : region_shape

    field_sim = run(sim, region_shape, [ω]; xres=xres, yres=yres, zres=yres, exclude_region=exclude_region)
    xy_mat = reshape(field_sim.x, (xres+1, yres+1))
    x_pixels = [x[1] for x in xy_mat[:,1]]
    y_pixels = [x[2] for x in xy_mat[1,:]]

    @series begin

        # Turn the responses (a big long vector) into a matrix, so that the heatmap will understand us
        response_mat = transpose(reshape(field(field_sim), (xres+1, yres+1)))
        seriestype --> :contour
        fill --> true
        grid --> false
        aspect_ratio := 1.0
        seriescolor --> :balance
        title --> "Field at ω=$ω"

        (x_pixels, y_pixels, field_apply.(response_mat))
    end

    if drawparticles
        @series begin
            sim.particles
        end
    end

end

"Plot the source field for a particular wavenumber"
@recipe function plot(s::RegularSource, ω)
    (FrequencySimulation(s), ω)
end


module CBF

import ..FileFormats
import MathOptInterface

const MOI = MathOptInterface

MOI.Utilities.@model(
    Model,
    (MOI.Integer,),
    (),
    (
        MOI.Reals,
        MOI.Zeros,
        MOI.Nonnegatives,
        MOI.Nonpositives,
        MOI.SecondOrderCone,
        MOI.RotatedSecondOrderCone,
        MOI.PositiveSemidefiniteConeTriangle,
        MOI.ExponentialCone,
        MOI.DualExponentialCone,
    ),
    (MOI.PowerCone, MOI.DualPowerCone),
    (),
    (),
    (MOI.VectorOfVariables,),
    (MOI.VectorAffineFunction,)
)

function MOI.supports_constraint(
    ::Model{T},
    ::Type{MOI.VariableIndex},
    ::Type{<:MOI.Utilities.SUPPORTED_VARIABLE_SCALAR_SETS{T}},
) where {T}
    return false
end

function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VariableIndex},
    ::Type{MOI.Integer},
)
    return true
end

function MOI.supports(
    ::Model,
    ::MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}},
)
    return false
end

"""
    Model()

Create an empty instance of `FileFormats.CBF.Model`.
"""
Model() = Model{Float64}()

Base.show(io::IO, ::Model) = print(io, "A Conic Benchmark Format (CBF) model")

include("read.jl")
include("write.jl")

end


# conversion routines between different sparse and dense storage formats

SparseArrays.sparse(::DenseCuArray, args...) = error("CUSPARSE supports multiple sparse formats, use specific constructors instead (e.g. CuSparseMatrixCSC)")

## CSR to CSC

# by flipping rows and columns, we can use that to get CSC to CSR too

for (fname,elty) in ((:cusparseScsr2csc, :Float32),
                     (:cusparseDcsr2csc, :Float64),
                     (:cusparseCcsr2csc, :ComplexF32),
                     (:cusparseZcsr2csc, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSC{$elty}(csr::CuSparseMatrixCSR{$elty}; inda::SparseChar='O')
            m,n = csr.dims
            colPtr = CUDA.zeros(Cint, n+1)
            rowVal = CUDA.zeros(Cint, nnz(csr))
            nzVal = CUDA.zeros($elty, nnz(csr))
            if version() >= v"10.2"
                # TODO: algorithm configuratibility?
                @workspace size=@argout(
                        cusparseCsr2cscEx2_bufferSize(handle(), m, n, nnz(csr), nonzeros(csr),
                            csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
                            $elty, CUSPARSE_ACTION_NUMERIC, inda,
                            CUSPARSE_CSR2CSC_ALG1, out(Ref{Csize_t}(1)))
                    )[] buffer->begin
                        cusparseCsr2cscEx2(handle(), m, n, nnz(csr), nonzeros(csr),
                            csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
                            $elty, CUSPARSE_ACTION_NUMERIC, inda,
                            CUSPARSE_CSR2CSC_ALG1, buffer)
                    end
            else
                $fname(handle(), m, n, nnz(csr), nonzeros(csr),
                    csr.rowPtr, csr.colVal, nzVal, rowVal,
                    colPtr, CUSPARSE_ACTION_NUMERIC, inda)
            end
            CuSparseMatrixCSC(colPtr,rowVal,nzVal,csr.dims)
        end

        function CuSparseMatrixCSR{$elty}(csc::CuSparseMatrixCSC{$elty}; inda::SparseChar='O')
            m,n    = csc.dims
            rowPtr = CUDA.zeros(Cint,m+1)
            colVal = CUDA.zeros(Cint,nnz(csc))
            nzVal  = CUDA.zeros($elty,nnz(csc))
            if version() >= v"10.2"
                # TODO: algorithm configuratibility?
                @workspace size=@argout(
                        cusparseCsr2cscEx2_bufferSize(handle(), n, m, nnz(csc), nonzeros(csc),
                            csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
                            $elty, CUSPARSE_ACTION_NUMERIC, inda,
                            CUSPARSE_CSR2CSC_ALG1, out(Ref{Csize_t}(1)))
                    )[] buffer->begin
                        cusparseCsr2cscEx2(handle(), n, m, nnz(csc), nonzeros(csc),
                            csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
                            $elty, CUSPARSE_ACTION_NUMERIC, inda,
                            CUSPARSE_CSR2CSC_ALG1, buffer)
                    end
            else
                $fname(handle(), n, m, nnz(csc), nonzeros(csc),
                    csc.colPtr, rowvals(csc), nzVal, colVal,
                    rowPtr, CUSPARSE_ACTION_NUMERIC, inda)
            end
            CuSparseMatrixCSR(rowPtr,colVal,nzVal,csc.dims)
        end
    end
end

## CSR to BSR and vice-versa

for (fname,elty) in ((:cusparseScsr2bsr, :Float32),
                     (:cusparseDcsr2bsr, :Float64),
                     (:cusparseCcsr2bsr, :ComplexF32),
                     (:cusparseZcsr2bsr, :ComplexF64))
    @eval begin
        function CuSparseMatrixBSR{$elty}(csr::CuSparseMatrixCSR{$elty}, blockDim::Integer;
                                          dir::SparseChar='R', inda::SparseChar='O',
                                          indc::SparseChar='O')
            m,n = csr.dims
            nnz_ref = Ref{Cint}(1)
            mb = div((m + blockDim - 1),blockDim)
            nb = div((n + blockDim - 1),blockDim)
            bsrRowPtr = CUDA.zeros(Cint,mb + 1)
            cudesca = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, inda)
            cudescc = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, indc)
            cusparseXcsr2bsrNnz(handle(), dir, m, n, cudesca, csr.rowPtr,
                                csr.colVal, blockDim, cudescc, bsrRowPtr, nnz_ref)
            bsrNzVal = CUDA.zeros($elty, nnz_ref[] * blockDim * blockDim )
            bsrColInd = CUDA.zeros(Cint, nnz_ref[])
            $fname(handle(), dir, m, n,
                   cudesca, nonzeros(csr), csr.rowPtr, csr.colVal,
                   blockDim, cudescc, bsrNzVal, bsrRowPtr,
                   bsrColInd)
            CuSparseMatrixBSR{$elty}(bsrRowPtr, bsrColInd, bsrNzVal, csr.dims, blockDim, dir, nnz_ref[])
        end
    end
end

for (fname,elty) in ((:cusparseSbsr2csr, :Float32),
                     (:cusparseDbsr2csr, :Float64),
                     (:cusparseCbsr2csr, :ComplexF32),
                     (:cusparseZbsr2csr, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSR{$elty}(bsr::CuSparseMatrixBSR{$elty};
                                          inda::SparseChar='O', indc::SparseChar='O')
            m,n = bsr.dims
            mb = div(m,bsr.blockDim)
            nb = div(n,bsr.blockDim)
            nnzVal = nnz(bsr) * bsr.blockDim * bsr.blockDim
            cudesca = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, inda)
            cudescc = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, indc)
            csrRowPtr = CUDA.zeros(Cint, m + 1)
            csrColInd = CUDA.zeros(Cint, nnzVal)
            csrNzVal  = CUDA.zeros($elty, nnzVal)
            $fname(handle(), bsr.dir, mb, nb,
                   cudesca, nonzeros(bsr), bsr.rowPtr, bsr.colVal,
                   bsr.blockDim, cudescc, csrNzVal, csrRowPtr,
                   csrColInd)
            CuSparseMatrixCSR(csrRowPtr, csrColInd, csrNzVal, bsr.dims)
        end
    end
end

## CSR to COO and vice-versa

function CuSparseMatrixCSR(coo::CuSparseMatrixCOO{Tv}, ind::SparseChar='O') where {Tv}
    m,n = coo.dims
    nnz = coo.nnz
    csrRowPtr = CUDA.zeros(Cint, nnz)
    cusparseXcoo2csr(handle(), coo.rowInd, nnz, m, csrRowPtr, ind)
    CuSparseMatrixCSR{Tv}(csrRowPtr, coo.colInd, coo.nzVal, coo.dims)
end

function CuSparseMatrixCOO(csr::CuSparseMatrixCSR{Tv}, ind::SparseChar='O') where {Tv}
    m,n = csr.dims
    nnz = csr.nnz
    cooRowInd = CUDA.zeros(Cint, nnz)
    cusparseXcsr2coo(handle(), csr.rowPtr, nnz, m, cooRowInd, ind)
    CuSparseMatrixCOO{Tv}(cooRowInd, csr.colVal, csr.nzVal, csr.dims, nnz)
end

## sparse to dense, and vice-versa

# TODO: use cusparseSparseToDense & cusparseDenseToSparse (3x faster, CUDA 11.1 Update 1)

for (cname,rname,elty) in ((:cusparseScsc2dense, :cusparseScsr2dense, :Float32),
                           (:cusparseDcsc2dense, :cusparseDcsr2dense, :Float64),
                           (:cusparseCcsc2dense, :cusparseCcsr2dense, :ComplexF32),
                           (:cusparseZcsc2dense, :cusparseZcsr2dense, :ComplexF64))
    @eval begin
        function CUDA.CuMatrix{$elty}(csr::CuSparseMatrixCSR{$elty}; ind::SparseChar='O')
            m,n = csr.dims
            denseA = CUDA.zeros($elty,m,n)
            lda = max(1,stride(denseA,2))
            cudesc = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, ind)
            $rname(handle(), m, n, cudesc, nonzeros(csr),
                   csr.rowPtr, csr.colVal, denseA, lda)
            denseA
        end
        function CUDA.CuMatrix{$elty}(csc::CuSparseMatrixCSC{$elty}; ind::SparseChar='O')
            m,n = csc.dims
            denseA = CUDA.zeros($elty,m,n)
            lda = max(1,stride(denseA,2))
            cudesc = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, ind)
            $cname(handle(), m, n, cudesc, nonzeros(csc),
                   rowvals(csc), csc.colPtr, denseA, lda)
            denseA
        end
    end
end

for (nname,cname,rname,elty) in ((:cusparseSnnz, :cusparseSdense2csc, :cusparseSdense2csr, :Float32),
                                 (:cusparseDnnz, :cusparseDdense2csc, :cusparseDdense2csr, :Float64),
                                 (:cusparseCnnz, :cusparseCdense2csc, :cusparseCdense2csr, :ComplexF32),
                                 (:cusparseZnnz, :cusparseZdense2csc, :cusparseZdense2csr, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSR(A::CuMatrix{$elty}; ind::SparseChar='O')
            m,n = size(A)
            lda = max(1, stride(A,2))
            cudesc = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL,
                                        CUSPARSE_FILL_MODE_LOWER,
                                        CUSPARSE_DIAG_TYPE_NON_UNIT, ind)
            nnzRowCol = CUDA.zeros(Cint, m)
            nnzTotal = Ref{Cint}(1)
            $nname(handle(),
                   'R', m, n, cudesc, A, lda, nnzRowCol,
                   nnzTotal)
            nzVal = CUDA.zeros($elty,nnzTotal[])

            rowPtr = CUDA.zeros(Cint,m+1)
            colInd = CUDA.zeros(Cint,nnzTotal[])
            $rname(handle(), m, n, cudesc, A,
                    lda, nnzRowCol, nzVal, rowPtr, colInd)
            return CuSparseMatrixCSR(rowPtr,colInd,nzVal,size(A))
        end

        function CuSparseMatrixCSC(A::CuMatrix{$elty}; ind::SparseChar='O')
            m,n = size(A)
            lda = max(1, stride(A,2))
            cudesc = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL,
                                        CUSPARSE_FILL_MODE_LOWER,
                                        CUSPARSE_DIAG_TYPE_NON_UNIT, ind)
            nnzRowCol = CUDA.zeros(Cint, n)
            nnzTotal = Ref{Cint}(1)
            $nname(handle(),
                   'C', m, n, cudesc, A, lda, nnzRowCol,
                   nnzTotal)
            nzVal = CUDA.zeros($elty,nnzTotal[])

            colPtr = CUDA.zeros(Cint,n+1)
            rowInd = CUDA.zeros(Cint,nnzTotal[])
            $cname(handle(), m, n, cudesc, A,
                    lda, nnzRowCol, nzVal, rowInd, colPtr)
            return CuSparseMatrixCSC(colPtr,rowInd,nzVal,size(A))
        end
    end
end

function CUDA.CuMatrix{T}(bsr::CuSparseMatrixBSR{T}; inda::SparseChar='O',
                          indc::SparseChar='O') where {T}
    CuMatrix{T}(CuSparseMatrixCSR{T}(bsr; inda, indc))
end

function CuSparseMatrixBSR(A::CuMatrix; ind::SparseChar='O')
    m,n = size(A)   # TODO: always let the user choose, or provide defaults for other methods too
    CuSparseMatrixBSR(CuSparseMatrixCSR(A; ind), gcd(m,n))
end


using Test
using RigidBodyDynamics
using StrandbeestRobot
using NBInclude

@testset "Initial configuration" begin
    @testset "floating base" begin
        mechanism = StrandbeestRobot.mechanism(floating=true)
        state = MechanismState(mechanism)
        cost = StrandbeestRobot.loop_joint_error(mechanism)
        StrandbeestRobot.solve_initial_configuration!(state)
        for joint in tree_joints(mechanism)
            @test RigidBodyDynamics.is_configuration_normalized(joint, configuration(state, joint); atol=1e-8)
        end
        @test cost(configuration(state)) <= 1e-9
    end

    @testset "fixed base" begin
        mechanism = StrandbeestRobot.mechanism(floating=false)
        state = MechanismState(mechanism)
        cost = StrandbeestRobot.loop_joint_error(mechanism)
        StrandbeestRobot.solve_initial_configuration!(state)
        @test cost(configuration(state)) <= 1e-9
    end
end

@testset "notebooks" begin
    notebook_dir = joinpath(@__DIR__, "..", "notebooks")
    @nbinclude(joinpath(notebook_dir, "Basic simulation.ipynb"); regex = r"^((?!\#NBSKIP).)*$"s)
    @nbinclude(joinpath(notebook_dir, "Passive walking.ipynb"); regex = r"^((?!\#NBSKIP).)*$"s)
    @nbinclude(joinpath(notebook_dir, "Powered walking.ipynb"); regex = r"^((?!\#NBSKIP).)*$"s)
end


# # GARTEUR SM-AG19 Testbed: Construction of the geometry

# Source code: [`garteur_geometry_tut.jl`](garteur_geometry_tut.jl)

# ## Description

# This virtual test application is based on the test article
# used by the GARTEUR Structures & Materials Action Group 19
# which organized a Round Robin exercise where 12 European laboratories
# tested a single structure between 1995 and 1997. The benchmark structure
# was a laboratory structure built to simulate the dynamic behaviour
# of an aeroplane. The structure was initially built for a benchmark
# study on experimental modal analysis conducted by the
# Structures and Materials Action Group (SM-AG19) of the Group
# for Aeronautical Research and Technology in EURope (GARTEUR).
# The test-bed was designed and manufactured by ONERA, France.

# ![](IMAC97photo.png)

# ### References

# - [GARTEUR] Ground Vibration Test Techniques, compiled by A Gravelle, GARTEUR
#   Structures & Materials Action Group 19 Technical report TP-115, 1999.
# - [BW] Etienne Balmes, Jan R. Wright, GARTEUR GROUP ON GROUND VIBRATION
#   TESTING | RESULTS FROM THE TEST OF A SINGLE STRUCTURE BY 12 LABORATORIES IN
#   EUROPE, Proceedings of DETC'97, 1997 ASME Design Engineering Technical
#   Conferences, September 14-17, 1997, Sacramento, California.

# ## Goals

# - Show how to construct model from multiple connected beams.
# - Demonstrate the use of massless connectors.
# - Visualize the structure interactively.
#

# ## Geometry of the testbed airplane.

# The aluminum testbed was a rather simple structure which was reasonably
# dynamically representative of a simple airplane structure [GARTEUR](@ref
# References). It was composed of several beams simulating a fuselage with wings
# and a tail. Wing tip drums allowed to adjust bending and torsion frequencies
# similarly to airplane ones, with some very close modal frequencies.

# ![](garteur-geom.png)

# The structural-element finite element code relies on the basic functionality
# implemented in this general FEM package.

using FinEtools

# This is the characteristic length. The dimensions of the aircraft frame are
# expressed in terms of multiples of this characteristic unit.
L = 0.1*phun("m");

# ## Cross-section

# Cross-sectional properties are incorporated in the cross-section property.
# There are several rectangular cross-sections in the model: the fuselage, the
# wing, the tail. There are also three massless connectors: connections between
# the fuselage and the wing, between the wing structure and the viscoelastic
# damping layer, and between the fuselage and the tail.

using FinEtoolsFlexStructures.CrossSectionModule: CrossSectionRectangle
# Body of the frame (fuselage).
cs_body = CrossSectionRectangle(s -> 1.5*L, s -> L/2, s -> [1.0, 0.0, 1.0]; label = 1)
# Wing beam.
cs_wing = CrossSectionRectangle(s -> L/10, s -> L, s -> [0.0, 0.0, 1.0]; label = 2)
# Wing drums.
cs_drum = CrossSectionRectangle(s -> L/10, s -> L, s -> [0.0, 0.0, 1.0]; label = 3)
# Vertical part of the tail.
cs_tailv = CrossSectionRectangle(s -> L, s -> L/10, s -> [1.0, 0.0, 1.0]; label = 4)
# Horizontal part of the tail.
cs_tailh = CrossSectionRectangle(s -> L/10, s -> L, s -> [0.0, 0.0, 1.0]; label = 5)
# Constraining plate on top of the viscoelastic tape.
cs_constrp = CrossSectionRectangle(s -> L*(1.1/100), s -> L*(76.2/100), s -> [0.0, 0.0, 1.0]; label = 6)
# Massless connectors of the structural parts of the wing: the main beam and the
# constraining plate.
cs_connw2c = CrossSectionRectangle(s -> L/2, s -> L/2, s -> [1.0, 0.0, 1.0]; label = 7)
# Massless connectors of the wing and the drums.
cs_connw2d = CrossSectionRectangle(s -> L/2, s -> L/2, s -> [1.0, 0.0, 1.0]; label = 8)
# Massless connector between the body and the wings.
cs_connb2w = CrossSectionRectangle(s -> L/5, s -> L, s -> [0.0, 1.0, 0.0]; label = 9)
# Massless connector between the body and the tail.
cs_conntb2t = CrossSectionRectangle(s -> L, s -> L/3, s -> [1.0, 0.0, 1.0]; label = 10)
# Massless connector between the structure and the sensors and point masses.
cs_connta2p = CrossSectionRectangle(s -> L/5, s -> L/5, s -> [1.0, 0.0, 1.0]; label = 11)

# ## Mesh

# We shall use the `frame_member` utility function to generate the mesh of the individual
# parts. This will result in a number of separate meshes for the members. These
# separate meshes will then be glued together (merged) based on the tolerance on
# the location of the nodes.
using FinEtoolsFlexStructures.MeshFrameMemberModule: frame_member
tolerance = L/10000;
# Number of intervals from 0.25*L to 8.5*L (the extent of the constraining plate).
nc = 8

# We start with an empty array, and for each part that we mesh, we append the
# mesh to the array. All the meshes in the array will then be considered for
# merging based on geometrical tolerances.
meshes = Tuple{FENodeSet, AbstractFESet}[]

# Define the constituent parts of the body of the aircraft.
push!(meshes, frame_member([-9*L 0 0; -8.5*L 0 0], 1, cs_body; label = cs_body.label))
push!(meshes, frame_member([-8.5*L 0 0; -8.0*L 0 0], 1, cs_body; label = cs_body.label))
push!(meshes, frame_member([-8.0*L 0 0; -2.0*L 0 0], 2, cs_body; label = cs_body.label))
push!(meshes, frame_member([-2.0*L 0 0; 0 0 0], 1, cs_body; label = cs_body.label))
push!(meshes, frame_member([0 0 0; 6*L 0 0], 2, cs_body; label = cs_body.label))

# Define the aluminum parts of the wings.
push!(meshes, frame_member([0 0 0.805*L;  0 0.25*L 0.805*L], 1, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 0 0.805*L;  0 -0.25*L 0.805*L], 1, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 0.25*L 0.805*L;  0 8.5*L 0.805*L], nc, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 -0.25*L 0.805*L;  0 -8.5*L 0.805*L], nc, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 8.5*L 0.805*L;  0 9.5*L 0.805*L], 1, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 -8.5*L 0.805*L;  0 -9.5*L 0.805*L], 1, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 9.5*L 0.805*L;  0 10.0*L 0.805*L], 1, cs_wing; label = cs_wing.label))
push!(meshes, frame_member([0 -9.5*L 0.805*L;  0 -10.0*L 0.805*L], 1, cs_wing; label = cs_wing.label))

# Define the drums at the ends of the wings.
push!(meshes, frame_member([0 +9.5*L +0.91*L; +2*L +9.5*L +0.91*L], 1, cs_drum; label = cs_drum.label))
push!(meshes, frame_member([0 +9.5*L +0.91*L; -2*L +9.5*L +0.91*L], 1, cs_drum; label = cs_drum.label))
push!(meshes, frame_member([0 -9.5*L +0.91*L; +2*L -9.5*L +0.91*L], 1, cs_drum; label = cs_drum.label))
push!(meshes, frame_member([0 -9.5*L +0.91*L; -2*L -9.5*L +0.91*L], 1, cs_drum; label = cs_drum.label))

# Define the horizontal and vertical parts of the tail.
push!(meshes, frame_member([-8*L 0 .75*L; -8*L 0 3.35*L], 2, cs_tailv; label = cs_tailv.label))
push!(meshes, frame_member([-8*L 0 3.35*L; -8*L 0 3.75*L], 2, cs_tailv; label = cs_tailv.label))
push!(meshes, frame_member([-8*L 0 3.8*L; -8*L 2*L 3.8*L], 2, cs_tailh; label = cs_tailh.label))
push!(meshes, frame_member([-8*L 0 3.8*L; -8*L -2*L 3.8*L], 2, cs_tailh; label = cs_tailh.label))

# Define the parts of the aluminum constraining plate for the viscoelastic layer.
push!(meshes, frame_member([-.119*L 0 0.8665*L;  -.119*L 0.25*L 0.8665*L], 1, cs_constrp; label = cs_constrp.label))
push!(meshes, frame_member([-.119*L 0 0.8665*L;  -.119*L -0.25*L 0.8665*L], 1, cs_constrp; label = cs_constrp.label))
push!(meshes, frame_member([-.119*L 0.25*L 0.8665*L;  -.119*L 8.5*L 0.8665*L], nc, cs_constrp; label = cs_constrp.label))
push!(meshes, frame_member([-.119*L -0.25*L 0.8665*L;  -.119*L -8.5*L 0.8665*L], nc, cs_constrp; label = cs_constrp.label))

# Define the massless connectors between:
# Wing - Wingdrum
push!(meshes, frame_member([0 +9.5*L +0.805*L;0 +9.5*L +0.91*L], 1, cs_connw2d; label = cs_connw2d.label))
push!(meshes, frame_member([0 -9.5*L +0.805*L;0 -9.5*L +0.91*L], 1, cs_connw2d; label = cs_connw2d.label))
# Body-Wing
push!(meshes, frame_member([0 0 0; 0 0.25*L .805*L], 1, cs_connb2w; label = cs_connb2w.label))
push!(meshes, frame_member([0 0 0; 0 -0.25*L .805*L], 1, cs_connb2w; label = cs_connb2w.label))
# Body-Tail
push!(meshes, frame_member([-8*L 0 0; -8*L 0 .75*L], 1, cs_conntb2t; label = cs_conntb2t.label))
# Tail-Taildrum
push!(meshes, frame_member([-8*L 0 3.75*L; -8*L 0 3.8*L], 1, cs_connw2d; label = cs_connw2d.label))

# Wing-Constraining plate for the viscoelastic layer
# Middle connector created individually
push!(meshes, frame_member([0 0 .805*L; -.119*L 0 0.8665*L], 1, cs_connw2c; label = cs_connw2c.label))
# Connectors alongside both wings
for i in 1:nc+1
    push!(meshes, frame_member([0 (0.25+(i-1)*8.25/nc)*L .805*L;  -.119*L (0.25+(i-1)*8.25/nc)*L 0.8665*L], 1, cs_connw2c; label = cs_connw2c.label))
    push!(meshes, frame_member([0 -(0.25+(i-1)*8.25/nc)*L .805*L;  -.119*L -(0.25+(i-1)*8.25/nc)*L 0.8665*L], 1, cs_connw2c; label = cs_connw2c.label))
end

# Massless Sensor Connectors
# Tail Sensors
push!(meshes, frame_member([-8*L 2*L 3.8*L; -(153/20)*L (37/20)*L 3.85*L], 1, cs_connta2p; label = cs_connta2p.label))# 303
push!(meshes, frame_member([-8*L -2*L 3.8*L; -(153/20)*L -(37/20)*L 3.85*L], 1, cs_connta2p; label = cs_connta2p.label))# 301

# Wingdrum Sensors
push!(meshes, frame_member([0 9.5*L .91*L ; 0 9.8*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# 101
push!(meshes, frame_member([-2*L 9.5*L .91*L ; -1.8*L 9.8*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# 112
push!(meshes, frame_member([2*L 9.5*L .91*L ; 1.8*L 9.8*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# 111

push!(meshes, frame_member([0 -9.5*L .91*L ; 0 -9.8*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# 1
push!(meshes, frame_member([-2*L -9.5*L .91*L ; -1.8*L -9.8*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# 12
push!(meshes, frame_member([2*L -9.5*L .91*L ; 1.8*L -9.8*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# 11

# Wingdrum complementary masses
push!(meshes, frame_member([2*L 9.5*L .91*L ; 1.8*L 9.2*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# added mass
push!(meshes, frame_member([2*L -9.5*L .91*L ; 1.8*L -9.2*L .96*L], 1, cs_connta2p; label = cs_connta2p.label))# added mass

# ## Merge all the individual members into one coherent structure

# Merge all the meshes of individual parts. This will glue together nodes which
# are in the "same" location. The parts of the mesh can be distinguished based
# on the label.
fens, fesa = mergenmeshes(meshes, tolerance)

# The result is the set of the finite element nodes, and an array of the finite element sets.
# The sets of the finite elements can be distinguished based on the label.

# The number of nodes in the mesh:
@show count(fens)

# The number of the finite element sets:
@show length(fesa)

# The labels of the finite element sets
@show [s.label[1] for s in fesa]'

# End of the tutorial
nothing


using Base:Float64

# concrete result type -
struct VLResult{T}
    value::T
end

abstract type VLAbstractAsset end
abstract type VLAbstractLattice end

mutable struct VLEquityAsset <: VLAbstractAsset

    # data -
    assetSymbol::String
    purchasePricePerShare::Float64
    numberOfShares::Int64
    purchaseDate::Date

    function VLEquityAsset(assetSymbol::String, purchasePricePerShare::Float64, numberOfShares::Int64, purchaseDate::Date)
        this = new(assetSymbol, purchasePricePerShare, numberOfShares, purchaseDate)
    end
end

mutable struct VLCallOptionContract <: VLAbstractAsset

    # data -
    assetSymbol::String
    strikePrice::Float64
    expirationDate::Date
    premimumValue::Float64
    numberOfContracts::Int64
    sense::Symbol
    contractMultiplier::Float64

    function VLCallOptionContract(assetSymbol::String, expirationDate::Union{Nothing,Date}, strikePrice::Float64, premimumValue::Float64,
        numberOfContracts::Int64; sense::Symbol=:buy, contractMultiplier::Float64=100.0)

        this = new(assetSymbol, strikePrice, expirationDate, premimumValue, numberOfContracts, sense, contractMultiplier)
    end
end

mutable struct VLPutOptionContract <: VLAbstractAsset

    # data -
    assetSymbol::String
    strikePrice::Float64
    expirationDate::Union{Nothing,Date}
    premimumValue::Float64
    numberOfContracts::Int64
    sense::Symbol
    contractMultiplier::Float64

    function VLPutOptionContract(assetSymbol::String, strikePrice::Float64, premimumValue::Float64;
        expirationDate::Union{Nothing,Date}=nothing, numberOfContracts::Int64=1, sense::Symbol=:buy, contractMultiplier::Float64=100.0)

        _ = new(assetSymbol, strikePrice, expirationDate, premimumValue, numberOfContracts, sense, contractMultiplier)
    end
end

mutable struct VLBinomialLattice <: VLAbstractLattice

    # data -
    μ::Float64
    σ::Float64
    𝝙t::Float64
    numberOfLevels::Int64

    function VLBinomialLattice(μ::Float64, σ::Float64, 𝝙t::Float64, numberOfLevels::Int64)
        _ = new(μ, σ, 𝝙t, numberOfLevels)
    end
end


" solution status (from ipopt) "
ApplicationReturnStatus = Dict(
0=>:Solve_Succeeded,
1=>:Solved_To_Acceptable_Level,
2=>:Infeasible_Problem_Detected,
3=>:Search_Direction_Becomes_Too_Small,
4=>:Diverging_Iterates,
5=>:User_Requested_Stop,
6=>:Feasible_Point_Found,
-1=>:Maximum_Iterations_Exceeded,
-2=>:Restoration_Failed,
-3=>:Error_In_Step_Computation,
-4=>:Maximum_CpuTime_Exceeded,
-5=>:Optimize_not_called,
-6=>:Method_not_defined,
-10=>:Not_Enough_Degrees_Of_Freedom,
-11=>:Invalid_Problem_Definition,
-12=>:Invalid_Option,
-13=>:Invalid_Number_Detected,
-100=>:Unrecoverable_Exception,
-102=>:Insufficient_Memory,
-199=>:Internal_Error)


type Risk
    underlying::FinancialAsset
    hedge::FinancialAsset
    nominalrisk::Float64
end


## softmax

const CUDNN_SOFTMAX_FAST = (UInt32)(0)
const CUDNN_SOFTMAX_ACCURATE = (UInt32)(1)
const CUDNN_SOFTMAX_LOG = (UInt32)(2)

const CUDNN_SOFTMAX_MODE_INSTANCE = (UInt32)(0)
const CUDNN_SOFTMAX_MODE_CHANNEL = (UInt32)(1)

## batchnorm

const CUDNN_BATCHNORM_PER_ACTIVATION = (UInt32)(0)
const CUDNN_BATCHNORM_SPATIAL = (UInt32)(1)
const CUDNN_BN_MIN_EPSILON = 1.0e-5


module CallbackTests

using Photon, Test
using Photon.Layers
using Photon.Callbacks
using Photon.Metrics
using Photon.Losses

function simple_conv_model()
    model = Sequential(
        Conv2D(4, 3, :relu),
        Conv2D(8, 3, :relu),
        MaxPool2D(),
        Dense(16, :relu),
        Dense(10)
    )
    return model
end

function getdata(s=28)
    [(
        KorA(randn(Float32,s,s,1,16)),
        KorA(randn(Float32,10,16))
    ) for i=1:10]
end

function test_callbacks()
    model = simple_conv_model()
    workout = Workout(model, MSELoss())

    data = getdata()
    filename = "myworkout.dat"
    rm(filename, force=true)

    train!(workout, data, epochs=2, cb=AutoSave(:loss, filename))
    @assert isfile(filename)
    rm(filename, force=true)

    train!(workout, data, epochs=2, cb=EpochSave(filename))
    @assert isfile(filename)
    rm(filename, force=true)

    val_data = getdata()
    workout = Workout(model, MSELoss())
    train!(workout, data, val_data, epochs=10, cb=EarlyStop(:val_loss))
    @assert workout.epochs < 10
end

@testset "Callbacks" begin
    test_callbacks()
end

end


@testset "Testing Concat" begin
    # x, y, z = randn.(Float32, (10, 20, 30), 10)
    # X, Y, Z = OneDNN.Memory.((x, y, z))

    # # vcat
    # expected = vcat(x, y, z)
    # arg = [X, Y, Z]
    # op = OneDNN.Concat(arg, 1)
    # result = op(arg)
    # @test isa(result, OneDNN.Memory)
    # @test result == expected

    # # hcat
    # x, y, z = randn.(Float32, 10, (10, 20, 30))
    # X, Y, Z = OneDNN.Memory.((x, y, z))

    # # vcat
    # expected = hcat(x, y, z)
    # arg = [X, Y, Z]
    # op = OneDNN.Concat(arg, 2)
    # result = op(arg)
    # @test result == expected

    #####
    ##### Pullbacks
    #####

    function tester(f, fargs, g, gargs; fkw...)
        expected, back_expected = Zygote._pullback(f, fargs...; fkw...)

        Δ = Float32(0.125) * randn(Float32, size(expected))
        grads_expected = back_expected(Δ)

        result, back = Zygote._pullback(g, gargs...)
        @test isa(result, OneDNN.Memory)
        @test OneDNN.materialize(result) == expected

        # two flavors, one with a `Memory` passed in, and one with just a normal julia
        # array.
        #
        # Note that the structures of the gradients are different because `vcat` is variadic,
        # and thus gets a tuple entry for each argument.
        #
        # OneDNN.concat, hoever, takes an array, so the sensitivity to the array input comes
        # back as an array.
        grads_result = back(OneDNN.Memory(Δ))
        @test grads_result[2][1] == grads_expected[2]
        @test grads_result[2][2] == grads_expected[3]
        @test grads_result[2][3] == grads_expected[4]
        @test grads_result[3] === nothing

        grads_result = back(Δ)
        @test grads_result[2][1] == grads_expected[2]
        @test grads_result[2][2] == grads_expected[3]
        @test grads_result[2][3] == grads_expected[4]
        @test grads_result[3] === nothing

        # Test inference.
        @inferred back(OneDNN.Memory(Δ))
        @inferred back(Δ)
    end

    # vcat
    x, y, z = randn.(Float32, (10, 20, 30), 10)
    X, Y, Z = OneDNN.Memory.((x, y, z))
    tester(vcat, (x, y, z), OneDNN.concat, ([X, Y, Z], 1))
    tester(vcat, (x, y, z), OneDNN.concat, ((X, Y, Z), 1))

    # hcat
    x, y, z = randn.(Float32, 10, (10, 20, 30))
    X, Y, Z = OneDNN.Memory.((x, y, z))
    tester(hcat, (x, y, z), OneDNN.concat, ([X, Y, Z], 2))
    tester(hcat, (x, y, z), OneDNN.concat, ((X, Y, Z), 2))

    ### 3d

    # vcat
    x, y, z = randn.(Float32, (10, 20, 30), 10, 20)
    X, Y, Z = OneDNN.Memory.((x, y, z))
    tester(vcat, (x, y, z), OneDNN.concat, ([X, Y, Z], 1))
    tester(vcat, (x, y, z), OneDNN.concat, ((X, Y, Z), 1))

    # hcat
    x, y, z = randn.(Float32, 10, (10, 20, 30), 20)
    X, Y, Z = OneDNN.Memory.((x, y, z))
    tester(hcat, (x, y, z), OneDNN.concat, ([X, Y, Z], 2))
    tester(hcat, (x, y, z), OneDNN.concat, ((X, Y, Z), 2))
end


push!(LOAD_PATH, joinpath(@__DIR__, ".."))

using JLD2
using BenchmarkTools
using Benchmarks
using Plots
pyplot()
# Benchmark parameters

decomposition = Slab()

ranks = (1, 2, 4, 8, 16)

grid_size(R, decomposition) = (256, 256, 256)

rank_size(R, ::Slab) = (1, R, 1)
rank_size(R, ::Pencil) = Int.((1, √R, √R))

# Run benchmarks

print_system_info()

for R in ranks
    Nx, Ny, Nz = grid_size(R, decomposition)
    Rx, Ry, Rz = rank_size(R, decomposition)
    @info "Benchmarking distributed incompressible model strong scaling with $(typeof(decomposition)) decomposition [N=($Nx, $Ny, $Nz), ranks=($Rx, $Ry, $Rz)]..."
    julia = Base.julia_cmd()
    run(`mpiexec -np $R $julia --project strong_scaling_incompressible_model_single.jl $(typeof(decomposition)) $Nx $Ny $Nz $Rx $Ry $Rz`)
end

# Collect and merge benchmarks from all ranks

suite = BenchmarkGroup(["size", "ranks"])

for R in ranks
    Nx, Ny, Nz = grid_size(R, decomposition)
    Rx, Ry, Rz = rank_size(R, decomposition)
    case = ((Nx, Ny, Nz), (Rx, Ry, Rz))

    for local_rank in 0:R-1
        filename = string("strong_scaling_incompressible_model_$(R)ranks_$(typeof(decomposition))_$local_rank.jld2")
        jldopen(filename, "r") do file
            if local_rank == 0
                suite[case] = file["trial"]
            else
                merged_trial = suite[case]
                local_trial = file["trial"]
                append!(merged_trial.times, local_trial.times)
                append!(merged_trial.gctimes, local_trial.gctimes)
            end
        end
    end
end

# Summarize benchmarks

plot_keys = collect(keys(suite))
sort!(plot_keys, by = v -> v[2][2])
plot_num = length(plot_keys)
rank_num = zeros(Int64, plot_num)
run_times = zeros(Float64, plot_num)
eff_ratio = zeros(Float64, plot_num)
for i in 1:plot_num
    rank_num[i] = plot_keys[i][2][2]
    run_times[i] = mean(suite[plot_keys[i]].times) / 1.0e6
    eff_ratio[i] = median(suite[plot_keys[1]].times) /(rank_num[i] * median(suite[plot_keys[i]].times))
end

plt = plot(rank_num, run_times, lw=4, xaxis=:log2, legend=:none,
          xlabel="Cores", ylabel="Times (ms)", title="Strong Scaling Incompressible Times")
display(plt)
savefig(plt, "ss_incompressible_times.png")

plt2 = plot(rank_num, eff_ratio, lw=4, xaxis=:log2, legend=:none, ylims=(0,1.1),
            xlabel="Cores", ylabel="Efficiency", title="Strong Scaling Incompressible Efficiency")
display(plt2)
savefig(plt2, "ss_incompressible_efficiency.png")

df = benchmarks_dataframe(suite)
sort!(df, :ranks)
benchmarks_pretty_table(df, title="Incompressible model strong scaling benchmark")

base_case = (grid_size(1, decomposition), rank_size(1, decomposition))
suite_Δ = speedups_suite(suite, base_case=base_case)
df_Δ = speedups_dataframe(suite_Δ, efficiency=:strong, base_case=base_case, key2rank=k->prod(k[2]))
sort!(df_Δ, :ranks)
benchmarks_pretty_table(df_Δ, title="Incompressible model strong scaling speedup")


export phase, unphase, findploidy

"""
    isbinary(filepath::String)
Returns `true` if the `filepath` is a binary file.
"""
function isbinary(filepath::String)
    !isvalid(String(Base.read(filepath, 1000)))
end

## Utilities for file reading/writing ##

#TODO change in docs
"""
    findploidy(genotypes::T where T<:AbstractVector)
Used internally in the `genepop` and `delimited` file parsers to scan the genotypes
of a sample and return the ploidy of the first non-missing locus.
"""
@inline function findploidy(genotypes::T) where T<:AbstractVector
    @inbounds Int8(length(first(skipmissing(genotypes))))
end

"""
    phase(loc::Union{String, Int}, type::DataType, digit::Int)
Takes a String of numbers or Integers and returns a typed locus appropriate for PopGen.jl as used in the
`genepop` and `csv` file parsers. Use `type` to specify output type (`Int8` or `Int16`),
and `digit` to specify the number of digits/characters used per allele in a locus.

## Examples
```
ph_locus = phase("128114", Int16, 3)
map(i -> phase(i, Int16, 3), ["112131", "211112", "001003", "516500"])
# or #
[phase(i, Int8, 2) for i in ["0101", "0103", "0202", "0103"]]
```
"""
@inline function phase(loc::T, type::DataType, digit::Int) where T<:AbstractString
    loc == "-9" || iszero(parse(Int, loc)) && return missing
    phased = map(i -> parse(type, join(i)), Iterators.partition(loc, digit))
    NTuple{length(phased), type}(sort(phased))
end

phase(loc::Missing, type::DataType, digit::Int) = missing

@inline function phase(loc::T, type::DataType, digits::T) where T<:Integer
    loc == -9 || iszero(loc) && return missing
    out = type[]
    units = 10^digits
    d,r = divrem(loc, units)
    out = type[r]
    #@inbounds push!(out, r)
    @inbounds while d != 0
        d,r = divrem(d, units)
        @inbounds push!(out, r)
    end
    return NTuple{length(out), type}(sort(out))
end

"""
    unphase(geno::T; digits::Int = 3, ploidy::Int = 2, miss::Int = 0) where T <: Genotype
Takes a `Genotype` (e.g. `(131, 94)`) and returns a string of concatenated
alleles padded with *n* number of zeroes, where *n* is given by `digits = `.
`missing` values are returned as either a string of 'digits × ploidy' zeroes (`miss = 0`)
or `"-9"` (`miss = -9`). The `ploidy` flag is only relevant for unphasing `missing` genotypes
and not used otherwise.

### Example
```
unphase((1,2,3,4), digits = 3)
"001002003004"

unphase(missing, digits = 2, ploidy = 2, miss = -9)
"-9"

unphase(missing, digits = 2, ploidy = 2, miss = 0)
"0000"
```
"""
function unphase(geno::T; digits::Int = 3, ploidy::Int = 2, miss::Int = 0) where T <: Genotype
    join(map(i -> lpad(i, digits, "0"), geno))
end

function unphase(geno::Missing; digits::Int = 3, ploidy::Int, miss::Int = 0)
    if miss == 0
        return "0"^(digits * ploidy)
    else
        return "$miss"
    end
end


abstract type DataFormat end

struct IteratorFormat end

struct ArrayFormat{D} end

###
function get_format(x::AbstractArray, condition, default=getdim_error)
    return getdim(x, condition, default)
end

to_format(::IteratorFormat, ::IteratorFormat, data) = data

# TODO
#function to_format(::ArrayFormat{D_dst}, src::IteratorFormat, data) where {D_dst}
#    return eachslice(data; dims=D_src)
#end

function to_format(::IteratorFormat, src::ArrayFormat{D_src}, data) where {D_src}
    return eachslice(data; dims=D_src)
end

to_format(::ArrayFormat{D}, ::ArrayFormat{D}, data) where {D} = data

function to_format(::ArrayFormat{D_dst}, ::ArrayFormat{D_src}, data::A) where {D_dst,D_src,A}
    perm = ntuple(ndims(A)) do ii
        if ii === D_dst
            return D_src
        elseif ii === D_src
            return D_dst
        else
            return ii
        end
    end
    return PermutedDimsArray(data, perm)
end

function to_format(dst, src, condition, default=getdim_error)
    return to_format(get_format(dst, condition, default),
                     get_format(src, condition, default),
                     src)
end