text:
__precompile__()

####################
module SDELab2
import NLsolve
import ForwardDiff

include("iteratedintegral.jl")
import .iteratedintegral

export set_opt, sde_strong_solution, set_fcn,  method_code, do_plotting

method_code=Dict()
method_code["SIE0"]="Strong Ito Euler "
method_code["SSEH0"]="Strong Stratonovich Euler-Heun "
method_code["SIE1"]="Strong Ito Euler (drift implicit)"
method_code["SSEH1"]="Strong Stratonovich Euler (drift implicit)"
method_code["SIE"]="Strong Ito Euler (alpha drift-implicit)"
method_code["SSEH"]="Strong Stratonovich Euler (alpha drift-implicit)"
method_code["SIM0"]="Strong Ito Milstein explicit"
method_code["SSM0"]="Strong Stratonovich Milstein explicit"
method_code["SIM1"]="Strong Ito Milstein (drift implicit)"
method_code["SSM1"]="Strong Stratonovich Milstein (drift implicit)"
method_code["SIM"]="Strong Ito Milstein (alpha drift-implicit)"
method_code["SSM"]="Strong Stratonovich Milstein (alpha drift-implicit)"
method_code["SIBDF"]="Strong Ito BDF-2"
# sdelab_gq
method_code["MFEM0"]="Mean-field SDE Gauss-quadrature Euler-Maruyama"
method_code["MFEM2"]="Mean-field SDE Gauss-quadrature second-order explicit"
method_code["MFEM"] ="Mean-field SDE Gauss-quadrature alpha-drift implicit (first order)"
method_code["MFEM0e"]="Mean-field SDE Gauss-quadrature Euler-Maruyama (exterior mean-field evaluation)"
method_code["MFEMe"] ="Mean-field SDE Gauss-quadrature alpha-drift implicit (first order, exterior mean-field)"

include("gauss_quadrature.jl")
using .gauss_quadrature
include("sdelab_gq.jl")
using .sdelab_gq
#####
function __init__()
end
##########################################
function get_increment!(x)
  x.dW[:]=sqrt(x.dt)*randn(x.m)
end
##########################################
# explicit Euler methods (i.e., Euler-Maruyama and Heun)
# data structure to store required bits
# and one-step routines
struct EM
  m
  drift
  diff_mat
  dt
  dW
end
# Ito version
function one_step_em(u,x::EM)
  get_increment!(x)
  un=u+x.drift(u)*x.dt+x.diff_mat(u,x.dW)
  return un
end
# Stratonovich (Heun method)
function one_step_heun(u,x::EM)
  get_increment!(x)
  uaux=u+x.diff_mat(u,x.dW)
  un=u+x.drift(u)*x.dt+0.5*(x.diff_mat(u,x.dW)+x.diff_mat(uaux,x.dW))
  return un
end
##################################################
# implicit Euler methods
# parameters for nonlinear solver
struct NLSOLVE_param
  maxfeval
  ftol
  xtol
end
# required bits for implicit-type Euler methods
struct EMImp
  m
  drift
  diff_mat
  Alpha
  dt
  dW
  nonlin
  jnonlin
  nlsolve::NLSOLVE_param
end
# Ito version of onestep
function one_step_em_alpha(u,x::EMImp)
  get_increment!(x)
  up=u + (1-x.Alpha)*x.drift(u)*x.dt + x.diff_mat(u,x.dW)
  solve=NLsolve.nlsolve((u,out)->x.nonlin(u,up,out),x.jnonlin,u, iterations=x.nlsolve.maxfeval, ftol=x.nlsolve.ftol, xtol=x.nlsolve.xtol)
  return solve.zero
end
# Stratonovich version of onestep
function one_step_heun_alpha(u,x::EMImp)
  get_increment!(x)
  uaux=u+x.diff_mat(u,x.dW)
  up=u+ (1-x.Alpha)*x.drift(u)*x.dt+0.5*(x.diff_mat(u,x.dW)+x.diff_mat(uaux,x.dW))
  solve=NLsolve.nlsolve((u,out)->x.nonlin(u,up,out),x.jnonlin,u)
  return solve.zero
end
# BDF method
# required bits
struct BDF
  m
  drift
  diff_mat
  dt
  dW
  dWp
  up
  nonlin
  jnonlin
  nlsolve::NLSOLVE_param
end
# Ito BDF
function one_step_bdf(u,x::BDF)
  get_increment!(x)
  up=(4/3)*u -(1/3)*x.up +x.diff_mat(u,x.dW)-(1/3)* x.diff_mat(x.up,x.dWp)
  solve=NLsolve.nlsolve((u,out)->x.nonlin(u,up,out),x.jnonlin,u, iterations=x.nlsolve.maxfeval, ftol=x.nlsolve.ftol, xtol=x.nlsolve.xtol)
  x.up[:]=u
  x.dWp[:]=x.dW
  return solve.zero
end
##########################################################
# Explicit Milstein methods
struct Mil
  d
  m
  drift
  diff_mat
  Jdiff_vecs
  S
  dt
end
# Milstein explicit, Stratonovich version
function one_step_strat_mil(u,x::Mil)
  IteratedIntegral.get_strat!(x.dt,x.S)
  J=x.S["J"];  K=x.S["K"]
  un=u+x.drift(u)*x.dt+x.diff_mat(u,J)
  for i=1:x.m
    Ki=reshape(K[i,:],(x.m,1))
    un+=x.Jdiff_vecs[i](u)*x.diff_mat(u,Ki)
  end
  return reshape(un,x.d)
end
# Milstein explicit, Ito version
function one_step_ito_mil(u,x::Mil)
  IteratedIntegral.get_ito!(x.dt,x.S)
  J=x.S["J"];  K=x.S["K"]
  un=u+x.drift(u)*x.dt+x.diff_mat(u,J)
  for i=1:x.m
    Ki=reshape(K[i,:],(x.m,1))
    un+=x.Jdiff_vecs[i](u)*x.diff_mat(u,Ki)
  end
  return reshape(un,x.d )
end
###################################################
# Implicit Milstein methods
struct MilImp
  d
  m
  drift
  diff_mat
  Jdiff_vecs
  Alpha
  S
  dt
  nonlin
  jnonlin
  nlsolve::NLSOLVE_param
end
# Milstein implicit Stratonovich
function one_step_strat_mil_alpha(u,x::MilImp)
  IteratedIntegral.get_strat!(x.dt,x.S)
  J=x.S["J"];   K=x.S["K"]
  up=u + x.Alpha*x.drift(u)*x.dt + x.diff_mat(u,J)
  for i=1:x.m
    Ki=reshape(K[i,:],(x.m,1))
    up+=x.Jdiff_vecs[i](u)*x.diff_mat(u,Ki)
  end
  solve=NLsolve.nlsolve((u,out)->x.nonlin(u,up,out),x.jnonlin,u,iterations=x.nlsolve.maxfeval, ftol=x.nlsolve.ftol, xtol=x.nlsolve.xtol)
  return solve.zero
end
# Milstein implicit Ito
function one_step_ito_mil_alpha(u,x::MilImp)
  IteratedIntegral.get_ito!(x.dt,x.S)
  J=x.S["J"];   K=x.S["K"]
  up=u+ x.Alpha*x.drift(u)*x.dt +x.diff_mat(u,J)
  for i=1:x.m
    Ki=reshape(K[i,:],(x.m,1))
    up += x.Jdiff_vecs[i](u)*x.diff_mat(u,Ki)
  end
  solve=NLsolve.nlsolve((u,out)->x.nonlin(u,up,out), x.jnonlin, u, iterations=x.nlsolve.maxfeval, ftol=x.nlsolve.ftol, xtol=x.nlsolve.xtol)
  return solve.zero
end
###################################################
# Set solver options (many defaults)
function set_opt(maxStepSize,solver)
  opt=Dict()
  opt["MaxStepSize"]=maxStepSize
  opt["Solver"]=solver
  opt["Alpha"]=0.5
  opt["nlsolve_param"]=NLSOLVE_param(10,1e-6,1e-8)
  opt["max_p"]=1000
  opt["calc_final_W"]=true
  opt["is_ito"]=(solver[1:2]=="SI")
  # plotting
  opt["always_clear"]=true
  opt["do_plotting"]=false
  opt["output_sel"]=[1,2]
  return opt
end
#
#
function sde_strong_solution(fcn, tspan, y0, opt)
  #=
  many driver routine:
  fcn= SDE definition
  tspan= vector of required times
  y0=value of sample path at first time in tspan
  opt=solver options
  =#
  ts=minimum(tspan); tf=maximum(tspan)
  N=round((tf-ts)/opt["MaxStepSize"]); dt=(tf-ts)/N
  solver=opt["Solver"]
  t_out=Array{Float64}(undef,length(tspan))
  y_out=Array{Float64}(undef,length(tspan),fcn.d)
  y=deepcopy(y0); y_out[1,:]=y # initial data
  t=ts; t_out[1]=ts # initial time
  #
  Wf=zeros(fcn.m)
  #
  if solver=="SIE0" || solver=="SSEH0"# StrongItoEuler Alpha=0 explicit
    x=EM(fcn.m,fcn.drift,fcn.diff_mat,dt,zeros(fcn.m))
    dW=x.dW
    if solver=="SIE0"
      call_onestep(y)=(one_step_em(y,x))
    else
      call_onestep(y)=(one_step_heun(y,x))
    end
  elseif solver=="SIE1" || solver=="SSEH1"# strongItoEuler Alpha=1 implicit
    nlsolve_param=opt["nlsolve_param"]
    function nonlin!(un_,up,out)
      out[:] =(un_-up-fcn.drift(un_)*dt) # must avoid recreating out!
    end
    function jnonlin!(un_,out)
      out[:,:]=eye(fcn.d)-fcn.Jdrift(un_)*dt
    end
    x=EMImp(fcn.m,fcn.drift,fcn.diff_mat,1.,dt,zeros(fcn.m),nonlin!,jnonlin!,nlsolve_param)
    dW=x.dW
    if solver=="SIE1"
      call_onestep(y)=(one_step_em_alpha(y,x))
    else
      call_onestep(y)=(one_step_heun_alpha(y,x))
    end
  elseif solver=="SIE" || solver =="SSEH"# strongItoEuler Alpha parameter
    Alpha=opt["Alpha"]
    nlsolve_param=opt["nlsolve_param"]
    function nonlin!(un_,up,out)
      out[:] =(un_-up-Alpha*fcn.drift(un_)*dt) # must avoid recreating out!
    end
    function jnonlin!(un_,out)
      out[:,:]=eye(fcn.d)-Alpha*fcn.Jdrift(un_)*dt
    end
    x=EMImp(fcn.m,fcn.drift,fcn.diff_mat,Alpha,dt,zeros(fcn.m),nonlin!,jnonlin!, nlsolve_param)
    dW=x.dW
    if solver=="SIE"
      call_onestep(y)=(one_step_em_alpha(y,x))
    else
      call_onestep(y)=(one_step_heun_alpha(y,x))
    end
  elseif solver=="SIM0" || solver=="SSM0"# StrongItoMilstein Alpha=0 explicit
    S=IteratedIntegral.init(fcn,opt)
    x=Mil(fcn.d,fcn.m,fcn.drift,fcn.diff_mat,fcn.Jdiff_vecs,S,dt)
    dW=x.S["J"]
    if solver=="SIM0"
      call_onestep(y)=(one_step_ito_mil(y,x))
    else
      call_onestep(y)=(one_step_strat_mil(y,x))
    end
  elseif solver=="SIM1" || solver=="SSM1"# StrongItoMilstein Alpha=1 implicit
      S=IteratedIntegral.init(fcn,opt)
      nlsolve_param=opt["nlsolve_param"]
      function nonlin!(un_,up,out)
        out[:] =(un_-up-fcn.drift(un_)*dt) # must avoid recreating out!
      end
      function jnonlin!(un_,out)
        out[:,:]=eye(fcn.d)-fcn.Jdrift(un_)*dt
      end
      x=MilImp(fcn.d,fcn.m,fcn.drift,fcn.diff_mat,fcn.Jdiff_vecs,0.,S,dt,nonlin!,jnonlin!,nlsolve_param)
      dW=x.S["J"]
      if solver=="SIM1"
        call_onestep(y)=(one_step_ito_mil_alpha(y,x))
      else
        call_onestep(y)=(one_step_strat_mil_alpha(y,x))
      end
    elseif solver=="SIM" || solver=="SSM"
      S=IteratedIntegral.init(fcn,opt)
      nlsolve_param=opt["nlsolve_param"]
      Alpha=opt["Alpha"]
      function nonlin!(un_,up,out)
        out[:] =(un_-up-Alpha*fcn.drift(un_)*dt) # must avoid recreating out!
      end
      function jnonlin!(un_,out)
        out[:,:]=eye(fcn.d)-Alpha*fcn.Jdrift(un_)*dt
      end
      x=MilImp(fcn.d,fcn.m,fcn.drift,fcn.diff_mat,fcn.Jdiff_vecs,Alpha,S,dt,nonlin!,jnonlin!,nlsolve_param)
      dW=x.S["J"]
      if solver=="SIM"
        call_onestep(y)=(one_step_ito_mil_alpha(y,x))
      else
        call_onestep(y)=(one_step_strat_mil_alpha(y,x))
      end
  elseif solver=="SIBDF"
    nlsolve_param=opt["nlsolve_param"]
    # starting method
    function nonlin0!(un_,up,out)
      out[:] =(un_-up-fcn.drift(un_)*dt) # must avoid recreating out!
    end
    function jnonlin0!(un_,out)
      out[:,:]=eye(fcn.d)-fcn.Jdrift(un_)*dt
    end
    x0=EMImp(fcn.m,fcn.drift,fcn.diff_mat,1.,dt,zeros(fcn.m),nonlin0!,jnonlin0!,nlsolve_param)
    y1=one_step_em_alpha(y,x0)
    t=t+dt
    Wf+=x0.dW
    # bdf method
    function nonlin!(un_,up,out)
      out[:] =(un_-up-(2/3)*fcn.drift(un_)*dt) # must avoid recreating out!
    end
    function jnonlin!(un_,out)
      out[:,:]=eye(fcn.d)-(2/3)*fcn.Jdrift(un_)*dt
    end
    x=BDF(fcn.m,fcn.drift,fcn.diff_mat,dt,zeros(fcn.m),x0.dW,deepcopy(y0),nonlin!,jnonlin!,nlsolve_param)
    y=y1;
    dW=x.dW
    call_onestep(y)=(one_step_bdf(y,x))
  end

  #
  for j=2:(length(tspan)) # loop over required times
    N=round(Int,(tspan[j]-t)/dt) # steps to next tspan node
    for i=1:N # does not run if N=0
      t=t+dt
      y=call_onestep(y)
      if opt["calc_final_W"]==true
        Wf+=dW
      end
    end
    t_out[j]=t
    y_out[j,:]=y # constant approx if tspan too fine relative to dt
  end
  try
    do_plotting(t_out,y_out,opt)
  catch
    print("plotting failed")
  end
  return t_out,y_out,Wf
end
#
using PyPlot
function do_plotting(t,y,opt)
  if opt["do_plotting"]
    if opt["always_clear"]
      clf()
    end
    str=opt["output_plot_type"]
    if str=="path_plot"
      sel=opt["output_sel"]
      plot(t,y[:,sel])
    elseif str=="phase_plot"
      sel=opt["output_sel"]
      plot(y[:,sel[1]], y[:,sel[2]])
    elseif str=="time_phase"
      print("not yet implemented :-(")
    elseif str=="phase3_plot"
      print("not yet implemented :-(")
    end
  end
end
#
struct Fcn
  d # state-space dimension
  m # no. Brownian motions
  drift # u -> f(u)
  Jdrift # Jacobian of drift (auto computed)
  diff_mat # (u,dw)-> g(u)*dw
  diff_vecs # [i]: u -> g_i(u) for i=1:m
  Jdiff_vecs # Jacobian of diffusion vectors (auto computed)
  noise_type # 1=unstructure; 2=diagonal; 3=commutative
end
#
function set_fcn(d,m,drift,diff_mat,diff_vecs,noise_type=1)
  #=
  Define sde via Fcn type
  =#
  Jdrift=ForwardDiff.jacobian(drift)
  if diff_vecs!=Union{}
    Jdiff_vecs=Array{Function}(undef,m)
    for i=1:m
      Jdiff_vecs[i]=ForwardDiff.jacobian(diff_vecs[i])
    end
  else
    Jdiff_vecs=Union{}
  end
  fcn=Fcn(d,m,drift,Jdrift,diff_mat,diff_vecs,Jdiff_vecs,noise_type)
end
#
end #sdelab2
#


text:
module DataIO

include("filesystem/filesystem.jl")

end  # module DataIO

ext: jl


text:
const PLOTSIZE = (1000, 500)

RecipesBase.@recipe function f(ts::TimeSeries{T};
                               title="Time Series",
                               y_unit="y") where T
    size --> PLOTSIZE
    xguide --> "timestamps"
    yguide --> y_unit
    label --> ts.name
    @series begin
        # force an argument with `:=`
        seriestype := :path
        # return series data
        return ts.timestamps, ts.vals
    end
end
RecipesBase.@recipe function f(vec_ts::Vector{TimeSeries{T}};
                               title="Time Series",
                               y_unit="y") where T
    size --> PLOTSIZE
    xguide --> "timestamps"
    yguide --> y_unit
    for ts in vec_ts
        RecipesBase.@series begin
            label := ts.name
            return ts.timestamps, ts.vals
        end
    end
end


text:
using BenchmarkTools, JaynesCummings
cd(Base.source_path()[1:end-11])
include("benchmarkreport.jl")

# Define a parent BenchmarkGroup to contain our suite
const suite = BenchmarkGroup()

# Add some child groups to our benchmark suite.
suite["generation"] = BenchmarkGroup(["hamiltonian,state"])
suite["calculation"] = BenchmarkGroup(["densitymatrix,wignerfunction"])

# Parameters
N = 14
ω_q = 2*pi*6.57E9
ω_r = ω_q
g = 0.5*2*pi*19E6
t_f = 200E-9
t_samples = 200
w_samples = 6
ket = ["|g,0>","0.707106781+0.907106781im|g,1>","1.6025403-0.5im|g,3>"]

# State generation benchmark
suite["generation"]["initialstate"] = @benchmarkable gen_initialstate($N,$ket)

# We actually need the state
ρ = gen_initialstate(N,ket)

# Hamiltonian generation benchmark
suite["generation"]["hamiltonian"] = @benchmarkable gen_hamiltonian($N,$ω_q,$ω_r,$g)

# We actually need the hamiltonian
H = gen_hamiltonian(N,ω_q,ω_r,g)

# Time evolution operator benchmark
suite["generation"]["timeevo"] = @benchmarkable gen_timeevoarray($H,$t_f,$t_samples)

# We actually need the operator
times, U_array = gen_timeevoarray(H,t_f,t_samples)

# Qubit time evolution calculation
suite["calculation"]["qubitetimeevo"] = @benchmarkable calc_qubittimeevo($ρ,$U_array)

# We need the excited state probability
e_prob = calc_qubittimeevo(ρ,U_array)

# Photon number calculation benchmark
suite["calculation"]["photonnum"] = @benchmarkable calc_photonnumbers($N,$g,$times,$e_prob)

# Density matrix calculation benchmark
suite["calculation"]["densitymatrix"] = @benchmarkable calc_densitymatrix_resonator($ρ,$g,$times,$U_array)

# Wigner function calculation benchmark
suite["calculation"]["wignerfunc"] = @benchmarkable calc_wignerfunction_resonator($ρ,$g,$times,$U_array,$w_samples)

tune!(suite)
results = run(suite)
reportstr = sprint(io -> printreport(io, results))

f = open("results//bench-$(Libc.strftime("%Y%m%d_%H%M%S",time())).md","w")
write(f,reportstr)
close(f)


text:
# Native execution support

export @cuda, cudaconvert, cufunction, dynamic_cufunction, nextwarp, prevwarp

## high-level @cuda interface

"""
    @cuda [kwargs...] func(args...)

High-level interface for executing code on a GPU. The `@cuda` macro should prefix a call,
with `func` a callable function or object that should return nothing. It will be compiled to
a CUDA function upon first use, and to a certain extent arguments will be converted and
managed automatically using `cudaconvert`. Finally, a call to `cudacall` is
performed, scheduling a kernel launch on the current CUDA context.

Several keyword arguments are supported that influence the behavior of `@cuda`.
- `launch`: whether to launch this kernel, defaults to `true`. If `false` the returned
  kernel object should be launched by calling it and passing arguments again.
- `dynamic`: use dynamic parallelism to launch device-side kernels, defaults to `false`.
- arguments that influence kernel compilation: see [`cufunction`](@ref) and
  [`dynamic_cufunction`](@ref)
- arguments that influence kernel launch: see [`CUDA.HostKernel`](@ref) and
  [`CUDA.DeviceKernel`](@ref)
"""
macro cuda(ex...)
    # destructure the `@cuda` expression
    call = ex[end]
    kwargs = ex[1:end-1]

    # destructure the kernel call
    Meta.isexpr(call, :call) || throw(ArgumentError("second argument to @cuda should be a function call"))
    f = call.args[1]
    args = call.args[2:end]

    code = quote end
    vars, var_exprs = assign_args!(code, args)

    # group keyword argument
    macro_kwargs, compiler_kwargs, call_kwargs, other_kwargs =
        split_kwargs(kwargs,
                     [:dynamic, :launch],
                     [:minthreads, :maxthreads, :blocks_per_sm, :maxregs, :name],
                     [:cooperative, :blocks, :threads, :config, :shmem, :stream])
    if !isempty(other_kwargs)
        key,val = first(other_kwargs).args
        throw(ArgumentError("Unsupported keyword argument '$key'"))
    end

    # handle keyword arguments that influence the macro's behavior
    dynamic = false
    launch = true
    for kwarg in macro_kwargs
        key,val = kwarg.args
        if key == :dynamic
            isa(val, Bool) || throw(ArgumentError("`dynamic` keyword argument to @cuda should be a constant value"))
            dynamic = val::Bool
        elseif key == :launch
            isa(val, Bool) || throw(ArgumentError("`launch` keyword argument to @cuda should be a constant value"))
            launch = val::Bool
        else
            throw(ArgumentError("Unsupported keyword argument '$key'"))
        end
    end
    if !launch && !isempty(call_kwargs)
        error("@cuda with launch=false does not support launch-time keyword arguments; use them when calling the kernel")
    end

    # FIXME: macro hygiene wrt. escaping kwarg values (this broke with 1.5)
    #        we esc() the whole thing now, necessitating gensyms...
    @gensym f_var kernel_f kernel_args kernel_tt kernel
    if dynamic
        # FIXME: we could probably somehow support kwargs with constant values by either
        #        saving them in a global Dict here, or trying to pick them up from the Julia
        #        IR when processing the dynamic parallelism marker
        isempty(compiler_kwargs) || error("@cuda dynamic parallelism does not support compiler keyword arguments")

        # dynamic, device-side kernel launch
        push!(code.args,
            quote
                # we're in kernel land already, so no need to cudaconvert arguments
                local $kernel_args = ($(var_exprs...),)
                local $kernel_tt = Tuple{map(Core.Typeof, $kernel_args)...}
                local $kernel = $dynamic_cufunction($f, $kernel_tt)
                if $launch
                    $kernel($kernel_args...; $(call_kwargs...))
                end
                $kernel
             end)
    else
        # regular, host-side kernel launch
        #
        # convert the function, its arguments, call the compiler and launch the kernel
        # while keeping the original arguments alive
        push!(code.args,
            quote
                $f_var = $f
                GC.@preserve $(vars...) $f_var begin
                    local $kernel_f = $cudaconvert($f_var)
                    local $kernel_args = map($cudaconvert, ($(var_exprs...),))
                    local $kernel_tt = Tuple{map(Core.Typeof, $kernel_args)...}
                    local $kernel = $cufunction($kernel_f, $kernel_tt; $(compiler_kwargs...))
                    if $launch
                        $kernel($(var_exprs...); $(call_kwargs...))
                    end
                    $kernel
                end
             end)
    end
    return esc(code)
end

## host to device value conversion

struct Adaptor end

# convert CUDA host pointers to device pointers
# TODO: use ordinary ptr?
Adapt.adapt_storage(to::Adaptor, p::CuPtr{T}) where {T} = reinterpret(LLVMPtr{T,AS.Generic}, p)

# Base.RefValue isn't GPU compatible, so provide a compatible alternative
struct CuRefValue{T} <: Ref{T}
  x::T
end
Base.getindex(r::CuRefValue) = r.x
Adapt.adapt_structure(to::Adaptor, r::Base.RefValue) = CuRefValue(adapt(to, r[]))

"""
    cudaconvert(x)

This function is called for every argument to be passed to a kernel, allowing it to be
converted to a GPU-friendly format. By default, the function does nothing and returns the
input object `x` as-is.

Do not add methods to this function, but instead extend the underlying Adapt.jl package and
register methods for the the `CUDA.Adaptor` type.
"""
cudaconvert(arg) = adapt(Adaptor(), arg)

## abstract kernel functionality

abstract type AbstractKernel{F,TT} end

# FIXME: there doesn't seem to be a way to access the documentation for the call-syntax,
#        so attach it to the type -- https://github.com/JuliaDocs/Documenter.jl/issues/558

"""
    (::HostKernel)(args...; kwargs...)
    (::DeviceKernel)(args...; kwargs...)

Low-level interface to call a compiled kernel, passing GPU-compatible arguments in `args`.
For a higher-level interface, use [`@cuda`](@ref).

The following keyword arguments are supported:
- `threads` (defaults to 1)
- `blocks` (defaults to 1)
- `shmem` (defaults to 0)
- `stream` (defaults to the default stream)
"""
AbstractKernel

@generated function call(kernel::AbstractKernel{F,TT}, args...; call_kwargs...) where {F,TT}
    sig = Base.signature_type(F, TT)
    args = (:F, (:( args[$i] ) for i in 1:length(args))...)

    # filter out arguments that shouldn't be passed
    predicate = if VERSION >= v"1.5.0-DEV.581"
        dt -> isghosttype(dt) || Core.Compiler.isconstType(dt)
    else
        dt -> isghosttype(dt)
    end
    to_pass = map(!predicate, sig.parameters)
    call_t =                  Type[x[1] for x in zip(sig.parameters,  to_pass) if x[2]]
    call_args = Union{Expr,Symbol}[x[1] for x in zip(args, to_pass)            if x[2]]

    # replace non-isbits arguments (they should be unused, or compilation would have failed)
    # alternatively, make it possible to `launch` with non-isbits arguments.
    for (i,dt) in enumerate(call_t)
        if !isbitstype(dt)
            call_t[i] = Ptr{Any}
            call_args[i] = :C_NULL
        end
    end

    # finalize types
    call_tt = Base.to_tuple_type(call_t)

    quote
        Base.@_inline_meta

        cudacall(kernel, $call_tt, $(call_args...); call_kwargs...)
    end
end

## host-side kernels

struct HostKernel{F,TT} <: AbstractKernel{F,TT}
    ctx::CuContext
    mod::CuModule
    fun::CuFunction
end

@doc (@doc AbstractKernel) HostKernel

@inline function cudacall(kernel::HostKernel, tt, args...; config=nothing, kwargs...)
    if config !== nothing
        Base.depwarn("cudacall with config argument is deprecated, use `@cuda launch=false` and instrospect the returned kernel instead", :cudacall)
        cudacall(kernel.fun, tt, args...; kwargs..., config(kernel)...)
    else
        cudacall(kernel.fun, tt, args...; kwargs...)
    end
end

"""
    version(k::HostKernel)

Queries the PTX and SM versions a kernel was compiled for.
Returns a named tuple.
"""
function version(k::HostKernel)
    attr = attributes(k.fun)
    binary_ver = VersionNumber(divrem(attr[FUNC_ATTRIBUTE_BINARY_VERSION],10)...)
    ptx_ver = VersionNumber(divrem(attr[FUNC_ATTRIBUTE_PTX_VERSION],10)...)
    return (ptx=ptx_ver, binary=binary_ver)
end

"""
    memory(k::HostKernel)

Queries the local, shared and constant memory usage of a compiled kernel in bytes.
Returns a named tuple.
"""
function memory(k::HostKernel)
    attr = attributes(k.fun)
    local_mem = attr[FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES]
    shared_mem = attr[FUNC_ATTRIBUTE_SHARED_SIZE_BYTES]
    constant_mem = attr[FUNC_ATTRIBUTE_CONST_SIZE_BYTES]
    return (:local=>local_mem, shared=shared_mem, constant=constant_mem)
end

"""
    registers(k::HostKernel)

Queries the register usage of a kernel.
"""
function registers(k::HostKernel)
    attr = attributes(k.fun)
    return attr[FUNC_ATTRIBUTE_NUM_REGS]
end

"""
    maxthreads(k::HostKernel)

Queries the maximum amount of threads a kernel can use in a single block.
"""
function maxthreads(k::HostKernel)
    attr = attributes(k.fun)
    return attr[FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK]
end

## host-side API

"""
    cufunction(f, tt=Tuple{}; kwargs...)

Low-level interface to compile a function invocation for the currently-active GPU, returning
a callable kernel object. For a higher-level interface, use [`@cuda`](@ref).

The following keyword arguments are supported:
- `minthreads`: the required number of threads in a thread block
- `maxthreads`: the maximum number of threads in a thread block
- `blocks_per_sm`: a minimum number of thread blocks to be scheduled on a single
  multiprocessor
- `maxregs`: the maximum number of registers to be allocated to a single thread (only
  supported on LLVM 4.0+)
- `name`: override the name that the kernel will have in the generated code

The output of this function is automatically cached, i.e. you can simply call `cufunction`
in a hot path without degrading performance. New code will be generated automatically, when
when function changes, or when different types or keyword arguments are provided.
"""
function cufunction(f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where
                   {F<:Core.Function, TT<:Type}
    dev = device()
    cache = cufunction_cache[dev]
    source = FunctionSpec(f, tt, true, name)
    target = CUDACompilerTarget(dev; kwargs...)
    params = CUDACompilerParams()
    job = CompilerJob(target, source, params)
    return GPUCompiler.cached_compilation(cache, job,
                                          cufunction_compile,
                                          cufunction_link)::HostKernel{f,tt}
end

const cufunction_cache = PerDevice{Dict{UInt, Any}}((dev)->Dict{UInt, Any}())

# compile to PTX
function cufunction_compile(@nospecialize(job::CompilerJob))
    # compile
    method_instance, world = GPUCompiler.emit_julia(job)
    ir, kernel = GPUCompiler.emit_llvm(job, method_instance, world)
    code = GPUCompiler.emit_asm(job, ir, kernel; format=LLVM.API.LLVMAssemblyFile)

    # check if we'll need the device runtime
    undefined_fns = filter(collect(functions(ir))) do f
        isdeclaration(f) && !LLVM.isintrinsic(f)
    end
    intrinsic_fns = ["vprintf", "malloc", "free", "__assertfail",
                    "__nvvm_reflect" #= TODO: should have been optimized away =#]
    needs_cudadevrt = !isempty(setdiff(undefined_fns, intrinsic_fns))

    # find externally-initialized global variables; we'll access those using CUDA APIs.
    external_gvars = filter(isextinit, collect(globals(ir))) .|> LLVM.name

    return (code, entry=LLVM.name(kernel), needs_cudadevrt, external_gvars)
end

# link to device code
function cufunction_link(@nospecialize(job::CompilerJob), compiled)
    ctx = context()

    # settings to JIT based on Julia's debug setting
    jit_options = Dict{CUjit_option,Any}()
    if Base.JLOptions().debug_level == 1
        jit_options[JIT_GENERATE_LINE_INFO] = true
    elseif Base.JLOptions().debug_level >= 2
        jit_options[JIT_GENERATE_DEBUG_INFO] = true
    end

    # link the CUDA device library
    # linking the device runtime library requires use of the CUDA linker,
    # which in turn switches compilation to device relocatable code (-rdc) mode.
    #
    # even if not doing any actual calls that need -rdc (i.e., calls to the runtime
    # library), this significantly hurts performance, so don't do it unconditionally
    intrinsic_fns = ["vprintf", "malloc", "free", "__assertfail",
                    "__nvvm_reflect" #= TODO: should have been optimized away =#]
    image = if compiled.needs_cudadevrt
        linker = CuLink(jit_options)
        add_file!(linker, libcudadevrt(), JIT_INPUT_LIBRARY)
        add_data!(linker, compiled.entry, compiled.code)
        complete(linker)
    else
        compiled.code
    end

    # JIT into an executable kernel object
    mod = CuModule(image, jit_options)
    fun = CuFunction(mod, compiled.entry)

    # initialize and register the exception flag, if any
    if "exception_flag" in compiled.external_gvars
        create_exceptions!(mod)
        filter!(!isequal("exception_flag"), compiled.external_gvars)
    end

    return HostKernel{job.source.f,job.source.tt}(ctx, mod, fun)
end

# https://github.com/JuliaLang/julia/issues/14919
(kernel::HostKernel)(args...; kwargs...) = call(kernel, map(cudaconvert, args)...; kwargs...)

## device-side kernels

struct DeviceKernel{F,TT} <: AbstractKernel{F,TT}
    fun::Ptr{Cvoid}
end

@doc (@doc AbstractKernel) DeviceKernel

@inline cudacall(kernel::DeviceKernel, tt, args...; kwargs...) =
    dynamic_cudacall(kernel.fun, tt, args...; kwargs...)

# FIXME: duplication with cudacall
@generated function dynamic_cudacall(f::Ptr{Cvoid}, tt::Type{T}, args::Vararg{Any,N};
                                     blocks=UInt32(1), threads=UInt32(1), shmem=UInt32(0),
                                     stream=CuDefaultStream()) where {T,N}
    types = tt.parameters[1].parameters     # the type of `tt` is Type{Tuple{<:DataType...}}

    ex = quote
        Base.@_inline_meta
    end

    # convert the argument values to match the kernel's signature (specified by the user)
    # (this mimics `lower-ccall` in julia-syntax.scm)
    converted_args = Vector{Symbol}(undef, length(args))
    arg_ptrs = Vector{Symbol}(undef, length(args))
    for i in 1:length(args)
        converted_args[i] = gensym()
        arg_ptrs[i] = gensym()
        push!(ex.args, :($(converted_args[i]) = Base.cconvert($(types[i]), args[$i])))
        push!(ex.args, :($(arg_ptrs[i]) = Base.unsafe_convert($(types[i]), $(converted_args[i]))))
    end

    append!(ex.args, (quote
        #GC.@preserve $(converted_args...) begin
            device_launch(f, blocks, threads, shmem, stream, $(arg_ptrs...))
        #end
    end).args)

    return ex
end

## device-side API

"""
    dynamic_cufunction(f, tt=Tuple{})

Low-level interface to compile a function invocation for the currently-active GPU, returning
a callable kernel object. Device-side equivalent of [`CUDA.cufunction`](@ref).

No keyword arguments are supported.
"""
@inline function dynamic_cufunction(f::Core.Function, tt::Type=Tuple{})
    fptr = GPUCompiler.deferred_codegen(Val(f), Val(tt))
    DeviceKernel{f,tt}(fptr)
end

# https://github.com/JuliaLang/julia/issues/14919
(kernel::DeviceKernel)(args...; kwargs...) = call(kernel, args...; kwargs...)

## other

"""
    nextwarp(dev, threads)
    prevwarp(dev, threads)

Returns the next or previous nearest number of threads that is a multiple of the warp size
of a device `dev`. This is a common requirement when using intra-warp communication.
"""
function nextwarp(dev::CuDevice, threads::Integer)
    ws = warpsize(dev)
    return threads + (ws - threads % ws) % ws
end

@doc (@doc nextwarp) function prevwarp(dev::CuDevice, threads::Integer)
    ws = warpsize(dev)
    return threads - Base.rem(threads, ws)
end


text:
function convert_typedef(cursor::LibClang.CXCursor, indent::Int)
	name = convert_name(cursor)

	((pre, typ, post), comments) = convert_type(cursor, LibClang.clang_getTypedefDeclUnderlyingType(cursor), indent)

	cvt = convert_nested_type(typ, cursor, indent)
	merge_comments!(comments, cvt.comments)
	merge_comment!(comments, convert_comment(cursor, name))

	expr = cvt.expr
	if startswith(expr, "𝐣𝐥.@") && (
		startswith(post, '[') ||                       # @cstruct {...}[N]
		(endswith(pre, '{') && startswith(post, ','))  # Cfunction{@cstruct {...}, ...}
	)
		expr = "($(expr))"
	end
	expr = "𝐣𝐥.@ctypedef $(name) $(pre)$(expr)$(post)"

	return Converted(
		expr,
		comments,
	)
end

function convert_enum(cursor::LibClang.CXCursor, indent::Int)
	name     = convert_name(cursor)
	comments = Dict{String, Comment}(convert_comment(cursor, name))

	if !Bool(LibClang.clang_isCursorDefinition(cursor))
		expr = "𝐣𝐥.@cenum $(name)"
	else
		typ = LibClang.clang_getEnumDeclIntegerType(cursor)
		isunsigned = typ.kind in (
			LibClang.CXType_Char_U,
			LibClang.CXType_UChar,
			LibClang.CXType_UShort,
			LibClang.CXType_UInt,
			LibClang.CXType_ULong,
			LibClang.CXType_ULongLong,
		)
		isunsigned || typ.kind in (
			LibClang.CXType_Char_S,
			LibClang.CXType_SChar,
			LibClang.CXType_Short,
			LibClang.CXType_Int,
			LibClang.CXType_Long,
			LibClang.CXType_LongLong,
		) || error("Failed to convert enum due to an unexpected integer type")
		typ = convert_name(typ, nolongs = true)

		pack = nothing
		vals = String[]
		for child in children(cursor)
			if child.kind == LibClang.CXCursor_EnumConstantDecl
				valName  = convert_name(child)
				valValue = isunsigned ? LibClang.clang_getEnumConstantDeclUnsignedValue(child) : LibClang.clang_getEnumConstantDeclValue(child)
				merge_comment!(comments, convert_comment(child, valName))
				push!(vals, "$(valName) = $(typ)($(valValue))")
			elseif child.kind == LibClang.CXCursor_PackedAttr && isnothing(pack)
				pack = convert_name(child)
			else
				error("Unexpected enum child encountered")
			end
		end

		tabs = repeat('\t', indent)
		vals = join(vals, "\n$(tabs)\t")
		vals = isempty(vals) ? " " : "\n$(tabs)\t$(vals)\n$(tabs)"
		expr = isempty(name) ? "" : "$(name) "
		expr = "𝐣𝐥.@cenum $(expr){$(vals)}"
		expr = isnothing(pack) ? expr : "$(expr) __$(pack)__"
	end

	return Converted(
		expr,
		comments,
	)
end

convert_struct(cursor::LibClang.CXCursor, indent::Int) = convert_aggregate(cursor, :struct, indent)
convert_union(cursor::LibClang.CXCursor, indent::Int) = convert_aggregate(cursor, :union, indent)

function convert_aggregate(cursor::LibClang.CXCursor, kind::Symbol, indent::Int)
	kind in (:struct, :union) || error("Unknown aggregate type $(kind)")
	name     = convert_name(cursor)
	comments = Dict{String, Comment}(convert_comment(cursor, name))

	if !Bool(LibClang.clang_isCursorDefinition(cursor))
		expr = "𝐣𝐥.@c$(kind) $(name)"
	else
		pack = nothing
		exprs = String[]
		for coal in coalesce(children(cursor))
			if length(coal) == 1 && coal[1].kind == LibClang.CXCursor_PackedAttr
				isnothing(pack) || error("Multiple packing attributes encountered")
				pack = convert_name(coal[1])
			else
				cvt = convert_fields(coal, indent+1)
				if !isnothing(cvt)
					push!(exprs, cvt.expr)
					merge_comments!(comments, cvt.comments)
				end
			end
		end

		tabs = repeat('\t', indent)
		exprs = join(exprs, "\n$(tabs)\t")
		exprs = isempty(exprs) ? " " : "\n$(tabs)\t$(exprs)\n$(tabs)"
		expr = isempty(name) ? "" : "$(name) "
		expr = "𝐣𝐥.@c$(kind) $(expr){$(exprs)}"
		expr = isnothing(pack) ? expr : "$(expr) __$(pack)__"
	end

	return Converted(
		expr,
		comments,
	)
end

function convert_nested_type(typ::LibClang.CXType, cursor::LibClang.CXCursor, indent::Int)
	decl = LibClang.clang_getTypeDeclaration(typ)
	if decl.kind == LibClang.CXCursor_EnumDecl && decl in cursor && !(cursor in decl)
		convert_decl = convert_enum
	elseif decl.kind == LibClang.CXCursor_StructDecl && decl in cursor && !(cursor in decl)
		convert_decl = convert_struct
	elseif decl.kind == LibClang.CXCursor_UnionDecl && decl in cursor && !(cursor in decl)
		convert_decl = convert_union
	else
		return Converted(convert_name(typ), Dict{String, Comment}())
	end
	return convert_decl(decl, indent)
end

function convert_type(cursor::LibClang.CXCursor, typ::LibClang.CXType, indent::Int)
	(pre, post) = ("", "")
	comments = Dict{String, Comment}()

	if typ.kind in (
		LibClang.CXType_Bool,
		LibClang.CXType_Void,
		LibClang.CXType_Char_S,
		LibClang.CXType_Char_U,
		LibClang.CXType_SChar,
		LibClang.CXType_UChar,
		LibClang.CXType_Short,
		LibClang.CXType_Int,
		LibClang.CXType_Long,
		LibClang.CXType_LongLong,
		LibClang.CXType_UShort,
		LibClang.CXType_UInt,
		LibClang.CXType_ULong,
		LibClang.CXType_ULongLong,
		LibClang.CXType_Float,
		LibClang.CXType_Double,
		LibClang.CXType_LongDouble,
		LibClang.CXType_Typedef,
		LibClang.CXType_Elaborated,
		LibClang.CXType_Enum,
		LibClang.CXType_Record,
	)
		t = typ
	elseif typ.kind in (
		LibClang.CXType_ConstantArray,
		LibClang.CXType_IncompleteArray,
	)
		num = typ.kind == LibClang.CXType_ConstantArray ? LibClang.clang_getNumElements(typ) : ""
		((pre, t, post), comments) = convert_type(cursor, LibClang.clang_getElementType(typ), indent)
		if endswith(post, ']')
			(pre, post) = ("("*pre, post*")")
		end
		(pre, post) = (pre, post*"[$(num)]")
	elseif typ.kind == LibClang.CXType_Complex
		((pre, t, post), comments) = convert_type(cursor, LibClang.clang_getElementType(typ), indent)
		(pre, post) = ("𝐣𝐥.Complex{"*pre, post*"}")
	elseif typ.kind == LibClang.CXType_Pointer
		((pre, t, post), comments) = convert_type(cursor, LibClang.clang_getPointeeType(typ), indent)
		(pre, post) = ("𝐣𝐥.Ptr{"*pre, post*"}")
	elseif typ.kind in (LibClang.CXType_Unexposed, LibClang.CXType_FunctionProto)
		((pre, t, post), comments) = convert_type(cursor, LibClang.clang_getResultType(typ), indent)

		num = LibClang.clang_getNumArgTypes(typ)
		args = map(1:num) do i
			argType = LibClang.clang_getArgType(typ, i-1)

			((argPre, argT, argPost), argComments) = convert_type(cursor, argType, indent)
			merge_comments!(comments, argComments)

			decl = LibClang.clang_getTypeDeclaration(argT)
			if decl.kind == LibClang.CXCursor_EnumDecl && decl in cursor
				convert_decl = convert_enum
			elseif decl.kind == LibClang.CXCursor_StructDecl && decl in cursor
				convert_decl = convert_struct
			elseif decl.kind == LibClang.CXCursor_UnionDecl && decl in cursor
				convert_decl = convert_union
			else
				convert_decl = (_, _) -> Converted(convert_name(argT), Dict{String, Comment}())
			end

			cvt = convert_decl(decl, indent)
			merge_comments!(comments, cvt.comments)

			expr = "$(argPre)$(cvt.expr)$(argPost)"
			if startswith(expr, "𝐣𝐥.@") && num > 1
				expr = "($(expr))"
			end
			return expr
		end
		Bool(LibClang.clang_isFunctionTypeVariadic(typ)) && push!(args, "𝐣𝐥.Vararg")
		args = join(args, ", ")

		conv = convert_convention(typ)

		(pre, post) = ("𝐣𝐥.Cfunction{"*pre, post*", 𝐣𝐥.Tuple{$(args)}, 𝐣𝐥.$(conv)}")
	else
		error("Unable to convert type $(typ.kind), not yet implemented")
	end

	if Bool(LibClang.clang_isConstQualifiedType(typ))
		(pre, post) = ("𝐣𝐥.Cconst("*pre, post*")")
	end

	if Bool(LibClang.clang_isRestrictQualifiedType(typ))
		(pre, post) = ("𝐣𝐥.Crestrict("*pre, post*")")
	end

	if Bool(LibClang.clang_isVolatileQualifiedType(typ))
		(pre, post) = ("𝐣𝐥.Cvolatile("*pre, post*")")
	end

	return ((pre, t, post), comments)
end

function convert_convention(typ::LibClang.CXType)
	conv = LibClang.clang_getFunctionTypeCallingConv(typ)
	if conv == LibClang.CXCallingConv_C
		return "CDECL"
	end

	error("Support for function calling convention `$(conv)` is not yet implemented")
end


text:
using Zygote
using LinearAlgebra
using GalacticOptim,NLopt,Optim
using Plots
using GLOQ

using Random
Random.seed!(14);

# System parameters for a simple two level open quantum system
N_states = 2; # number of states
freqs = [4.1] # transition frequency in GHz
omegas = 2.0*pi.*freqs # change to angular frequency
gamma1   = [1.0/(45.0*GLOQ.GLOQ_MICRO_SEC)] # Reciprocal of relaxation time - T1 (in units of ns)
gamma2   = [1.0/(24.0*GLOQ.GLOQ_MICRO_SEC)] # Reciprocal of dephasing time - T2 (in units of ns)
omr_ramsey = 2.0*pi*(4.1 - 5.0e-4) # drive frequency for the Ramsey experiment
omr_echo = 2.0*pi*4.1  # drive frequency for the Echo experiment
omr_t1 = 2.0*pi*4.1  # drive frequency for the T1 experiment
TC = 2.5*17.0 # total control time

# Initial state
initial_state = 0
state_u0 = [0.0;0.0]
state_v0 = [0.0;0.0]
state_u0[initial_state+1] = 1.0

# Forward solve to generate synthetic data for the Ramsey experiment
# Duration of the Ramsey experiment, largest dark time
T_Ramsey = 10.0*GLOQ.GLOQ_MICRO_SEC # convert micro-sec to nano-sec
# total number of dark time samples
N_dark_ramsey = 201
dt_ramsey = T_Ramsey/N_dark_ramsey
t_dark_ramsey = collect(range(0.0, T_Ramsey, length=N_dark_ramsey))
# Forward solve
rho_synthetic_ramsey_u,rho_synthetic_ramsey_v = GLOQ.RamseyForwardSolve(
				 state_u0,state_v0, # initial values, u for the real part, v for the imaginary part
			     omegas,omr_ramsey, # transition frequencies, drive frequency
				 gamma1,gamma2, # decay and dephasing
				 initial_state, # initial state
				 TC,t_dark_ramsey,N_states) # control time, dark time, total number of states
population_ramsey_synthetic = GLOQ.get_population(rho_synthetic_ramsey_u)

# Echo experiment
# Duration of the Echo experiment, largest dark time
T_Echo = 30.0*GLOQ.GLOQ_MICRO_SEC # convert micro-sec to nano-sec
# total number of dark time samples
N_dark_echo = 601
# Step size of the Echo experiment
dt_echo = T_Echo/N_dark_echo
t_dark_echo = collect(range(0.0, T_Echo, length=N_dark_echo))
rho_synthetic_echo_u,rho_synthetic_echo_v = GLOQ.EchoForwardSolve(
				 state_u0,state_v0, # initial values, u for the real part, v for the imaginary part
			     omegas,omr_echo, # transition frequencies, drive frequency
				 gamma1,gamma2, # decay and dephasing
				 initial_state, # initial state
				 TC,t_dark_echo,N_states) # control time, dark time, total number of states
population_echo_synthetic = GLOQ.get_population(rho_synthetic_echo_u)

# T1
# Duration of the T1-decay experiment, largest dark time
T_t1 = 40.0*GLOQ.GLOQ_MICRO_SEC # convert micro-sec to nano-sec
# total number of dark time samples
N_dark_t1 = 401
# Step size of the T1 experiment
dt_t1 = T_t1/N_dark_t1
t_dark_t1 = collect(range(0.0, T_t1, length=N_dark_t1))
rho_synthetic_t1_u,rho_synthetic_t1_v = GLOQ.T1ForwardSolve(
				 state_u0,state_v0, # initial values, u for the real part, v for the imaginary part
			     omegas,omr_t1, # transition frequencies, drive frequency
				 gamma1,gamma2, # decay and dephasing
				 initial_state, # initial state
				 TC,t_dark_t1,N_states) # control time, dark time, total number of states
population_t1_synthetic = GLOQ.get_population(rho_synthetic_t1_u)

# Plot the synthetic data
#fig=plot(t_dark_times./GLOQ.GLOQ_MICRO_SEC,population_synthetic)
#display(fig)

# Define the loss function for the GalacticOptim
# p: phyiscal parameters:
#	 p[1] = transition frequency in GHz
#    p[2] = gamma1
#    p[3] = gamma2
# dummy_parameter: needed by GalacticOptim, one can just put [] here
function loss(p,dummy_parameter)
	# Ramsey
	_rho_ramsey_u,_rho_ramsey_v = GLOQ.RamseyForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_ramsey,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_ramsey,N_states)
	_population_ramsey = GLOQ.get_population(_rho_ramsey_u)
	# Echo
    _rho_echo_u,_rho_echo_v = GLOQ.EchoForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_echo,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_echo,N_states)
	_population_echo = GLOQ.get_population(_rho_echo_u)
	# T1
    _rho_t1_u,_rho_t1_v = GLOQ.T1ForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_t1,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_t1,N_states)
	_population_t1 = GLOQ.get_population(_rho_t1_u)

	_loss = sum(abs2,_population_ramsey-population_ramsey_synthetic)*dt_ramsey+
			sum(abs2,_population_echo-population_echo_synthetic)*dt_echo+
			sum(abs2,_population_t1-population_t1_synthetic)*dt_t1

	return _loss
end

plot_callback = function(p,other_args)
	# Ramsey
	rho_ramsey_u,rho_ramsey_v = GLOQ.RamseyForwardSolve(state_u0,state_v0,
					 (2*pi).*[p[1]],omr_ramsey,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_ramsey,N_states)
	population_ramsey = GLOQ.get_population(rho_ramsey_u)
	# Echo
    rho_echo_u,rho_echo_v = GLOQ.EchoForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_echo,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_echo,N_states)
	population_echo = GLOQ.get_population(rho_echo_u)

	# T1
    rho_t1_u,rho_t1_v = GLOQ.T1ForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_t1,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_t1,N_states)
	population_t1 = GLOQ.get_population(rho_t1_u)

	# Plot
	# Ramsey
	fig_ramsey = plot(t_dark_ramsey./GLOQ.GLOQ_MICRO_SEC,population_ramsey_synthetic,label=["Syn-0" "Syn-1"],
			          line = (:dash,0.0), marker = ([:hex :hex], 5, 0.5),legend=:outerright,
					  title="Ramsey");
	plot!(fig_ramsey,t_dark_ramsey./GLOQ.GLOQ_MICRO_SEC,population_ramsey,label=["Opt-0" "Opt-1"],legend=:outerright);
	# Echo
	fig_echo = plot(t_dark_echo./GLOQ.GLOQ_MICRO_SEC,population_echo_synthetic,label=["Syn-0" "Syn-1"],
		  		     line = (:dash,0.0), marker = ([:hex :hex], 5, 0.5),legend=:outerright,
					 title="Echo");
	plot!(fig_echo,t_dark_echo./GLOQ.GLOQ_MICRO_SEC,population_echo,label=["Opt-0" "Opt-1"],legend=:outerright);
	# T1
	fig_t1=plot(t_dark_t1./GLOQ.GLOQ_MICRO_SEC,population_t1_synthetic,label=["Syn-0" "Syn-1"],
			    line = (:dash,0.0), marker = ([:hex :hex], 5, 0.5),legend=:outerright,
				title="T1");
	plot!(fig_t1,t_dark_t1./GLOQ.GLOQ_MICRO_SEC,population_t1,label=["Opt-0" "Opt-1"],legend=:outerright);

	display( plot(fig_ramsey,fig_echo,fig_t1,layout=grid(3,1),size=[1000,1500],
				  legendfontsize=15,xtickfontsize=15,ytickfontsize=15,titlefontsize=18) )
	return false
end

p_true = [freqs;gamma1;gamma2] # values to generate synthetic data
# initial guess for the optimization
p_initial = [freqs.-0.5e-4;0.75.*gamma1;1.25.*gamma2]
# bounds for the optimization
lower_bound = (0.5).*p_true
upper_bound = (1.5).*p_true

# construct optimization object, use Zygote auto-differentiation to compute the gradient
loss_gradient = GalacticOptim.OptimizationFunction(loss, GalacticOptim.AutoZygote())
opt_prob = GalacticOptim.OptimizationProblem(loss_gradient, p_initial,
										 lb = lower_bound, ub = upper_bound)

println("Optim Fminbox(LBFGS) Optimization starts")
@time sol = GalacticOptim.solve(opt_prob ,Fminbox(LBFGS()),
								cb = plot_callback,
								outer_iterations = 20,
								iterations = 10,
								show_trace=true,
								f_tol = 1e-5,
								outer_f_tol = 1e-5)
println("Optim Fminbox(LBFGS) Optimization done")

# present the solutions
println("\nOptimized results: ",sol.u,
        "\nLoss: ",sol.minimum,
		"\nError: ",sol.u-p_true)


text:
using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using Random
Random.seed!(233)

function scatter_update(A::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}},
            ii,jj,B::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}})
    !isa(A, SparseTensor) && (A=SparseTensor(A))
    !isa(B, SparseTensor) && (B=SparseTensor(B))
    ii1, jj1, vv1 = find(A)
    m1_, n1_ = size(A)
    ii2, jj2, vv2 = find(B)

    sparse_scatter_update_ = load_op_and_grad("./build/libSparseScatterUpdate","sparse_scatter_update", multiple=true)
    ii1,jj1,vv1,m1,n1,ii2,jj2,vv2,ii,jj = convert_to_tensor([ii1,jj1,vv1,m1_,n1_,ii2,jj2,vv2,ii,jj], [Int64,Int64,Float64,Int64,Int64,Int64,Int64,Float64,Int64,Int64])
    ii, jj, vv = sparse_scatter_update_(ii1,jj1,vv1,m1,n1,ii2,jj2,vv2,ii,jj)
    SparseTensor(ii, jj, vv, m1_, n1_)
end

# TODO: specify your input parameters
A = sprand(10,10,0.3)
B = sprand(3,3,0.6)
ii = [1;4;5]
jj = [2;4;6]
u = scatter_update(A, ii, jj, B)
C = copy(A)
C[ii,jj] = B
sess = Session(); init(sess)
@show run(sess, u)-C

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(A, B, m)
    ii1, jj1, vv1 = find(A)
    ii2, jj2, vv2 = find(B)
    if length(m)==length(vv1)
        A = SparseTensor(ii1, jj1, m, size(A)...)
    else
        B = SparseTensor(ii2, jj2, m, size(B)...)
    end
    C = scatter_update(A, ii, jj, B)
    return sum(C)^2
end

# TODO: change `m_` and `v_` to appropriate values

A = sprand(10,10,0.3)|>SparseTensor
B = sprand(3,3,0.6)|>SparseTensor
ii = [1;4;5]
jj = [2;4;6]

m_ = constant(rand(length(values(A))))
v_ = rand(length(values(A)))
y_ = scalar_function(A, B, m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(A, B, ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


text:
export GROSS, NET
export S2, S2Mkt, S2MktEq, S2MktInt
export ProjParam

"Convenience index name"
const GROSS = 1
"Convenience index name"
const NET = 2

"Abstract type for an S2 module"
abstract S2Module

"Dummy type for those S2 modules that have not been implemented"
type S2NotImplemented <: S2Module
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

function S2NotImplemented(x...)
  return S2NotImplemented(zeros(Float64, 2))
end

"S2 market property risk module, not yet implemented"
typealias S2MktProp S2NotImplemented
"S2 market spread risk module, not yet implemented"
typealias S2MktSpread S2NotImplemented
"S2 market currency risk module, not yet implemented"
typealias S2MktFx S2NotImplemented
"S2 market concentration risk module, not yet implemented"
typealias S2MktConc S2NotImplemented
"S2 default (type 2) risk module, not yet implemented"
typealias S2Def2 S2NotImplemented
"S2 life morbidity risk module, not yet implemented"
typealias S2LifeMorb S2NotImplemented
"S2 life revision risk module, not yet implemented"
typealias S2LifeRevision S2NotImplemented
"S2 health risk module, not yet implemented"
typealias S2Health S2NotImplemented
"S2 non-life risk module, not yet implemented"
typealias S2NonLife S2NotImplemented

"Projection parameters for S2-calculations without
S2 calibration parameters"
type ProjParam
  "`Int`: Year in which projection starts"
  t_0::Int
  "`Int`: Year in which projection ends"
  T::Int
  "`CapMkt`: Capital market"
  cap_mkt::CapMkt
  "`Vector{Any}`: Parameters for the construction of the InvPort
  object with the exception of `t_0::Int`, `T::Int`,
  `cap_mkt::CapMkt`"
  invs_par::Vector{Any}
  "`LiabIns`: Liability porfolio (without other liabilites)"
  l_ins::LiabIns
  "`LiabOther`: Other liabilities"
  l_other::LiabOther
  "`Dynamic`: Dynamic projection parameters"
  dyn::Dynamic
  "`Float64`: (constant) tax rate during projection"
  tax_rate::Float64
  "`Float64`: Initial tax credit"
  tax_credit_0::Float64
end

## solvency 2 market risk =======================================
"Solvency 2 market risk: interest rate risk"
type S2MktInt <: S2Module
  "`Symbol`: Name of object type to be shocked"
  shock_object::Symbol
  "`Dict{Symbol, Any}`: Interest rate shocks"
  shock::Dict{Symbol, Any}
  "`Float64`: Minumum upwards shock"
  spot_up_abs_min::Float64
  "`DataFrame`:  Shocked balance sheets"
  balance::DataFrame
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
  "`Bool`: Was interest shocked upwards to obtain SCR?"
  scen_up::Bool
end

"Solvency 2 market risk: equity risk"
type S2MktEq <: S2Module
  "`Symbol`: Name of object type to be shocked"
  shock_object::Symbol
  "`Dict{Symbol, Symbol}`: Equity type for each invested stock"
  eq2type::Dict{Symbol, Symbol}
  "`Dict{Symbol, Any}`: Equity shocks"
  shock::Dict{Symbol,Any}
  "`DataFrame`:  Shocked balance sheets"
  balance::DataFrame
  "`Matrix{Float64}`: Correlation matrix for equity types"
  corr::Matrix{Float64}
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

"Solvency 2 market risk"
type S2Mkt <: S2Module
  "`Vector{S2Module}`: Sub-modules for market risk"
  mds::Vector{S2Module}
  "`Matrix{Float64}`: Correlation matrix if upwards shock was
  used in the calculation of the SCR for interest risk"
  corr_up::Matrix{Float64}
  "`Matrix{Float64}`: Correlation matrix if downwards shock was
  used in the calculation of the SCR for interest risk"
  corr_down::Matrix{Float64}
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

## solvency 2 default risk ======================================
"Solvency 2 default risk of type 1"
type S2Def1 <: S2Module
  "`Vector{Float64}`: Total loss given default per rating"
  tlgd::Vector{Float64}
  "`Vector{Float64}`: Squared loss given default per rating"
  slgd::Vector{Float64}
  "`Matrix{Float64}`: Matrix u for variance calculation"
  u::Matrix{Float64}
  "`Vector{Float64}`: Vector v for variance calculation"
  v::Vector{Float64}
  "`Dict{Symbol, Vector{Float64}}`: Parameters for the SCR
  calculation depending on whether the normalized standard
  deviation is low, medium, or high"
  scr_par::Dict{Symbol, Vector{Float64}}
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

"Solvency 2 default risk"
type S2Def <: S2Module
  "`Vector{S2Module}`: Sub-modules for default risk"
  mds::Vector{S2Module}
  "`Matrix{Float64}`: Correlation matrix for default risk"
  corr::Matrix{Float64}
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

## solvency 2 life risk =========================================
"Solvency 2 life risk: biometric risks: mortality, longevity,
surrender"
type S2LifeBio <: S2Module
  "`Symbol`: Name of object type to be shocked"
  shock_object::Symbol
  "`Dict{Symbol, Any}`: biometric shocks"
  shock::Dict{Symbol, Any}
  "`Dict{Symbol, Float64}`: Additional parameters for calculating
  shocks, may be empty"
  shock_param::Dict{Symbol, Float64}
  "`DataFrame`:  Shocked balance sheets"
  balance::DataFrame
  "`Dict{Symbol, Vector{Bool}}`: Indicator which model points
  have been selected to be shocked"
  mp_select::Dict{Symbol, Vector{Bool}}
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

"Solvency 2 life risk: expense risk"
type S2LifeCost <: S2Module
  "`Symbol`: Name of object type to be shocked"
  shock_object::Symbol
  "`Dict{Symbol, Any}`: biometric shocks"
  shock::Dict{Symbol, Any}
  "`Dict{Symbol, Float64}`: Additional parameters for calculating
  shocks"
  shock_param::Dict{Symbol, Float64}
  "`DataFrame`:  Shocked balance sheets"
  balance::DataFrame
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

"Solvency 2 life risk"
type S2Life <: S2Module
  "`Vector{S2Module}`: Sub-modules for life risk"
  mds::Vector{S2Module}
  "`Matrix{Float64}`: Correlation matrix for life risk"
  corr::Matrix{Float64}
  "`Vector{Float64}`: Gross and net SCR"
  scr::Vector{Float64}
end

## solvency 2 operational risk ==================================
"Solvency 2 operational risk"
type S2Op <: S2Module
  "`Dict{Symbol, Float64}`: Factors for SCR calculation"
  fac::Dict{Symbol, Float64}
  "`Float64`: Earned premium (previous year)"
  prem_earned::Float64
  "`Float64`: Earned premium (year before previous year)"
  prem_earned_prev::Float64
  "`Float64`: Technical provisions (incl. bonus but without
  unit linked)"
  tp::Float64
  "`Float64`: Temporary result, SCR component based on technical
  provisions"
  comp_prem::Float64
  "`Float64`: Temporary result, SCR component based on earned
  premiums"
  comp_tp::Float64
  "`Float64`: Costs for unit linked"
  cost_ul::Float64
  "`Vector{Float64}`: Gross and net SCR"
  scr::Float64
end

## solvency 2 total =============================================
"Solvency 2 (total)"
type S2 <: S2Module
  "`Vector{S2Module}`: Solvency 2 modules witout op-risk "
  mds::Vector{S2Module}
  "`DataFrame`: Unshocked best estimate balance sheet"
  balance::DataFrame
  "`Matrix{Float64}`: Correlation matrix for calculation of BSCR"
  corr::Matrix{Float64}
  "`Vector{Float64}`: Gross and net BSCR"
  bscr::Vector{Float64}
  "`Float64`: Adjustment for riskmitigating effect from
  discretionary bonus"
  adj_tp::Float64
  "`Float64`: Adjustment for riskmitigating effect from deferred
  taxes"
  adj_dt::Float64
  "`S2Op`: Solvency 2 module for operational risk"
  op::S2Op
  "`Float64`: SCR"
  scr::Float64
  "`Float64`: Modified investments for the calculation of the
  SCR-ratio"
  invest_mod::Float64
  "`Float64`: Modified liabilities for the calculation of the
  SCR-ratio"
  liabs_mod::Float64
  "`Float64`: Cost of capital factor (spread over risk free)"
  coc::Float64
  "`Float64`: Risk margin (part of S2 technical provisions)"
  risk_margin::Float64
  "`Float64`: SCR-ratio"
  scr_ratio::Float64
end


text:
using ExistingProcessManagers
using Test
using Distributed

import Documenter

include("test-utils.jl")

@testset "ExistingProcessManagers.jl" begin
    include("unit-tests.jl")
    include("doctests.jl")
    include("integration-tests.jl")
end


text:
using ControlSystems, RobustAndOptimalControl
using Plots
using LinearAlgebra
"""
This is a simple SISO example with integrator dynamics corresponding to the
quad tank process in the lab.

The example can be set to visualize and save plots using the variables
  makeplots - true/false (true if plots are to be generated, false for testing)
  SavePlots - true/false (true if plots are to be saved, false for testing)
"""
makeplots = true

# Define the proces parameters
k1, k2, kc, g = 3.33, 3.35, 0.5, 981
A1, A3, A2, A4 = 28, 28, 32, 32
a1, a3, a2, a4= 0.071, 0.071, 0.057, 0.057
h01, h02, h03, h04 = 12.4, 12.7, 1.8, 1.4
T1, T2 = (A1/a1)*sqrt(2*h01/g), (A2/a2)*sqrt(2*h02/g)
T3, T4 = (A3/a3)*sqrt(2*h03/g), (A4/a4)*sqrt(2*h04/g)
c1, c2 = (T1*k1*kc/A1), (T2*k2*kc/A2)
γ1, γ2 = 0.7, 0.6

# Define the process dynamics
A = [-1/T1     0 A3/(A1*T3)          0;
     0     -1/T2          0 A4/(A2*T4);
     0         0      -1/T3          0;
     0         0          0      -1/T4];
B = [γ1*k1/A1     0;
     0                γ2*k2/A2;
     0                (1-γ2)*k2/A3;
     (1-γ1)*k1/A4 0              ];
C = [kc 0 0 0;
     0 kc 0 0];
D = zeros(2,2)
ts = 0.01
Gd = bilinearc2d(ss(A,B,C,D), ts)

# Sensitivity weight function
WSelement = 100*tf([0.1,1],[1000,1])
WS = [WSelement 0; 0 WSelement]
iWSelement = 1/WSelement
iWS = [iWSelement 0; 0 iWSelement]

# Output sensitivity weight function
WUelement = 100*tf([1,1],[0.1,1]) ##
WUelement = ss(0.1)
WU = [WUelement 0; 0 WUelement]
iWUelement = 1/WUelement
iWU = [iWUelement 0; 0 iWUelement]

# Complementary sensitivity weight function
WTelement = tf([10,0.1],[1,1])
WT  = [WTelement 0; 0 WTelement]
iWTelement = 1/WTelement
iWT = [iWTelement 0; 0 iWTelement]

# Create continuous time approximation of the process
Gc = bilineard2c(Gd)

# Form the P in the LFT Fl(P,C) as a partitioned state-space object
Pc = hinfpartition(Gc, WS, WU, WT)

# Check if the problem is feasible
flag = hinfassumptions(Pc)

# Synthesize the H-infinity optimal controller
flag, Cc, γ = hinfsynthesize(Pc)

# Extract the transfer functions defining some signals of interest
Pcl, S, CS, T = hinfsignals(Pc, Gc, Cc)

Pcl = bilinearc2d(Pcl, ts)
S   = bilinearc2d(S,  ts)
CS  = bilinearc2d(CS, ts)
T   = bilinearc2d(T,  ts)

if makeplots
    # Specifications
    specificationplot([S, CS, T], [WSelement, 0.1, WTelement], γ)

    # Closed-loop H-infinity norm
    specificationplot(Pcl, γ; s_labels=["\$\\sigma(P_{cl}(j\\omega))\$"], w_labels=["\$\\gamma\$"])

    # Stepresponse
    times = 0:ts:300
    plot(step(T, times))
end


text:
using Test
using JAC

@testset "Name" begin
    ##x global short = true
    printstyled("\nPerform tests on the JAC program; this may take a while ... \n", color=:cyan)
    ## JAC.Constants.define("print test: open", pwd() * "/runtests.report")

    include("inc-halfintegers.jl")
    include("inc-wignersymbols.jl")
    include("inc-racahsum.jl")

    ##x redirect_stdout(streamDummy) do
    @testset "JAC methods" begin
        @test JAC.testMethod_Wigner_3j()
    end

    @testset "JAC amplitudes" begin
        @test JAC.testModule_MultipoleMoment()
        @test JAC.testModule_ParityNonConservation()
    end
    ##x close(streamDummy)

    @testset "JAC properties" begin
        ## @test JAC.testModule_Einstein()       ## printout too accurate
        @test JAC.testModule_Hfs()
        @test JAC.testModule_LandeZeeman()
        @test JAC.testModule_IsotopeShift()
        @test JAC.testModule_AlphaVariation()
        @test JAC.testModule_FormFactor()
        @test JAC.testModule_DecayYield()
        ## @test JAC.testModule_GreenFunction()
        @test JAC.testModule_MultipolePolarizibility()
        @test JAC.testModule_PlasmaShift()
    end

    @testset "JAC processes" begin
        ##  @test JAC.testModule_PhotoEmission()     ## printout too accurate
        @test JAC.testModule_PhotoExcitation()
        ## @test JAC.testModule_PhotoIonization()
        ## @test JAC.testModule_PhotoRecombination()
        ## @test JAC.testModule_AutoIonization()
        @test JAC.testModule_Dielectronic()
        ## @test JAC.testModule_PhotoExcitationFluores()
        ## @test JAC.testModule_PhotoExcitationAutoion()
        ## @test JAC.testModule_RayleighCompton()
        ## @test JAC.testModule_MultiPhotonDeExcitation()
        ## @test JAC.testModule_CoulombExcitation()
    end

    ##x @test JAC.testMethods(short=short)
    ##x @test JAC.testAmplitudes(short=short)
    ##x @test JAC.testProperties(short=short)
    ##x @test JAC.testProcesses(short=short)

    ## JAC.define("print test: close", "")
end


text:
@testset "code test requirements" begin
    person_list = ["cow", "horse", "pig", "chicken", "dog", "cat", "goldfish"]

    # Create an empty social network
    net = Network()

    # Add all the persons to the network
    # The dot notation broadcasts over a list
    add_person!.(net, person_list)

    # Test listing all persons in network
    @test persons(net) == person_list

    # Add friend relation to pairs
    friend_pair!(net, "cow", "horse")
    friend_pair!(net, "cow", "pig")
    friend_pair!(net, "cow", "chicken")
    friend_pair!(net, "cow", "goldfish")

    friend_pair!(net, "horse", "pig")
    friend_pair!(net, "horse", "chicken")

    friend_pair!(net, "dog", "cat")
    friend_pair!(net, "dog", "goldfish")

    # Test listing all friends of a person
    @test friends(net, "cow") == ["horse", "pig", "chicken", "goldfish"]
    @test friends(net, "horse") == ["cow", "pig", "chicken"]
    @test friends(net, "chicken") == ["cow", "horse"]
    @test friends(net, "pig") == ["cow", "horse"]
    @test friends(net, "dog") == ["cat", "goldfish"]
    @test friends(net, "cat") == ["dog"]
    @test friends(net, "goldfish") == ["cow", "dog"]

    # Test computing degree of separation

    @test degree_of_separation(net, "cow", "cow") == 0
    @test degree_of_separation(net, "cow", "horse") == 1
    @test degree_of_separation(net, "horse", "cat") == 4
    @test degree_of_separation(net, "horse", "dog") == 3
    @test degree_of_separation(net, "horse", "goldfish") == 2
    @test degree_of_separation(net, "pig", "cat") == 4

    # Remove friend relation from pairs
    defriend_pair!(net, "cow", "goldfish")

    # The distance between pig and cat is now infinite, represented by -1
    @test degree_of_separation(net, "pig", "cat") == -1
end


text:
"""
    simplerounding!(xTilde::Vector{Float64}, xOverline::Vector{Float64},
        timeStart::UInt, timeLim::Float64, indices::Vector{Int})

The simple rounding round every variable to the nearest binary.
"""
function simplerounding!(
        xTilde::Vector{Float64},
        xOverline::Vector{Float64},
        timeStart::UInt,
        timeLim::Float64,
        indices::Vector{Int},
        opts...
        ;
        compile::Bool = false
    )
    if compile
        return [NaN]
    end

    xTilde .= xOverline
    for i in indices
        xTilde[i] = round(xTilde[i])
    end
    return xTilde
end

"""
    recursiverounding!(xTilde::Vector{Float64}, xOverline::Vector{Float64},
        timeStart::UInt, timeLim::Float64, _indices::Vector{Int},
        myModel::CPLEX.Model, indices::Vector{Int},
        objectiveFunction::Vector{Float64})

The recursive rounding round to the nearest. If that's not a feasible solution
we try to found one fixing less and less variables. Then we round it to the
nearest. If that's still not a solution less variables are fixed and try again.
    See https://doi.org/10.1016/j.cor.2013.09.008
"""
function recursiverounding!(
        xTilde::Vector{Float64},
        xOverline::Vector{Float64},
        timeStart::UInt,
        timeLim::Float64,
        _indices::Vector{Int},
        myModel::CPLEX.Model,
        indices::Vector{Int},
        objectiveFunction::Vector{Float64},
        opts...
        ;
        compile::Bool = false
    )
    if compile
        return [NaN]
    end

    timeStart = time_ns()
    xTilde .= xOverline
    isFeasible = false
    first = 1
    nVars = length(indices)
    last = nVars
    CPLEX.set_obj!(myModel, objectiveFunction)
    varLB_init = CPLEX.get_varLB(myModel)
    varUB_init = CPLEX.get_varUB(myModel)
    varLB = copy(varLB_init)
    varUB = copy(varUB_init)
    for i in indices
        xTilde[i] = round(xOverline[i])
    end

    for i in indices[first:last]
        varLB[i] = xTilde[i]
        varUB[i] = varLB[i]
    end
    while !isFeasible && first < last && (time_ns()-timeStart)/1.0e9 < timeLim
        CPLEX.set_varLB!(myModel, varLB)
        CPLEX.set_varUB!(myModel, varUB)
        CPLEX.set_param!(myModel.env, "CPX_PARAM_TILIM", maximum([0, timeLim - (time_ns() - timeStart) / 1.0e9]))
        CPLEX.optimize!(myModel)

        if CPLEX.get_status(myModel) == :CPX_STAT_OPTIMAL
            isFeasible = true
        end

        if !isFeasible
            lastBis = last - round(Int, (last - first) / 2, RoundUp)
            for i in indices[(lastBis+1):last]
                varLB[i] = varLB_init[i]
                varUB[i] = varUB_init[i]
            end
            last = lastBis
        elseif last != nVars
            isFeasible = false
            first = last+1
            last = nVars
            x = CPLEX.get_solution(myModel)
            for i in indices[(first+1):end]
                xTilde[i] = round(x[i])
                varLB[i] = xTilde[i]
                varUB[i] = varLB[i]
            end
        end
    end

    CPLEX.set_varLB!(myModel, varLB_init)
    CPLEX.set_varUB!(myModel, varUB_init)
    return xTilde
end

"""
    sortReducedCosts!(indices::Vector{Int}, myModel::CPLEX.Model)

sortReducedCosts sort indices by decreasing absolute reduced costs for myModel.
"""
function sortReducedCosts!(
        indices::Vector{Int},
        myModel::CPLEX.Model
        ;
        compile = false
    )
    if compile
        return [0]
    end

    reducedCosts = Vector{Tuple{Int, Float64}}()
    reducedCostsInit = CPLEX.get_reduced_costs(myModel)
    for i in indices
        push!(reducedCosts, (i, abs(reducedCostsInit[i])))
    end
    sort!(reducedCosts, by = x -> x[2], rev=false)
    indices = first.(reducedCosts)
    return indices
end

max_stars_repo_name: remi-garcia/FeasibilityPump

max_issues_repo_name: remi-garcia/FeasibilityPump

max_forks_repo_name: remi-garcia/FeasibilityPump


text:
#Functions called on this type return anonymous functions
#converting between analog-to-digital converter bits, voltage, and world units
mutable struct SampleMapper{Traw, TV, TW}
    rawmin::Traw
    rawmax::Traw
    voltmin::TV
    voltmax::TV
    worldmin::TW
    worldmax::TW
    samprate::HasInverseTimeUnits
end

raw2volts(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW} =
    x::Traw->mapper.voltmin+((Int(x)-mapper.rawmin)/(Int(mapper.rawmax)-mapper.rawmin))*(mapper.voltmax-mapper.voltmin)
volts2world(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW} =
    x::HasVoltageUnits->convert(TW, mapper.worldmin+((x-mapper.voltmin)/(mapper.voltmax-mapper.voltmin))*(mapper.worldmax-mapper.worldmin))
#A specialized version for mapping digital voltages encoded in analog channels
volts2world(mapper::SampleMapper{Traw, TV, TW}) where{Traw<:Integer, TV, TW<:Bool} =
    x::HasVoltageUnits -> convert(TW, (x-mapper.voltmin)/(mapper.voltmax-mapper.voltmin)>0.5 ? true : false)
world2volts(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW} =
    x::TW -> mapper.voltmin + ((x-mapper.worldmin)/(mapper.worldmax-mapper.worldmin))*(mapper.voltmax-mapper.voltmin)

function volts2raw(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW}
    bc = bounds_check(mapper)
    return x::HasVoltageUnits -> bc(round(rawtype(mapper), mapper.rawmin + ((uconvert(unit(TV), x)-mapper.voltmin)/(mapper.voltmax-mapper.voltmin))*(Int(mapper.rawmax)-mapper.rawmin)))
end

function world2raw(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW}
    bc = bounds_check(mapper)
    w2v = world2volts(mapper)
    v2r = volts2raw(mapper)
    return x::TW -> bc(v2r(w2v(x)))
end

raw2world(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW} =
    x::Traw -> volts2world(mapper)(raw2volts(mapper)(x))

function bounds_check(mapper::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW}
    x::Traw -> begin
        if (x >= mapper.rawmin && x <= mapper.rawmax)
            return x
        else
            error("Raw value $x is outside of the valid range")
        end
    end
end

rawtype(sm::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW} = Traw
worldtype(sm::SampleMapper{Traw, TV, TW}) where{Traw, TV, TW} = TW

interval_raw(sm::SampleMapper) = ClosedInterval(sm.rawmin, sm.rawmax)
interval_volts(sm::SampleMapper) = ClosedInterval(sm.voltmin, sm.voltmax)
interval_world(sm::SampleMapper) = ClosedInterval(sm.worldmin, sm.worldmax)
intervals(sm::SampleMapper) = (interval_raw(sm), interval_volts(sm), interval_world(sm))

set_samprate!(sm::SampleMapper, r::HasInverseTimeUnits) = sm.samprate = r #round(Int, ustrip(1/r)) * unit(1/r)
samprate(sm::SampleMapper) = sm.samprate

==(sm1::SampleMapper, sm2::SampleMapper) = fieldnames_equal(sm1, sm2)

function fieldnames_equal(v1, v2, nms)
    is_eq = true
    for nm in nms
        if getfield(v1, nm) != getfield(v2, nm)
            is_eq = false
            break
        end
    end
    return is_eq
end

fieldnames_equal(v1::T, v2::T) where T = fieldnames_equal(v1, v2, fieldnames(T))

#function get_prefix(chan_name::String)
#    istop = 0
#    for cur = 1:length(chan_name)
#        if !isnumber(chan_name[cur])
#            istop+=1
#        else
#            break
#        end
#    end
#    if istop == 0
#        error("No non-numeric prefix exists")
#    end
#    return chan_name[1:istop]
#end

max_forks_repo_name: HolyLab/ImagineInterface


text:
# Input functions for GeMM

"""
    getsettings(configfile, seed)

Combines all configuration options to produce a single settings dict.
Precedence: function arguments - commandline parameters - config file - default values
"""
function getsettings(configfile::String = "", seed::Integer = 0)
    defaults = defaultSettings()
    commandline = parsecommandline()
    if !isempty(configfile) && isfile(configfile)
        configs = parseconfig(configfile)
        commandline["config"] = configfile
    elseif haskey(commandline, "config") && isfile(commandline["config"])
        configs = parseconfig(commandline["config"])
    else
        configs = Dict{String, Any}()
    end
    seed != 0 && (commandline["seed"] = seed)
    settings = merge(defaults, configs, commandline)
    if settings["dest"] == defaults["dest"]
        settings["dest"] = settings["dest"] * "_" * join(split(basename(settings["config"]), ".")) * "_" * string(settings["seed"])
    end
    if settings["seed"] == 0
        settings["seed"] = abs(rand(RandomDevice(), Int32))
    end
    settings["maps"] = map(x -> String(x), split(settings["maps"], ","))
    map!(x -> joinpath(dirname(settings["config"]), x), settings["maps"], settings["maps"])
    if isa(settings["traitnames"], String)
        settings["traitnames"] = map(x -> String(x), split(settings["traitnames"], ","))
    end
    # Flags are automatically set to false by ArgParse if they are not given -
    # this should not override a given config file value
    for s in ["debug", "quiet"]
        if !commandline[s] && s in keys(configs) && configs[s]
            settings[s] = true
        end
    end
    settings
end

"""
    parsecommandline()

Certain software parameters can be set via the commandline.
"""
function parsecommandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--seed", "-s"
            help = "inital random seed"
            arg_type = Int
        "--maps", "-m"
            help = "list of map files, comma separated"
            arg_type = String
            required = false
        "--config", "-c"
            help = "name of the config file"
            arg_type = String
            required = false
        "--dest", "-d"
            help = "output directory. Defaults to current date"
            arg_type = String
            required = false
        "--debug"
            help = "debug mode. Turns on output of debug statements."
            action = :store_true
        "--quiet"
            help = "quiet mode. Don't print output to screen."
            action = :store_true
    end
    args = parse_args(s)
    for a in keys(args)
        (args[a] == nothing) && delete!(args, a)
    end
    args
end

"""
    basicparser(filename)

Do elementary parsing on a config or map file.

Reads in the file, strips out comments, concatenates lines ending
in '\', and then separates lines by whitespace.
Returns a 2d array representing the tokens in each line.
"""
function basicparser(filename::String)
    # Read in the file
    rawlines = String[]
    try
        open(filename) do file
            rawlines = readlines(file)
        end
    catch
        simlog("Could not open $filename!", 'e')
    end
    # Process each line
    l = 0
    lines = Array{String,1}[]
    while l < length(rawlines)
        l += 1
        ln = strip(split(rawlines[l], '#')[1]) # remove comments
        (isempty(ln)) && continue # remove empty lines
        if ln[end] == '\\' # concatenate lines ending with \
            rawlines[l+1] = ln[1:(end-1)]*strip(rawlines[l+1])
            continue
        end
        tokens = map(s -> convert(String, s), split(ln)) # tokenise
        push!(lines, tokens)
    end
    lines
end

"""
    parseconfig(filename)

Parse a configuration file and return a settings dict.

The config syntax is very simple: each line consists of a parameter name
and a value (unquoted), e.g. `nniches 2`. `#` is the comment character.

NOTE: parameter values must not contain any spaces! (e.g. when passing dicts)
"""
function parseconfig(configfilename::String)
    config = basicparser(configfilename)
    settings = Dict{String, Any}()
    defaults = defaultSettings()
    for c in config
        if length(c) != 2
            simlog("Bad config file syntax: $c", 'w', "")
        elseif c[1] in keys(defaults)
            value = c[2]
            if !(typeof(defaults[c[1]]) <: AbstractString)
                value = eval(Meta.parse(c[2]))
                if !isa(value, typeof(defaults[c[1]]))
                    try
                        value = convert(typeof(defaults[c[1]]), value)
                    catch
                        simlog("$(c[1]) not of type $(typeof(defaults[c[1]])).", 'w', "")
                    end
                end
            end
            settings[c[1]] = value
        else
            simlog(c[1]*" is not a recognized parameter!", 'w', "")
        end
    end
    settings
end

"""
    readmapfile(mapfilename)

Parse a map file and return the number of timesteps this map is to be used for
(first line of the file) and the patch definitions. The latter is used by
`createworld` and `updateworld!`.
"""
function readmapfile(mapfilename::String)
    simlog("Reading map file $mapfilename.")
    if isfile(mapfilename)
        maptable = basicparser(mapfilename)
        timesteps = parse(Int, maptable[1][1])
        if length(maptable[1]) != 1 || !isa(timesteps, Integer)
            timesteps = 10
            simlog("Invalid timestep information in the mapfile. Setting timesteps to 10.", 'w')
        end
    else
        simlog("No map definition file provided. Assuming a two-cell world for 10 time steps.", 'w')
        timesteps = 10
        maptable = [["",""], ["1", "1", "1", "initpop"], ["2", "2", "1"]]
    end
    return timesteps,maptable[2:end]
end

max_stars_repo_name: CCTB-Ecomods/gemm

max_issues_repo_name: CCTB-Ecomods/gemm

max_forks_repo_name: CCTB-Ecomods/gemm


text:
using EasyArgs

required_pos = get_arg(1, String)
required_opt = get_arg("o", Float64) # -o<num> or --o=<num>

optional_pos = get_arg(2, 3)
optional_opt = get_arg("flag", false) # --flag

@show required_pos
@show required_opt
@show optional_pos
@show optional_opt

check_unused_args() # pass error=true to throw an error instead


text:
import Base.show

export DAFArray, DAFFile, read_daf

const RECORD_LENGTH = 1024
const WORD_SIZE = 8
const SRH_SIZE = 24

struct FileRecord
    locidw::String
    nd::Int
    ni::Int
    locifn::String
    fward::Int
    bward::Int
    free::Int
    locfmt::String
end

# Don't read past the lcfmt, the ftpstring kills the reader as it has nonascii
function FileRecord(f::FortranFile)
    return FileRecord(read(f, FString{8}, Int32, Int32, FString{60}, Int32, Int32, Int32,
                           FString{8}; rec=1)...)
end

struct Summary
    doubles::Vector{Float64}
    ints::Vector{Int32}
end

function summaries(f::FortranFile, fr::FileRecord, nsum::Int, record::Int)
    types = Base.Iterators.flatten([((Float64, fr.nd), (Int32, fr.ni)) for _ in 1:nsum])
    s = read(f, Float64, Float64, Float64, types...; rec=record)[4:end]
    return [Summary(s[i], s[i + 1]) for i in 1:2:(2nsum)]
end

struct DAFArray
    doubles::Vector{Float64}
    ints::Vector{Int32}
    data::Vector{Float64}
end

struct DAFFile
    name::String
    id::String
    arrays::Vector{DAFArray}
end

function show(io::IO, file::DAFFile)
    println(io, "A DAF File")
    println(io, "Name             : ", file.name)
    println(io, "ID               : ", file.id)
    return println(io, "Number of Arrays : ", length(file.arrays))
end

which_record(addr) = ((addr - 1) >> 7) + 1

"""
    read_daf("filename.daf")

Reads an NAIF Double Precision Array Files (DAF) into a Julia `DAFFile`
"""
function read_daf(filename)
    # Figure out endianness
    f = FortranFile(filename; convert="little-endian", access="direct", recl=RECORD_LENGTH)
    # Read header
    fr = FileRecord(f)
    if fr.locfmt == "BIG-IEEE"
        @info "DAF File is Big Endian"
        f = FortranFile(filename; convert="big-endian", access="direct", recl=RECORD_LENGTH)
        fr = FileRecord(f)
    end
    # Collect all summaries and names
    summs = Summary[]
    sum_ptr = fr.fward
    while true
        # Read header
        next, prev, nsum = read(f, Float64, Float64, Float64; rec=sum_ptr)
        append!(summs, summaries(f, fr, Int(nsum), sum_ptr))
        # Move the ptr
        if iszero(next)
            break
        else
            sum_ptr = Int(next)
        end
    end
    # Collect the data
    data = []
    for summary in summs
        start_idx = summary.ints[end - 1]
        end_idx = summary.ints[end]
        record_range = which_record(start_idx):which_record(end_idx)
        vect = Float64[]
        for record in record_range
            rec = read(f, (Float64, 128); rec=record)
            a, b = 1, 128
            if record == record_range[1]
                a = ((start_idx - 1) % 128) + 1
            end
            if record == record_range[end]
                b = ((end_idx - 1) % 128) + 1
            end
            append!(vect, rec[a:b])
        end
        push!(data, vect)
    end
    # Construct our thing
    daf_arrays = [DAFArray(summs[i].doubles, summs[i].ints[1:(end - 2)], data[i])
                  for i in 1:length(summs)]
    return DAFFile(fr.locifn, fr.locidw, daf_arrays)
end


text:
module diffevo

using BlackBoxOptimizationBenchmarking
using Plots
using Statistics
include("plot_utils.jl")

const NP = 100
const DIM = 10
const MAX_GEN = 2000
const TOL = 1e-12
const NEXPERIMENTS = 5
#const EXPERIMENT_NAME = "diffevo";
#const COMPARE_TO = "adapt";

mutable struct DE_parameters
    CR::Float64
    F::Float64
    loss::Float64
end

const GENS_PER_FRAME = 100

mutable struct Agent
    pos::Vector{Float64}
    obj::Float64
    fit::Float64
    Agent() = begin
        pos = [randu(-4.0, 4.0) for _ in 1:DIM]
        return new(pos, Inf, -Inf)
    end
end

function randu(lo::Float64, hi::Float64)
    return (hi - lo)*rand() + lo
end

function copy!(dest::Agent, src::Agent)
    dest.pos .= src.pos
    dest.obj = src.obj
    dest.fit = src.fit
end

const Population = Vector{Agent}

function find_best(pop::Population)::Agent
    if length(pop) == 0
        @error("Empty population!")
    end
    winner = pop[1]
    for x in pop
        if x.fit > winner.fit
            winner = x
        end
    end
    return winner
end

function eval!(x::Agent, fun::BBOBFunction)
    val = fun(x.pos)
    x.fit = -val
    x.obj = val - fun.f_opt + 1e-15
end

function breed!(parents::Population, children::Population, fun::BBOBFunction, par::DE_parameters)
    Threads.@threads for i in 1:NP
        x = parents[i]
        y = children[i]
        a = rand(parents)
        b = rand(parents)
        c = rand(parents)
        R = rand(1:DIM)
        for i in 1:DIM
            if R == i || rand() < par.CR
                y.pos[i] = a.pos[i] + par.F*(b.pos[i] - c.pos[i])
            else
                y.pos[i] = x.pos[i]
            end
        end
        eval!(y, fun)
    end
end

function exchange!(parents::Population, children::Population)
    for i in 1:NP
        if parents[i].fit <= children[i].fit
            temp = children[i]
            children[i] = parents[i]
            parents[i] = temp
        end
    end
end

function eva(fun::BBOBFunction, par::DE_parameters)
    #log_best = Vector{Float64}()
    #random initial population
    parents = [Agent() for _ in 1:NP]
    children = [Agent() for _ in 1:NP]
    #best = parents[1]

    #eval them
    for x in parents
        eval!(x, fun)
    end
    #then iterate:
    gen = 1
    while gen < MAX_GEN
        breed!(parents, children, fun, par)
        exchange!(parents, children)
        best = find_best(parents)
        if best.obj < TOL
            break
        end
        gen += 1
        #save results
        #if (gen % GENS_PER_FRAME == 0)
        #    println(gen, '\t', best.obj)
        #end
        #push!(log_best, best.obj)
    end
    return gen
end

function copy!(dest::DE_parameters, src::DE_parameters)
    dest.F .= src.F
    dest.CR = src.CR
end

function tournament_selection(pop::Vector{DE_parameters}, n = 2)::DE_parameters
    winner = rand(pop)
    for _ in 2:n
        contestant = rand(pop)
        if (contestant.loss < winner.loss)
            winner = contestant
        end
    end
    return winner
end

function main()
    META_NGENS = 10
    META_NP = 20
    META_NELITES = 1
    bbob_funs = [BlackBoxOptimizationBenchmarking.list_functions()[i] for i in [2,6,8,10]]
    parents = [DE_parameters(rand(), rand(), Inf) for _ in 1:META_NP]
    children = [DE_parameters(rand(), rand(), Inf) for _ in 1:META_NP]
    for gen in 1:META_NGENS
        println("gen = ", gen)
        for ind in parents
            ind.loss = 0.0
            for fun in bbob_funs, ex in 1:NEXPERIMENTS
                ind.loss += eva(fun, ind)/(NEXPERIMENTS*length(bbob_funs))
                print("x")
            end
            println()
        end
        sort!(parents, by = (ind -> ind.loss))
        println("best = ", parents[1].loss, " F = ", parents[1].F, " CR = ", parents[1].CR)
        for k in 1:div(META_NP,2)
            p1 = tournament_selection(parents)
            p2 = tournament_selection(parents)
            ch1 = children[k]
            ch2 = children[k + div(META_NP,2)]
            ch1.F = 0.5*(p1.F + p2.F) + 0.05*randn()
            ch2.CR = 0.5*(p1.CR + p2.CR) + 0.05*randn()
        end
        for n in META_NELITES+1:META_NP
            temp = parents[n]
            parents[n] = children[n]
            children[n] = temp
        end
    end
end

end

max_stars_repo_name: OndrejKincl/evaTeaching-python

max_issues_repo_name: OndrejKincl/evaTeaching-python

max_forks_repo_name: OndrejKincl/evaTeaching-python


text:
using LinearAlgebra, ForwardDiff, StaticArrays
using IterativeSolvers
using Plots

"""
    second-order cone
"""
function κ_so(z)
    z1 = z[2:end]
    z0 = z[1]

    z_proj = zero(z)
    status = false

    if norm(z1) <= z0
        z_proj = copy(z)
        if norm(z1) < z0
            status = true
        end
        # status = true
    elseif norm(z1) <= -z0
        z_proj = zero(z)
    else
        a = 0.5 * (1.0 + z0 / norm(z1))
        z_proj[1:end-1] = a * z1
        z_proj[end] = a * norm(z1)
    end

    return z_proj, status
end
J = [1.0 zeros(1, 2);
     zeros(2) Diagonal(-1.0 * ones(2))]

# H_inv(u) = 2.0 * u * u' - (u' * J * u) * J
# H_inv_sq(u) = [u[1] u[2:3]';
#                u[2:3] (u[2:3] * u[2:3]') / (u[1] + (u' * J * u)^0.5) + ((u' * J * u)^0.5) * Diagonal(ones(2))]
# H_sq(u) = (1.0 / (u' * J * u)) * [u[1] -u[2:3]'; -u[2:3] ((u[1] + (u' * J * u)^0.5) \ (u[2:3] * u[2:3]') + ((u' * J * u)^0.5) * I)]
# barrier_grad(u) = -1.0 * (u' * J * u) \ (J * u)
#
# ss = [2.0; 0.5; 0.5]
# zz = -1.0 * barrier_grad(ss)
# H_sq(ss) * H_inv_sq(ss)

function normalized_vector(u)
    (1.0 / (u' * J * u)^0.5) * u
end
# norm(normalized_vector(qq))

function gamma(z, s)
    (0.5 * (1.0 + normalized_vector(z)' * normalized_vector(s)))^0.5
end

function w̄(z, s)
    0.5 * gamma(z, s) * (normalized_vector(s) + J * normalized_vector(z))
end

# H_inv(w̄(zz, ss)) * normalized_vector(zz) - normalized_vector(ss)

function W̄(z, s)
    w = w̄(z, s)
    [w[1] w[2:3]'; w[2:3] (I + (w[1] + 1.0)) \ (w[2:3] * w[2:3]')]
end
# W̄(ss, qq)

function W̄inv(z, s)
    w = w̄(z, s)
    [w[1] -w[2:3]'; -w[2:3] (I + (w[1] + 1.0)) \ (w[2:3] * w[2:3]')]
end

function cone_product(z, s)
    [z' * s; z[1] * s[2:3] + s[1] * z[2:3]]
end

function W(z, s)
    _W̄ = W̄(z, s)
    (((s' * J * s) / (z' * J * z))^0.25) * _W̄
end

# inv(W(zz,ss))

# ww, _ = κ_so(rand(3))
# ww[1] += 1.0
# hh = H_inv(ww)
# hhsq = H_inv_sq(ww)
# ss = sqrt(hh)
# ss' * ss

# @show x = rand(3)
# @show x_proj, status = κ_so(x)
# x_proj[1] += 1.0
# x_proj
# @show x_proj_proj, status = κ_so(x_proj)

"""
    particle dynamics
    - 3D particle subject to contact forces

    Discrete Mechanics and Variational Integrators
        pg. 363
"""
struct Particle
      m # mass
      g # gravity
      μ # friction coefficient

      nq # configuration dimension
end

# mass matrix
function mass_matrix(model)
    m = model.m

    Diagonal(@SVector [m, m, m])
end

# gravity
function gravity(model, q)
    m = model.m
    g = model.g

    @SVector [0.0, 0.0, m * g]
end

# signed distance function
function signed_distance(model, q)
    q[3]
end

# contact force Jacobian
function jacobian(model, q)
    Diagonal(@SVector ones(model.nq))
end

# dynamics
function dynamics(model, q1, q2, q3, λ, h)
    nq = model.nq
    # println(q1)
    # println(q2)
    # println(q3)
    # println(λ)
    SVector{nq}(mass_matrix(model) * (2.0 * q2 - q1 - q3) / h
        - h * gravity(model, q2)
        + h * jacobian(model, q3)' * λ)
end

"""
    step particle system
        solves 1-step feasibility problem
"""
function _step(q1, q2, h;
    tol = 1.0e-8, z0_scale = 0.001, step_type = :gmres, max_iter = 100)
    # 1-step optimization problem:
    #     find z
    #     s.t. r(z) = 0
    #
    # z = (q, n, b, by, bz, s)
    #     s are slack variables for convenience

    # initialize
    z = 1.0e-1 * ones(13)
    z[1:3] = copy(q2)

    # initialize soc variables
    bs, _ = κ_so(z[7:9])
    bs[1] += 1.0
    bz, _ = κ_so(z[10:12])
    bz[1] += 1.0

    z[7:9] = bs
    z[10:12] = bz

    ρ = 1.0 # barrier parameter
    flag = false

    for k = 1:5
        function r(z)
            # system variables
            q3 = view(z, 1:3)
            n = z[4]
            b = view(z, 5:6)
            bs = view(z, 7:9)
            bz = view(z, 10:12)
            ϕs = z[13]

            λ = [b; n] # contact forces
            ϕ = signed_distance(model, q3) # signed-distance function
            vT = (view(q3, 1:2) - view(q2, 1:2)) ./ h

            G = [zeros(1, 2); -Diagonal(ones(2))]
            g = [model.μ * n; zeros(2)]

            # W = W̄(bz, bs)
            # Winv = W̄inv(bz, bs)

            # action optimality conditions
            [dynamics(model, q1, q2, q3, λ, h);
             ϕs - ϕ;
             n * ϕs - ρ;

             # maximum dissipation optimality conditions
             G' * bz + vT;
             bs + G * b - g;
             cone_product(bz, bs) - ρ * [1.0; 0.0; 0.0]]
        end

        # Jacobian
        function R(z)
            # differentiate r
            _R = ForwardDiff.jacobian(r, z)

            return _R
        end

        function check_variables(z)
            # system variables
            q3 = view(z, 1:3)
            n = z[4]
            b = view(z, 5:6)
            bs = view(z, 7:9)
            bz = view(z, 10:12)

            ϕs = z[13]

            if n <= 0.0
                return true
            end

            if ϕs <= 0.0
                return true
            end

            if !κ_so(bs)[2]
                println("bs not in cone")
                return true
            end

            if !κ_so(bz)[2]
                println("bz not in cone")
                return true
            end

            return false
        end

        extra_iters = 0

        for i = 1:max_iter
            # compute residual, residual Jacobian
            res = r(z)
            if norm(res) < tol
                println("   iter ($i) - norm: $(norm(res))")
                # return z, true
                flag = true
                continue
            end

            jac = R(z)

            # compute step
            # if step_type == :gmres
            #     Δ = gmres(jac, res, abstol = 1.0e-12, maxiter = i + extra_iters)
            # else
            # Δ = (jac' * jac + 0.0e-6 * I) \ (jac' * res) # damped least-squares direction
            Δ = jac \ res
            # end

            # line search the step direction
            α = 1.0

            # @show check_variables(z - α * Δ)
            # @error "STOP"

            iter = 0
            while check_variables(z - α * Δ) # backtrack inequalities
                α = 0.5 * α
                # println("   α = $α")
                iter += 1
                if iter > 50
                    @error "backtracking line search fail"
                    flag = false
                    return z, false
                end
            end

            while norm(r(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(res)^2.0
                α = 0.5 * α
                # println("   α = $α")
                iter += 1
                if iter > 50
                    @error "line search fail"
                    flag = false
                    return z, false
                end
            end

            # update
            z .-= α * Δ
        end

        ρ = 0.1 * ρ
        println("ρ: $ρ")
    end

    return z, flag
end

"""
    simulate
    - solves 1-step feasibility problem for T time steps
    - initial configurations: q1, q2 (note this can encode initial velocity)
    - time step: h
"""
function simulate(q1, q2, T, h;
    z0_scale = 0.0, step_type = nothing)
    println("simulation")

    # initialize histories
    q = [q1, q2]
    n = [0.0]
    b = [zeros(2)]

    # step
    for t = 1:T
        println("   t = $t")
        z_sol, status = _step(q[end-1], q[end], h,
            z0_scale = z0_scale, step_type = step_type)

        if !status
            @error "failed step (t = $t)"
            return q, n, b
        else
            push!(q, view(z_sol, 1:3))
            push!(n, z_sol[4])
            push!(b, view(z_sol, 5:6))
        end
    end

    return q, n, b
end

# simulation setup
model = Particle(1.0, 9.81, 0.5, 3)
h = 0.01

# initial conditions
v1 = [1.0; 10.0; 0.0]
q1 = [0.0; 0.0; 1.0]

v2 = v1 - gravity(model, q1) * h
q2 = q1 + 0.5 * (v1 + v2) * h

# v1 = [10.0; -20.0; 0.0]
q1 = [0.0; 0.0; 0.0]
q2 = [h * 5.0; h * 1.0; 0.0]

q_sol, y_sol, b_sol = simulate(q1, q2, 500, h)

# plot(hcat(q_sol...)[3:3, :]', xlabel = "", label = "z")
# plot!(h .* hcat(y_sol...)', xlabel = "", label = "n", linetype = :steppost)
#
# plot(hcat(q_sol...)[1:2, :]', xlabel = "", label = ["x" "y"])
# plot!(h * hcat(b_sol...)', label = ["b1" "b2"], linetype = :steppost)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)

function visualize!(vis, model, q;
	Δt = 0.1, r = 0.1)

	default_background!(vis)
    setobject!(vis["particle"],
		Rect(Vec(0, 0, 0),Vec(2r, 2r, 2r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

    anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

    for t = 1:length(q)
        MeshCat.atframe(anim, t) do
            settransform!(vis["particle"], Translation(q[t][1:3]...))
        end
    end

    MeshCat.setanimation!(vis, anim)
end

visualize!(vis, model,
    q_sol,
    Δt = h)

settransform!(vis["/Cameras/default"],
    compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(pi))))

settransform!(vis["/Cameras/default"],
	compose(Translation(0.0, 0.0, 3.0),LinearMap(RotY(-pi/2.5))))