text:
module SKPreprocessors

using PyCall

# standard included modules
using DataFrames
using Random
using AutoMLPipeline.AbsTypes
using AutoMLPipeline.Utils

import AutoMLPipeline.AbsTypes: fit!, transform!
export fit!, transform!
export SKPreprocessor, skpreprocessors

function __init__()
  global DEC=pyimport_conda("sklearn.decomposition","scikit-learn")
  global FS=pyimport_conda("sklearn.feature_selection","scikit-learn")
  global IMP=pyimport_conda("sklearn.impute","scikit-learn")
  global PREP=pyimport_conda("sklearn.preprocessing","scikit-learn")

  # Available scikit-learn learners.
  global preprocessor_dict = Dict(
     "DictionaryLearning" => DEC.DictionaryLearning,
     "FactorAnalysis" => DEC.FactorAnalysis,
     "FastICA" => DEC.FastICA,
     "IncrementalPCA" => DEC.IncrementalPCA,
     "KernelPCA" => DEC.KernelPCA,
     "LatentDirichletAllocation" => DEC.LatentDirichletAllocation,
     "MiniBatchDictionaryLearning" => DEC.MiniBatchDictionaryLearning,
     "MiniBatchSparsePCA" => DEC.MiniBatchSparsePCA,
     "NMF" => DEC.NMF,
     "PCA" => DEC.PCA,
     "SparsePCA" => DEC.SparsePCA,
     "SparseCoder" => DEC.SparseCoder,
     "TruncatedSVD" => DEC.TruncatedSVD,
     "dict_learning" => DEC.dict_learning,
     "dict_learning_online" => DEC.dict_learning_online,
     "fastica" => DEC.fastica,
     "non_negative_factorization" => DEC.non_negative_factorization,
     "sparse_encode" => DEC.sparse_encode,
     "GenericUnivariateSelect" => FS.GenericUnivariateSelect,
     "SelectPercentile" => FS.SelectPercentile,
     "SelectKBest" => FS.SelectKBest,
     "SelectFpr" => FS.SelectFpr,
     "SelectFdr" => FS.SelectFdr,
     "SelectFromModel"  => FS.SelectFromModel,
     "SelectFwe" => FS.SelectFwe,
     "RFE" => FS.RFE,
     "RFECV" => FS.RFECV,
     "VarianceThreshold"  => FS.VarianceThreshold,
     "chi2" => FS.chi2,
     "f_classif"  => FS.f_classif,
     "f_regression" => FS.f_regression,
     "mutual_info_classif" => FS.mutual_info_classif,
     "mutual_info_regression" => FS.mutual_info_regression,
     "SimpleImputer" => IMP.SimpleImputer,
     #"IterativeImputer" => IMP.IterativeImputer,
     #"KNNImputer" => IMP.KNNImputer,
     "MissingIndicator" => IMP.MissingIndicator,
     "Binarizer" => PREP.Binarizer,
     "FunctionTransformer" => PREP.FunctionTransformer,
     "KBinsDiscretizer" => PREP.KBinsDiscretizer,
     "KernelCenterer" => PREP.KernelCenterer,
     "LabelBinarizer" => PREP.LabelBinarizer,
     "LabelEncoder" => PREP.LabelEncoder,
     "MultiLabelBinarizer" => PREP.MultiLabelBinarizer,
     "MaxAbsScaler" => PREP.MaxAbsScaler,
     "MinMaxScaler" => PREP.MinMaxScaler,
     "Normalizer" => PREP.Normalizer,
     "OneHotEncoder" => PREP.OneHotEncoder,
     "OrdinalEncoder" => PREP.OrdinalEncoder,
     "PolynomialFeatures" => PREP.PolynomialFeatures,
     "PowerTransformer" => PREP.PowerTransformer,
     "QuantileTransformer" => PREP.QuantileTransformer,
     "RobustScaler" => PREP.RobustScaler,
     "StandardScaler" => PREP.StandardScaler
     #"add_dummy_feature" => PREP.add_dummy_feature,
     #"binarize" => PREP.binarize,
     #"label_binarize" => PREP.label_binarize,
     #"maxabs_scale" => PREP.maxabs_scale,
     #"minmax_scale" => PREP.minmax_scale,
     #"normalize" => PREP.normalize,
     #"quantile_transform" => PREP.quantile_transform,
     #"robust_scale" => PREP.robust_scale,
     #"scale" => PREP.scale,
     #"power_transform" => PREP.power_transform
    )
end

"""
    SKPreprocessor(preprocessor::String,args::Dict=Dict())

A wrapper for Scikitlearn preprocessor functions.
Invoking `skpreprocessors()` will list the acceptable
and supported functions. Please check Scikitlearn
documentation for arguments to pass.

Implements `fit!` and `transform!`.
"""
mutable struct SKPreprocessor <: Transformer
  name::String
  model::Dict
  args::Dict

  function SKPreprocessor(args=Dict())
    default_args=Dict(
                      :name => "skprep",
                      :preprocessor => "PCA",
                      :impl_args => Dict(),
                      :autocomponent=>false
                     )
    cargs = nested_dict_merge(default_args, args)
    cargs[:name] = cargs[:name]*"_"*randstring(3)
    prep = cargs[:preprocessor]
    if !(prep in keys(preprocessor_dict))
      println("$prep is not supported.")
      println()
      skpreprocessors()
      error("Argument keyword error")
    end
    new(cargs[:name],Dict(),cargs)
  end
end

function SKPreprocessor(prep::String,args::Dict=Dict())
  SKPreprocessor(Dict(:preprocessor => prep,:name=>prep,args...))
end

function skpreprocessors()
  processors = keys(preprocessor_dict) |> collect |> x-> sort(x,lt=(x,y)->lowercase(x)<lowercase(y))
  println("syntax: SKPreprocessor(name::String, args::Dict=Dict())")
  println("where *name* can be one of:")
  println()
  [print(processor," ") for processor in processors]
  println()
  println()
  println("and *args* are the corresponding preprocessor's initial parameters.")
  println("Note: Please consult Scikitlearn's online help for more details about the preprocessor's arguments.")
end

function fit!(skp::SKPreprocessor, x::DataFrame, y::Vector=[])
  features = x |> Array
  impl_args = copy(skp.args[:impl_args])
  autocomp = skp.args[:autocomponent]
  if autocomp == true
    cols = ncol(x)
    ncomponents = 1
    if cols > 0
      ncomponents = round(sqrt(cols),digits=0) |> Integer
      push!(impl_args,:n_components => ncomponents)
    end
  end
  preprocessor = skp.args[:preprocessor]
  py_preprocessor = preprocessor_dict[preprocessor]

  # Train model
  preproc = py_preprocessor(;impl_args...)
  preproc.fit(features)
  skp.model = Dict(
                   :skpreprocessor => preproc,
                   :impl_args => impl_args
                  )
end

function transform!(skp::SKPreprocessor, x::DataFrame)
  features = deepcopy(x) |> Array
  model=skp.model[:skpreprocessor]
  return collect(model.transform(features)) |> DataFrame
end

end


text:
trunc_above(v::T,thr) where {T} = ifelse(v<thr || isnan(v), v, T(thr))
nanz(x) = ifelse(isnan(x), zero(x), x)
#Assume that float values are also normed
correctbias(pix::T, bias=100/2^16) where T<:Union{Normed{UInt16, 16}, Float32, Float64}  =
    T(max(0.0, pix-bias))
correctbias(img::AbstractArray{T, N}, bias=100/2^16) where {T<:Union{Normed{UInt16, 16},Float32,Float64},N} =
    mappedarray(x->correctbias(x), img)
sqrtimg(img::AbstractArray{T,N}) where {T,N} = mappedarray(x->sqrt(x), img)

function ypad(img, padded_size::Int; flip_y = true, fillval = eltype(img)(NaN), pad_side=:both)
    @assert ndims(img) == 2
    @assert padded_size >= size(img,2)
    npix = div(padded_size - size(img,2), 2)
    if flip_y
        img = flipy(img)
    end
    first_idx = (1,1)
    if pad_side == :both
        first_idx = (1,npix+1)
    elseif pad_side == :top
        first_idx = (1,2*npix+1)
    elseif pad_side != :bottom
        error("Unrecognized padding argument $pad_side")
    end
    return PaddedView(fillval, img, (size(img,1),padded_size), first_idx)
end

flipy(img::AbstractArray{T,2}) where {T} = view(img, :, reverse(last(axes(img))))

#Crops the image to the minimal rectangle without val along any edge
function crop_vals(img::AbstractArray{T,2}, val) where {T}
    xmax = ymax = 1
    xmin = size(img,1)
    ymin = size(img,2)
    for y = 1:size(img,2)
        for x = 1:size(img,1)
            if !isequal(img[x,y], val)
                xmin = min(x,xmin)
                xmax = max(x,xmax)
                ymin = min(y,ymin)
                ymax = max(y,ymax)
            end
        end
    end
    if xmax < xmin || ymax < ymin
        error("All elements of image seem to be $val")
    end
    return view(img, xmin:xmax, ymin:ymax)
end

function nanplus!(prealloc, img1, img2)
    for i in eachindex(img1)
        val1 = img1[i]
        val2 = img2[i]
        isn1 = isnan(val1)
        if isn1 || isnan(val2)
            prealloc[i] = ifelse(isn1, val2 , val1)
        else
            prealloc[i] = val1 + val2
        end
    end
    return prealloc
end

nanplus(img1::AbstractArray{T}, img2::AbstractArray{T}) where {T<:AbstractFloat} = nanplus!(similar(img1), img1, img2)
nanplus(img1, img2) = img1.+img2

##Not used
##Crops the image to the specified y size (symmetrically crops top and bottom)
#function crop_y(img::AbstractArray{T,2}, ysz::Int) where {T}
#    @assert iseven(ysz)
#    cursz = size(img,2)
#    @assert iseven(cursz) && cursz >= ysz
#    halfcrop = div(cursz-ysz, 2)
#    return view(img, :, halfcrop+1:cursz-halfcrop)
#end

ext: jl


text:
using Test, FrameFun, DomainSets, BasisFunctions

@testset "FrameFunPlatforms" begin
    plat1 = FourierPlatform()
    plat2 = FourierExtensionPlatform(0.0..0.5)
    plat3 = ChebyshevPlatform()
    plat4 = OPSExtensionFramePlatform(ChebyshevU(10),0.0..0.5)

    measure(plat1) == FourierWeight()
    measure(plat2) == submeasure(FourierWeight(),0.0..0.5)
    measure(plat3) == ChebyshevWeight()
    measure(plat4) == submeasure(ChebyshevWeight(),0.0..0.5)

    @test SamplingStyle(plat1) == InterpolationStyle()
    @test SamplingStyle(plat2) == OversamplingStyle()
    @test SamplingStyle(plat3) == InterpolationStyle()
    @test SamplingStyle(plat4) == OversamplingStyle()

    @test DictionaryStyle(plat1) == BasisStyle()
    @test DictionaryStyle(plat2) == FrameStyle()
    @test DictionaryStyle(plat3) == BasisStyle()
    @test DictionaryStyle(plat4) == FrameStyle()

    @test SolverStyle(plat1,InterpolationStyle()) == InverseStyle()
    @test SolverStyle(plat2,OversamplingStyle()) == AZStyle()
    @test SolverStyle(plat3,InterpolationStyle()) == InverseStyle()
    @test SolverStyle(plat4,OversamplingStyle()) == AZStyle()

    @test dictionary(plat1,10) == Fourier(10)
    @test iscompatible(dictionary(plat2,10), extensionframe(Fourier(10),0.0..0.5))
    @test dictionary(plat3,10) == ChebyshevT(10)
    @test iscompatible(dictionary(plat4,10), extensionframe(ChebyshevU(10),0.0..0.5))

    @test operator(dualdictionary(plat1,10,FourierWeight())) isa DiagonalOperator
    @test dualdictionary(plat2,10,measure(plat2)) isa ExtensionFrame
    @test operator(dualdictionary(plat3,10,ChebyshevWeight())) isa DiagonalOperator
    @test dualdictionary(plat4,10,measure(plat4)) isa ExtensionFrame
end


text:
"""
Quad tree like division

|---|---|
| 2 | 4 |
|---|---|
| 1 | 3 |
|---|---|

"""
function split(box::Box{2})
    X, Y = box.X, box.Y
    M = 0.5(X + Y)
    [Box(Point(X[1], X[2]), Point(M[1], M[2])),
     Box(Point(X[1], M[2]), Point(M[1], Y[2])),
     Box(Point(M[1], X[2]), Point(Y[1], M[2])),
     Box(Point(M[1], M[2]), Point(Y[1], Y[2]))]
end

"""
Four segments that are the boundary

 |------|
 |   2  |
 | 1   3|
 |   4  |
 |------|

"""
function surface(box::Box{2})
    X, Y = box.X, box.Y
    [Segment(Point(X[1], X[2]), Point(X[1], Y[2])),
     Segment(Point(X[1], Y[2]), Point(Y[1], Y[2])),
     Segment(Point(Y[1], Y[2]), Point(Y[1], X[2])),
     Segment(Point(Y[1], X[2]), Point(X[1], X[2]))]
end

"""Do 2 segments collide?"""
function collides(s1::Segment{2}, s2::Segment{2})
    # Nope if bbox don't
    !(bbox_collides(s1.box, s2.box)) && return false

    c = dot(s1.t, s2.t)
    DO = s2.O - s1.O
    #  Possible coplanar
    if(abs(c - 1) < EPS || abs(c + 1) < EPS)
        # Same line
        norm(DO) < EPS && return true
        # If they lie on a same line then bbox colision implier collision
        return abs(dot(normalize(DO), s1.t) - 1) < EPS ||
               abs(dot(normalize(DO), s1.t) + 1) < EPS
    end

    #A = [1 -c; c -1]
    #b = [dot(DO, s1.t), dot(DO, s2.t)]
    #x = A\b
    #println((x, s1(x[1]), s2(x[2])))

    # We solve the 2x2 system
    α = (-dot(s1.t, DO) + c*dot(s2.t, DO))/(c^2-1)
    # Point would not lie on the first line; [0, L] is in side
    !(-EPS < α < s1.length + EPS) && return false

    β = (-c*dot(s1.t, DO) + dot(s2.t, DO))/(c^2-1)
    !(-EPS < β < s2.length + EPS) && return false

    #@show (α, β)

    true
end

"""See if segment collides with a box"""
function collides(segment::Segment{2}, box::Box{2})
    is_A_in = segment.O ∈ box
    is_B_in = segment.B ∈ box

    # At least on is inside
    (is_A_in || is_B_in) && return true

    # Check collisions with cases --[---]--
    for line in surface(box)
        collides(segment, line) && return true
    end

    false
end

max_stars_repo_name: MiroK/fractional-dim


text:
module AdDynamics
    using DifferentialEquations
    using ForwardDiff
    using DiffResults
    using StaticArrays
    using ..ProbInfo, ..LinPoint
    struct LinRes
        endpoint::SArray{Tuple{14},Float64,1,14}
        derivative::SArray{Tuple{14,21},Float64,2,294}
    end

    const mass_idx = 1
    const r_idx = SVector(2,3,4)
    const v_idx = SVector(5,6,7)
    const qbi_idx = SVector(8,9,10,11)
    const omb_idx = SVector(12,13,14)

    struct IntegratorParameters{T}
        dt :: Float64
        uk :: SArray{Tuple{3}, T, 1, 3}
        up :: SArray{Tuple{3}, T, 1, 3}
        sigma :: T
        pinfo :: ProbInfo
    end

    function DCM(quat::SArray{Tuple{4}, T, 1, 4} where T)
        q0 = quat[1]
        q1 = quat[2]
        q2 = quat[3]
        q3 = quat[4]

        return transpose(@SMatrix [(1-2*(q2^2 + q3^2))         (2*(q1 * q2 + q0 * q3))  (2*(q1*q3 - q0 * q2));
                        (2*(q1 * q2 - q0 * q3))     (1-2*(q1^2 + q3^2))      (2*(q2*q3 + q0 * q1));
                        (2*(q1 * q3 + q0 * q2))     (2*(q2*q3 - q0 * q1))    (1-2*(q1^2 + q2^2)) ])
    end

    function Omega{T}(omegab::SArray{Tuple{3}, T, 1, 3})
        z = zero(T)
        return @SMatrix [z        -omegab[1] -omegab[2] -omegab[3];
                        omegab[1] z           omegab[3] -omegab[2];
                        omegab[2] -omegab[3]  z          omegab[1];
                        omegab[3]  omegab[2] -omegab[1]  z         ]
    end

    function dx(state::SArray{Tuple{14}, T, 1, 14} where T, u::SArray{Tuple{3}, T, 1, 3} where T, info::ProbInfo)
        qbi = state[qbi_idx]
        omb = state[omb_idx]
        thr_acc = DCM(qbi) * u/state[mass_idx]
        grv_acc = @SVector [-info.g0,0,0]
        acc = thr_acc + grv_acc
        rot_vel = 0.5*Omega(omb)*qbi
        rot_acc = info.jBi*(cross(info.rTB,u) - cross(omb,info.jB*omb))
        return @SVector [-info.a*norm(u),
                state[5],state[6],state[7],
                acc[1], acc[2], acc[3],
                rot_vel[1], rot_vel[2], rot_vel[3], rot_vel[4],
                rot_acc[1], rot_acc[2], rot_acc[3]]
    end

    function dynamics_sim(state::SArray{Tuple{14}, T, 1, 14} where T, p::IntegratorParameters, t)
        lkm = (p.dt-t)/p.dt
        lkp = t/p.dt
        u = p.uk*lkm + p.up*lkp
        p.sigma*dx(state,u,p.pinfo)
    end

    macro genIdx(start,last)
        return :(@SVector [i for i=$start:$last])
    end
    const stateC = @genIdx(1,14)
    const ukC = @genIdx(15,17)
    const upC = @genIdx(18,20)
    const sigmaC = 21
    function fdiff(pinfo,dt)
        function f(p)
            state = p[stateC]
            uk = p[ukC]
            up = p[upC]
            sigma = p[sigmaC]
            prob = ODEProblem(dynamics_sim, state, eltype(p).((0.0,dt)), IntegratorParameters(dt,uk,up,sigma,pinfo))
            DifferentialEquations.solve(prob, Tsit5(),reltol=1e-8,abstol=1e-8,force_dtmin=true,dtmin=0.0001,save_everystep=false,saveat=[dt])[end]
        end
        return f
    end

    function linearize_segment(fun, sigma_lin::Float64, istate::SVector{14}, c1::SVector{3}, c2::SVector{3})
        ip = SVector{21}(vcat(istate, c1, c2, sigma_lin))
        res = DiffResults.JacobianResult((@MVector zeros(14)),MVector(ip))
        ForwardDiff.jacobian!(res, fun, ip)
        jcb = DiffResults.jacobian(res)
        return LinRes(DiffResults.value(res),SMatrix{14,21}(jcb))
    end

    function linearize_dynamics(states::Array{LinPoint,1}, sigma_lin::Float64, dt::Float64, info::ProbInfo)
        results = Array{LinRes,1}(length(states)-1)
        fun = fdiff(info,dt)
        for i=1:length(states)-1
            cstate = states[i]
            nstate = states[i+1]
            results[i] = linearize_segment(fun, sigma_lin, cstate.state, cstate.control, nstate.control)
        end
        return results
    end
    function next_step(dynam, ab, abn, state, control_k, control_kp, sigma, sigHat, relax)
        ctrl = vcat(state-ab.state, control_k-ab.control, control_kp-abn.control,sigma-sigHat)
        return dynam.derivative *ctrl + dynam.endpoint + relax
    end
    export linearize_dynamics, next_step, LinRes
end

max_stars_repo_name: BenChung/SuccessiveConvexification

max_issues_repo_name: BenChung/SuccessiveConvexification

max_forks_repo_name: BenChung/SuccessiveConvexification


text:
using Revise
using PDDL

domain = load_domain("test/pddl/domain.pddl")
problem = load_problem("test/pddl/problem.pddl")
1
# inits = collect(initstate(domain, problem).facts)


text:
const MOVE_STRAIGHT_INDEX = 1
const TURN_RIGHT_INDEX = 2
const TURN_LEFT_INDEX = 3

next_lane(lane::Lane, roadway::Roadway, direction::Int) = roadway[lane.exits[clamp(direction,1,length(lane.exits))].target.tag]

prev_lane(lane::Lane, roadway::Roadway, direction::Int) = roadway[lane.entrances[clamp(direction,1,length(lane.entrances))].target.tag]

next_lane_point(lane::Lane, roadway::Roadway, direction::Int) = roadway[lane.exits[clamp(direction,1,length(lane.exits))].target]

prev_lane_point(lane::Lane, roadway::Roadway, direction::Int) = roadway[lane.entrances[clamp(direction,1,length(lane.entrances))].target]

function move_along_with_direction(roadind::RoadIndex, roadway::Roadway, Δs::Float64; direction::Int = 1)

    lane = roadway[roadind.tag]
    curvept = lane[roadind.ind, roadway]

    if curvept.s + Δs < 0.0
        if has_prev(lane)
            pt_lo = prev_lane_point(lane, roadway, direction)
            pt_hi = lane.curve[1]
            s_gap = norm(VecE2(pt_hi.pos - pt_lo.pos))

            if curvept.s + Δs < -s_gap
                lane_prev = prev_lane(lane, roadway, direction)
                curveind = curveindex_end(lane_prev.curve)
                roadind = RoadIndex(curveind, lane_prev.tag)
                return move_along_with_direction(roadind, roadway, Δs + curvept.s + s_gap, direction=direction)
            else # in the gap between lanes
                t = (s_gap + curvept.s + Δs) / s_gap
                curveind = CurveIndex(0, t)
                RoadIndex(curveind, lane.tag)
            end

        else # no prev lane, return the beginning of this one
            curveind = CurveIndex(1, 0.0)
            return RoadIndex(curveind, roadind.tag)
        end
    elseif curvept.s + Δs > lane.curve[end].s
        if has_next(lane)
            pt_lo = lane.curve[end]
            pt_hi = next_lane_point(lane, roadway, direction)
            s_gap = norm(VecE2(pt_hi.pos - pt_lo.pos))

            if curvept.s + Δs ≥ pt_lo.s + s_gap # extends beyond the gap
                curveind = lane.exits[min(length(lane.exits),direction)].target.ind
                roadind = RoadIndex(curveind, lane.exits[min(length(lane.exits),direction)].target.tag)
                return move_along_with_direction(roadind, roadway, Δs - (lane.curve[end].s + s_gap - curvept.s),direction=direction)
            else # in the gap between lanes
                t = (Δs - (lane.curve[end].s - curvept.s)) / s_gap
                curveind = CurveIndex(0, t)
                RoadIndex(curveind, lane.exits[clamp(direction,1,length(lane.exits))].target.tag)
            end
        else # no next lane, return the end of this lane
            curveind = curveindex_end(lane.curve)
            return RoadIndex(curveind, roadind.tag)
        end
    else
        if Δs >= 0.0
            if roadind.ind.i == 0
                ind = get_curve_index(CurveIndex(1,0.0), lane.curve, curvept.s+Δs)
            elseif roadind.ind.i == length(lane.curve)
                ind = get_curve_index(curveindex_end(lane.curve), lane.curve, curvept.s+Δs)
            else
                ind = get_curve_index(roadind.ind, lane.curve, Δs)
            end
            RoadIndex(ind, roadind.tag)
        else
            if roadind.ind.i == 0
                ind = get_curve_index(CurveIndex(1,0.0), lane.curve, curvept.s+Δs)
            elseif roadind.ind.i == length(lane.curve)
                ind = get_curve_index(curveindex_end(lane.curve), lane.curve, Δs)
            else
                ind = get_curve_index(roadind.ind, lane.curve, Δs)
            end
            RoadIndex(ind, roadind.tag)
        end
    end
end

function in_lanes(posG::VecSE2, roadway::Roadway)
    projections = Array{RoadProjection}(undef, 0)
    for seg in roadway.segments
        for lane in seg.lanes
            roadproj = proj(posG, lane, roadway, move_along_curves=false)
            targetlane = roadway[roadproj.tag]
            footpoint = targetlane[roadproj.curveproj.ind, roadway]
            dist = normsquared(VecE2(posG - footpoint.pos))
            if dist < targetlane.width/2.0
                push!(projections,roadproj)
            end
        end
    end

    return projections
end


text:
# This is an example on training the SCAMPy implementation of the EDMF
# scheme with data generated using PyCLES.
#
# The example seeks to find the optimal values of the entrainment and
# detrainment parameters of the EDMF scheme to replicate the LES profiles
# of the BOMEX experiment.
#
# This example is fully parallelized and can be run in the cluster with
# the included script.

# Import modules to all processes
@everywhere using Pkg
@everywhere Pkg.activate("../..")
@everywhere using Distributions
@everywhere using StatsBase
@everywhere using LinearAlgebra
# Import EKP modules
@everywhere using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
@everywhere using EnsembleKalmanProcesses.Observations
@everywhere using EnsembleKalmanProcesses.ParameterDistributionStorage
@everywhere include(joinpath(@__DIR__, "../../src/helper_funcs.jl"))
using JLD2
using NPZ

#########
#########  Define the parameters and their priors
#########

# Define the parameters that we want to learn
param_names = ["entrainment_factor", "detrainment_factor"]
n_param = length(param_names)

# Prior information: Define transform to unconstrained gaussian space
constraints = [ [bounded(0.01, 0.3)],
                [bounded(0.01, 1.0)]]
# All vars are standard Gaussians in unconstrained space
prior_dist = [Parameterized(Normal(0.0, 1.0))
                for x in range(1, n_param, length=n_param) ]
priors = ParameterDistribution(prior_dist, constraints, param_names)

#########
#########  Retrieve true LES samples from PyCLES data and transform
#########

# Define observation window (s)
ti = [14400.0]
tf = [21600.0]
# Define variables considered in the loss function
y_names = Array{String, 1}[]
push!(y_names, ["thetal_mean", "ql_mean", "qt_mean", "total_flux_h", "total_flux_qt"])

# Define preconditioning and regularization of inverse problem
perform_PCA = true # Performs PCA on data

# Define name of PyCLES simulation to learn from
sim_names = ["Bomex"]
sim_suffix = [".may18"]

# Init arrays
yt = zeros(0)
yt_var_list = []
yt_big = zeros(0)
yt_var_list_big = Array{Float64, 2}[]
P_pca_list = []
pool_var_list = []
for (i, sim_name) in enumerate(sim_names)
    les_dir = string("/groups/esm/ilopezgo/Output.", sim_name, sim_suffix[i])
    # Get SCM vertical levels for interpolation
    z_scm = get_profile(string("Output.", sim_name, ".00000"), ["z_half"])
    # Get (interpolated and pool-normalized) observations, get pool variance vector
    yt_, yt_var_, pool_var = obs_LES(y_names[i], les_dir, ti[i], tf[i], z_scm = z_scm)
    push!(pool_var_list, pool_var)
    if perform_PCA
        yt_pca, yt_var_pca, P_pca = obs_PCA(yt_, yt_var_)
        append!(yt, yt_pca)
        push!(yt_var_list, yt_var_pca)
        push!(P_pca_list, P_pca)
    else
        append!(yt, yt_)
        push!(yt_var_list, yt_var_)
        global P_pca_list = nothing
    end
    # Save full dimensionality (normalized) output for error computation
    append!(yt_big, yt_)
    push!(yt_var_list_big, yt_var_)
end
d = length(yt) # Length of data array

# Construct global observational covariance matrix, TSVD
yt_var = zeros(d, d)
vars_num = 1
for (k,config_cov) in enumerate(yt_var_list)
    vars = length(config_cov[1,:])
    yt_var[vars_num:vars_num+vars-1, vars_num:vars_num+vars-1] = config_cov
    global vars_num = vars_num+vars
end

yt_var_big =cov_from_cov_list(yt_var_list_big)

n_samples = 1
samples = zeros(n_samples, length(yt))
samples[1,:] = yt
Γy = yt_var

#########
#########  Calibrate: Ensemble Kalman Inversion
#########

algo = Inversion() # Sampler(vcat(get_mean(priors)...), get_cov(priors))
N_ens = 20 # number of ensemble members
N_iter = 10 # number of EKP iterations.
Δt = 1.0 # Artificial time stepper of the EKI.
println("NUMBER OF ENSEMBLE MEMBERS: ", N_ens)
println("NUMBER OF ITERATIONS: ", N_iter)

initial_params = construct_initial_ensemble(priors, N_ens, rng_seed=rand(1:1000))
ekobj = EnsembleKalmanProcess(initial_params, yt, Γy, algo )
scm_dir = "SCAMPy/"

# Define caller function
@everywhere g_(x::Array{Float64,1}) = run_SCAMPy(x, $param_names,
   $y_names, $scm_dir, $ti, $tf, P_pca_list = $P_pca_list,
   norm_var_list = $pool_var_list, scampy_handler = "call_BOMEX.sh")

# Create output dir
prefix = "results_"
prefix = typeof(algo) == Sampler{Float64} ? string(prefix, "eks_") : string(prefix, "eki_")
prefix = Δt ≈ 1 ? prefix : string(prefix, "dt", Δt, "_")
outdir_path = string(prefix, "p", n_param,"_e", N_ens, "_i", N_iter, "_d", d)
println("Name of outdir path for this EKP, ", outdir_path)
command = `mkdir $outdir_path`
try
    run(command)
catch e
    println("Output directory already exists. Output may be overwritten.")
end

# EKP iterations
g_ens = zeros(N_ens, d)
norm_err_list = []
g_big_list = []
for i in 1:N_iter
    # Note that the parameters are transformed when used as input to SCAMPy
    params_cons_i = deepcopy(transform_unconstrained_to_constrained(priors,
        get_u_final(ekobj)) )
    params = [row[:] for row in eachrow(params_cons_i')]
    @everywhere params = $params
    array_of_tuples = pmap(g_, params) # Outer dim is params iterator
    (g_ens_arr, g_ens_arr_pca) = ntuple(l->getindex.(array_of_tuples,l),2) # Outer dim is G̃, G
    println(string("\n\nEKP evaluation ",i," finished. Updating ensemble ...\n"))
    for j in 1:N_ens
      g_ens[j, :] = g_ens_arr_pca[j]
    end
    # Get normalized error for full dimensionality output
    push!(norm_err_list, compute_errors(g_ens_arr, yt_big))
    # Store full dimensionality output
    push!(g_big_list, g_ens_arr)
    # Get normalized error
    if typeof(algo) != Sampler{Float64}
        update_ensemble!(ekobj, Array(g_ens') , Δt_new=Δt)
    else
        update_ensemble!(ekobj, Array(g_ens') )
    end
    println("\nEnsemble updated. Saving results to file...\n")
    # Convert to arrays
    phi_params = Array{Array{Float64,2},1}(transform_unconstrained_to_constrained(priors, get_u(ekobj)))
    phi_params_arr = zeros(i+1, n_param, N_ens)
    for (k,elem) in enumerate(phi_params)
      phi_params_arr[k,:,:] = elem
    end

    # Save EKP information to JLD2 file
    save(string(outdir_path,"/ekp.jld2"),
        "ekp_u", transform_unconstrained_to_constrained(priors, get_u(ekobj)),
        "ekp_g", get_g(ekobj),
        "truth_mean", ekobj.obs_mean,
        "truth_cov", ekobj.obs_noise_cov,
        "ekp_err", ekobj.err,
        "truth_mean_big", yt_big,
        "truth_cov_big", yt_var_big,
        "g_big", g_big_list,
        "norm_err", norm_err_list,
        "P_pca", P_pca_list,
        "pool_var", pool_var_list,
        "phi_params", phi_params_arr,
        )
    # Convert to arrays
    phi_params = Array{Array{Float64,2},1}(transform_unconstrained_to_constrained(priors, get_u(ekobj)))
    phi_params_arr = zeros(i+1, n_param, N_ens)
    g_big_arr = zeros(i, N_ens, length(yt_big))
    for (k,elem) in enumerate(phi_params)
      phi_params_arr[k,:,:] = elem
      if k < i + 1
        g_big_arr[k,:,:] = hcat(g_big_list[k]...)'
      end
    end
    norm_err_arr = hcat(norm_err_list...)' # N_iter, N_ens
    # Or you can also save information to numpy files with NPZ
    npzwrite(string(outdir_path,"/y_mean.npy"), ekobj.obs_mean)
    npzwrite(string(outdir_path,"/Gamma_y.npy"), ekobj.obs_noise_cov)
    npzwrite(string(outdir_path,"/ekp_err.npy"), ekobj.err)
    npzwrite(string(outdir_path,"/phi_params.npy"), phi_params_arr)
    npzwrite(string(outdir_path,"/y_mean_big.npy"), yt_big)
    npzwrite(string(outdir_path,"/Gamma_y_big.npy"), yt_var_big)
    npzwrite(string(outdir_path,"/norm_err.npy"), norm_err_arr)
    npzwrite(string(outdir_path,"/g_big.npy"), g_big_arr)
    for (l, P_pca) in enumerate(P_pca_list)
      npzwrite(string(outdir_path,"/P_pca_",sim_names[l],".npy"), P_pca)
      npzwrite(string(outdir_path,"/pool_var_",sim_names[l],".npy"), pool_var_list[l])
    end
end

# EKP results: Has the ensemble collapsed toward the truth?
println("\nEKP ensemble mean at last stage (original space):")
println( mean( transform_unconstrained_to_constrained(priors, get_u_final(ekobj)), dims=2) ) # Parameters are stored as columns


text:
using Documenter
using Literate
using RHEOSafm

"""
    docprepare()
Convert all source documentation (.jl and .md files)
into markdown, ready for documentation build, and
relevant .jl files into Jupyter notebooks ready for
deployment to notebooks branch.
"""
function docprepare()
    # remove remnants of previous build and recreate staging dir
    rm("docs/staging-docs", force=true, recursive=true)
    mkdir("docs/staging-docs")

    # copy readme to staging-docs
    # cp("README.md", "docs/staging-docs/index.md")

    # copy logo to staging-docs
    cp("logo.png", "docs/staging-docs/logo.png")

    # copy assets to staging directory
    cp("docs/src/assets", "docs/staging-docs/assets")

    # iterate through src and convert/copy as appropriate
    for file in readdir("docs/src")
        if endswith(file, "md")
            cp("docs/src/$file", "docs/staging-docs/$file")
        elseif endswith(file, "jl")
            Literate.markdown("docs/src/$file", "docs/staging-docs/"; documenter=true)
        end
    end
end

function notebookprepare()
    # create notebook staging dir
    mkdir("docs/staging-docs/notebooks")

    # copy assets to notebooks staging directory
    cp("docs/src/assets", "docs/staging-docs/notebooks/assets")

    for file in readdir("docs/src")
        if endswith(file, "jl")
            Literate.notebook("docs/src/$file", "docs/staging-docs/notebooks/")
        end
    end
end

function maindocbuilder()
    # prepare doc files from source
    docprepare()

    # prepare notebook files from source
    notebookprepare()

    # build docs from staging area
    makedocs(modules=[RHEOSafm],
            doctest= false, clean=true,
            format = Documenter.HTML(),
            sitename=" ",
            source = "staging-docs",
            authors="Alessandra Bonfanti, Louis Kaplan, Alexandre Kabla",
            pages = ["Home" => "index.md",
                     "Examples" => ["Relaxation experiment" => "example_01.md",
                                    "Indentation experiment" => "example_02.md"],
                     "API" => "API.md"
                    ]
            )
    deploydocs(repo = "github.com/JuliaRheology/RHEOSafm.jl.git",
               deps = nothing,
               make = nothing,
               target = "build")

end

maindocbuilder()


text:
# Copyright (c) 2014: SCS.jl contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

const MOI = MathOptInterface

include("scaled_psd_cone_bridge.jl")

MOI.Utilities.@product_of_sets(
    _Cones,
    MOI.Zeros,
    MOI.Nonnegatives,
    MOI.SecondOrderCone,
    ScaledPSDCone,
    MOI.ExponentialCone,
    MOI.DualExponentialCone,
    MOI.PowerCone{T},
    MOI.DualPowerCone{T},
)

struct _SetConstants{T}
    b::Vector{T}
    power_coefficients::Dict{Int,T}
    _SetConstants{T}() where {T} = new{T}(T[], Dict{Int,T}())
end

function Base.empty!(x::_SetConstants)
    empty!(x.b)
    empty!(x.power_coefficients)
    return x
end

Base.resize!(x::_SetConstants, n) = resize!(x.b, n)

function MOI.Utilities.load_constants(x::_SetConstants, offset, f)
    MOI.Utilities.load_constants(x.b, offset, f)
    return
end

function MOI.Utilities.load_constants(
    x::_SetConstants{T},
    offset,
    set::Union{MOI.PowerCone{T},MOI.DualPowerCone{T}},
) where {T}
    x.power_coefficients[offset+1] = set.exponent
    return
end

function MOI.Utilities.function_constants(x::_SetConstants, rows)
    return MOI.Utilities.function_constants(x.b, rows)
end

function MOI.Utilities.set_from_constants(x::_SetConstants, S, rows)
    return MOI.Utilities.set_from_constants(x.b, S, rows)
end

function MOI.Utilities.set_from_constants(
    x::_SetConstants{T},
    ::Type{S},
    rows,
) where {T,S<:Union{MOI.PowerCone{T},MOI.DualPowerCone{T}}}
    @assert length(rows) == 3
    return S(x.power_coefficients[first(rows)])
end

const OptimizerCache{T} = MOI.Utilities.GenericModel{
    Cdouble,
    MOI.Utilities.ObjectiveContainer{Cdouble},
    MOI.Utilities.VariablesContainer{Cdouble},
    MOI.Utilities.MatrixOfConstraints{
        Cdouble,
        MOI.Utilities.MutableSparseMatrixCSC{
            Cdouble,
            T,
            MOI.Utilities.ZeroBasedIndexing,
        },
        _SetConstants{Cdouble},
        _Cones{Cdouble},
    },
}

function _to_sparse(
    ::Type{T},
    A::MOI.Utilities.MutableSparseMatrixCSC{
        Cdouble,
        T,
        MOI.Utilities.ZeroBasedIndexing,
    },
) where {T}
    return -A.nzval, A.rowval, A.colptr
end

mutable struct MOISolution
    primal::Vector{Float64}
    dual::Vector{Float64}
    slack::Vector{Float64}
    ret_val::Int
    raw_status::String
    objective_value::Float64
    dual_objective_value::Float64
    objective_constant::Float64
    solve_time_sec::Float64
    iterations::Int
end

function MOISolution()
    return MOISolution(
        Float64[],
        Float64[],
        Float64[],
        0, # SCS_UNFINISHED
        "",
        NaN,
        NaN,
        NaN,
        0.0,
        0,
    )
end

"""
    Optimizer()

Create a new SCS optimizer.
"""
mutable struct Optimizer <: MOI.AbstractOptimizer
    cones::Union{Nothing,_Cones{Cdouble}}
    sol::MOISolution
    silent::Bool
    options::Dict{Symbol,Any}

    Optimizer() = new(nothing, MOISolution(), false, Dict{Symbol,Any}())
end

function MOI.get(::Optimizer, ::MOI.Bridges.ListOfNonstandardBridges)
    return [ScaledPSDConeBridge{Cdouble}]
end

MOI.get(::Optimizer, ::MOI.SolverName) = "SCS"

MOI.is_empty(optimizer::Optimizer) = optimizer.cones === nothing

function MOI.empty!(optimizer::Optimizer)
    optimizer.cones = nothing
    optimizer.sol = MOISolution()
    return
end

###
### MOI.RawOptimizerAttribute
###

function MOI.set(optimizer::Optimizer, param::MOI.RawOptimizerAttribute, value)
    return optimizer.options[Symbol(param.name)] = value
end

function MOI.get(optimizer::Optimizer, param::MOI.RawOptimizerAttribute)
    return optimizer.options[Symbol(param.name)]
end

###
### MOI.Silent
###

MOI.supports(::Optimizer, ::MOI.Silent) = true

function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    optimizer.silent = value
    return
end

MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

###
### MOI.AbstractModelAttribute
###

function MOI.supports(
    ::Optimizer,
    ::Union{
        MOI.ObjectiveSense,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}},
    },
)
    return true
end

###
### MOI.AbstractVariableAttribute
###

function MOI.supports(
    ::Optimizer,
    ::MOI.VariablePrimalStart,
    ::Type{MOI.VariableIndex},
)
    return true
end

###
### MOI.AbstractConstraintAttribute
###

function MOI.supports(
    ::Optimizer,
    ::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
    ::Type{<:MOI.ConstraintIndex},
)
    return true
end

function MOI.supports_constraint(
    ::Optimizer,
    ::Type{MOI.VectorAffineFunction{Cdouble}},
    ::Type{
        <:Union{
            MOI.Zeros,
            MOI.Nonnegatives,
            MOI.SecondOrderCone,
            ScaledPSDCone,
            MOI.ExponentialCone,
            MOI.DualExponentialCone,
            MOI.PowerCone{Cdouble},
            MOI.DualPowerCone{Cdouble},
        },
    },
)
    return true
end

function _map_sets(f, ::Type{T}, sets, ::Type{S}) where {T,S}
    F = MOI.VectorAffineFunction{Cdouble}
    cis = MOI.get(sets, MOI.ListOfConstraintIndices{F,S}())
    return T[f(MOI.get(sets, MOI.ConstraintSet(), ci)) for ci in cis]
end

function MOI.optimize!(
    dest::Optimizer,
    src::MOI.Utilities.UniversalFallback{OptimizerCache{T}},
) where {T}
    linear_solver = get(dest.options, :linear_solver, DirectSolver)
    if T != scsint_t(linear_solver)
        cache = MOI.Utilities.UniversalFallback(
            OptimizerCache{scsint_t(linear_solver)}(),
        )
        MOI.copy_to(cache, src)
        return MOI.optimize!(dest, cache)
    end
    # The real stuff starts here.
    MOI.empty!(dest)
    index_map = MOI.Utilities.identity_index_map(src)
    Ab = src.model.constraints
    A = Ab.coefficients

    for (F, S) in keys(src.constraints)
        throw(MOI.UnsupportedConstraint{F,S}())
    end

    model_attributes = MOI.get(src, MOI.ListOfModelAttributesSet())
    objective_function_attr =
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Cdouble}}()
    for attr in model_attributes
        if attr == MOI.ObjectiveSense() || attr == objective_function_attr
            continue  # These are handled later
        elseif attr == MOI.Name()
            continue  # This can be skipped without consequence
        end
        throw(MOI.UnsupportedAttribute(attr))
    end
    max_sense = false
    if MOI.ObjectiveSense() in model_attributes
        max_sense = MOI.get(src, MOI.ObjectiveSense()) == MOI.MAX_SENSE
    end
    objective_constant = 0.0
    c = zeros(A.n)
    if objective_function_attr in model_attributes
        obj = MOI.get(src, objective_function_attr)
        objective_constant = MOI.constant(obj)
        for term in obj.terms
            c[term.variable.value] += (max_sense ? -1 : 1) * term.coefficient
        end
    end
    # `model.primal` contains the result of the previous optimization.
    # It is used as a warm start if its length is the same, e.g.
    # probably because no variable and/or constraint has been added.
    if A.n != length(dest.sol.primal)
        resize!(dest.sol.primal, A.n)
        fill!(dest.sol.primal, 0.0)
    end
    if A.m != length(dest.sol.dual)
        resize!(dest.sol.dual, A.m)
        fill!(dest.sol.dual, 0.0)
    end
    if A.m != length(dest.sol.slack)
        resize!(dest.sol.slack, A.m)
        fill!(dest.sol.slack, 0.0)
    end
    # Set starting values and throw error for other variable attributes
    has_warm_start = false
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    for attr in MOI.get(src, MOI.ListOfVariableAttributesSet())
        if attr == MOI.VariableName()
            # Skip
        elseif attr == MOI.VariablePrimalStart()
            has_warm_start = true
            for (i, x) in enumerate(vis_src)
                dest.sol.primal[i] = something(MOI.get(src, attr, x), 0.0)
            end
        else
            throw(MOI.UnsupportedAttribute(attr))
        end
    end
    # Set starting values and throw error for other constraint attributes
    for (F, S) in MOI.get(src, MOI.ListOfConstraintTypesPresent())
        cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
        for attr in MOI.get(src, MOI.ListOfConstraintAttributesSet{F,S}())
            if attr == MOI.ConstraintName()
                # Skip
            elseif attr == MOI.ConstraintPrimalStart()
                has_warm_start = true
                for ci in cis_src
                    start = MOI.get(src, attr, ci)
                    if start !== nothing
                        rows = MOI.Utilities.rows(Ab, ci)
                        dest.sol.slack[rows] .= start
                    end
                end
            elseif attr == MOI.ConstraintDualStart()
                has_warm_start = true
                for ci in cis_src
                    start = MOI.get(src, attr, ci)
                    if start !== nothing
                        rows = MOI.Utilities.rows(Ab, ci)
                        dest.sol.dual[rows] .= start
                    end
                end
            else
                throw(MOI.UnsupportedAttribute(attr))
            end
        end
    end
    options = copy(dest.options)
    if dest.silent
        options[:verbose] = 0
    end
    dest.cones = deepcopy(Ab.sets)
    sol = scs_solve(
        linear_solver,
        A.m,
        A.n,
        A,
        SparseArrays.spzeros(A.n, A.n), # placeholder: P
        Ab.constants.b,
        c,
        Ab.sets.num_rows[1],
        Ab.sets.num_rows[2] - Ab.sets.num_rows[1],
        Float64[], # # placeholder: bu
        Float64[], # # placeholder: bl
        _map_sets(MOI.dimension, T, Ab, MOI.SecondOrderCone),
        _map_sets(MOI.side_dimension, T, Ab, ScaledPSDCone),
        div(Ab.sets.num_rows[5] - Ab.sets.num_rows[4], 3),
        div(Ab.sets.num_rows[6] - Ab.sets.num_rows[5], 3),
        vcat(
            _map_sets(set -> set.exponent, Float64, Ab, MOI.PowerCone{Cdouble}),
            _map_sets(
                set -> -set.exponent,
                Float64,
                Ab,
                MOI.DualPowerCone{Cdouble},
            ),
        ),
        dest.sol.primal,
        dest.sol.dual,
        dest.sol.slack;
        warm_start = has_warm_start,
        options...,
    )
    # If the solution is an infeasibility certificate, the objective values are
    # left as infinite, not the value corresponding to the ray.
    if !isfinite(sol.info.dobj)
        sol.info.dobj = -Ab.constants.b' * sol.y
    end
    if !isfinite(sol.info.pobj)
        sol.info.pobj = c' * sol.x
    end
    if sol.ret_val == -4  # SCS_FAILED
        dest.sol = MOISolution(
            sol.x,
            sol.y,
            sol.s,
            sol.ret_val,
            "",
            (max_sense ? -1 : 1) * sol.info.pobj,
            (max_sense ? -1 : 1) * sol.info.dobj,
            objective_constant,
            NaN,
            -1,
        )
    else
        dest.sol = MOISolution(
            sol.x,
            sol.y,
            sol.s,
            sol.ret_val,
            raw_status(sol.info),
            (max_sense ? -1 : 1) * sol.info.pobj,
            (max_sense ? -1 : 1) * sol.info.dobj,
            objective_constant,
            (sol.info.setup_time + sol.info.solve_time) / 1000,
            sol.info.iter,
        )
    end
    return index_map, false
end

function MOI.optimize!(dest::Optimizer, src::MOI.ModelLike)
    linear_solver = get(dest.options, :linear_solver, DirectSolver)
    T = scsint_t(linear_solver)
    cache = MOI.Utilities.UniversalFallback(OptimizerCache{T}())
    index_map = MOI.copy_to(cache, src)
    MOI.optimize!(dest, cache)
    return index_map, false
end

function MOI.get(optimizer::Optimizer, ::MOI.SolveTimeSec)
    return optimizer.sol.solve_time_sec
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.sol.raw_status
end

"""
    ADMMIterations()

The number of ADMM iterations completed during the solve.
"""
struct ADMMIterations <: MOI.AbstractModelAttribute end

MOI.is_set_by_optimize(::ADMMIterations) = true

function MOI.get(optimizer::Optimizer, ::ADMMIterations)
    return optimizer.sol.iterations
end

# Implements getter for result value and statuses
# SCS returns one of the following integers:
# -7 SCS_INFEASIBLE_INACCURATE
# -6 SCS_UNBOUNDED_INACCURATE
# -5 SCS_SIGINT
# -4 SCS_FAILED
# -3 SCS_INDETERMINATE
# -2 SCS_INFEASIBLE  : primal infeasible, dual unbounded
# -1 SCS_UNBOUNDED   : primal unbounded, dual infeasible
#  0 SCS_UNFINISHED  : never returned, used as placeholder
#  1 SCS_SOLVED
#  2 SCS_SOLVED_INACCURATE
function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    s = optimizer.sol.ret_val
    @assert -7 <= s <= 2
    if s == -7
        return MOI.ALMOST_INFEASIBLE
    elseif s == -6
        return MOI.ALMOST_DUAL_INFEASIBLE
    elseif s == 2
        return MOI.ALMOST_OPTIMAL
    elseif s == -5
        return MOI.INTERRUPTED
    elseif s == -4
        return MOI.INVALID_MODEL
    elseif s == -3
        return MOI.SLOW_PROGRESS
    elseif s == -2
        return MOI.INFEASIBLE
    elseif s == -1
        return MOI.DUAL_INFEASIBLE
    elseif s == 1
        return MOI.OPTIMAL
    else
        @assert s == 0
        return MOI.OPTIMIZE_NOT_CALLED
    end
end

function MOI.get(optimizer::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    value = optimizer.sol.objective_value
    if !MOI.Utilities.is_ray(MOI.get(optimizer, MOI.PrimalStatus()))
        value += optimizer.sol.objective_constant
    end
    return value
end

function MOI.get(optimizer::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    value = optimizer.sol.dual_objective_value
    if !MOI.Utilities.is_ray(MOI.get(optimizer, MOI.DualStatus()))
        value += optimizer.sol.objective_constant
    end
    return value
end

function MOI.get(optimizer::Optimizer, attr::MOI.PrimalStatus)
    if attr.result_index > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    elseif optimizer.sol.ret_val in (-3, 1, 2)
        return MOI.FEASIBLE_POINT
    elseif optimizer.sol.ret_val in (-6, -1)
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif optimizer.sol.ret_val == -4
        return MOI.NO_SOLUTION
    end
    return MOI.INFEASIBLE_POINT
end

function MOI.get(
    optimizer::Optimizer,
    attr::MOI.VariablePrimal,
    vi::MOI.VariableIndex,
)
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.sol.primal[vi.value]
end

function MOI.get(
    optimizer::Optimizer,
    attr::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.sol.slack[MOI.Utilities.rows(optimizer.cones, ci)]
end

function MOI.get(optimizer::Optimizer, attr::MOI.DualStatus)
    if attr.result_index > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    elseif optimizer.sol.ret_val in (-3, 1, 2)
        return MOI.FEASIBLE_POINT
    elseif optimizer.sol.ret_val in (-7, -2)
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif optimizer.sol.ret_val == -4
        return MOI.NO_SOLUTION
    end
    return MOI.INFEASIBLE_POINT
end

function MOI.get(
    optimizer::Optimizer,
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.sol.dual[MOI.Utilities.rows(optimizer.cones, ci)]
end

MOI.get(optimizer::Optimizer, ::MOI.ResultCount) = 1


text:
struct cmpo
    Q::TensorMap{ComplexSpace, 1, 1}
    L::TensorMap{ComplexSpace, 2, 1}
    R::TensorMap{ComplexSpace, 2, 1}
    P::TensorMap{ComplexSpace, 2, 2}
end

@inline get_phy(T::cmpo) = get_chi(T.R)
@inline get_vir(T::cmpo) = get_d(T.R) + 1

@non_differentiable get_phy(T::cmpo)
@non_differentiable get_vir(T::cmpo)


text:
# dv[SingleItemIndex, SingleItemIndex)
function Base.getindex(d::DataVector,
                       i::SingleIndex,
                       j::SingleIndex)
    if j != 1
        throw(ArgumentError("Second index must be 1"))
    end
    if d.na[i]
        return NA
    else
        return d.data[i]
    end
end

head(dv::AbstractDataVector) = dv[1:min(6, length(dv))]
tail(dv::AbstractDataVector) = dv[max(length(dv) - 6, 1):length(dv)]

# Container operations

# TODO: Macroize these definitions

function Base.push!{T}(dv::DataVector{T}, v::NAtype)
    push!(dv.data, baseval(T))
    push!(dv.na, true)
    return v
end

function Base.push!{S, T}(dv::DataVector{S}, v::T)
    push!(dv.data, v)
    push!(dv.na, false)
    return v
end

function Base.pop!(dv::DataVector)
    d, m = pop!(dv.data), pop!(dv.na)
    if m
        return NA
    else
        return d
    end
end

function Base.unshift!{T}(dv::DataVector{T}, v::NAtype)
    unshift!(dv.data, baseval(T))
    unshift!(dv.na, true)
    return v
end

function Base.unshift!{S, T}(dv::DataVector{S}, v::T)
    unshift!(dv.data, v)
    unshift!(dv.na, false)
    return v
end

function Base.shift!{T}(dv::DataVector{T})
    d, m = shift!(dv.data), shift!(dv.na)
    if m
        return NA
    else
        return d
    end
end

# TODO: should this be an AbstractDataVector, so it works with PDV's?
function Base.map(f::Function, dv::DataVector)
    n = length(dv)
    res = DataArray(Any, n)
    for i in 1:n
        res[i] = f(dv[i])
    end
    return res
end

function Base.push!{T,R}(pdv::PooledDataVector{T,R}, v::NAtype)
    push!(pdv.refs, zero(R))
    return v
end

function Base.push!{S,R,T}(pdv::PooledDataVector{S,R}, v::T)
    v = convert(S,v)
    push!(pdv.refs, getpoolidx(pdv, v))
    return v
end

Base.pop!(pdv::PooledDataVector) = pdv.pool[pop!(pdv.refs)]

function Base.unshift!{T,R}(pdv::PooledDataVector{T,R}, v::NAtype)
    unshift!(pdv.refs, zero(R))
    return v
end

function Base.unshift!{S,R,T}(pdv::PooledDataVector{S,R}, v::T)
    v = convert(S,v)
    unshift!(pdv.refs, getpoolidx(pdv, v))
    return v
end

Base.shift!(pdv::PooledDataVector) = pdv.pool[shift!(pdv.refs)]

Base.reverse(x::AbstractDataVector) = x[end:-1:1]

# Pad a vector with NA's

function padNA(dv::AbstractDataVector,
               front::Integer,
               back::Integer)
    n = length(dv)
    res = similar(dv, front + n + back)
    for i in 1:n
        res[i + front] = dv[i]
    end
    return res
end


text:
#-----------------------------------------------------------------------
#
#   BROYDEN ALGORITHM
#
#-----------------------------------------------------------------------

#-----------------------------------------------------------------------
"""
    JiSMU( ΔF::RealVec , Δx::RealVec , Ji::RealMtx )

The function `JiSMU` implements the Sherman-Morrison update formula,
returning an updated value of the Jacobian matrix `Ji`, given the
respective differences `ΔF` and `Δx` for the function ``F(x)`` and its
argument ``x``.

"""
function JiSMU( ΔF::RealVec , Δx::RealVec , Ji::RealMtx )
    ΔxTJi = transpose(Δx)*Ji
    return Ji + ( ( Δx - Ji*ΔF )/( ΔxTJi*ΔF ) )*( ΔxTJi )
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
"""
    bsolve( F::Function , J::RealMtx , f0::RealVec , x0::RealVec ,
            nb::Int=24 )

The function `bsolve` implements the Broyden algorithm; in particular,
it finds the roots of the function `F(x)`, given an initial Jacobian
matrix `J` and the initial guesses `f0` and `x0`. The parameter `nb`
specifies the maximum number of iterations.

"""
function bsolve( F::Function , J::RealMtx , f0::RealVec , x0::RealVec ,
                 nb::Int=24 )
    tpfl = typeof(x0[1])

    Ji = inv(J)
    F0 = f0

    p = dot(F0,F0)

    ( b , k ) = ( true , 1 )

    if nb > 0
        Δx = - Ji*f0
        xi = x0 + Δx
        xs = xi
        if nb > 1
            Fi = F(xi)
            ΔF = Fi - F0
            F0 = Fi
            p  = norm(F0)
            ps = p
            i  = 2
            while i <= nb && b
                Ji = JiSMU( ΔF , Δx , Ji )
                Δx = - Ji*F0
                xi = xi + Δx
                Fi = F(xi)
                ( ΔF , F0 ) = ( Fi - F0 , Fi )
                pc = norm(F0)
                if pc<p && pc<ps
                    ( p , ps , xs ) = ( pc , pc , xi )
                elseif pc>p && k>=4
                    b = false
                else
                    p  = pc
                    k += 1
                end
                i += 1
            end
        end
    end

    return xs
end     #---------------------------------------------------------------