const GAEVF = JuMP.GenericAffExpr{Float64, VariableRef}
const GQEVF = JuMP.GenericQuadExpr{Float64, VariableRef}

function moi_tests(
    op_problem::OperationsProblem,
    params::Bool,
    vars::Int,
    interval::Int,
    lessthan::Int,
    greaterthan::Int,
    equalto::Int,
    binary::Bool,
)
    JuMPmodel = PSI.get_jump_model(op_problem)
    @test (:ParameterJuMP in keys(JuMPmodel.ext)) == params
    @test JuMP.num_variables(JuMPmodel) == vars
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.Interval{Float64}) == interval
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.LessThan{Float64}) == lessthan
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.GreaterThan{Float64}) == greaterthan
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.EqualTo{Float64}) == equalto
    @test ((JuMP.VariableRef, MOI.ZeroOne) in JuMP.list_of_constraint_types(JuMPmodel)) ==
          binary

    return
end

function psi_constraint_test(
    op_problem::OperationsProblem,
    constraint_names::Vector{Symbol},
)
    constraints = PSI.get_constraints(op_problem)
    for con in constraint_names
        @test !isnothing(get(constraints, con, nothing))
    end
    return
end

function psi_checkbinvar_test(
    op_problem::OperationsProblem,
    bin_variable_names::Vector{Symbol},
)
    container = PSI.get_optimization_container(op_problem)
    for variable in bin_variable_names
        for v in PSI.get_variable(container, variable)
            @test JuMP.is_binary(v)
        end
    end
    return
end

function psi_checkobjfun_test(op_problem::OperationsProblem, exp_type)
    model = PSI.get_jump_model(op_problem)
    @test JuMP.objective_function_type(model) == exp_type
    return
end

function moi_lbvalue_test(op_problem::OperationsProblem, con_name::Symbol, value::Number)
    for con in PSI.get_constraints(op_problem)[con_name]
        @test JuMP.constraint_object(con).set.lower == value
    end
    return
end

function psi_checksolve_test(op_problem::OperationsProblem, status)
    model = PSI.get_jump_model(op_problem)
    JuMP.optimize!(model)
    @test termination_status(model) in status
end

function psi_checksolve_test(
    op_problem::OperationsProblem,
    status,
    expected_result,
    tol = 0.0,
)
    res = solve!(op_problem)
    model = PSI.get_jump_model(op_problem)
    @test termination_status(model) in status
    obj_value = JuMP.objective_value(model)
    @test isapprox(obj_value, expected_result, atol = tol)
end

function psi_ptdf_lmps(res::ProblemResults, ptdf)
    duals = get_duals(res)
    λ = convert(Array, duals[:CopperPlateBalance][:, :var])

    nf_duals = collect(keys(duals))[occursin.(PSI.NETWORK_FLOW, string.(keys(duals)))]
    flow_duals =
        hcat([duals[k][:, propertynames(duals[k]) .!== :DateTime] for k in nf_duals]...)
    μ = Matrix(flow_duals[:, ptdf.axes[1]])

    buses = get_components(Bus, get_system(res))
    lmps = OrderedDict()
    for bus in buses
        lmps[get_name(bus)] = μ * ptdf[:, get_number(bus)]
    end
    lmp = λ .- DataFrames.DataFrame(lmps)
    return lmp[!, sort(propertynames(lmp))]
end

function check_variable_unbounded(op_problem::OperationsProblem, var_name)
    psi_cont = PSI.get_optimization_container(op_problem)
    variable = PSI.get_variable(psi_cont, var_name)
    for var in variable
        if JuMP.has_lower_bound(var) || JuMP.has_upper_bound(var)
            return false
        end
    end
    return true
end

function check_variable_bounded(op_problem::OperationsProblem, var_name)
    psi_cont = PSI.get_optimization_container(op_problem)
    variable = PSI.get_variable(psi_cont, var_name)
    for var in variable
        if !JuMP.has_lower_bound(var) || !JuMP.has_upper_bound(var)
            return false
        end
    end
    return true
end

function check_flow_variable_values(
    op_problem::OperationsProblem,
    var_name::Symbol,
    device_name::String,
    limit::Float64,
)
    psi_cont = PSI.get_optimization_container(op_problem)
    variable = PSI.get_variable(psi_cont, var_name)
    for var in variable[device_name, :]
        if !(JuMP.value(var) <= (limit + 1e-2))
            return false
        end
    end
    return true
end

function check_flow_variable_values(
    op_problem::OperationsProblem,
    var_name::Symbol,
    device_name::String,
    limit_min::Float64,
    limit_max::Float64,
)
    psi_cont = PSI.get_optimization_container(op_problem)
    variable = PSI.get_variable(psi_cont, var_name)
    for var in variable[device_name, :]
        if !(JuMP.value(var) <= (limit_max + 1e-2)) ||
           !(JuMP.value(var) >= (limit_min - 1e-2))
            return false
        end
    end
    return true
end

function check_flow_variable_values(
    op_problem::OperationsProblem,
    pvar_name::Symbol,
    qvar_name::Symbol,
    device_name::String,
    limit_min::Float64,
    limit_max::Float64,
)
    psi_cont = PSI.get_optimization_container(op_problem)
    time_steps = PSI.model_time_steps(psi_cont)
    pvariable = PSI.get_variable(psi_cont, pvar_name)
    qvariable = PSI.get_variable(psi_cont, qvar_name)
    for t in time_steps
        fp = JuMP.value(pvariable[device_name, t])
        fq = JuMP.value(qvariable[device_name, t])
        flow = sqrt((fp)^2 + (fq)^2)
        if !(flow <= (limit_max + 1e-2)^2) || !(flow >= (limit_min - 1e-2)^2)
            return false
        end
    end
    return true
end

function check_flow_variable_values(
    op_problem::OperationsProblem,
    pvar_name::Symbol,
    qvar_name::Symbol,
    device_name::String,
    limit::Float64,
)
    psi_cont = PSI.get_optimization_container(op_problem)
    time_steps = PSI.model_time_steps(psi_cont)
    pvariable = PSI.get_variable(psi_cont, pvar_name)
    qvariable = PSI.get_variable(psi_cont, qvar_name)
    for t in time_steps
        fp = JuMP.value(pvariable[device_name, t])
        fq = JuMP.value(qvariable[device_name, t])
        flow = sqrt((fp)^2 + (fq)^2)
        if !(flow <= (limit + 1e-2)^2)
            return false
        end
    end
    return true
end

function PSI._jump_value(int::Int)
    @warn("This is for testing purposes only.")
    return int
end

function _test_plain_print_methods(list::Array)
    for object in list
        normal = repr(object)
        io = IOBuffer()
        show(io, "text/plain", object)
        grabbed = String(take!(io))
        @test !isnothing(grabbed)
    end
end

function _test_html_print_methods(list::Array)
    for object in list
        normal = repr(object)
        io = IOBuffer()
        show(io, "text/html", object)
        grabbed = String(take!(io))
        @test !isnothing(grabbed)
    end
end


import Base: length, first, last, getindex

abstract IndexRepr

immutable LineRepr <: IndexRepr
    id::NTuple{2, Int}
end
interior_indices(::LineRepr) = ()

immutable CircleArcRepr <: IndexRepr
    id::NTuple{3, Int}
end
interior_indices(::CircleArcRepr) = (2,)

immutable EllipseArcRepr <: IndexRepr
    id::NTuple{3, Int}
end
interior_indices(::EllipseArcRepr) = (2,)

for (T, size) in zip((:LineRepr, :CircleArcRepr, :EllipseArcRepr), (2, 3, 3))
    eval(quote
            first(r::$T) = first(r.id)
            last(r::$T) = last(r.id)
            getindex(r::$T, i::Int) = r.id[i]
         end)
end

# Convert to primitives  -----------------------------------------------------------------

function primitives(l::Line)
    points = [l.p0, l.p1]
    indices = LineRepr((1, 2))
    (points, indices)
end

function primitives(c::CircleArc)
    points = [c.p0, c.center, c.p1]
    indices = CircleArcRepr((1, 2, 3))
    (points, indices)
end

function primitives(c::EllipseArc)
    points = [c.p0, c.center, c.p1]
    indices = EllipseArcRepr((1, 2, 3))
    (points, indices)
end

function primitives(shape::BoundingBox)
    dx = shape.ur.x - shape.ll.x
    dy = shape.ur.y - shape.ll.y

    points = [shape.ll, shape.ll + Point(dx, 0), shape.ur, shape.ll + Point(0, dy)]
    primitives(ClosedPolygon(points))
end

function primitives(shape::Square)
    points = [shape.ll,
              shape.ll + Point(shape.size, 0),
              shape.ll + Point(shape.size, shape.size),
              shape.ll + Point(0, shape.size)]

    primitives(ClosedPolygon(points))
end

function primitives(shape::Rectangle)
    points = [shape.ll,
              shape.ll + Point(shape.size_x, 0),
              shape.ll + Point(shape.size_x, shape.size_y),
              shape.ll + Point(0, shape.size_y)]

    primitives(ClosedPolygon(points))
end

# Physical points, curves defined in terms of LOCAL point numbering, orientation
function primitives(shape::ClosedPolygon)
    # Point part
    points = shape.points

    npoints = length(points)
    pindices = 1:npoints  # Local
    curves = map(LineRepr, zip([pindices[2:end]..., first(pindices)],
                               [pindices[1:end-1]..., last(pindices)]))

    (points, [(c, 1) for c in curves])
end

function primitives(shape::Circle)
    points = [shape.center,
              shape.center + Point(shape.radius, 0),
              shape.center + Point(0, shape.radius),
              shape.center + Point(-shape.radius, 0),
              shape.center + Point(0, -shape.radius)]

    curves = map(CircleArcRepr, [(2, 1, 3), (3, 1, 4), (4, 1, 5), (5, 1, 2)])

    (points, [(c, 1) for c in curves])
end

function primitives(shape::Ellipse)
    if shape.size_x > shape.size_y
        points = [shape.center,
                  shape.center + Point(shape.size_x, 0),
                  shape.center + Point(0, shape.size_y),
                  shape.center + Point(-shape.size_x, 0),
                  shape.center + Point(0, -shape.size_y)]
    else
        points = [shape.center,
                  shape.center + Point(0, shape.size_y),
                  shape.center + Point(shape.size_x, 0),
                  shape.center + Point(0, -shape.size_y),
                  shape.center + Point(-shape.size_x, 0)]
    end
    curves = map(EllipseArcRepr, [(2, 1, 3), (4, 1, 3), (4, 1, 5), (2, 1, 5)])
    orientation = [1, -1, 1, -1]

    (points, collect(zip(curves, orientation)))
end

function primitives(shape::Loop)
    # NOTE: loop as it comes here has curves defined in terms of points which can be not can
    # coincide. Here we try to keep only the uniqe points - the check for duplicates,
    # hoever, involves only the interior points of curves.

    curves = shape.curves
    orientation = shape.orientation
    @assert first(orientation) == 1

    points, indices = primitives(first(curves))
    # First curve for sure can add all its points
    all_points = Vector{Point}(points)
    # And also its indices
    all_indices = Vector{IndexRepr}([indices])
    # The connection over last point

    link = last(points)
    link_index = last(indices)
    for (curve, sign) in zip(curves[2:end-1], orientation[2:end-1])
        counter = length(all_points)
        # Connection over first
        points, indices = primitives(curve)
        # First, interior, Last
        if sign == 1
            # First, is known
            global_indices = [link_index]
            # Need to check interior points
            for interior_id in interior_indices(indices)
                point = points[interior_id]
                index = findfirst(q -> q == point, all_points)
                # Not seen
                if index == 0
                    push!(all_points, point)

                    counter += 1
                    index = counter
                end
                # Seen with index
                push!(global_indices, index)
            end
            # The last gut is new
            link = last(points)
            push!(all_points, link)
            counter += 1

            link_index = counter
            push!(global_indices, link_index)

            repr = typeof(indices)(tuple(global_indices...))
        else
            # Last, is known so remember to add it later
            # First is new
            link = first(points)
            push!(all_points, link)
            counter +=1

            global_indices = [counter]
            # Need to check interior points
            for interior_id in interior_indices(indices)
                point = points[interior_id]
                index = findfirst(q -> q == point, all_points)
                # Not seen
                if index == 0
                    push!(all_points, point)

                    counter += 1
                    index = counter
                end
                # Seen with index
                push!(global_indices, index)
            end
            # The last point is there already as link_index
            push!(global_indices, link_index)

            repr = typeof(indices)(tuple(global_indices...))
            link_index = first(repr)
        end
        push!(all_indices, repr)
    end

    # The last curve closes the loop so if 1/-1 oriented we have seen its last/first point
    # That point was seen as 1, last is given by link_index
    counter = length(all_points)
    curve, sign = last(curves), last(orientation)
    points, indices = primitives(curve)

    global_indices = []
    fi, li = (sign == 1) ? (link_index, 1) : (1, link_index)

    push!(global_indices, fi)
    # Interior guys
    for interior_id in interior_indices(indices)
        point = points[interior_id]
        index = findfirst(q -> q == point, all_points)
        # Not seen
        if index == 0
            push!(all_points, point)

            counter += 1
            index = counter
        end
        # Seen with index
        push!(global_indices, index)
    end
    push!(global_indices, li)

    repr = typeof(indices)(tuple(global_indices...))
    push!(all_indices, repr)

    (all_points, collect(zip(all_indices, orientation)))
end

# Convert to gmsh ------------------------------------------------------------------------

function mesh(canvas::Canvas, size::Vector, file::AbstractString)
    base, ext = splitext(file)
    for (index, size_i) in enumerate(size)
        file_i = base * "$(index)" * ext
        mesh(canvas, size_i, file_i)
    end
end

function mesh(canvas::Canvas, size::Real, file::AbstractString)
    base, ext = splitext(file)

    geo_file = "$(base).geo"
    if ext == ".msh"
        @assert gmsh_script(canvas, "$(geo_file)", size)
        gmsh = success(`gmsh -2 $(geo_file) -o $(file)`)
        return gmsh && success(`rm $(geo_file)`)
    end
    # XML mesh, mesh itself & physical & facet
    msh_file = "$(base).msh"
    if ext == ".xml"
        @assert mesh(canvas, size, msh_file)
        convert = success(`dolfin-convert $(msh_file) $(file)`)
        return convert && success(`rm $(msh_file)`)
    end
    # HDF5
    xml_file = "$(base).xml"
    if ext == ".h5"
        @assert mesh(canvas, size, xml_file)

const python = "
from dolfin import Mesh, MeshFunction, mpi_comm_world, HDF5File;
import os;

xml_mesh = '$(xml_file)'
base, ext = os.path.splitext(xml_mesh);
assert os.path.exists(xml_mesh);

xml_volumes = '$(base)_physical_region.xml';
assert os.path.exists(xml_volumes);

xml_facets = '$(base)_facet_region.xml';
has_facet_regions = os.path.exists(xml_facets);

mesh = Mesh(xml_mesh);
cell_f = MeshFunction('size_t', mesh, xml_volumes);

h5_file = '$(base).h5';
out = HDF5File(mesh.mpi_comm(), h5_file, 'w');
out.write(mesh, '/mesh');
out.write(cell_f, '/cell_markers');

if has_facet_regions: out.write(MeshFunction('size_t', mesh, xml_facets), '/facet_markers')
#"
        convert = success(`python -c "$(python)"`)
        xmls = [xml_file, "$(base)_physical_region.xml", "$(base)_facet_region.xml"]
        return convert && all(success(`rm $(file)`) for file in filter(isfile, xmls))
    end
end

gmsh_script(canvas::Canvas, size::Real) = write(STDOUT, gmsh_code(canvas, size))

function gmsh_script(canvas::Canvas, file::AbstractString, size::Real)
    @assert last(splitext(file)) == ".geo"
    f = open(file, "w")
    count = write(f, gmsh_code(canvas, size))
    close(f)
    count > 0
end

function gmsh_code(canvas::Canvas, size::Real)
    counter = 1  # Of any GMSH entities

    code = Vector{AbstractString}(["SIZE = $(size);\n"])
    loops = Vector{Int}()  # Each shape is a closed loop, their collection together with the
                           # canvas loop defines the outer domain

    # We first write shapes
    for (tag, shape) in enumerate(canvas.shapes)
        points, curves = primitives(shape)
        # Will need to convert entities from local to global
        local_to_global = collect(counter:counter+length(points)-1)
        # Code for points
        points_gmsh, counter = gmsh_code(points, local_to_global, counter)
        # In case of lines we want a loop as well; for loops and to define physical ...
        curves_gmsh, shape_loop, counter = gmsh_code(curves, local_to_global, counter)
        # Line Loop is another entity
        # NOTE, all stuff below is what gmsh_code of a curve should do .... Anyways, there
        # will be rewrite of this because of 3d
        loop_gmsh = "Line Loop($(counter)) = {$(join(map(x -> "$(x)", shape_loop), ", "))};"
        counter += 1
        # Plane surface in terms of line loop
        surf_gmsh = "Plane Surface($(counter)) = {$(counter-1)};"
        # Physical, not entities
        physurf_gmsh = "Physical Surface($(tag)) = {$(counter)};"
        physline_gmsh = "Physical Line($(tag)) = {$(join(map(x-> "$(x)", shape_loop), ", "))};"

        loops = vcat(loops, counter-1)  # Add to 'inner' boundary
        # Prepare for next round
        counter += 1  # Counter so that it is ready to use
        shape_code = join([points_gmsh,
                           curves_gmsh,
                           loop_gmsh,
                           surf_gmsh,
                           physurf_gmsh,
                           physline_gmsh,
                           "//--------------------\n"], "\n")  # For readability
        push!(code, shape_code)
    end
    # Now the canvas
    tag = length(canvas.shapes) + 1

    points, curves = primitives(canvas.bbox)
    local_to_global = collect(counter:counter+length(points)-1)

    points_gmsh, counter = gmsh_code(points, local_to_global, counter)

    curves_gmsh, shape_loop, counter = gmsh_code(curves, local_to_global, counter)
    loop_gmsh = "Line Loop($(counter)) = {$(join(map(x -> "$(x)", shape_loop), ", "))};"
    counter += 1

    surf_gmsh = "Plane Surface($(counter)) = {$(counter-1), $(join(map(x -> "$(x)", loops), ", "))};"
    physurf_gmsh = "Physical Surface($(tag)) = {$(counter)};"
    physline_gmsh = "Physical Line($(tag)) = {$(join(map(x-> "$(x)", shape_loop), ", "))};"

    canvas_code = join([points_gmsh,
                        curves_gmsh,
                        loop_gmsh,
                        surf_gmsh,
                        physurf_gmsh,
                        physline_gmsh,
                        "//--------------------\n"], "\n")
    push!(code, canvas_code)

    # Done
    join(code, "\n")
end

function gmsh_code(points::Vector{Point}, local_to_global::Vector{Int}, counter::Int)
    # Code for points
    points_gmsh = Vector{AbstractString}()

    for (local_index, point) in enumerate(points)
        str = "Point($(local_to_global[local_index])) = {$(point.x), $(point.y), 0., SIZE};"
        push!(points_gmsh, str)
        counter +=1
    end
    points_gmsh = join(points_gmsh, "\n")

    (points_gmsh, counter)
end

function gmsh_code(line::Tuple{LineRepr, Int}, local_to_global::Vector{Int}, counter::Int)
    line, orientation = line
    v0, v1 = local_to_global[first(line)], local_to_global[last(line)]
    line_gmsh = "Line($(counter)) = {$(v0), $(v1)};"
    loop_piece = orientation*counter

    counter += 1
    (line_gmsh, loop_piece, counter)
end

function gmsh_code(curve::Tuple{CircleArcRepr, Int}, local_to_global::Vector{Int}, counter::Int)
    curve, orientation = curve
    v0, v1 = local_to_global[first(curve)], local_to_global[last(curve)]
    center = local_to_global[curve[2]]
    curve_gmsh = "Circle($(counter)) = {$(v0), $(center), $(v1)};"
    loop_piece = orientation*counter

    counter += 1
    (curve_gmsh, loop_piece, counter)
end

function gmsh_code(curve::Tuple{EllipseArcRepr, Int}, local_to_global::Vector{Int}, counter::Int)
    curve, orientation = curve
    v0, center, v1 = local_to_global[curve[1]], local_to_global[curve[2]], local_to_global[curve[3]]
    curve_gmsh = "Ellipse($(counter)) = {$(v0), $(center), $(v0), $(v1)};"
    loop_piece = orientation*counter

    counter += 1
    (curve_gmsh, loop_piece, counter)
end

function gmsh_code{T<:IndexRepr}(curves::Vector{Tuple{T, Int}}, local_to_global::Vector{Int}, counter::Int)
    curves_gmsh = Vector{AbstractString}()
    shape_loop = Vector{Int}()

    for curve in curves
        curve_gmsh, loop_piece, counter = gmsh_code(curve, local_to_global, counter)
        push!(curves_gmsh, curve_gmsh)
        push!(shape_loop, loop_piece)
    end
    curves_gmsh = join(curves_gmsh, "\n")

    (curves_gmsh, shape_loop, counter)
end


#mostly legacy code used in testing

#the below funciton is now only for testing purposes
function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterLevel,
  RegressionType::Val) where {TM<:Matrix,TV<:Vector,T}
  @unpack G,A = Θ

  #initialize A₁.==1 and all other As to the amount implied by G
  #print("prod(Θ.G): f0=$(prod(Θ.G|>Matrix)) x64=$(prod(Float64.(Θ.G |> Matrix) )) |")
  prevA₁ = A[:,1]
  factored::TM = initializeÃandfactor(Θ,Xv)

  #run the regression for A₁
  try
    A[:,1] .= reg(factored, Xv.v, RegressionType)
  catch err
    (err == InterruptException()) && throw(err)
    errmessage = "$err"
    errmessage = length(errmessage > 1000) ? errmessage[1:1000] : errmessage
    all((isfinite.(factored))) || error("Factored includes nonfinite values!
      prevA₁: $(prevA₁)
      ****************\nG: $G
      ****************\nfactored[1:1000,:]: $(factored[1:1000,:])")

    A[:,1] .= reg(factored, Xv.v, Val{:fallbackreg}())
    @warn "regression method $RegressionType failed due to $errmessage. Using fallbackreg fallback val=$(A[:,1])"
  end

  #println("A[:,1]: $(A[:,1])")

  #absμ .= vec(sum(absμₖ,dims=2))
  absμ = factored * (x->Float64(abs(x))).(A[:,1])

  return absμ
end

#primary prediction alogrithm for the demeaned version
function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterDemean,  RegressionType::Val
    ) where {TM,TV,T}
  @unpack G,A = Θ

  #initialize A₁.==1 and all other As to the amount implied by G
  prevA₁ = A[:,1]
  factored::TM = Xv.M * initializeÃandfactor(Θ,Xv)

  #run the regression for A₁
  A[:,1] .= reg(factored, Xv.ṽ, RegressionType)
  #@eval Main demeanbeta = $(A[:,1] |> Vector)

  absμ = factored * (x->Float64(abs(x))).(A[:,1])

  return absμ
end

(Θ::VolumePartsIter)(Xv::AbstractXYIter{TM,TV,T},
    ::Val{:leveldemean}) where {TM<:CuMatrix,TV<:CuVector,T} = error(
    "Do not use CUDA for the demeaning version")

function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterLevel, RegressionType::Val
  ) where {TM<:CuMatrix,TV<:CuVector,T}

  @unpack G,A = Θ

  #initialize A₁.==1 and all other As to the amount implied by G
  #print("prod(Θ.G): f0=$(prod(Θ.G|>Matrix)) x64=$(prod(Float64.(Θ.G |> Matrix) )) |")
  prevA₁ = A[:,1]
  factored::TM = initializeÃandfactor(Θ,Xv)

  #run the regression for A₁
  try
    A[:,1] .= reg(factored, Xv.v, RegressionType)
    #@info "culevel A[:,1]: $(A[:,1])"
    #@info "culevel factored[20,:]: $(factored[20,:])"
    #@info "culevel A[:,20]: $(A[:,20])"
  catch err #NOTE: below may be broken, could try to fix if needed
    (err == InterruptException()) && throw(err)
    errmessage = "$err"
    errmessage = length(errmessage) > 1000 ? errmessage[1:1000] : errmessage
    all((isfinite.(factored))) || @warn("Factored includes nonfinite values!
      prevA₁: $(prevA₁)
      ****************\nG: $G
      ****************\nfactored[1:1000,:]: $(factored[1:1000,:])
      attempting correction")
    factored = (x->min(x,LARGE_VAL)).(factored)

    A[:,1] .= reg(factored, Xv.v, Val{:fallbackreg}())
    @warn "regression method $RegressionType failed due to $errmessage. Using fallback val=$(A[:,1])"
  end

  #println("A[:,1]: $(A[:,1])")

  absμ = factored * (x->Float64(abs(x))).(A[:,1])

  return absμ
end

#NOTE: this is only for testing purposes, but is it even needed for that???
#handles the intercept case
function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterLevel,
  ::Union{Val{:interceptzygote}, Val{:intercept}},

  ) where {TM<:CuMatrix,TV<:CuVector,T}

  @unpack G,A = Θ
  factored::TM = hcat(fill!(similar(Xv.v), T(1.0)), initializeÃandfactor(Θ,Xv))

  #run the regression for A₁
  A₁ = reg(factored, Xv.v, Val{:choleskyzygote}())
  absμ = factored * A₁

  return absμ
end

#not sure why I need to write this again...but it throws a method error otherwise
function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterLevel,
  ::Union{Val{:interceptzygote},Val{:intercept}},
  ) where {TM<:Matrix,TV<:Vector,T}

  @unpack G,A = Θ
  factored::TM = hcat(fill!(similar(Xv.v), T(1.0)), initializeÃandfactor(Θ,Xv))
  #run the regression for A₁
  A₁ = reg(factored, Xv.v, Val{:choleskyzygote}())
  absμ = factored * A₁

  return absμ
end

function (Θ::VolumePartsIter)(Xv::XYIterLevel{TM,TV,T}, ::Val{:intercept}
  ) where {TM<:CuMatrix,TV<:CuVector,T}

  @unpack G,A = Θ

  #initialize A₁.==1 and all other As to the amount implied by G
  #print("prod(Θ.G): f0=$(prod(Θ.G|>Matrix)) x64=$(prod(Float64.(Θ.G |> Matrix) )) |")
  prevA₁ = A[:,1]
  factored::TM = hcat(Xv.tv1s, initializeÃandfactor(Θ,Xv))
  local b::TV
  #expectations = mean(factored, dims=1)
  #v̄ = mean(Xv.v)

  #run the regression for A₁
  try
    b = reg(factored, Xv.v, Val{:cholesky}())
    #@info "culevel A[:,1]: $(A[:,1])"
    #@info "culevel factored[20,:]: $(factored[20,:])"
    #@info "culevel A[:,20]: $(A[:,20])"
  catch err
    (err == InterruptException()) && throw(err)
    errmessage = "$err"
    errmessage = length(errmessage)  > 1000 ? errmessage[1:1000] : errmessage
    all((isfinite.(factored))) || @warn("Factored includes nonfinite values!
      prevA₁: $(prevA₁)
      ****************\nG: $G
      ****************\nfactored[1:1000,:]: $(factored[1:1000,:])
      attempting correction")
    factored .= (x->min(x,LARGE_VAL)).(factored)
    #expectations .= mean(factored, dims=1)

    b = reg(factored, Xv.v, Val{:fallbackreg}())
    @warn "regression method $RegressionType failed due to $errmessage. Using fallback val=$(A[:,1])"
  end

  #intercept = v̄ - dot(vec(expectations), A[:,1])

  #println("A[:,1]: $(A[:,1])")
  A[:,1] .= b[2:end]
  absμ = factored * b

  return absμ
end

function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T}, Xv::XYIterLevel,
  RegressionType::Val) where {TM,TV,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2), Nexpanded=size(Xv.ws,1))

  #special version for zygote
  #expand a matrix for A
  (RegressionType === Val{:intercept}()) && error("Don't use the pre-compute method with an intercept regression")

  Ã = hcat(ones(T, dims.K, 1), cumprodbyrow(G))
  Ãexpanded = view(Ã, Θ.expand.index.Ã...)
  LÃexpanded = view(Ã, Θ.expand.index.LÃ...)
  factored = genabsμₖ.(Ãexpanded, LÃexpanded, Xv.ws', Xv.Rtws', Xv.RLws')' #0.4ms

  absA₁ = (x->Float64(abs(x))).(reg(factored, Xv.v, RegressionType))#0.5ms
  #println("CPU absAₜ: $absAₜ")

  #sum across signals
  absμ = factored * absA₁
  #absμ = sum(Aexpanded, dims=1) |> vec
  return absμ
end

#used in testing
function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T}, Xv::XYIterDemean,
  RegressionType::Val) where {TM,TV,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2), Nexpanded=size(Xv.ws,1))

  #special version for zygote
  #expand a matrix for A
  (RegressionType === Val{:intercept}()) && error("Don't use the pre-compute method with an intercept regression")

  Ã = hcat(ones(T, dims.K, 1), cumprodbyrow(G))
  Ãexpanded = view(Ã, Θ.expand.index.Ã...)
  LÃexpanded = view(Ã, Θ.expand.index.LÃ...)
  factored = Xv.M * genabsμₖ.(Ãexpanded', LÃexpanded', Xv.ws, Xv.Rtws, Xv.RLws) #0.4ms

  absA₁ = (x->Float64(abs(x))).(reg(factored, Xv.ṽ, Val{:choleskyzygote}()))#0.5ms
  #@info ("absA₁ (:zygoteleveldemean): $(absA₁)")

  #sum across signals
  absμ = factored * absA₁
  #absμ = sum(Aexpanded, dims=1) |> vec
  return absμ
end

function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T}, Xv::XYIterLevel,
  RegressionType::Val) where {TM<:CuMatrix,TV<:CuVector,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2), Nexpanded=size(Xv.ws,1))

  someones = ones(T, dims.K, 1)
  prodmatG = cumprodbyrow(G)
  Ã = hcat(someones, prodmatG)

  Ãexpanded = CuMatrix{T}(Ã[Θ.expand.index.Ã...])
  LÃexpanded = CuMatrix{T}(Ã[Θ.expand.index.LÃ...])

  factored = genabsμₖ.(Ãexpanded, LÃexpanded, Xv.ws', Xv.Rtws', Xv.RLws')' #0.4ms

  #run the regression (use a spacial non-inplace version)

  #println(sum(cholesky(factored'*factored)\(factored'*Xv.v)))
  #println("Cu absAₜT: $absAₜT")

  absA₁ = reg(factored, Xv.v, RegressionType)

  absμ = factored * absA₁

  #=return absμ=#

  return absμ
end

#verion for regressions with intercept
function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T}, Xv::XYIterLevel,
  RegressionType::Val{:interceptzygote}) where {TM<:CuMatrix,TV<:CuVector,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2), Nexpanded=size(Xv.ws,1))

  someones = ones(T, dims.K, 1)
  prodmatG = cumprodbyrow(G)
  Ã = hcat(someones, prodmatG)

  Ãexpanded = CuMatrix{T}(Ã[Θ.expand.index.Ã...])
  LÃexpanded = CuMatrix{T}(Ã[Θ.expand.index.LÃ...])
  factored = CUDA.hcat(
    Xv.tv1s, genabsμₖ.(Ãexpanded, LÃexpanded, Xv.ws', Xv.Rtws', Xv.RLws')') #0.4ms

  absA₁ =reg(factored, Xv.v, Val{:choleskyzygote}())
  absμ = factored * absA₁
  #absμ = sum(Aexpanded, dims=1) |> vec
  return absμ
end
#cpu version of above
function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T}, Xv::XYIterLevel,
  RegressionType::Val{:interceptzygote}) where {TM<:Matrix,TV<:Vector,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2), Nexpanded=size(Xv.ws,1))

  someones = ones(T, dims.K, 1)
  prodmatG = cumprodbyrow(G)
  Ã = hcat(someones, prodmatG)

  Ãexpanded = view(Ã, Θ.expand.index.Ã...)
  LÃexpanded = view(Ã, Θ.expand.index.LÃ...)
  factored = hcat(Xv.tv1s, genabsμₖ.(Ãexpanded, LÃexpanded, Xv.ws', Xv.Rtws', Xv.RLws')') #0.4ms

  absA₁ =reg(factored, Xv.v, Val{:choleskyzygote}())
  absμ = factored * absA₁
  #absμ = sum(Aexpanded, dims=1) |> vec
  return absμ
end

function lossforzygote(G::Matrix{T}, Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel;
  RegressionType::Val=Val{PARAM[:iterregressiontypezygote]}()) where {TM,TV,T}

  #ε = factored * Aₜ .- Xv.ṽ
  Σε2 = sum((Θ(G, Xv, RegressionType) .- Xv.v).^2) #0.3ms

  return Σε2

end

function lossforzygote(G::Matrix{T}, Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterDemean;
  RegressionType::Val=Val{PARAM[:iterregressiontypezygote]}()) where {TM,TV,T}

  #ε = factored * Aₜ .- Xv.ṽ
  Σε2 = sum((Θ(G, Xv, RegressionType) .- Xv.ṽ).^2) #0.3ms

  return Σε2
end

#computs the predicted values for each strategy k and time t and demeans the result
#via the projection matrix
function projectvolumefromA(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterDemean) where {TM,TV,T}

  @unpack A, LA = Θ.expand
  @unpack ws, Rtws, RLws, M = Xv
  #first map each coefficient to its appropriate

  absμₖ = genabsμₖ.(A', LA', ws, Rtws, RLws)

  multiply!(M, absμₖ)
  return absμₖ
end

#same as the above without demeaning
function projectvolumefromA(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel) where {TM,TV,T}

  @unpack A, LA = Θ.expand
  @unpack ws, Rtws, RLws= Xv
  #first map each coefficient to its appropriate

  absμₖ =genabsμₖ.(A', LA', ws, Rtws, RLws)

  return absμₖ
end

function projectvolumefromA(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel
    ) where {TM<:CuMatrix, TV<:CuVector, T}

  @unpack A, LA = Θ.expand
  @unpack ws, Rtws, RLws = Xv
  #first map each coefficient to its appropriate

  dims = (K=size(Θ.A,1), T = size(Θ.A,2))

  genabsμₖ.(A', LA', ws, Rtws, RLws)
  #absμₖ .= cudagenabsμₖ.(A', LA', ws, RLws)

  return absμₖ
end

##############Iter loss methods
#again, mostly legacy

function iterloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel, RegressionType::Val;
  ) where {TM<:CuMatrix, TV<:CuVector, T}

  Δ = Θ(Xv, RegressionType) .- Xv.v
  λ = CUDA.sum(Δ .* Δ)# / Ncleanpanel

  return λ
end

iterloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel, RegressionType::Val;
  ) where {TM, TV, T} = sum(
    (Θ(Xv, RegressionType) .- Xv.v).^2)

iterloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterDemean, RegressionType::Val;
  ) where {TM, TV, T} = sum(
    (Θ(Xv, RegressionType) .- Xv.ṽ).^2)

function currentloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterDemean, absμₖ::TM) where {TM, TV, T}

  someones::Vector{T} = ones(T, size(Θ.A,1))
  return sum((absμₖ*someones .- Xv.ṽ).^2)
end

function currentloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterDemean,
   absμₖ::TM) where {TM<:CuMatrix, TV<:CuVector, T}

  someones::CuVector{T} = CUDA.ones(T, size(Θ.A,1))
  return sum((absμₖ*someones .- Xv.ṽ).^2)
end

function currentloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel,
   absμₖ::TM) where {TM<:CuMatrix, TV<:CuVector, T}

  someones::CuVector{T} = CUDA.ones(T, size(Θ.A,1))
  return sum((absμₖ*someones .- Xv.v).^2)
end

function testiterloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterLevel;
    RegressionType::Val=Val(PARAM[:iterregressiontype])) where {TM, TV, T}

  vhattest = iterloss(deepcopy(Θ), Xv, RegressionType)

  vhatver = sum((Θ(Xv, RegressionType) .- Xv.v) .^ 2)

  @assert vhatver ≈ vhattest
end

function testiterloss(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterDemean;
  RegressionType::Val=Val(:choleskyzygote)) where {TM, TV, T}

  vhattest = iterloss(deepcopy(Θ), Xv, RegressionType)

  vhatver = sum((Θ(Xv, RegressionType) .- Xv.ṽ) .^ 2)

  @assert vhatver ≈ vhattest
end


using PowerSimulationsDynamics
using Sundials

"""
Case 8:
This case study a 19-state virtual synchronous machine against an infinite bus located at bus 1, with VSM located at bus 2.
The perturbation increase the reference power (analogy for mechanical power) from 0.5 to 0.7.
"""

##################################################
############### LOAD DATA ########################
##################################################

include(joinpath(TEST_FILES_DIR, "data_tests/test08.jl"))

##################################################
############### SOLVE PROBLEM ####################
##################################################

#PSCAD benchmark data
csv_file = joinpath(TEST_FILES_DIR, "benchmarks/pscad/Test08/Test08_omega.csv")
t_offset = 9.0

#Define Fault using Callbacks
Pref_change = ControlReferenceChange(1.0, case_inv, :P_ref, 0.7)

@testset "Test 08 VSM Inverter Infinite Bus ResidualModel" begin
    path = (joinpath(pwd(), "test-08"))
    !isdir(path) && mkdir(path)
    try
        # Define Simulation Problem
        sim = Simulation(
            ResidualModel,
            omib_sys, # system
            path,
            (0.0, 4.0),
            Pref_change,
        )

        # Test Initial Condition
        diff_val = [0.0]
        res = get_init_values_for_comparison(sim)
        for (k, v) in test08_x0_init
            diff_val[1] += LinearAlgebra.norm(res[k] - v)
        end
        @test (diff_val[1] < 1e-3)

        # Obtain small signal results for initial conditions
        small_sig = small_signal_analysis(sim)
        eigs = small_sig.eigenvalues
        @test small_sig.stable

        # Test Eigenvalues
        @test LinearAlgebra.norm(eigs - test08_eigvals) < 1e-3

        # Solve problem
        @test execute!(sim, IDA(), dtmax = 0.005, saveat = 0.005) ==
              PSID.SIMULATION_FINALIZED
        results = read_results(sim)

        # Obtain frequency data
        series = get_state_series(results, ("generator-102-1", :ω_oc))
        t = series[1]
        ω = series[2]

        # Should return zeros and a warning
        series3 = get_field_current_series(results, "generator-102-1")
        series4 = get_field_voltage_series(results, "generator-102-1")

        # Obtain PSCAD benchmark data
        M = get_csv_data(csv_file)
        t_pscad = M[:, 1] .- t_offset
        ω_pscad = M[:, 2]

        power = PSID.get_activepower_series(results, "generator-102-1")
        rpower = PSID.get_reactivepower_series(results, "generator-102-1")
        @test isa(power, Tuple{Vector{Float64}, Vector{Float64}})
        @test isa(rpower, Tuple{Vector{Float64}, Vector{Float64}})
        @test LinearAlgebra.norm(ω - ω_pscad) <= 1e-4
        @test LinearAlgebra.norm(t - round.(t_pscad, digits = 3)) == 0.0

    finally
        @info("removing test files")
        rm(path, force = true, recursive = true)
    end
end

@testset "Test 08 VSM Inverter Infinite Bus MassMatrixModel" begin
    path = (joinpath(pwd(), "test-08"))
    !isdir(path) && mkdir(path)
    try
        #Define Simulation Problem
        sim = Simulation(
            MassMatrixModel,
            omib_sys, # system
            path,
            (0.0, 4.0),
            Pref_change,
        )

        # Test Initial Condition
        diff_val = [0.0]
        res = get_init_values_for_comparison(sim)
        for (k, v) in test08_x0_init
            diff_val[1] += LinearAlgebra.norm(res[k] - v)
        end
        @test (diff_val[1] < 1e-3)

        # Obtain small signal results for initial conditions
        small_sig = small_signal_analysis(sim)
        eigs = small_sig.eigenvalues
        @test small_sig.stable

        # Test Eigenvalues
        @test LinearAlgebra.norm(eigs - test08_eigvals) < 1e-3

        # Solve problem
        @test execute!(sim, Rodas5(), dtmax = 0.005, saveat = 0.005) ==
              PSID.SIMULATION_FINALIZED
        results = read_results(sim)

        # Obtain frequency data
        series = get_state_series(results, ("generator-102-1", :ω_oc))
        t = series[1]
        ω = series[2]

        # Obtain PSCAD benchmark data
        M = get_csv_data(csv_file)
        t_pscad = M[:, 1] .- t_offset
        ω_pscad = M[:, 2]

        power = PSID.get_activepower_series(results, "generator-102-1")
        rpower = PSID.get_reactivepower_series(results, "generator-102-1")
        @test isa(power, Tuple{Vector{Float64}, Vector{Float64}})
        @test isa(rpower, Tuple{Vector{Float64}, Vector{Float64}})
        @test LinearAlgebra.norm(ω - ω_pscad) <= 1e-4
        @test LinearAlgebra.norm(t - round.(t_pscad, digits = 3)) == 0.0

    finally
        @info("removing test files")
        rm(path, force = true, recursive = true)
    end
end


using Test, Qsosed

@testset "test parse parameters" begin
    config_path = String(@__DIR__) * "/../configs/config_example.yaml"
    parameters = Parameters(config_path)
    @test parameters.M == 1e8
    @test parameters.mdot == 0.5
    @test parameters.mdot == 0.5
    @test parameters.spin == 0.0
    @test parameters.mu_nucleon == 0.61
    @test parameters.reprocessing == true
    @test parameters.hard_xray_fraction == 0.02
    @test parameters.corona_electron_energy == 100
    @test parameters.warm_electron_energy == 0.2
    @test parameters.warm_photon_index == 2.5
    @test parameters.reflection_albedo == 0.3
    @test parameters.min_energy == 1e-4 # kev
    @test parameters.max_energy == 200
    @test parameters.uv_min_energy == 0.00387
    @test parameters.uv_max_energy == 0.06
    @test parameters.xray_min_energy == 0.06
    @test parameters.xray_max_energy == 1
end


module C

const NAMES = [
    "V1"
    "Km1"
    "V2"
    "Km2"
    "V3"
    "Km3"
    "V4"
    "Km4"
    "V5"
    "Km5"
    "V6"
    "Km6"
    "KimERK"
    "KexERK"
    "KimpERK"
    "KexpERK"
    "KimppERK"
    "KexppERK"
    "V10"
    "Km10"
    "n10"
    "p11"
    "p12"
    "p13"
    "V14"
    "Km14"
    "V15"
    "Km15"
    "p16"
    "p17"
    "KimDUSP"
    "KexDUSP"
    "KimpDUSP"
    "KexpDUSP"
    "V20"
    "Km20"
    "V21"
    "Km21"
    "p22"
    "p23"
    "V24"
    "Km24"
    "V25"
    "Km25"
    "KimRSK"
    "KexRSK"
    "V27"
    "Km27"
    "V28"
    "Km28"
    "V29"
    "Km29"
    "V30"
    "Km30"
    "V31"
    "Km31"
    "n31"
    "p32"
    "p33"
    "p34"
    "V35"
    "Km35"
    "V36"
    "Km36"
    "V37"
    "Km37"
    "p38"
    "p39"
    "KimFOS"
    "KexFOS"
    "KimpcFOS"
    "KexpcFOS"
    "V42"
    "Km42"
    "V43"
    "Km43"
    "V44"
    "Km44"
    "p45"
    "p46"
    "p47"
    "m47"
    "p48"
    "p49"
    "m49"
    "p50"
    "p51"
    "m51"
    "p52"
    "m52"
    "p53"
    "p54"
    "m54"
    "p55"
    "p56"
    "m56"
    "V57"
    "Km57"
    "n57"
    "p58"
    "p59"
    "p60"
    "p61"
    "KimF"
    "KexF"
    "p63"
    "KF31"
    "nF31"
    #
    "a"
    "Vn"
    "Vc"
    "Ligand"
    "EGF"
    "HRG"
    "no_ligand"
]

for (idx,name) in enumerate(NAMES)
    eval(Meta.parse("const $name = $idx"))
end

const NUM = length(NAMES)

end  # module


"""
    $(SIGNATURES)

Creates an axis array with given range and units, as well as spin polarization if
`polarized` is `true`. The spin dimension is the last. The arguments pairs indicate the name
and ranges (or categories) of the axes.
"""
function zeros_axisarray(T::Type, polarized::Bool, axes::Vararg{Tuple{Symbol, Any}};
                         kwargs...)
    if polarized
        zeros_axisarray(T, axes..., kwargs..., (:spin, [:↑, :↓ ]))
    else
        zeros_axisarray(T, axes...; kwargs...)
    end
end
function zeros_axisarray(T::Type, axes::Vararg{Tuple{Symbol, Any}}; kwargs...)
    @argcheck all(typeof(axis[1]) <: Symbol for axis in axes)
    underlying = zeros(T, ((length(axis[2]) for axis in chain(axes, kwargs))...))
    AxisArray(underlying, (Axis{n}(r) for (n, r) in chain(axes, kwargs))...)
end

"""
    $(SIGNATURES)

Creates an array with named axis and with correct units for an electronic density. The first
argument can be either an actual quantity with the correct dimensions, or a unitless subtype
of number. In the latter case, the units default to Hartree atomic units. In the former
case, one should take care to give a fully qualified quantity, including actual units (at
least until triangular dispatch makes its way into Julia).
"""
function density_array{Q <: Number}(T::Type{Q}, spin::Bool, args...; kwargs...)
    @argcheck dimension(Q) == NoDims
    density_array(HartreeUnits.ρ{Q}, spin, args...; kwargs...)
end
function density_array{Q <: DFTUnits.Ρ}(T::Type{Q}, spin::Bool, args...; kwargs...)
    zeros_axisarray(T, spin, args...; kwargs...)
end

@traitfn function similar_potential_from_rho(U::Type, ρ::::(!HasSpinDim))
    similar(ρ, U)
end
@traitfn function similar_potential_from_rho(U::Type, ρ::::HasSpinDim)
    similar(view(ρ, Axis{:spin}(:↑)), U)
end

"""
    $(SIGNATURES)

Creates an axis arrays for the potential, guessing the exact type from the type of ρ.
"""
function similar_potential{Q <: DFTUnits.Ρ}(ρ::AxisArray{Q})
    @argcheck :radius ∈ axisnames(ρ)
    const i = findfirst(axisnames(ρ), :radius)
    const axis = axisvalues(ρ)[i]
    T = typeof(one(eltype(ρ)) / one(eltype(axis)))
    U = HartreeUnits.ϵ{T}
    fill!(similar_potential_from_rho(U, ρ), zero(U))
end


module SplitGraph

using TreeKnit
using TreeTools
using ProgressMeter

include("objects.jl")
include("tools.jl")
include("energy.jl")
include("likelihood.jl")

let verbose::Bool = false, vverbose::Bool = false
	global v() = verbose
	global vv() = vverbose
	global set_verbose(v) = (verbose = v)
	global set_vverbose(v) = (vverbose = v)
end

# function opttrees(γ, Trange, M, seq_lengths, t...)
# 	opttrees!(γ, Trange, M, seq_lengths, [copy(x, TreeTools.EmptyData) for x in t]...)
# end

# function opttrees(t...; kwargs...)
# 	opttrees!([copy(x, TreeTools.EmptyData) for x in t]...; kwargs...)
# end
"""
	opttrees!(t... ; kwargs...)

Return a list of MCCs for input trees.
Output:
1.
"""
function opttrees(t...;
	γ=2,
	seq_lengths=1000 * ones(Int64, length(t)),
	Trange=reverse(0.001:0.01:1.),
	M = 10,
	likelihood_sort=true,
	resolve=true,
	sa_rep = 1,
	verbose=false
)
	opttrees!(
		γ, Trange, M, seq_lengths, [copy(convert(Tree{TreeTools.EmptyData}, x)) for x in t]...;
		likelihood_sort, resolve, sa_rep, verbose
	)
end
function opttrees!(
	γ, Trange, M, seq_lengths, t::Vararg{Tree};
	likelihood_sort=true, resolve=true, sa_rep=1, verbose=false,
)
	set_verbose(verbose)

	treelist = t
	mcc = naive_mccs(treelist...)
	if length(mcc) == 1
		return mcc, 0, 0., Int64[], Float64[], Union{Missing,Float64}[]
	end
	mcc_names = TreeKnit.name_mcc_clades!(treelist, mcc)
	for (i,t) in enumerate(treelist)
		TreeKnit.reduce_to_mcc!(t, mcc)
	end
	g = trees2graph(treelist)

	# SA - Optimization
	oconfs, F, nfound = sa_opt(g, γ=γ, Trange=Trange, M=M, rep=sa_rep, resolve=resolve)
	# Computing likelihoods
	if length(oconfs) != 1
		v() && @info "Sorting $(length(oconfs)) topologically equivalent configurations."
		vv() && @info "Configurations\n $oconfs"
		vv() && @info g.labels
		oconf, L = sortconf(oconfs, treelist, g, seq_lengths, mcc_names, likelihood_sort, false)
	else
		oconf = oconfs[1]
		L = Union{Missing,Float64}[]
	end
	vv() && @info "Final configuration for this iteration: $oconf."
	vv() && @info "MCCs removed: $([mcc_names[x] for x in g.labels[.!oconf]])"
	return [mcc_names[x] for x in g.labels[.!oconf]], compute_energy(oconf,g), compute_F(oconf, g, γ), L
end

function sortconf(oconfs, trees, g::Graph, seq_lengths, mcc_names, likelihood_sort, E_sort)
	if E_sort # Only considering configurations of lowest energies
		E = [compute_energy(conf,g) for conf in oconfs]
		Emin = minimum(E)
		oconfs_ = oconfs[findall(x->x==Emin, E)]
		v() && @info "Removing ", length(oconfs) - length(oconfs_), " configurations using energy."
	else # Removing configurations where nothing is removed
		oconfs_ = oconfs[findall(c->sum(c)<length(c), oconfs)]
	end

	# Sorting the remaining configurations using Poisson likelihood ratios
	if length(oconfs_) == 1
		return oconfs_[1], Union{Missing,Float64}[]
	elseif !likelihood_sort
		v() && @info "No likelihood sort: picking a random configuration among remaining ones"
		return rand(oconfs_), Union{Missing,Float64}[]
	else
		v() && @info "Comparing $(length(oconfs_)) configurations using likelihood"
		L = Union{Missing,Float64}[]
		for conf in oconfs_
			push!(L, conf_likelihood(conf, g, seq_lengths, trees, mode=:time, v=vv()))
		end
		vv() && println("Confs: ", [[mcc_names[x] for x in g.labels[.!conf]] for conf in oconfs_])
		vv() && @info "Likelihoods: $L"
		Lmax = maximum(L)
		ismissing(Lmax) && @warn "Maximum likelihood is `missing`"
		oconfs_ = oconfs_[findall(isequal(Lmax), L)]
		if length(oconfs_) != 1 # Final sort by energy if more than one most likely conf
			E = [compute_energy(conf,g) for conf in oconfs_]
			Emin = minimum(E)
			oconfs_ = oconfs_[findall(isequal(Emin), E)]
		end
		return rand(oconfs_), L
	end
end

end # module


module Distances

include("DistanceTypes.jl")

export
    # Function
    GetDistance,

    # Enum
    DistanceType,
    Euclidean,
    Cityblock,
    TotalVariation,
    Chebyshev,
    Jaccard,
    BrayCurtis,
    SpanNormDist
end


const MSB = 0x80
const MASK7 = 0x7f
const MASK8 = 0xff

const _wfbuf = (Vector{UInt8}(undef, 1), Vector{UInt8}(undef, 2), Vector{UInt8}(undef, 4), Vector{UInt8}(undef, 8), Vector{UInt8}(undef, 16))

const TIO = Union{IO, TTransport}

function _write_fixed(io::TIO, ux::T, bigendian::Bool) where T <: Unsigned
    N = sizeof(ux)
    _write_fixed(io, ux, _wfbuf[Int(log2(N))+1], bigendian ? (N:-1:1) : (1:N))
end

function _write_fixed(io::TIO, ux::T, a::Vector{UInt8}, r::R) where {T <: Unsigned, R <: AbstractRange}
    for n in r
        a[n] = UInt8(ux & MASK8)
        ux >>>= 8
    end
    write(io, a)
end

function _read_fixed(io::TIO, ret::T, N::Int, bigendian::Bool) where T <: Unsigned
    r = bigendian ? ((N-1):-1:0) : (0:1:(N-1))
    for n in r
        byte = convert(T, read(io, UInt8))
        ret |= (byte << (8*n))
    end
    ret
end

function _write_uleb(io::TIO, x::T) where T <: Integer
    nw = 0
    cont = true
    while cont
        byte = x & MASK7
        if (x >>>= 7) != 0
            byte |= MSB
        else
            cont = false
        end
        nw += write(io, UInt8(byte))
    end
    nw
end

function _read_uleb(io::TIO, typ::Type{T}) where T <: Integer
    res = convert(typ, 0)
    n = 0
    byte = UInt8(MSB)
    while (byte & MSB) != 0
        byte = read(io, UInt8)
        res |= (convert(typ, byte & MASK7) << (7*n))
        n += 1
    end
    res
end

function _write_zigzag(io::TIO, x::T) where T <: Integer
    nbits = 8*sizeof(x)
    zx = xor((x << 1), (x >> (nbits-1)))
    _write_uleb(io, zx)
end

function _read_zigzag(io::TIO, typ::Type{T}) where T <: Integer
    zx = _read_uleb(io, UInt64)
    # result is positive if zx is even
    convert(typ, iseven(zx) ? (zx >>> 1) : -((zx+1) >>> 1))
end


"""
    FermiFS{N,M,S} <: SingleComponentFockAddress

Address type that represents a Fock state of `N` fermions of the same spin in `M` modes
by wrapping a bitstring of type `S <: BitString`.

# Constructors

* `FermiFS{N,M}(bs::BitString)`: Unsafe constructor. Does not check whether the number of
  particles in `bs` is equal to `N`.

* `FermiFS(::BitString)`: Automatically determine `N` and `M`. This constructor is not type
  stable!

* `FermiFS{[N,M,S]}(onr)`: Create `FermiFS{N,M}` from [`onr`](@ref) representation. This is
  efficient as long as at least `N` is provided.

See also: [`SingleComponentFockAddress`](@ref), [`BoseFS`](@ref), [`BitString`](@ref).
"""
struct FermiFS{N,M,S<:BitString{M}} <: SingleComponentFockAddress{N,M}
    bs::S

    FermiFS{N,M,S}(bs::S) where {N,M,S<:BitString{M}} = new{N,M,S}(bs)
end

function FermiFS{N,M,S}(onr::Union{SVector{M},NTuple{M}}) where {N,M,C,T,S<:BitString{M,C,T}}
    @boundscheck sum(onr) == N && all(in((0, 1)), onr) || error("Invalid ONR: may only contain 0s and 1s.")
    result = zero(SVector{C,T})
    for mode in 1:M
        iszero(onr[mode]) && continue
        minus_j, offset = fldmod(mode - 1, 64)
        j = C - minus_j
        new = result[j] | T(1) << T(offset)
        result = setindex(result, new, j)
    end
    return FermiFS{N,M,S}(S(SVector(result)))
end

function FermiFS{N,M}(onr::Union{AbstractVector,Tuple}) where {N,M}
    S = typeof(BitString{M}(0))
    return FermiFS{N,M,S}(SVector{M}(onr))
end
function FermiFS{N}(onr::Union{SVector{M},Tuple{M}}) where {N,M}
    return FermiFS{N,M}(onr)
end
function FermiFS(onr::Union{AbstractVector,Tuple})
    M = length(onr)
    N = sum(onr)
    return FermiFS{N,M}(onr)
end

function Base.show(io::IO, f::FermiFS{N,M}) where {N,M}
    print(io, "FermiFS{$N,$M}(", tuple(onr(f)...), ")")
end

Base.bitstring(a::FermiFS) = bitstring(a.bs)
Base.isless(a::FermiFS, b::FermiFS) = isless(a.bs, b.bs)
Base.hash(a::FermiFS,  h::UInt) = hash(a.bs, h)
Base.:(==)(a::FermiFS, b::FermiFS) = a.bs == b.bs
num_occupied_modes(::FermiFS{N}) where {N} = N

function near_uniform(::Type{FermiFS{N,M}}) where {N,M}
    return FermiFS([fill(1, N); fill(0, M - N)])
end

onr(a::FermiFS) = SVector(m_onr(a))

"""
    FermiFSIndex

Struct used for indexing and performing [`excitation`](@ref)s on a [`FermiFS`](@ref).

## Fields:

* `occnum`: the occupation number.
* `mode`: the index of the mode.

"""
struct FermiFSIndex<:FieldVector{2,Int}
    occnum::Int
    mode::Int
end

function Base.show(io::IO, i::FermiFSIndex)
    @unpack occnum, mode = i
    print(io, "FermiFSIndex(occnum=$occnum, mode=$mode)")
end
Base.show(io::IO, ::MIME"text/plain", i::FermiFSIndex) = show(io, i)

find_mode(f::FermiFS, i) = FermiFSIndex(Int(is_occupied(f, i)), i)
find_mode(f::FermiFS, is::Tuple) = map(i -> find_mode(f, i), is)

"""
    FermiOccupiedModes{N,S<:BitString}

Iterator over occupied modes in address. `N` is the number of fermions. See [`occupied_modes`](@ref).
"""
struct FermiOccupiedModes{N,S}
    bs::S
end

occupied_modes(a::FermiFS{N,<:Any,S}) where {N,S} = FermiOccupiedModes{N,S}(a.bs)

Base.length(::FermiOccupiedModes{N}) where {N} = N
Base.eltype(::FermiOccupiedModes) = FermiFSIndex

function Base.iterate(o::FermiOccupiedModes)
    c = 0
    chunk = o.bs.chunks[end]
    while iszero(chunk)
        c += 1
        chunk = o.bs.chunks[end - c]
    end
    zeros = trailing_zeros(chunk % Int)
    return iterate(o, (chunk >> (zeros % UInt64), c * 64 + zeros, c))
end
function Base.iterate(o::FermiOccupiedModes, st)
    chunk, index, c = st
    while iszero(chunk)
        c += 1
        c == num_chunks(o.bs) && return nothing
        chunk = o.bs.chunks[end - c]
        index = c * 64
    end
    zeros = trailing_zeros(chunk % Int)
    index += zeros
    chunk >>= zeros
    return FermiFSIndex(1, index + 1), (chunk >> 1, index + 1, c)
end

function Base.iterate(o::FermiOccupiedModes{<:Any,<:BitString{<:Any,1,T}}) where {T}
    chunk = o.bs.chunks[end]
    zeros = trailing_zeros(chunk % Int)
    return iterate(o, (chunk >> (zeros % T), zeros))
end
function Base.iterate(o::FermiOccupiedModes{<:Any,<:BitString{<:Any,1,T}}, st) where {T}
    chunk, index = st
    iszero(chunk) && return nothing
    chunk >>= 0x1
    index += 1
    zeros = trailing_zeros(chunk % Int)
    return FermiFSIndex(1, index), (chunk >> (zeros % T), index + zeros)
end

function find_occupied_mode(a::FermiFS, i::Integer)
    for k in occupied_modes(a)
        i -= 1
        i == 0 && return k
    end
    return FermiFSIndex(0, 0)
end

@inline function m_onr(a::FermiFS{<:Any,M}) where {M}
    result = zero(MVector{M,Int32})
    @inbounds for (_, mode) in occupied_modes(a)
        result[mode] = 1
    end
    return result
end

function is_occupied(a::FermiFS{<:Any,M,S}, mode) where {M,T,S<:BitString{<:Any,1,T}}
    @boundscheck 1 ≤ mode ≤ M || throw(BoundsError(a, mode))
    return a.bs.chunks[1] & (T(1) << (mode - 1) % T) > 0
end
function is_occupied(a::FermiFS{<:Any,M}, mode) where {M}
    @boundscheck 1 ≤ mode ≤ M || throw(BoundsError(a, mode))
    j, i = fldmod1(mode, 64)
    return a.bs.chunks[end + 1 - j] & (UInt(1) << UInt(i - 1)) > 0
end

"""
    _flip_and_count(bs::BitString, k)

Count the number of ones before the `k`-th mode, flip the `k`th bit. Return the new
bitstring, the count, and the value of the bit after the flip.
"""
@inline function _flip_and_count(bs::BitString{<:Any,1,T}, k::Unsigned) where {T}
    chunk = bs.chunks[1]
    # highlights the k-th bit
    kmask = one(T) << k

    count = count_ones((kmask - 0x1) & chunk)
    chunk = chunk ⊻ kmask
    val = chunk & kmask > 0
    return typeof(bs)(chunk), count, val
end
@inline function _flip_and_count(bs::BitString, k::Unsigned)
    j, i = fldmod(k % Int, UInt(64))
    j = length(bs.chunks) - j
    chunk = bs.chunks[j]

    kmask = one(UInt64) << i

    count = count_ones((kmask - 0x1) & chunk)
    chunk = chunk ⊻ kmask
    val = chunk & kmask > 0

    count += mapreduce(count_ones, +, bs.chunks[j + 1:end]; init=0)
    return typeof(bs)(setindex(bs.chunks, chunk, j)), count, val
end

function excitation(a::FermiFS, creations::NTuple{N}, destructions::NTuple{N}) where {N}
    bs = a.bs
    count = 0
    for i in N:-1:1
        d = destructions[i].mode
        bs, x, val = _flip_and_count(bs, UInt(d - 0x1))
        val && return a, 0.0
        count += x
    end
    for i in N:-1:1
        c = creations[i].mode
        bs, x, val = _flip_and_count(bs, UInt(c - 0x1))
        !val && return a, 0.0
        count += x
    end
    return typeof(a)(bs), ifelse(iseven(count), 1.0, -1.0)
end


#Distribute a optigraph among workers.  Each worker should have the same first stage model.  Each worker will be allocated some of the nodes in the original modelgraph
function distribute_optigraph(graph::OptiGraph,to_workers::Vector{Int64};remote_name = :graph)
    #NOTE: Does not yet support subgraphs.  Need to aggregate first.
    to_workers = sort(to_workers)
    n_nodes = num_nodes(graph)
    n_workers = length(to_workers)
    nodes_per_worker = Int64(floor(n_nodes/n_workers))
    nodes = all_nodes(graph)
    node_indices = [getindex(graph,node) for node in nodes]
    user_pips_data = PipsNLP._setup_pips_nlp_data!(graph)

    #Allocate optinodes onto workers.  Split it up evenly
    allocations = []
    node_indices = []
    j = 1
    while  j <= n_nodes
        if j + nodes_per_worker > n_nodes
            push!(allocations,nodes[j:end])
            push!(node_indices,[getindex(graph,node) for node in nodes[j:end]])
        else
            push!(allocations,nodes[j:j+nodes_per_worker - 1])
            push!(node_indices,[getindex(graph,node) for node in nodes[j:j+nodes_per_worker - 1]])
        end
        j += nodes_per_worker
    end

    println("Distributing optigraph among workers: $to_workers")
    channel = RemoteChannel(1)
    remote_references = []
    #Fill channel with sets of nodes to send
    #TODO: Make this runs in parallel with an @async?
    @sync begin
        for (i,worker) in enumerate(to_workers)
            @spawnat(1, put!(channel, allocations[i]))
            ref1 = @spawnat worker begin
                Core.eval(Main, Expr(:(=), :nodes, take!(channel)))
            end
            wait(ref1)
            ref2 = @spawnat worker Core.eval(Main, Expr(:(=), remote_name, PipsNLP._create_worker_optigraph(getfield(Main,:nodes),
            node_indices[i],
            user_pips_data,
            n_nodes)))
            push!(remote_references,ref2)
        end
    end
    #graph.ext[:allocations] = allocations
    return remote_references
end

function fill_solution!(graph::OptiGraph,remote_name::Symbol,worker::Int64)
    for i = 1:num_all_nodes(graph)
        x_sol = @fetchfrom worker getnode(getfield(Main,remote_name),i).ext[:pips_data].x_sol
        status = @fetchfrom worker termination_status(getnode(getfield(Main,remote_name),i))

        node = getnode(graph,i)
        Plasmo.set_node_primals(node,all_variables(node),x_sol)
        Plasmo.set_node_status(node,status)
    end
    return nothing
end

function _create_worker_optigraph(optinodes::Vector{OptiNode},
    node_indices::Vector{Int64},
    user_pips_data,
    n_nodes)

    graph = OptiGraph()
    graph.node_idx_map = Dict{OptiNode,Int64}()
    graph.ext[:user_pips_data] = user_pips_data

    #Add nodes to worker's graph.  Each worker should have the same number of nodes, but some will be empty.
    for i = 1:n_nodes
        node = add_node!(graph)
        node.ext[:pips_data] = PIPSNLPData()  #all nodes need this
    end
    for (i,index) in enumerate(node_indices)
        graph.optinodes[index] = optinodes[i]
    end

    return graph
end


using SeisMain
using Test

#From USGS, National Petroleum Reserve - Alaska Data Archive
#https://energy.usgs.gov/GeochemistryGeophysics/SeismicDataProcessingInterpretation/NPRASeismicDataArchive.aspx#3862174-data-

download("http://certmapper.cr.usgs.gov/data/NPRA/seismic/1976/125_76/PROCESSED/125_76_PT2_PR.SGY","125_76.SGY")

SegyToSeis("125_76.SGY","125_76")
d,h,ext = SeisRead("125_76")
imx = SeisMain.ExtractHeader(h,"imx")
println("size(d)=",size(d))
println("imx=",imx)
SeisWrite("125_76_copy",d,h,ext)


module VarintTranslator # LoRDeckCodes

using ..LoRDeckCodes: ArgumentException

const AllButMSB = 0x7f
const JustMSB   = 0x80

function pop_varint(stream::IOBuffer)::Int
    result::UInt = 0
    currentShift = 0
    bytesPopped = 0
    while !eof(stream)
        (byte,) = read(stream, 1)
        current = byte & AllButMSB
        result |= current << currentShift
        (byte & JustMSB) != JustMSB && return Int(result)
        currentShift += 7
    end
    throw(ArgumentException("Byte array did not contain valid varints."))
end

function get_varint(value::Int)::Vector{UInt8}
    buff = zeros(UInt8, 10)
    currentIndex = 1
    iszero(value) && return UInt8[0]
    while !iszero(value)
        byteVal = value & AllButMSB
        value >>= 7
        if !iszero(value)
            byteVal |= JustMSB
        end
        buff[currentIndex] = byteVal
        currentIndex += 1
    end
    buff[1:currentIndex-1]
end

end # module LoRDeckCodes.VarintTranslator


function test_fitness(fitness::Function, cont_f::Function;
                      maxfit::Float64=1.0, minfit::Float64=0.0)
    fit = fitness(cont_f; seed=0)
    @test fit <= maxfit
    @test fit >= minfit
    fit2, t, bytes, gctime, memallocs = @timed fitness(cont_f; seed=1234)
    println("Fitness of $fit2 in $t s, $bytes bytes")
    @test fit2 <= maxfit
    @test fit2 >= minfit
end

rand_cont(x::Array{Float64}) = rand(1:4)

@testset "Coverage fitness" begin
    test_fitness(RoboGrid.coverage_fitness, rand_cont)
end

@testset "Water search fitness" begin
    test_fitness(RoboGrid.water_search_fitness, rand_cont)
end

@testset "Water memorize fitness" begin
    test_fitness(RoboGrid.water_memorize_fitness, rand_cont; minfit=-7.0)
end

@testset "Cross search fitness" begin
    test_fitness(RoboGrid.cross_search_fitness, rand_cont)
end

@testset "Cross memorize fitness" begin
    test_fitness(RoboGrid.cross_memorize_fitness, rand_cont; minfit=-7.0)
end

@testset "Cross strategy fitness" begin
    test_fitness(RoboGrid.cross_strategy_fitness, rand_cont; minfit=-7.0)
end


"""
    reverse(A; dims) => AbstractDimArray
    reverse(dim::Dimension) => Dimension

Reverse the array order, and update the dim to match.
"""
Base.reverse(A::AbstractDimArray; dims=1) =
    _reverse(IndexOrder, _reverse(ArrayOrder, A, dims), dims)
Base.reverse(s::AbstractDimStack; dims=1) = map(A -> reverse(A; dims=dims), s)
Base.reverse(ot::Type{<:SubOrder}, x; dims) = _reverse(ot, x, dims)
_reverse(ot::Type{<:SubOrder}, x, lookup) = set(x, reverse(ot, dims(x, lookup)))
_reverse(ot::Type{<:ArrayOrder}, x, lookup) = begin
    newdims = reverse(ArrayOrder, dims(x, lookup))
    newdata = _reversedata(x, dimnum(x, lookup))
    setdims(rebuild(x, newdata), newdims)
end
# Dimension
Base.reverse(ot::Type{<:SubOrder}, dims::DimTuple) = map(d -> reverse(ot, d), dims)
Base.reverse(ot::Type{<:SubOrder}, dim::Dimension) = reverse(ot, mode(dim), dim)
Base.reverse(ot::Type{<:SubOrder}, mode::IndexMode, dim::Dimension) =
    rebuild(dim; mode=reverse(ot, mode))
# Reverse the index
Base.reverse(ot::Type{<:IndexOrder}, mode::IndexMode, dim::Dimension) =
    rebuild(dim, reverse(index(dim)), reverse(ot, mode))
Base.reverse(ot::Type{<:IndexOrder}, mode::IndexMode, dim::Dimension{<:Val{I}}) where I =
    rebuild(dim, Val(reverse(I)), reverse(ot, mode))
Base.reverse(ot::Type{<:IndexOrder}, mode::NoIndex, dim::Dimension) = dim
Base.reverse(dim::Dimension) = reverse(IndexOrder, dim)
# Mode
Base.reverse(ot::Type{<:SubOrder}, mode::NoIndex) = mode
Base.reverse(ot::Type{<:SubOrder}, mode::IndexMode) =
    rebuild(mode; order=reverse(ot, order(mode)))
Base.reverse(ot::Type{<:SubOrder}, mode::AbstractSampled) =
    rebuild(mode; order=reverse(ot, order(mode)), span=reverse(ot, span(mode)))
# Order
Base.reverse(::Type{<:IndexOrder}, o::Ordered) =
    Ordered(reverse(indexorder(o)), arrayorder(o), reverse(relation(o)))
Base.reverse(::Type{<:Union{ArrayOrder,Relation}}, o::Ordered) =
    Ordered(indexorder(o), reverse(arrayorder(o)), reverse(relation(o)))
Base.reverse(::Type{<:SubOrder}, o::Unordered) = Unordered(reverse(relation(o)))
# SubOrder
Base.reverse(::ReverseIndex) = ForwardIndex()
Base.reverse(::ForwardIndex) = ReverseIndex()
Base.reverse(::ReverseArray) = ForwardArray()
Base.reverse(::ForwardArray) = ReverseArray()
Base.reverse(::ReverseRelation) = ForwardRelation()
Base.reverse(::ForwardRelation) = ReverseRelation()
# Span
Base.reverse(::Type{<:SubOrder}, span::Span) = span
Base.reverse(::Type{<:IndexOrder}, span::Irregular) = span
Base.reverse(::Type{<:IndexOrder}, span::Span) = reverse(span)
Base.reverse(span::Regular) = Regular(-step(span))
Base.reverse(span::Explicit) = Explicit(reverse(val(span), dims=2))

_reversedata(A::AbstractDimArray, dimnum) = reverse(parent(A); dims=dimnum)
_reversedata(s::AbstractDimStack, dimnum) =
    map(a -> reverse(parent(a); dims=dimnum), data(s))

"""
    fliparray(Order, A, dims) => AbstractDimArray
    fliparray(dim::Dimension) => Dimension

`Flip` the array order without changing any data.
"""
function flip end

flip(ot::Type{<:SubOrder}, x; dims) = flip(ot, x, dims)
flip(ot::Type{<:SubOrder}, x, lookupdims) = set(x, flip(ot, dims(x, lookupdims)))
flip(ot::Type{<:SubOrder}, dims::DimTuple) = map(d -> flip(ot, d), dims)
flip(ot::Type{<:SubOrder}, dim::Dimension) = _set(dim, flip(ot, mode(dim)))
flip(ot::Type{<:SubOrder}, mode::IndexMode) = _set(mode, flip(ot, order(mode)))
flip(ot::Type{<:SubOrder}, mode::NoIndex) = mode
flip(ot::Type{<:SubOrder}, o::Order) = _set(o, reverse(ot, o))

"""
    reorder(::Order, A::AbstractDimArray) => AbstractDimArray
    reorder(::Order, A::Dimension) => AbstractDimArray
    reorder(A::AbstractDimArray, order::Union{Order,Dimension{<:Order},Tuple}) => AbstractDimArray
    reorder(A::AbstractDimArray, order::Pair{<:Dimension,<:SubOrder}...) => AbstractDimArray

Reorder every dims index/array/relation to `order`, or reorder index for
the the given dimension(s) to the `Order` they wrap.

Reordering `Relation` will reverse the array, not the dimension index.

`order` can be an [`Order`](@ref), a single [`Dimension`](@ref)
or a `Tuple` of `Dimension`.

If no axis reversal is required the same objects will be returned, without allocation.
"""
function reorder end

reorder(x, args...; kw...) = reorder(x, (args..., _kwdims(kw)...))
reorder(x, nt::NamedTuple) = reorder(x, _kwdims(nt))
reorder(x, p::Pair, ps::Vararg{<:Pair}) = reorder(x, (p, ps...))
reorder(x, ps::Tuple{Vararg{<:Pair}}) = reorder(x, _pairdims(ps...))
# Reorder specific dims.
reorder(x, dimwrappers::Tuple) = _reorder(x, dimwrappers)
# Reorder all dims.
reorder(x, ot::Union{SubOrder,Type{<:SubOrder}}) =
    _reorder(x, map(d -> basetypeof(d)(ot), dims(x)))
reorder(dim::Dimension, ot::Union{SubOrder,Type{<:SubOrder}}) =
    ot == basetypeof(order(ot, dim)) ? dim : reverse(ot, dim)

# Recursive reordering. x may be reversed here
_reorder(x, orderdims::DimTuple) = _reorder(reorder(x, orderdims[1]), tail(orderdims))
_reorder(x, orderdims::Tuple{}) = x

reorder(x, orderdim::Dimension) = _reorder(val(orderdim), x, dims(x, orderdim))

_reorder(neworder::SubOrder, x, orderdim::DimOrDimType) = _reorder(basetypeof(neworder), x, orderdim)
# Reverse the dimension index
_reorder(ot::Type{<:IndexOrder}, x, orderdim::DimOrDimType) =
    ot == basetypeof(order(ot, orderdim)) ? x : set(x, reverse(ot, orderdim))
# If either ArrayOrder or Relation are reversed, we reverse the array
_reorder(ot::Type{<:Union{ArrayOrder,Relation}}, x, dim::DimOrDimType) =
    ot == basetypeof(order(ot, dim)) ? x : reverse(ArrayOrder, x; dims=dim)

"""
    modify(f, A::AbstractDimArray) => AbstractDimArray
    modify(f, s::AbstractDimStack) => AbstractDimStack
    modify(f, dim::Dimension) => Dimension
    modify(f, x, lookupdim::Dimension) => typeof(x)

Modify the parent data, rebuilding the object wrapper without
change. `f` must return a `AbstractArray` of the same size as the original.

This method is mostly useful as a way of swapping the parent array type of
an object.

## Example

If we have a previously-defined `DimArray`, we can copy it to an Nvidia GPU with:

```julia
A = DimArray(rand(100, 100), (X, Y))
modify(CuArray, A)
```

This also works for all the data layers in a `DimStack`.
"""
function modify end
modify(f, s::AbstractDimStack) = map(a -> modify(f, a), s)
function modify(f, A::AbstractDimArray)
    newdata = f(parent(A))
    size(newdata) == size(A) || error("$f returns an array with a different size")
    rebuild(A, newdata)
end
modify(f, x, dim::DimOrDimType) = set(x, modify(f, dims(x, dim)))
function modify(f, dim::Dimension)
    newindex = f(index(dim))
    size(newindex) == size(dim) || error("$f returns a vector with a different size")
    rebuild(dim, newindex)
end
function modify(f, dim::Dimension{<:Val{Index}}) where Index
    newindex = f(Index)
    length(newindex) == length(dim) || error("$f returns a Tuple with a different size")
    rebuild(dim, Val(newindex))
end

"""
    dimwise(f, A::AbstractDimArray{T,N}, B::AbstractDimArray{T2,M}) => AbstractDimArray{T3,N}

Dimension-wise application of function `f` to `A` and `B`.

## Arguments

- `a`: `AbstractDimArray` to broacast from, along dimensions not in `b`.
- `b`: `AbstractDimArray` to broadcast from all dimensions. Dimensions must be a subset of a.

This is like broadcasting over every slice of `A` if it is
sliced by the dimensions of `B`.
"""
function dimwise(f, A::AbstractDimArray, B::AbstractDimArray)
    dimwise!(f, similar(A, promote_type(eltype(A), eltype(B))), A, B)
end

"""
    dimwise!(f, dest::AbstractDimArray{T1,N}, A::AbstractDimArray{T2,N}, B::AbstractDimArray) => dest

Dimension-wise application of function `f`.

## Arguments

- `dest`: `AbstractDimArray` to update
- `A`: `AbstractDimArray` to broacast from, along dimensions not in `b`.
- `B`: `AbstractDimArray` to broadcast from all dimensions. Dimensions must be a subset of a.

This is like broadcasting over every slice of `A` if it is
sliced by the dimensions of `B`, and storing the value in `dest`.
"""
function dimwise!(
    f, dest::AbstractDimArray{T,N}, a::AbstractDimArray{TA,N}, b::AbstractDimArray{TB,NB}
) where {T,TA,TB,N,NB}
    N >= NB || error("B-array cannot have more dimensions than A array")
    comparedims(dest, a)
    common = commondims(a, dims(b))
    od = otherdims(a, common)
    # Lazily permute B dims to match the order in A, if required
    if !dimsmatch(common, dims(b))
        b = PermutedDimsArray(b, common)
    end
    # Broadcast over b for each combination of dimensional indices D
    if length(od) == 0
        dest .= f.(a, b)
    else
        map(DimIndices(od)) do D
            dest[D...] .= f.(a[D...], b)
        end
    end
    return dest
end

"""
    basetypeof(x) => Type

Get the "base" type of an object - the minimum required to
define the object without it's fields. By default this is the full
`UnionAll` for the type. But custom `basetypeof` methods can be
defined for types with free type parameters.

In DimensionalData this is primariliy used for comparing `Dimension`s,
where `Dim{:x}` is different from `Dim{:y}`.
"""
@inline basetypeof(x::T) where T = basetypeof(T)
@generated function basetypeof(::Type{T}) where T
    if T isa Union
        T
    else
        getfield(parentmodule(T), nameof(T))
    end
end

# Left pipe operator for cleaning up brackets
@deprecate f <| x f(x)

# Unwrap Val
unwrap(::Val{X}) where X = X
unwrap(::Type{Val{X}}) where X = X
unwrap(x) = x

# Get a tuple of unique keys for DimArrays. If they have the same
# name we call them layerI.
function uniquekeys(das::Tuple{AbstractDimArray,Vararg{<:AbstractDimArray}})
    uniquekeys(Symbol.(map(name, das)))
end
function uniquekeys(keys::Tuple{Symbol,Vararg{<:Symbol}})
    ids = ntuple(x -> x, length(keys))
    map(keys, ids) do k, id
        count(k1 -> k == k1, keys) > 1 ? Symbol(:layer, id) : k
    end
end

# Produce a 2 * length(dim) matrix of interval bounds from a dim
function dim2boundsmatrix(dim::Dimension)
    samp = sampling(dim)
    samp isa Intervals || error("Cannot create a bounds matrix for $(nameof(typeof(samp)))")
    _dim2boundsmatrix(locus(dim), span(dim), dim)
end

_dim2boundsmatrix(::Locus, span::Explicit, dim) = val(span)
_dim2boundsmatrix(::Locus, span::Regular, dim) =
    vcat(permutedims(_shiftindexlocus(Start(), dim)), permutedims(_shiftindexlocus(End(), dim)))
@noinline _dim2boundsmatrix(::Center, span::Regular{Dates.TimeType}, dim) =
    error("Cannot convert a Center TimeType index to Explicit automatically: use a bounds matrix e.g. Explicit(bnds)")
@noinline _dim2boundsmatrix(::Start, span::Irregular, dim) =
    error("Cannot convert Irregular to Explicit automatically: use a bounds matrix e.g. Explicit(bnds)")

#=
Shift the index from the current locus to the new locus. We only actually
shift Regular Intervals, and do this my multiplying the offset of
-1, -0.5, 0, 0.5 or 1 by the absolute value of the span.
=#
function shiftlocus(locus::Locus, dim::Dimension)
    samp = sampling(dim)
    samp isa Intervals || error("Cannot shift locus of $(nameof(typeof(samp)))")
    rebuild(dim; val=_shiftindexlocus(locus, dim), mode=DD._set(mode(dim), locus))
end

_shiftindexlocus(locus::Locus, dim::Dimension) = _shiftindexlocus(locus::Locus, mode(dim), dim)
_shiftindexlocus(locus::Locus, mode::IndexMode, dim::Dimension) = index(dim)
_shiftindexlocus(locus::Locus, mode::AbstractSampled, dim::Dimension) =
    _shiftindexlocus(locus, span(mode), sampling(mode), dim)
_shiftindexlocus(locus::Locus, span::Span, sampling::Sampling, dim::Dimension) = index(dim)
_shiftindexlocus(destlocus::Locus, span::Regular, sampling::Intervals, dim::Dimension) =
    index(dim) .+ (abs(step(span)) * _offset(locus(sampling), destlocus))
_shiftindexlocus(::Start, span::Explicit, sampling::Intervals, dim::Dimension) = val(span)[1, :]
_shiftindexlocus(::End, span::Explicit, sampling::Intervals, dim::Dimension) = val(span)[2, :]
_shiftindexlocus(destlocus::Center, span::Explicit, sampling::Intervals, dim::Dimension) =
    _shiftindexlocus(destlocus, locus(dim), span, sampling, dim)
_shiftindexlocus(::Center, ::Center, span::Explicit, sampling::Intervals, dim::Dimension) = index(dim)
_shiftindexlocus(::Center, ::Locus, span::Explicit, sampling::Intervals, dim::Dimension) =
    view(val(span), 2, :)  .- view(val(span), 1, :)

_offset(::Start, ::Center) = 0.5
_offset(::Start, ::End) = 1
_offset(::Center, ::Start) = -0.5
_offset(::Center, ::End) = 0.5
_offset(::End, ::Start) = -1
_offset(::End, ::Center) = -0.5
_offset(::T, ::T) where T<:Locus = 0

maybeshiftlocus(locus::Locus, dim::Dimension) = _maybeshiftlocus(locus, sampling(dim), dim)

_maybeshiftlocus(locus::Locus, sampling::Intervals, dim::Dimension) = shiftlocus(locus, dim)
_maybeshiftlocus(locus::Locus, sampling::Sampling, dim::Dimension) = dim

_astuple(t::Tuple) = t
_astuple(x) = (x,)


#-------------------------------------------------------------------------#
#                                                                         #
#                               Simulating_nets                           #
#                                                                         #
#-------------------------------------------------------------------------#

# V1 18 may 2020
#-------------------------------------------------------------------------#
# OBSERVATIONS:

# 1. Bla bla bla
# 2. Bla bla bla
# 3. Bla bla bla
#
#-------------------------------------------------------------------------#
# using PhyloNetworks

## FUNCTIONS

using Random

"""
    testing if a one_taxa string is valid
        This is a function in progress
"""
function is_valid_1taxa(one_taxa::String)
    test = true
    if(tryparse(Float64, SubString(one_taxa, 1,1)) !== nothing)
        error("taxa should not start with a number")
    end
    return test
end

"""
    Simulation of a phylogenetic tree topology
        We start with a two taxa tree and start to paste subtrees by choosing
        randomly taxa to paste.
        Currently the branch lengths are uniformly distributed (0,1)
        Returs a the topology in Newick format
"""
function random_tree_topology(taxa)
    shuffle!(taxa)
    # Starting the tree
    x = Any["(",taxa[1],",",taxa[2],")",";"]

    for i in 3:length(taxa)
        # taxa already in the tree and taxa not in the tree
        used_taxa = taxa[1:(i-1)]
        left_taxa = taxa[i:end]
        # choosing a taxa in the tree to add new branches
        taxa_to_change = rand(used_taxa)
        # possition to add branches
        position_of_change = findfirst(x.==taxa_to_change)
        # updating tree
        insert!(x,position_of_change+1,")")
        insert!(x,position_of_change,",")
        insert!(x,position_of_change,taxa[i])
        insert!(x,position_of_change,"(")
    end
    return x
end

"""
    Adds random branch lengths to a phylogenetic tree topology
        Currently the branch lengths are uniformly distributed (0,1)
        Returs a the topology in Newick format
"""
function add_random_bl(x,taxa)
    # We get the index of the taxa and right parenthesis (to add the branch lengths)
    idx_right_parenthesis = findall(x.==")")
    idx_taxa = zeros(Int64,length(taxa))
    for i in 1:length(taxa)
        idx_taxa[i] = findfirst(x.==taxa[i])
    end
    idx_taxa_and_right_p = vcat(idx_taxa, idx_right_parenthesis)
    # To paste the new branch length we need to sort in decreasing order
    sort!(idx_taxa_and_right_p, rev=true)
    # Random uniform branch length
    bl = rand(length(idx_taxa_and_right_p))
    # Adding the branch length
    for i in 1:length(idx_taxa_and_right_p)
        insert!(x,idx_taxa_and_right_p[i]+1,bl[i])
        insert!(x,idx_taxa_and_right_p[i]+1,":")
    end
    return x
end

"""
    Given a topology and branch lengths, this function change the branch lengths
    in a way that the tree is now (insert word to say that all the paths from a
    leaf to the root are of the same length)
        Only keeps the topology and with a stick break methodology gives the
        new branch lengths.
        Returs a the topology in Newick format
"""
function balance_a_tree(x,total_bl)
    # Case of taxa-colon-branch length-semicolon; eg A:0.32;
    if length(x)==4
        x[3] = total_bl[1]
        return x
    end

    # Case of two taxa
    if length(x)==12
      partition = rand(1)
      # bl = total_bl.*partition
      x[[4,8]] .= total_bl.*partition
      aux = total_bl.*(Float64[1]-partition)
      x[11] = aux[1]
      return x
    end

    # We need to find the place of the comma that divides the tree
    aux = (x.==")") - (x.=="(")
    pop!(aux)
    popfirst!(aux)
    central_comma_place = findfirst(cumsum(aux).==0) +4

    # Stuff of the left tree
    left_tree = vcat(x[2:(central_comma_place-1)],";")
    partition = rand(1)
    left_tree = balance_a_tree(left_tree,total_bl.*partition)

    # Stuff of the right tree
    right_tree = vcat(x[(central_comma_place+1):(length(x)-4)],";")
    right_tree[length(right_tree)-1] = left_tree[length(left_tree)-1]
    right_tree = balance_a_tree(right_tree,total_bl.*partition)

    # Pasting the left and right trees
    pop!(left_tree)
    pop!(right_tree)
    x = vcat("(",left_tree,",",right_tree,")",":",total_bl.*(Float64[1]-partition),";")

    return x
end

"""
    Addition of a random hybrid node to a phylogenetic tree
        We start a tree "x" with taxa "taxa" and select two taxa to add the hybrid
        node.
        Currently the branch length of the new hybrid and its gamma parameter are
        uniformly distributed (0,1). The first node that is selected we currently
        TRANSFORM it to a hybrid node. For the second node we only paste the
        other hybrid node.
        For instance, if we have the tree "((a,b),(c,d));" and select "b" and "c",
        we get "((a,#H1),((c)#H1,d));", without considering the branch lengths and
        the gamma parameter.
        Currently, sometimes the hybrid is in the same branch.
        Returs the tree with the added hybrid, in Newick format.

        Future improvements:
            - Extra variable with the name of the hybrid, now is H1
            - Correcting that in some simulations the hybrid nodes are in the same
            path an this probuces a weird hybrid that follows the same path as
            the original tree
"""
function add_random_hybrid(x,taxa)
   # sampling two taxa to insert the hybrid
   taxa_to_change1,taxa_to_change2 = taxa[randperm(length(taxa))[1:2]]
   # We get the branch length and the gamma parameter for the new hybrid
   bl_hybrid = rand(1)
   gamma = rand(1)
   # Finding the index of the taxa to change
   idx_node1 = findfirst(x.==taxa_to_change1)
   idx_node2 = findfirst(x.==taxa_to_change2)
   # We allways prefere the index of node 1 to be smaller
   if idx_node1 > idx_node2
      idx_node1,idx_node2 = idx_node2,idx_node1
   end
   # Now we can make the changes on the array x
   x[idx_node2] = "#H1"

   insert!(x,idx_node1+3,gamma[1])
   insert!(x,idx_node1+3,":")
   insert!(x,idx_node1+3,":")
   insert!(x,idx_node1+3,bl_hybrid[1])
   insert!(x,idx_node1+3,":")
   insert!(x,idx_node1+3,"#H1")
   insert!(x,idx_node1+3,")")
   insert!(x,idx_node1,"(")
   return x
end

"""
    Simulating a phylogenetic tree/network
        Returs a the topology in Newick format

        Future improvements:
            - To add the hybrid a node is eliminated, this could be solved by
            since the begining ask the name of the hybrid and past it as part of
            the tree, but remember that this change also affects the function
            add_random_hybrid; in particular, since this function chose randomly
            which node of the tree is transformed to hybrid.
"""
function Sim_tree(taxa,total_bl,is_network)
    # Getting the topology
    x = random_tree_topology(taxa)
    # Adding some branch lengths
    x = add_random_bl(x,taxa)
    # Fixing the branch lengths to be "equal"
    x = balance_a_tree(x,total_bl)
    # Adding a random hybrid
    if is_network
        x = add_random_hybrid(x,taxa)
    end
    return x
end

## Funcitons of njWillems2014

function Julia_willems2014(x,name_file_dist_txt,
                           name_file_Input_Matrix_Distances,
                           name_file_Parameters_Distances,
                           name_file_new_Network_Distances,
                           folder_output_distances_willems)
   # Folder to save the output of willems
   output_willems = folder_output_distances_willems*"/"*name_file_new_Network_Distances
   # Getting the netwokr from the Newick format
   net = readTopology(join(x))
   # Distance matrix
   Sigma = sharedPathMatrix(net)
   Vy = Sigma[:Tips]

   # From Cecil
   M = pairwiseTaxonDistanceMatrix(net, keepInternal=false)
   tipLabels(net)

   # Generating file to run njWillems2014
   ind = parse.(Int, tipLabels(net))
   # ind =  tipLabels(net)

   Mext = hcat(ind,M)
   n = length(tipLabels(net))

   df = DataFrame(Mext)
   df[!,:x1] = convert.(Int,df[!,:x1])

   CSV.write(name_file_dist_txt, df, writeheader=false, delim=' ')

   # the first row with the number of tips
   f=readlines(name_file_dist_txt)
   g = open(name_file_Input_Matrix_Distances,"w")
   write(g,"$n \n")

   for l in f
      write(g,l)
      write(g,"\n")
   end
   close(g)

   # This "run" generates the file "Network_Distances.txt"
   text_to_run = `./njwillems2014 $name_file_Input_Matrix_Distances $name_file_Parameters_Distances`
   run(text_to_run)

   # We copy the file "Network_Distances.txt", that generates njwillems2014
   cp("Network_Distances.txt", output_willems, force=true)

   net_dist = readdlm("Network_Distances.txt")
   # rm(name_file_new_Network_Distances)

   # txt_file = open("Network_Distances.txt","r")
   # readline(txt_file)
   # close(txt_file)

   # print(join(net_dist[1:10]))
   # println(net_dist)

   idx_Main_text = findfirst(net_dist .== "Main")
   println(net_dist[CartesianIndex(1, 1):(idx_Main_text)])

   # println((net_dist[1:20]))
   return net_dist
end
## End of functions


mutable struct SROCK1ConstantCache{zType,tType} <: StochasticDiffEqConstantCache
  ms::SVector{10,Int}
  mη::SVector{10,tType}
  zprev::zType
  mdeg::Int
  optimal_η::tType
end
@cache struct SROCK1Cache{uType,rateType,noiseRateType,C<:SROCK1ConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  gₘ₋₁::noiseRateType
  gₘ₋₂::noiseRateType
  tmp::uType
  k::rateType
  fsalfirst::rateType
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::SROCK1,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  SROCK1ConstantCache{typeof(t)}(u)
end

function alg_cache(alg::SROCK1,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  k = zero(rate_prototype)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  gₘ₋₁ = zero(noise_rate_prototype)
  gₘ₋₂ = zero(noise_rate_prototype)
  tmp  = uᵢ₋₂             # these 3 variables are dummied to use same memory
  fsalfirst = k
  atmp = zero(rate_prototype)
  constantcache = SROCK1ConstantCache{typeof(t)}(u)
  SROCK1Cache(u,uprev,uᵢ₋₁,uᵢ₋₂,gₘ₋₁,gₘ₋₂,tmp,k,fsalfirst,atmp,constantcache)
end

mutable struct SROCK2ConstantCache{zType,T} <: StochasticDiffEqConstantCache
  ms::SVector{46,Int}
  recf::Vector{T}
  mσ::SVector{46,T}
  mτ::SVector{46,T}
  recf2::Vector{T}
  mα::SVector{46,T}
  zprev::zType
  mdeg::Int
  deg_index::Int
  start::Int
end

@cache struct SROCK2Cache{uType,rateType,noiseRateType,T,C<:SROCK2ConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ::uType
  uₓ::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  Gₛ::noiseRateType
  Gₛ₁::noiseRateType
  WikRange::T
  vec_χ::T
  tmp::uType
  k::rateType
  fsalfirst::rateType
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::SROCK2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  SROCK2ConstantCache{typeof(t)}(u)
end

function alg_cache(alg::SROCK2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  k = zero(rate_prototype)
  uᵢ = zero(u)
  uₓ = zero(u)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  Gₛ = zero(noise_rate_prototype)
  Gₛ₁ = zero(noise_rate_prototype)
  WikRange = false .* vec(ΔW)
  vec_χ = false .* vec(ΔW)
  tmp  = uᵢ₋₂            # these 2 variables are dummied to use same memory
  fsalfirst = k
  atmp = zero(rate_prototype)
  constantcache = SROCK2ConstantCache{typeof(t)}(u)
  SROCK2Cache(u,uprev,uᵢ,uₓ,uᵢ₋₁,uᵢ₋₂,Gₛ,Gₛ₁,WikRange,vec_χ,tmp,k,fsalfirst,atmp,constantcache)
end

mutable struct SROCKEMConstantCache{zType,tType} <: StochasticDiffEqConstantCache
  ms::SVector{10,Int}
  mη::SVector{10,tType}
  zprev::zType
  mdeg::Int
  optimal_η::tType
end

@cache struct SROCKEMCache{uType,rateType,noiseRateType,T,C<:SROCKEMConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  Gₛ::noiseRateType
  Gₛ₁::noiseRateType
  tmp::uType
  k::rateType
  fsalfirst::rateType
  WikRange::T
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::SROCKEM,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  SROCKEMConstantCache{typeof(t)}(u)
end

function alg_cache(alg::SROCKEM,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  k = zero(rate_prototype)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  Gₛ = zero(noise_rate_prototype)
  if (!alg.strong_order_1 || is_diagonal_noise(prob) || typeof(ΔW) <: Number || length(ΔW) == 1)
    Gₛ₁ = Gₛ
  else
    Gₛ₁ = zero(noise_rate_prototype)
  end
  WikRange = false .* vec(ΔW)
  tmp  = zero(u)             # these 3 variables are dummied to use same memory
  fsalfirst = k
  atmp = zero(rate_prototype)
  constantcache = SROCKEMConstantCache{typeof(t)}(u)
  SROCKEMCache(u,uprev,uᵢ₋₁,uᵢ₋₂,Gₛ,Gₛ₁,tmp,k,fsalfirst,WikRange,atmp,constantcache)
end

mutable struct SKSROCKConstantCache{zType,T} <: StochasticDiffEqConstantCache
  mc::Vector{T}
  mα::Vector{T}
  zprev::zType
end

@cache struct SKSROCKCache{uType,rateType,noise_rate_prototype,T,C<:SKSROCKConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  Gₛ::noise_rate_prototype
  tmp::uType
  k::rateType
  fsalfirst::rateType
  WikRange::T
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::SKSROCK,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  SKSROCKConstantCache{typeof(t)}(u)
end

function alg_cache(alg::SKSROCK,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  k = zero(rate_prototype)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  Gₛ = zero(noise_rate_prototype)
  tmp  = uᵢ₋₂             # Dummmy variables
  fsalfirst = k
  WikRange = false .* vec(ΔW)
  atmp = zero(rate_prototype)
  constantcache = SKSROCKConstantCache{typeof(t)}(u)
  SKSROCKCache(u,uprev,uᵢ₋₁,uᵢ₋₂,Gₛ,tmp,k,fsalfirst,WikRange,atmp,constantcache)
end

mutable struct TangXiaoSROCK2ConstantCache{zType,T} <: StochasticDiffEqConstantCache
  ms::SVector{46,Int}
  recf::Vector{T}
  mσ::SVector{46,T}
  mτ::SVector{46,T}
  recf2::Vector{T}
  mα::SVector{5, T}
  mn̂::SVector{5, Int}
  c1::SVector{13, T}
  c2::SVector{13, T}
  zprev::zType
  mdeg::Int
  deg_index::Int
  start::Int
  start_mcs::Int
end

@cache struct TangXiaoSROCK2Cache{uType,rateType,noiseRateType,C<:TangXiaoSROCK2ConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ::uType
  uₓ::uType
  Û₁::uType
  Û₂::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  Gₛ::noiseRateType
  Gₛ₁::noiseRateType
  tmp::uType
  k::rateType
  fsalfirst::rateType
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::TangXiaoSROCK2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  TangXiaoSROCK2ConstantCache{typeof(t)}(u)
end

function alg_cache(alg::TangXiaoSROCK2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  k = zero(rate_prototype)
  uᵢ = zero(u)
  uₓ = zero(u)
  Û₁ = zero(u)
  Û₂ = zero(u)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  Gₛ = zero(noise_rate_prototype)
  if typeof(ΔW) <: Number || length(ΔW) == 1 || is_diagonal_noise(prob)
    Gₛ₁ = Gₛ
  else
    Gₛ₁ = zero(noise_rate_prototype)
  end
  tmp  = uᵢ₋₂            # these 2 variables are dummied to use same memory
  fsalfirst = k
  atmp = zero(rate_prototype)
  constantcache = TangXiaoSROCK2ConstantCache{typeof(t)}(u)
  TangXiaoSROCK2Cache(u,uprev,uᵢ,uₓ,Û₁,Û₂,uᵢ₋₁,uᵢ₋₂,Gₛ,Gₛ₁,tmp,k,fsalfirst,atmp,constantcache)
end

mutable struct KomBurSROCK2ConstantCache{zType,T} <: StochasticDiffEqConstantCache
  ms::SVector{46,Int}
  recf::Vector{T}
  mσ::SVector{46,T}
  mτ::SVector{46,T}
  mδ::Vector{T}
  zprev::zType
  mdeg::Int
  deg_index::Int
  start::Int
end

@cache struct KomBurSROCK2Cache{uType,rateType,noiseRateType,T,C<:KomBurSROCK2ConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  utmp::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  k::rateType
  yₛ₋₁::rateType
  yₛ₋₂::rateType
  yₛ₋₃::rateType
  SXₛ₋₁::uType
  SXₛ₋₂::uType
  SXₛ₋₃::uType
  Gₛ::noiseRateType
  Xₛ₋₁::noiseRateType
  Xₛ₋₂::noiseRateType
  Xₛ₋₃::noiseRateType
  vec_χ::T
  tmp::uType
  fsalfirst::rateType
  WikRange::T
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::KomBurSROCK2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  KomBurSROCK2ConstantCache{typeof(t)}(u)
end

function alg_cache(alg::KomBurSROCK2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  utmp = zero(u)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  k = zero(rate_prototype)
  yₛ₋₁ = k
  yₛ₋₂ = zero(rate_prototype)
  yₛ₋₃ = zero(rate_prototype)
  Xₛ₋₁ = zero(noise_rate_prototype)
  Xₛ₋₂ = zero(noise_rate_prototype)
  Xₛ₋₃ = zero(noise_rate_prototype)
  vec_χ = false .* vec(ΔW)
  WikRange = false .* vec(ΔW)
  if typeof(ΔW) <: Number || length(ΔW) == 1 || is_diagonal_noise(prob)
    Gₛ = Xₛ₋₁
    SXₛ₋₁ = utmp
    SXₛ₋₂ = utmp
    SXₛ₋₃ = utmp
  else
    Gₛ = zero(noise_rate_prototype)
    SXₛ₋₁ = zero(uprev)
    SXₛ₋₂ = zero(uprev)
    SXₛ₋₃ = zero(uprev)
  end
  tmp  = uᵢ₋₂            # these 3 variables are dummied to use same memory
  fsalfirst = k
  atmp = yₛ₋₂
  constantcache = KomBurSROCK2ConstantCache{typeof(t)}(u)
  KomBurSROCK2Cache(u,uprev,utmp,uᵢ₋₁,uᵢ₋₂,k,yₛ₋₁,yₛ₋₂,
                    yₛ₋₃,SXₛ₋₁,SXₛ₋₂,SXₛ₋₃,Gₛ,Xₛ₋₁,Xₛ₋₂,Xₛ₋₃,vec_χ,
                    tmp,fsalfirst,WikRange,atmp,constantcache)
end

mutable struct SROCKC2ConstantCache{zType,T} <: StochasticDiffEqConstantCache
  ms::SVector{46,Int}
  recf::Vector{T}
  mσ::SVector{46,T}
  mτ::SVector{46,T}
  zprev::zType
  mdeg::Int
  deg_index::Int
  start::Int
end

@cache struct SROCKC2Cache{uType,rateType,noiseRateType,T,C<:SROCKC2ConstantCache} <: StochasticDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  Gₛ::noiseRateType
  Gₛ₁::noiseRateType
  WikRange::T
  tmp::uType
  k::rateType
  fsalfirst::rateType
  atmp::rateType
  constantcache::C
end

function alg_cache(alg::SROCKC2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
  SROCKC2ConstantCache{typeof(t)}(u)
end

function alg_cache(alg::SROCKC2,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
  k = zero(rate_prototype)
  uᵢ = zero(u)
  uᵢ₋₁ = zero(u)
  uᵢ₋₂ = zero(u)
  Gₛ = zero(noise_rate_prototype)
  Gₛ₁ = zero(noise_rate_prototype)
  WikRange = false .* vec(ΔW)
  tmp  = zero(u)
  fsalfirst = k           # this variables are dummied to use same memory
  atmp = zero(rate_prototype)
  constantcache = SROCKC2ConstantCache{typeof(t)}(u)
  SROCKC2Cache(u,uprev,uᵢ,uᵢ₋₁,uᵢ₋₂,Gₛ,Gₛ₁,WikRange,tmp,k,fsalfirst,atmp,constantcache)
end