using ClimateMachine.Mesh.Grids: _x1, _x2, _x3, VerticalDirection

abstract type AbstractSimulation end

struct Simulation{𝒯, 𝒰, 𝒱, 𝒲, 𝒳, 𝒴, L} <: AbstractSimulation
    model::𝒯
    grid::L
    timestepper::𝒰
    time::𝒱
    callbacks::𝒲
    rhs::𝒳
    state::𝒴
end

function Simulation(model::ModelSetup; grid, timestepper, time, callbacks)
    rhs = DGModel(
        model,
        grid.numerical,
        model.numerics.flux,
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )

    FT = eltype(rhs.grid.vgeo)
    state = init_ode_state(rhs, FT(0); init_on_cpu = true)

    return Simulation(model, grid, timestepper, time, callbacks, rhs, state)
end

function Simulation(model::DryAtmosModel; grid, timestepper, time, callbacks)
    rhs = ESDGModel(
        model,
        grid.numerical,
        surface_numerical_flux_first_order = model.numerics.flux,
        volume_numerical_flux_first_order = KGVolumeFlux(),
    )

    FT = eltype(rhs.grid.vgeo)
    state = init_ode_state(rhs, FT(0); init_on_cpu = true)

    return Simulation(model, grid, timestepper, time, callbacks, rhs, state)
end

function Simulation(model::Tuple; grid, timestepper, time, callbacks)
    rhs = []
    for item in model
        if typeof(item) <: DryAtmosModel
            tmp = ESDGModel(
                item,
                grid.numerical,
                surface_numerical_flux_first_order = item.numerics.flux,
                volume_numerical_flux_first_order = KGVolumeFlux(),
            )
            push!(rhs, tmp)
        elseif typeof(item) <: DryAtmosLinearModel
            tmp = DGModel(
                item,
                grid.numerical,
                item.numerics.flux,
                CentralNumericalFluxSecondOrder(),
                CentralNumericalFluxGradient();
                direction = item.numerics.direction,
            )
            push!(rhs, tmp)
        end
    end
    rhs = Tuple(rhs)

    FT = eltype(rhs[1].grid.vgeo)
    state = init_ode_state(rhs[1], FT(0); init_on_cpu = true)

    return Simulation(model, grid, timestepper, time, callbacks, rhs, state)
end

struct CplSimulation{𝒯, 𝒰, 𝒱, 𝒲, 𝒳, 𝒴, L, B, CBV} <: AbstractSimulation
    model::𝒯
    grid::L
    odesolver::𝒰
    time::𝒱
    callbacks::𝒲
    rhs::𝒳
    state::𝒴
    nsteps::Int
    boundary::B
    cbvector::CBV
end

function CplSimulation(model; grid, timestepper, time, boundary_z = nothing, nsteps, callbacks)
    rhs = DGModel(
        model,
        grid.numerical,
        model.numerics.flux,
        model.numerics.flux_second_order,
        CentralNumericalFluxGradient(),
        direction = model.numerics.direction,
    )

    FT = eltype(rhs.grid.vgeo)
    state = init_ode_state(rhs, FT(0); init_on_cpu = true)

    xc = grid.numerical.vgeo[:, _x1:_x1, :]
    yc = grid.numerical.vgeo[:, _x2:_x2, :]
    zc = grid.numerical.vgeo[:, _x3:_x3, :]

    boundary(boundary_z, xc, yc, zc) = isnothing(boundary_z) ? zc .^ 2 .< eps(FT) : boundary_z(param_set, xc, yc, zc)
    boundary = boundary(boundary_z, xc, yc, zc)

    simulation = Simulation(model, grid, timestepper, time, callbacks, rhs, state)
    t0 = simulation.time.start
    Δt = timestepper.timestep

    # Instantiate time stepping method
    odesolver = timestepper.method(rhs, state, dt = Δt, t0 = t0)

    # Make callbacks from callbacks tuple
    cbvector = create_callbacks(simulation, odesolver)

    return CplSimulation(model, grid, odesolver, time, callbacks, rhs, state, nsteps, boundary, cbvector)
end

function CplSimulation(model::Tuple; grid, timestepper, time, boundary_z = nothing, nsteps, callbacks)
    rhs = []
    for item in model
        if typeof(item) <: DryAtmosModel
            tmp = ESDGModel(
                item,
                grid.numerical,
                surface_numerical_flux_first_order = item.numerics.flux,
                volume_numerical_flux_first_order = KGVolumeFlux(),
            )
            push!(rhs, tmp)
        elseif typeof(item) <: DryAtmosLinearModel
            tmp = DGModel(
                item,
                grid.numerical,
                item.numerics.flux,
                CentralNumericalFluxSecondOrder(),
                CentralNumericalFluxGradient();
                direction = item.numerics.direction,
            )
            push!(rhs, tmp)
        end
    end
    rhs = Tuple(rhs)

    xc = grid.numerical.vgeo[:, _x1:_x1, :]
    yc = grid.numerical.vgeo[:, _x2:_x2, :]
    zc = grid.numerical.vgeo[:, _x3:_x3, :]

    boundary(boundary_z, xc, yc, zc) = isnothing(boundary_z) ? zc .^ 2 .< eps(FT) : boundary_z(param_set, xc, yc, zc)
    boundary = boundary(boundary_z, xc, yc, zc)

    FT = eltype(rhs[1].grid.vgeo)
    state = init_ode_state(rhs[1], FT(0); init_on_cpu = true)

    simulation = Simulation(model[1], grid, timestepper, time, callbacks, rhs, state)
    t0 = simulation.time.start
    tend = simulation.time.finish
    Δt = timestepper.timestep

    # Instantiate time stepping method
    odesolver = timestepper.method(
        rhs[1],
        rhs[2],
        LinearBackwardEulerSolver(ManyColumnLU(); isadjustable = false),
        state;
        dt = Δt,
        t0 = t0,
        split_explicit_implicit = false,
    )

    # Make callbacks from callbacks tuple
    cbvector = create_callbacks(simulation, odesolver)

    return CplSimulation(model, grid, odesolver, time, callbacks, rhs, state, nsteps, boundary, cbvector)
end

function initialize!(simulation::Simulation; overwrite = false)
    if overwrite
        simulation = Simulation(
            model = simulation.model,
            timestepper = simulation.timestepper,
            time = simulation.time,
            callbacks = simulation.callbacks,
        )
    end

    return nothing
end

function evolve!(simulation::Simulation; refDat = ())
    # Unpack everything we need in this routine here
    timestepper = simulation.timestepper
    state = simulation.state
    rhs = simulation.rhs
    grid = simulation.grid

    npoly = convention(grid.resolution.polynomial_order, Val(ndims(grid.domain)))

    t0 = simulation.time.start
    tend = simulation.time.finish
    Δt = timestepper.timestep

    # Set up overintegration and polynomial order-based spatial staggering
    if haskey(grid.resolution, :overintegration_order)
        nover = convention(grid.resolution.overintegration_order, Val(ndims(grid.domain)))
    else
        nover = (0, 0, 0)
    end
    staggering = get(simulation.model.numerics, :staggering, false)

    # Perform overintegration filtering (only works if nover > 0)
    overintegration_filter!(state, rhs, npoly, nover)

    # Make right-hand side function with built-in filters
    rhs = rhs_closure(rhs, npoly, nover, staggering = staggering)

    # Instantiate time stepping method
    odesolver = timestepper.method(rhs, state, dt = Δt, t0 = t0)

    # Make callbacks from callbacks tuple
    cbvector = create_callbacks(simulation, odesolver)

    # Perform evolution of simulations
    if isempty(cbvector)
        solve!(state, odesolver; timeend = tend)
    else
        solve!(state, odesolver; timeend = tend, callbacks = cbvector)
    end

    # Check results against reference if StateCheck callback is used
    # TODO: TB: I don't think this should live within this function
    if any(typeof.(simulation.callbacks) .<: StateCheck)
        check_inds = findall(typeof.(simulation.callbacks) .<: StateCheck)
        @assert length(check_inds) == 1 "Only use one StateCheck in callbacks!"

        ClimateMachine.StateCheck.scprintref(cbvector[check_inds[1]])
        if length(refDat) > 0
            @test ClimateMachine.StateCheck.scdocheck(cbvector[check_inds[1]], refDat)
        end
    end

    return nothing
end

function evolve!(cpl_solver, numberofsteps; refDat = ())

    # Perform evolution of simulations
    solve!(
        nothing,
        cpl_solver;
        numberofsteps = numberofsteps,
        callbacks = cpl_solver.component_list.domainAtmos.component_model.cbvector,
    )

    # Check results against reference if StateCheck callback is used
    # TODO: TB: I don't think this should live within this function
    # if any(typeof.(simulation.callbacks) .<: StateCheck)
    #   check_inds = findall(typeof.(simulation.callbacks) .<: StateCheck)
    #   @assert length(check_inds) == 1 "Only use one StateCheck in callbacks!"

    #   ClimateMachine.StateCheck.scprintref(cbvector[check_inds[1]])
    #   if length(refDat) > 0
    #     @test ClimateMachine.StateCheck.scdocheck(cbvector[check_inds[1]], refDat)
    #   end
    # end

    return nothing
end

function rhs_closure(rhs, npoly, nover; staggering = false)
    if staggering
        function rhs_staggered(state_array, args...; kwargs...)
            rhs(state_array, args...; kwargs...)
            overintegration_filter!(state_array, rhs, npoly, nover)
            staggering_filter!(state_array, rhs, npoly, nover)
            return nothing
        end

        rhs_filtered = rhs_staggered
    else
        function rhs_unstaggered(state_array, args...; kwargs...)
            rhs(state_array, args...; kwargs...)
            overintegration_filter!(state_array, rhs, npoly, nover)
            return nothing
        end

        rhs_filtered = rhs_unstaggered
    end

    return rhs_filtered
end # returns a closure

function evolve!(simulation::Simulation{<:Tuple}; refDat = ())
    # Unpack everything we need in this routine here
    model = simulation.model[1]
    state = simulation.state
    rhs = simulation.rhs
    grid = simulation.grid.numerical
    timestepper = simulation.timestepper
    t0 = simulation.time.start
    tend = simulation.time.finish
    Δt = timestepper.timestep

    # Instantiate time stepping method
    odesolver = timestepper.method(
        rhs[1],
        rhs[2],
        LinearBackwardEulerSolver(ManyColumnLU(); isadjustable = false),
        state;
        dt = Δt,
        t0 = t0,
        split_explicit_implicit = false,
    )

    # Make callbacks from callbacks tuple
    cbvector = create_callbacks(simulation, odesolver)

    # Perform evolution of simulations
    if isempty(cbvector)
        solve!(state, odesolver; timeend = tend, adjustfinalstep = false)
    else
        solve!(state, odesolver; timeend = tend, callbacks = cbvector, adjustfinalstep = false)
    end

    # Check results against reference if StateCheck callback is used
    # TODO: TB: I don't think this should live within this function
    if any(typeof.(simulation.callbacks) .<: StateCheck)
        check_inds = findall(typeof.(simulation.callbacks) .<: StateCheck)
        @assert length(check_inds) == 1 "Only use one StateCheck in callbacks!"

        ClimateMachine.StateCheck.scprintref(cbvector[check_inds[1]])
        if length(refDat) > 0
            @test ClimateMachine.StateCheck.scdocheck(cbvector[check_inds[1]], refDat)
        end
    end

    return nothing
end

function overintegration_filter!(state_array, rhs, npoly, nover)
    if sum(nover) > 0
        cutoff_order = npoly .+ 1 # yes this is confusing
        cutoff = MassPreservingCutoffFilter(rhs.grid, cutoff_order)
        num_state_prognostic = number_states(rhs.balance_law, Prognostic())
        filterstates = 1:num_state_prognostic
        ClimateMachine.Mesh.Filters.apply!(state_array, filterstates, rhs.grid, cutoff)
    end

    return nothing
end

function staggering_filter!(state_array, rhs, npoly, nover)
    if sum(nover) > 0
        cutoff_order = npoly .+ 0 # one polynomial order less than scalars
        cutoff = MassPreservingCutoffFilter(rhs.grid, cutoff_order)
        num_state_prognostic = number_states(rhs.balance_law, Prognostic())
        filterstates = 2:4
        ClimateMachine.Mesh.Filters.apply!(state_array, filterstates, rhs.grid, cutoff)

    end

    return nothing
end


include("./game.jl")

using Distributions: Dirichlet
using LinearAlgebra: normalize!
using StatsBase
using Sockets

## MCTS setting

const Setting = (
        MAX_SIMULATION = 800,
        T              = 1.,
        C              = √3,
        ϵ              = 0.35,
        α              = 0.3,
        max_parall     = 100,
        PORT           = 8877
        )

function makeClients()
    Clients = Channel{TCPSocket}(Setting.max_parall)
    @simd for i=1:Setting.max_parall
        client = connect(Setting.PORT)
        put!(Clients, client)
    end
    Clients
end

function closeClients(Clients::Channel{TCPSocket})
    while isready(Clients)
        client = take!(Clients)
        close(client)
    end
    nothing
end

mutable struct Node
    N::Int
    W::Float16
    P::Float16
    vloss::Int
    parent::Union{Node, Nothing}
    children::Dict{Integer, Node}
    action::Union{Int8, String}
    isLeaf::Bool
    function Node(action=false, parent=nothing, P=Float16(0))
        new(Int(0), Float16(0), Float16(P), Int(0), parent, Dict{Integer, Node}(), action, true)
    end
end

function Base.show(io::IO, self::Node)
    "print out the node"
    children = keys(self.children)
    parent = isnothing(self.parent) ? "root" : self.parent.action
    print(io, "Node(action: $(self.action), parent: $parent, $children, ")
    p = self.N >0 ? self.W/self.N : 0
    print(io, "N/W/Q/P/VL: $(self.N)/$(self.W)/$(p)/$(self.P)/$(self.vloss)\n")
end

function bestChild(self::Node)
    best_child = nothing
    best_score = -Inf
    for child in values(self.children)
        Q = child.N > 0 ? (child.W-child.vloss)/child.N : 0
        curr = Q + Setting.C*child.P*√(self.N) / (1+child.N)
        if curr > best_score
            best_score = curr
            best_child = child
        end
    end
    return best_child
end

function add_noise!(probs::Array{Float32}, validmoves::Array{Int8,1})
    " Add Dirichlet noise in the root node "
    dirch = rand(Dirichlet(length(validmoves), Setting.α))
    @inbounds @. probs[validmoves] = probs[validmoves] * (1-Setting.ϵ) + dirch *Setting.ϵ
    nothing
end

function expand(self::Node, probs::Array{Float32}, moves::Array{Int8})
    for child in moves
        self.children[child] = Node(child, self, probs[child])
    end
    self.isLeaf = false
end

mutable struct MCTS
    root::Union{Nothing, Node}
    leaves::Int
    Clients::Channel{TCPSocket}
    function MCTS()
        new(nothing, 1, makeClients())
    end
end

function Base.show(io::IO, self::MCTS)
    "print out the node"
    compact = get(io, :compact, false)
    if !compact
        showroot = isnothing(self.root) ? "nothing" : self.root.action
        print(io, "MCTS(", showroot, ", ", self.leaves, ")")
    end
end

function backpropgation(root::Union{Node, Nothing}, q::Float16)
    while !isnothing(root.parent)
        root.vloss -= 1
        root.W += q
        root = root.parent
    end
end

function predict(game::Game, invaMoves::Array{Int8}, Clients)
    client = take!(Clients)
    println(client, game.state)
    flush(client)
    message = readuntil(client, "\n")
    put!(Clients, client)
    probs, vals = eval(Meta.parse(message))
    !isempty(invaMoves) && @inbounds probs[invaMoves] .= 0.
    normalize!(probs, 1)
    return probs, vals
end

function root_add_noise(game::Game, self::MCTS)
    vaMoves = validMoves(game)
    probs, _ = predict(game, setdiff(ALL_MOVES[1], vaMoves), self.Clients)
    add_noise!(probs, vaMoves)
    expand(self.root, probs, vaMoves)
end

function search!(self::MCTS, initState::String)
    root = self.root
    game = Game(initState)
    # selection
    curr_player = game.nextPlayer
    root.N += 1
    while !root.isLeaf
        root = bestChild(root)
        place(game, root.action)
        root.N += 1; root.vloss +=1
        if !iszero(game.winner)
            backpropgation(root, Float16(game.winner*curr_player))
            break
        end
    end
    # expansion
    if iszero(game.winner)
        vaMoves = validMoves(game)
        #bottleneck
        @inbounds probs, q = predict(game, setdiff(ALL_MOVES[1], vaMoves), self.Clients)
        game = nothing
        self.leaves += -1+length(vaMoves)
        expand(root, probs, vaMoves)
        # backpropgation
        backpropgation(root, q)
    end
end

function (self::MCTS)(game::Game; test::Bool=false, counts::Int = Setting.MAX_SIMULATION)
    "return move, Pi"

    initState = game.state
    if isnothing(self.root)
        self.root = Node(initState)
    elseif @inbounds self.root.action == initState[1:end-1]
        for child in values(self.root.children)
            if @inbounds child.action == int(initState[end])
                self.root = child
                self.root.parent = nothing
                break
            end
        end
    end
    total = 0
    if self.leaves < 9
        root_add_noise(game, self)
        total += 1
    end
    while total < counts
        parall = min(self.leaves, counts-total, Setting.max_parall)
        @sync for i=1:parall
            @async search!(self, initState)
        end
        total += parall
    end

    moves = self.root.children
    Pi = zeros(Float32, 9)
    if test
        # choose the action from the current state with greatest N
        for i in keys(moves)
            @inbounds Pi[i] = moves[i].N
        end
        normalize!(Pi, 1)
        move = argmax(Pi)
    else
        for i in keys(moves)
            @inbounds Pi[i] = moves[i].N ^ (1. / Setting.T)
        end
        normalize!(Pi, 1)
        @inbounds move = sample(ALL_MOVES[1], Weights(Pi))
    end
    # cut off the subtree
    self.root = moves[move]
    self.root.action = "$(game.state)$move"
    self.root.parent = nothing
    return move, Pi
end


#= Create ODEProblem from the selected Model and Simulation options, and also the selected
    vehicle configurations
=#
import JPLEphemeris: SPK, Dates, position, velocity, state
include("..\\Vehicle\\Spacecraft.jl")
include("..\\Environment\\Environment.jl")
include("State.jl")

#Create EOM from the vehicle configurations
function EOM(vehicle::Type{<:AbstractSpacecraft}, env::Environment)

    # Gravity (Environment-based)
    #

    # VehicleShapeAerodynamics (Vehicle properties, combined with environmental properties)

    # VehicleActuators (Vehicle properties, combined with environmental properties)

    eomfunc = function (y::SimState, t::Float64)

    end #eomfunc
    return eomfunc
end #EOM


include("gamess_custom_basis_sets/631+gd_split.jl")
include("gamess_custom_basis_sets/6311++g2d2p_split.jl")
include("gamess_custom_basis_sets/6311++gdp-J.jl")

const atom_to_atomic_number_mapping = Dict(
  "H" => 1.0,
  "He" => 2.0,
  "Li" => 3.0,
  "Be" => 4.0,
  "B" => 5.0,
  "C" => 6.0,
  "N" => 7.0,
  "O" => 8.0,
  "F" => 9.0,
  "Ne" => 10.0,
  "Na" => 11.0,
  "Mg" => 12.0,
  "Al" => 13.0,
  "Si" => 14.0,
  "P" => 15.0,
  "S" => 16.0,
  "Cl" => 17.0,
  "Ar" => 18.0,
  "K" => 19.0,
  "Ca" => 20.0,
  "Sc" => 21.0,
  "Ti" => 22.0,
  "V" => 23.0,
  "Cr" => 24.0,
  "Mn" => 25.0,
  "Fe" => 26.0,
  "Co" => 27.0,
  "Ni" => 28.0,
  "Cu" => 29.0,
  "Zn" => 30.0,
  "Ga" => 31.0,
  "Ge" => 32.0,
  "As" => 33.0,
  "Se" => 34.0,
  "Br" => 35.0,
  "Kr" => 36.0,
  "Rb" => 37.0,
  "Sr" => 38.0,
  "Y" => 39.0,
  "Zr" => 40.0,
  "Nb" => 41.0,
  "Mo" => 42.0,
  "Tc" => 43.0,
  "Ru" => 44.0,
  "Rh" => 45.0,
  "Pd" => 46.0,
  "Ag" => 47.0,
  "Cd" => 48.0,
  "In" => 49.0,
  "Sn" => 50.0,
  "Sb" => 51.0,
  "Te" => 52.0,
  "I" => 53.0,
  "Xe" => 54.0
)

function create_input(input)
  xyz_array::Vector{String} = []
  xyz = open(input, "r") do file
    fragment = readlines(file)
    xyz_array = fragment[3:end]
  end
  natoms = length(xyz_array)

  output = open(joinpath(@__DIR__, input[1:(end-4)]*".inp"), "w") do file
    write(file, " \$scf dirscf=.t. damp=.t. shift=.f.\n")
    write(file, "   soscf=.f. diis=.t. fdiff=.f. extrap=.f. \$end\n")
    write(file, " \$system mwords=500 memddi=500 \$end\n")
    write(file, " \$intgrl schwrz=.t. intomp=1 shfock=.f. \$end\n")
    write(file, " \$guess guess=hcore \$end\n")
    write(file, " \$contrl scftyp=rhf runtyp=energy nprint=-5\n")
    write(file, "     icharg=0 ispher=0 inttyp=eric \n")
    write(file, "     maxit=100 \$end\n")
    write(file, " \$end\n")
    write(file, " \$data\n")
    write(file, "$(splitdir(input[1:(end-4)])[end]) \n")
    write(file, "C1\n")
    for line in xyz_array
      atom = rstrip(line[1:2])
      newline = line[1:2]*"   $(atom_to_atomic_number_mapping[atom])   "*
        line[3:end]
      write(file, "$newline\n")

      newline = six_three_one_one_2plus_2d_2p_split[atom]
      #newline = six_three_one_one_2plus_d_p_J[atom]
      #newline = six_three_one_plus_g_d_split[atom]
      write(file, "$newline\n")
    end
    write(file, " \$end\n")
  end
end


export PIBasis

# ---------------------- Implementation of the PIBasisSpec

"""
`struct PIBasisSpec`
"""
struct PIBasisSpec
   orders::Vector{Int}     # order (length) of ith basis function
   iAA2iA::Matrix{Int}     # where in A can we find the ith basis function
end

==(B1::PIBasisSpec, B2::PIBasisSpec) = _allfieldsequal(B1, B2)

Base.length(spec::PIBasisSpec) = length(spec.orders)

maxcorrorder(spec::PIBasisSpec) = size(spec.iAA2iA, 2)

function _get_pibfcn(spec0, Aspec, vv)
   vv1 = vv[2:end]
   vvnz = vv1[findall(vv1 .!= 0)]
   return (spec0[vv[1]], Aspec[vvnz])
end

function _get_pibfcn(Aspec, vv)
   vvnz = vv[findall(vv .!= 0)]
   return Aspec[vvnz]
end

# TODO: maybe instead of property == nothing, there should be a
#       generic property with no symmetry attached to it.

function PIBasisSpec( basis1p::OneParticleBasis,
                      symgrp::SymmetryGroup,
                      Bsel::DownsetBasisSelector;
                      property = nothing,
                      filterfun = _->true,
                      init1pbasis = true )

   # we initialize the 1p-basis here; to prevent this it must be manually
   # avoided by passing in init1pbasis = false
   if init1pbasis
      init1pspec!(basis1p, Bsel)
   end

   # get the basis spec of the one-particle basis
   #  Aspec[i] described the basis function that will get written into A[i]
   Aspec = get_spec(basis1p)

   # we assume that `Aspec` is sorted by degree, but best to double-check this
   # since the notion of degree used to construct `Aspec` might be different
   # from the one used to construct AAspec.
   if !issorted(Aspec; by = b -> level(b, Bsel, basis1p))
      error("""PIBasisSpec : AAspec construction failed because Aspec is not
               sorted by degree. This could e.g. happen if an incompatible
               notion of degree was used to construct the 1-p basis spec.""")
   end
   # An AA basis function is given by a tuple 𝒗 = vv. Each index 𝒗ᵢ = vv[i]
   # corresponds to the basis function Aspec[𝒗ᵢ] and the tuple
   # 𝒗 = (𝒗₁, ...) to a product basis function
   #   ∏ A_{vₐ}
   tup2b = vv -> _get_pibfcn(Aspec, vv)

   #  degree or level of a basis function ↦ is it admissible?
   admissible = bb -> (level(bb, Bsel, basis1p) <= maxlevel(bb, Bsel, basis1p))

   if property != nothing
      filter1 = bb -> filterfun(bb) && filter(bb, Bsel, basis1p) && filter(property, symgrp, bb)
   else
      filter1 = bb -> filterfun(bb) && filter(bb, Bsel, basis1p)
   end

   # we can now construct the basis specification; the `ordered = true`
   # keyword signifies that this is a permutation-invariant basis
   maxord = maxorder(Bsel)
   AAspec = gensparse(; NU = maxorder(Bsel),
                        tup2b = tup2b,
                        admissible = admissible,
                        ordered = true,
                        maxvv = [length(Aspec) for _=1:maxord],
                        filter = filter1)

   return PIBasisSpec(AAspec)
end

function PIBasisSpec(AAspec)
   orders = zeros(Int, length(AAspec))
   iAA2iA = zeros(Int, (length(AAspec), length(AAspec[1])))
   for (iAA, vv) in enumerate(AAspec)
      # we use reverse because gensparse constructs the indices in
      # ascending order, but we want descending here.
      # (I don't remember why though)
      iAA2iA[iAA, :] .= reverse(vv)
      orders[iAA] = length( findall( vv .!= 0 ) )
   end
   return PIBasisSpec(orders, iAA2iA)
end

get_spec(AAspec::PIBasisSpec, i::Integer) = AAspec.iAA2iA[i, 1:AAspec.orders[i]]

# ------------------ PISpec sparsification

"""
returns a new `PIBasisSpec` constructed from the old one, but keeping only
the basis indices `Ikeep`. This maintains the order of the basis functions.
"""
sparsify(spec::PIBasisSpec, Ikeep::AbstractVector{<: Integer}) =
      PIBasisSpec(spec.orders[Ikeep], spec.iAA2iA[Ikeep, :])

function _fix_A_indices!(spec::PIBasisSpec, new_inds::AbstractVector{<: Integer})
   for iAA = 1:size(spec.iAA2iA, 1)
      for α = 1:spec.orders[iAA]
         vα = spec.iAA2iA[iAA, α]
         new_vα = new_inds[vα]
         @assert new_vα > 0
         spec.iAA2iA[iAA, α] = new_vα
      end
   end
   return nothing
end

# --------------------------------- PIBasis implementation

"""
`mutable struct PIBasis:` implementation of a permutation-invariant
basis based on the density projection trick.

The standard constructor is
```
PIBasis(basis1p, N, D, maxdeg)
```
* `basis1p` : a one-particle basis
* `N` : maximum interaction order
* `D` : an abstract degee specification, e.g., SparsePSHDegree
* `maxdeg` : the maximum polynomial degree as measured by `D`
"""
mutable struct PIBasis{BOP, REAL} <: ACEBasis
   basis1p::BOP             # a one-particle basis
   spec::PIBasisSpec
   real::REAL     # could be `real` or `identity` to keep AA complex
end               # TODO -> chain??

cutoff(basis::PIBasis) = cutoff(basis.basis1p)

==(B1::PIBasis, B2::PIBasis) =
      ( (B1.basis1p == B2.basis1p) &&
        (B1.spec == B2.spec) &&
        (B1.real == B2.real) )

Base.length(basis::PIBasis) = length(basis.spec)

# default symmetry group
PIBasis(basis1p, Bsel::AbstractBasisSelector; kwargs...) =
   PIBasis(basis1p, O3(), Bsel; kwargs...)

PIBasis(basis1p, symgrp, Bsel::AbstractBasisSelector;
        isreal = false, kwargs...) =
   PIBasis(basis1p,
           PIBasisSpec(basis1p, symgrp, Bsel; kwargs...),
           isreal ? Base.real : Base.identity )

get_spec(pibasis::PIBasis) =
   [ get_spec(pibasis, i) for i = 1:length(pibasis) ]

get_spec(pibasis::PIBasis, i::Integer) =
      get_spec.( Ref(pibasis.basis1p), get_spec(pibasis.spec, i) )

setreal(basis::PIBasis, isreal::Bool) =
   PIBasis(basis.basis1p, basis.spec, isreal)

maxcorrorder(basis::PIBasis) = maxcorrorder(basis.spec)

# ------------------ sparsification

function sparsify!(basis::PIBasis, Ikeep::AbstractVector{<: Integer})
   basis.spec = sparsify(basis.spec, Ikeep)
   return basis
end

"""
This should allow the 1p basis to sparsify itself, then feed back to the
pibasis what the correct indices are.
"""
function clean_1pbasis!(basis::PIBasis)
   spec = get_spec(basis)
   B1p = basis.basis1p
   spec1p = NamedTuple[]
   for bb in spec
      append!(spec1p, bb)
   end
   identity.(unique!(spec1p))
   # sparsify the product 1p basis
   _, new_inds = sparsify!(basis.basis1p, spec1p)
   # now fix the indexing of the PIBasis specification
   _fix_A_indices!(basis.spec, new_inds)
   return basis
end

# -------------------

_scaling_absvalue(x::Number) = abs(x)
_scaling_absvalue(x::Symbol) = 0

# TODO: this is a hack; cf. #68
function scaling(pibasis::PIBasis, p)
   ww = zeros(Float64, length(pibasis))
   bspec = get_spec(pibasis)
   for i = 1:length(pibasis)
      for b in bspec[i]
         # TODO: revisit how this should be implemented for a general basis
         ww[i] += sum(x -> _scaling_absvalue(x)^p, b)  #  abs.(values(b)).^p
      end
   end
   return ww
end

# -------------------------------------------------
# FIO codes

write_dict(basis::PIBasis) =
   Dict(  "__id__" => "ACE_PIBasis",
         "basis1p" => write_dict(basis.basis1p),
            "spec" => write_dict(basis.spec),
            "real" => basis.real == Base.real ? true : false )

read_dict(::Val{:ACE_PIBasis}, D::Dict) =
   PIBasis( read_dict(D["basis1p"]),
            read_dict(D["spec"]),
            D["real"] ? Base.real : Base.identity )

write_dict(spec::PIBasisSpec) =
   Dict( "__id__" => "ACE_PIBasisSpec",
         "orders" => spec.orders,
         "iAA2iA" => write_dict(spec.iAA2iA) )

read_dict(::Val{:ACE_PIBasisSpec}, D::Dict) =
   PIBasisSpec( D["orders"], read_dict(D["iAA2iA"]) )

# -------------------------------------------------
# Evaluation codes

function evaluate!(AA, basis::PIBasis, config::UConfig)
   A = evaluate(basis.basis1p, config)
   evaluate!(AA, basis, A)
   release!(A)
   return AA
end

function evaluate!(AA, basis::PIBasis, A::AbstractVector{<: Number})
   fill!(AA, 1)
   for iAA = 1:length(basis)
      aa = one(eltype(A))
      for t = 1:basis.spec.orders[iAA]
         aa *= A[ basis.spec.iAA2iA[ iAA, t ] ]
      end
      AA[iAA] = basis.real(aa)
   end
   return AA
end

_valtype(basis::PIBasis, A::AbstractVector{<: Number}) =
      basis.real(eltype(A))

# draft of defining bases via chains
function evaluate(basis::PIBasis, A::AbstractVector{<: Number})
   VT = _valtype(basis, A)
   AA = Vector{VT}(undef, length(basis))
   evaluate!(AA, basis, A)
   return AA
end

# draft of defining bases via chains
function evaluate(basis::PIBasis, config::UConfig)
   A = evaluate(basis.basis1p, config)
   AA = evaluate(basis, A)
   release!(A)
   return AA
end

# -------------------------------------------------
# gradients

evaluate_d(basis::PIBasis, cfg::UConfig, args...) =
         evaluate_ed(basis, cfg, args...)[2]

function evaluate_ed(basis::PIBasis, cfg::UConfig, args...)
   A, dA = evaluate_ed(basis.basis1p, cfg, args...)
   AA, dAA = evaluate_ed(basis, A, dA)
   release!(A)
   release!(dA)
   return AA, dAA
end

_gradtype(basis::PIBasis, A, dA) =
      basis.real( promote_type(eltype(A), eltype(dA)) )

# this is really an frule I think
function evaluate_ed(basis::PIBasis, A, dA)
   VT = _valtype(basis, A)
   GT = _gradtype(basis, A, dA)
   AA = Vector{VT}(undef, length(basis))
   dAA = Matrix{GT}(undef, length(basis), size(dA, 2))
   evaluate_ed!(AA, dAA, basis, A, dA)
   release!(A)
   release!(dA)
   return AA, dAA
end

function evaluate_ed!(AA, dAA, basis::PIBasis,
                      cfg::UConfig, args...)
   A, dA = evaluate_ed(basis.basis1p, cfg, args...)
   evaluate_ed!(AA, dAA, basis, A, dA)
   release!(dA)
   release!(A)
   return AA, dAA
end

function _AA_local_adjoints!(dAAdA, A, iAA2iA, iAA, ord, _real)
   if ord == 1
      return _AA_local_adjoints_1!(dAAdA, A, iAA2iA, iAA, ord, _real)
   elseif ord == 2
      return _AA_local_adjoints_2!(dAAdA, A, iAA2iA, iAA, ord, _real)
   else
      return _AA_local_adjoints_x!(dAAdA, A, iAA2iA, iAA, ord, _real)
   end
end

function _AA_local_adjoints_1!(dAAdA, A, iAA2iA, iAA, ord, _real)
   @inbounds dAAdA[1] = 1
   @inbounds A1 = A[iAA2iA[iAA, 1]]
   return _real(A1)
end

function _AA_local_adjoints_2!(dAAdA, A, iAA2iA, iAA, ord, _real)
   @inbounds A1 = A[iAA2iA[iAA, 1]]
   @inbounds A2 = A[iAA2iA[iAA, 2]]
   @inbounds dAAdA[1] = A2
   @inbounds dAAdA[2] = A1
   return _real(A1 * A2)
end

function _AA_local_adjoints_x!(dAAdA, A, iAA2iA, iAA, ord, _real)
   @assert length(dAAdA) >= ord
   @assert ord >= 2
   # TODO - optimize a bit more? can move one operation out of the loop
   # Forward pass:
   @inbounds A1 = A[iAA2iA[iAA, 1]]
   @inbounds A2 = A[iAA2iA[iAA, 2]]
   @inbounds dAAdA[1] = 1
   @inbounds dAAdA[2] = A1
   @inbounds AAfwd = A1 * A2
   @inbounds for a = 3:ord-1
      dAAdA[a] = AAfwd
      AAfwd *= A[iAA2iA[iAA, a]]
   end
   @inbounds dAAdA[ord] = AAfwd
   @inbounds Aend = A[iAA2iA[iAA, ord]]
   aa = _real(AAfwd * Aend)
   # backward pass
   @inbounds AAbwd = Aend
   @inbounds for a = ord-1:-1:3
      dAAdA[a] *= AAbwd
      AAbwd *= A[iAA2iA[iAA, a]]
   end
   dAAdA[2] *= AAbwd
   AAbwd *= A2
   dAAdA[1] *= AAbwd

   return aa
end

_acquire_dAAdA!(basis::PIBasis, A) = Vector{eltype(A)}(undef, maxcorrorder(basis))

function evaluate_ed!(AA, dAA, basis::PIBasis,
                      A::AbstractVector, dA::AbstractMatrix)
   dAAdA = _acquire_dAAdA!(basis, A)
   _evaluate_ed!(AA, dAA, basis, A, dA, dAAdA)
end

function _evaluate_ed!(AA, dAA, basis::PIBasis,
                       A::AbstractVector, dA::AbstractMatrix, dAAdA)
   orders = basis.spec.orders
   iAA2iA = basis.spec.iAA2iA

   # Must treat the constants separately. This is not so elegant and could
   # maybe be improved?
   if orders[1] == 0  # SHOULD BE THE ONLY ONE with ord=0!!
      iAAinit = 2
      AA[1] = 1.0
      dAA[1, :] .= Ref(zero(eltype(dAA)))
   else
      iAAinit = 1
   end

   for iAA = iAAinit:length(basis)
      ord = orders[iAA]

      # ----- compute the local adjoints dAA / dA
      # dAAdA[a] ← ∏_{t ≂̸ a} A_{v_t}
      AA[iAA] = _AA_local_adjoints!(dAAdA, A, iAA2iA, iAA, orders[iAA], basis.real)

      # ----- now convert them into dAA / dX
      for j = 1:size(dA, 2)
         dAA[iAA, j] = sum(dAAdA[a] * dA[iAA2iA[iAA, a], j]
                           for a = 1:ord) |> basis.real
      end
   end

   return AA, dAA
end


"""
    struct MultiFieldFEFunction <: CellField
      # private fields
    end
"""
struct MultiFieldFEFunction{T<:MultiFieldCellField} <: CellField
  single_fe_functions::Vector{<:SingleFieldFEFunction}
  free_values::AbstractArray
  cell_dof_values::AbstractArray
  fe_space::MultiFieldFESpace
  multi_cell_field::T

  function MultiFieldFEFunction(
    free_values::AbstractVector,
    space::MultiFieldFESpace,
    single_fe_functions::Vector{<:SingleFieldFEFunction})

    multi_cell_field = MultiFieldCellField(map(i->i.cell_field,single_fe_functions))
    T = typeof(multi_cell_field)

    f(i) = get_cell_axes(get_fe_space(i))
    cell_axes = create_array_of_blocked_axes(map(f,single_fe_functions)...)
    blocks = Tuple(map(get_cell_values,single_fe_functions))
    blockids = [(i,) for i in 1:length(single_fe_functions)]
    cell_dof_values = VectorOfBlockArrayCoo(blocks,blockids,cell_axes)

    new{T}(
      single_fe_functions,
      free_values,
      cell_dof_values,
      space,
      multi_cell_field)
  end
end

#function MultiFieldFEFunction(
#  space::MultiFieldFESpace,
#  blocks::Vector{<:SingleFieldFEFunction})
#  fv = zero_free_values(space)
#  xh0 = MultiFieldFEFunction(fv,space,blocks)
#end

Arrays.get_array(a::MultiFieldFEFunction) = get_array(a.multi_cell_field)

CellData.get_memo(a::MultiFieldFEFunction) = get_memo(a.multi_cell_field)

CellData.get_cell_map(a::MultiFieldFEFunction) = get_cell_map(a.multi_cell_field)

CellData.get_cell_axes(a::MultiFieldFEFunction) = get_cell_axes(a.multi_cell_field)

CellData.RefStyle(::Type{<:MultiFieldFEFunction{T}}) where T = RefStyle(T)

CellData.MetaSizeStyle(::Type{<:MultiFieldFEFunction{T}}) where T = MetaSizeStyle(T)

FESpaces.FEFunctionStyle(::Type{<:MultiFieldFEFunction}) = Val{true}()

FESpaces.get_free_values(f::MultiFieldFEFunction) = f.free_values

FESpaces.get_fe_space(f::MultiFieldFEFunction) = f.fe_space

Base.length(f::MultiFieldFEFunction) = length(f.multi_cell_field)

"""
    num_fields(m::MultiFieldFEFunction)
"""
num_fields(m::MultiFieldFEFunction) = length(m.single_fe_functions)

Base.iterate(m::MultiFieldFEFunction) = iterate(m.single_fe_functions)

Base.iterate(m::MultiFieldFEFunction,state) = iterate(m.single_fe_functions,state)

Base.getindex(m::MultiFieldFEFunction,field_id::Integer) = m.single_fe_functions[field_id]

function Geometry.restrict(a::MultiFieldFEFunction,trian::Triangulation)
  f = (ai) -> restrict(ai,trian)
  blocks = map(f,a.single_fe_functions)
  blocks
end

function FESpaces.get_cell_values(f::MultiFieldFEFunction)
  f.cell_dof_values
end


module Fmincon

using MATLAB

function fmincon(fundef::String, fun::String, x0::AbstractArray{<:Real,1},
    lb::AbstractArray{<:Real,1}, ub::AbstractArray{<:Real,1};
    options::AbstractDict=Dict(),
    A::AbstractArray{<:Real,2}=zeros(0,0), b::AbstractArray{<:Real,1}=zeros(0),
    Aeq::AbstractArray{<:Real,2}=zeros(0,0), beq::AbstractArray{<:Real,1}=zeros(0),
    gradients::Bool=false, printfile::String="")

    # get the default matlab session
    msession = MATLAB.get_default_msession()

    # add matlab optimization function to the MATLAB path
    moptimpath = dirname(pathof(Fmincon))
    eval_string(msession, "addpath('$moptimpath');")

    # make sure inputs are compatible with MATLAB
    x0_m = Float64.(x0)
    A_m = Float64.(A)
    b_m = Float64.(b)
    Aeq_m = Float64.(Aeq)
    beq_m = Float64.(beq)
    lb_m = Float64.(lb)
    ub_m = Float64.(ub)
    options_m = Dict{String, Any}()
    for key in keys(options)
        if isa(options[key], Number)
            options_m[key] = Float64.(options[key])
        else
            options_m[key] = options[key]
        end
    end
    gradients_m = Bool(gradients)

    # # copy input variables to MATLAB session
    put_variable(msession, :x0, x0_m)
    put_variable(msession, :A, A_m)
    put_variable(msession, :b, b_m)
    put_variable(msession, :Aeq, Aeq_m)
    put_variable(msession, :beq, beq_m)
    put_variable(msession, :lb, lb_m)
    put_variable(msession, :ub, ub_m)
    put_variable(msession, :options, options_m)
    put_variable(msession, :gradients, gradients_m)

    # load objective/constraint function into embedded Julia
    if isfile(fundef)
        eval_string(msession, "jl.include('$fundef');")
    else
        eval_string(msession, "jl.eval('$fundef');")
    end

    nargout = gradients ? 6 : 3

    # make user function MEX-like
    mexusrfun =
        """
        jl.eval(...
            ['function mexusrfun(prhs)' newline...
             '    args = vec(jarray(prhs[1]))' newline...
             '    outs = $fun(args)' newline...
             '    return [Float64.(outs[i]) for i=1:$nargout]' newline...
             'end']);
        """
    eval_string(msession, mexusrfun)

    # create a function handle for the user function
    eval_string(msession, "mexusrfun = jl.wrapmex('mexusrfun', $nargout);")

    # clear outputs
    eval_string(msession, "clear xopt fopt exitflag output")

    # perform optimization
    eval_string(msession, "[xopt, fopt, exitflag, output] = optimize(mexusrfun, x0, A, b, Aeq, beq, lb, ub, options, gradients, '$printfile');")

    # check if optimization was performed succesfully
    eval_string(msession, "fail = exist('xopt') == 0;")
    fail = jvalue(get_mvariable(msession, :fail))

    # throw error if it didn't, and redirect user to other error messages
    if fail
        error("An error occured during call to fmincon, see error messages from MATLAB and embedded Julia")
    end

    # copy output variables from MATLAB session
    xopt = vec(jarray(get_mvariable(msession, :xopt)))
    fopt = jvalue(get_mvariable(msession, :fopt))
    exitflag = jvalue(get_mvariable(msession, :exitflag))
    output = jvalue(get_mvariable(msession, :output))

    return xopt, fopt, exitflag, output

end

end


module LiveSlice

using Interact
using Blink
using CMDimData
using CMDimData.EasyPlot
using MDDatasets

struct SliceParam
	sweep::PSweep
	choosevalue::Observable
	sweepindex::Observable
end

#==Helper functions
===============================================================================#
_throt(x) = throttle(0.1, x) #Default throttling of inputs for this tool.

#Sweep value string
_swvalstr(x::Int) = string(x)
_swvalstr(x::Float64) = EasyPlot.SI(x)

#==Main autoslice() function
===============================================================================#
"""
    autoslice(refreshfn, refparamlist::Vector{PSweep}, slicewfrmlist::Vector{Waveform})

Automatically updates provided waveforms, and calls provided function to refresh plots

#Arguments
 - refreshfn: Function to be called when caller should refresh plots
 - refparamlist: List of parameters to use to slice input data
 - slicewfrmlist: List of DataHR-waveforms to be sliced along different parameter dimensions
"""
function autoslice(refreshfn::Function, refparamlist::Vector{PSweep}, slicewfrmlist::Vector{EasyPlot.Waveform})
	origdatalist = DataHR[wfrm.data for wfrm in slicewfrmlist] #Only supports DataHR
	paramlist = SliceParam[]
	obslist = Observable[] #Index into each parameter sweep value
	widgetlist = [] #To be displayed in Blink

	#Build observables & sliders for each parameter provided by caller:
	for p in refparamlist
		choosevalue = Observable(true)
		sweepindex = Observable(1)
		push!(paramlist, SliceParam(p, choosevalue, sweepindex))
		push!(obslist, choosevalue)
		push!(obslist, _throt(sweepindex)) #Throttle observables to reduce load
		cb = checkbox(choosevalue)
		sld = slider(1:length(p.v), value=sweepindex, label=p.id)
		push!(widgetlist, hbox(cb, sld))
	end

	#Build up an expression string to see which parameter values we are selecting:
	_paramlist_str = map(obslist...) do varlist...
		isfirst = true
		paramstr = ""

		for p in paramlist
			if !isfirst
				paramstr *= " / "
			end
			valstr = "ALL"
			if p.choosevalue[]
				i = p.sweepindex[]
				valstr = _swvalstr(p.sweep.v[i])
			end
			paramstr *= p.sweep.id * "=" * valstr
			isfirst=false
		end
		return HTML(paramstr)
	end
	push!(widgetlist, _paramlist_str)

	#Swap out waveform data with a sliced version (dep. on slider values):
	_updater = map(obslist...) do varlist...
		for (od, wfrm) in zip(origdatalist, slicewfrmlist)
			idxlen = length(size(od.elem))
			Δidxlen = idxlen - length(paramlist)
			idxlist = fill!(Array{Any}(undef, idxlen), :)
			swlist = PSweep[]
			for (i, p) in enumerate(paramlist)
				if p.choosevalue[] #Want a specific index, not all swept values
					idxlist[i] = p.sweepindex[]
				else
					push!(swlist, p.sweep)
				end
			end
			append!(swlist, od.sweeps[(end+1-Δidxlen):end])
			newelem = od.elem[idxlist...]
			wfrm.data = DataHR(swlist, newelem)
		end
		refreshfn() #Ask user to refresh plots, etc
	end

	#Display main window:
	wnd = Window()
	body!(wnd, vbox()) #first call to body! does not always work.
	body!(wnd, hbox(vbox(widgetlist...), nothing))
	return
end

end #module


import PeriodicGraphs

function export_dataline(f, x)
    inbetween = '\t'
    if any(isspace, x)
        if any(==('\n'), x)
            x = ';'*x*"\n;"
            inbetween = '\n'
        else
            @assert x[1]!= '\''
            x = '\''*x*'\''
        end
    end
    println(f, inbetween*x)
end

function export_vtf(file, c::Union{Crystal{Nothing},CrystalNet}, repeatedges=1, colorname=false)
    mkpath(splitdir(file)[1])
    n = length(c.types)
    if c isa CrystalNet
        @assert length(c.pos) == n
    else
        @assert size(c.pos)[2] == n
    end
    invcorres = [PeriodicVertex3D(i) for i in 1:n]
    corres = Dict{PeriodicVertex3D,Int}([invcorres[i]=>i for i in 1:n])
    encounteredtypes = Dict{Symbol,String}()
    numencounteredtypes = 0
    open(file, write=true) do f
        println(f, """
        ###############################
        # written by PeriodicGraphs.jl
        ###############################
        """)

        for i in 1:n
            ty = c.types[i]
            sty = ty === Symbol("") ? string(i) : string(ty)
            name = colorname ? get!(encounteredtypes, ty) do
                string(numencounteredtypes+=1)
            end : string(i)
            println(f, "atom $(i-1) type $sty name $name resid $(i-1)")
        end
        j = n + 1
        for _ in 1:repeatedges
            jmax = j - 1
            for i in 1:jmax
                vertex = invcorres[i]
                for x in neighbors(c.graph, vertex.v)
                    y = PeriodicVertex3D(x.v, x.ofs .+ vertex.ofs)
                    if get!(corres, y, j) == j
                        j += 1
                        push!(invcorres, y)
                    end
                end
            end
        end
        for i in n+1:length(invcorres)
            v = invcorres[i].v
            ofs = invcorres[i].ofs
            ty = c.types[v]
            sty = ty === Symbol("") ? string(i) : string(ty)
            name = colorname ? encounteredtypes[ty] : string(v)
            println(f, "atom $(i-1) type $sty name $name resid $v")
        end
        println(f)

        for (i,x) in enumerate(invcorres)
            for neigh in neighbors(c.graph, x.v)
                j = get(corres, PeriodicVertex3D(neigh.v, neigh.ofs .+ x.ofs), nothing)
                isnothing(j) && continue
                if i < j
                    println(f, "bond ", i - 1, ':', j - 1)
                end
            end
        end
        println(f)

        _a, _b, _c, α, β, γ = cell_parameters(c.cell)
        println(f, "pbc $_a $_b $_c $α $β $γ\n")

        println(f, "ordered")
        for x in invcorres
            coord = c isa CrystalNet ?
                c.cell.mat * (widen.(c.pos[x.v]) .+ x.ofs) :
                c.pos[:,x.v] + c.cell.mat * x.ofs
            println(f, join(round.(Float64.(coord); digits=15), ' '))
        end
    end
end

function export_vtf(file, cif::CIF, repeatedges=1)
    export_vtf(file, CrystalNet(cif), repeatedges, true)
end

function export_cif(file, c::Union{Crystal, CIF})
    mkpath(splitdir(file)[1])
    info = c isa CIF ? copy(c.cifinfo) : Dict{String,Union{String,Vector{String}}}("data"=>last(splitdir(file)))
    loops = Dict{Int, Tuple{Vector{String}, Vector{Vector{String}}}}()
    open(file, write=true) do f
        print(f, "data_"); println(f, pop!(info, "data")*"\n\n")
        for (id, data) in info
            if data isa String
                print(f, '_'); print(f, id)
                export_dataline(f, data)
            else
                @assert data isa Vector{String}
                n = length(data)
                if haskey(loops, n)
                    push!(loops[n][1], id)
                    push!(loops[n][2], data)
                else
                    loops[n] = (String[id], Vector{String}[data])
                end
            end
        end

        #scale_factor::Float64 = unique!(sort(c.types)) == [:Si] && c.pos[:,1] != [0,0,0] ? 1.2 : 1.0
        _a, _b, _c, α, β, γ = Float64.(cell_parameters(c.cell))

        println(f, """

        _cell_length_a\t\t$_a
        _cell_length_b\t\t$_b
        _cell_length_c\t\t$_c
        _cell_angle_alpha\t\t$α
        _cell_angle_beta\t\t$β
        _cell_angle_gamma\t\t$γ

        _symmetry_space_group_name_H-M\t'$(c.cell.spacegroup)'""")
        if c.cell.tablenumber != 0
            println(f, "_symmetry_Int_Tables_number\t$(c.cell.tablenumber)")
        end
        if c.cell.latticesystem !== Symbol("")
            println(f, "_symmetry_cell_setting\t$(c.cell.latticesystem)")
        end

        println(f, """

        loop_
        _symmetry_equiv_pos_as_xyz
        x,y,z""")
        for eq in c.cell.equivalents
            println(f, eq)
        end
        println(f)

        n = c isa CIF ? length(c.ids) : length(c.types)
        if !haskey(loops, n)
            loops[n] = (String[], Vector{String}[])
        end
        append!(loops[n][1],
                ["atom_site_label", "atom_site_type_symbol",
                 "atom_site_fract_x", "atom_site_fract_y", "atom_site_fract_z"])
        labels = String[string(c.types[c isa CIF ? c.ids[i] : i])*string(i) for i in 1:n]
        pos = string.(round.(c.pos; sigdigits=6))
        append!(loops[n][2], [labels, string.(c.types[x] for x in (c isa CIF ? c.ids : collect(1:n))),
                pos[1,:], pos[2,:], pos[3,:]])

        if c isa Crystal
            bonds = edges(c.graph)
            src = String[]
            dst = String[]
            mult = Int[]
            last_bond = (0, 0)
            n = 0
            for e in bonds
                if (e.src, e.dst.v) == last_bond
                    mult[end] += 1
                else
                    push!(src, string(labels[e.src]))
                    push!(dst, string(labels[e.dst.v]))
                    push!(mult, 1)
                    last_bond = (e.src, e.dst.v)
                    n += 1
                end
            end
            if !haskey(loops, n)
                loops[n] = (String[], Vector{String}[])
            end
            append!(loops[n][1], ["geom_bond_atom_site_label_1", "geom_bond_atom_site_label_2", "geom_bond_multiplicity"])
            append!(loops[n][2], [src, dst, string.(mult)])
        end

        for (n, (ids, datas)) in loops
            println(f, "loop_")
            for id in ids
                print(f, '_'); println(f, id)
            end
            for i in 1:n
                for j in 1:length(datas)
                    print(f, datas[j][i]); print(f, '\t')
                end
                println(f)
            end
            println(f)
        end
        println(f, "\n#END")
    end
    nothing
end

function export_cgd(file, c::Crystal)
    mkpath(splitdir(file)[1])
    invmat = Float64.(inv(c.cell.mat))
    open(file, write=true) do f
        println(f, "CRYSTAL\n")
        margin = 1
        println(f, "\tNAME\t", last(splitdir(file)))
        #scale_factor::Float64 = unique!(sort(c.types)) == [:Si] && c.pos[:,1] != [0,0,0] ? 1.2 : 1.0
        _a, _b, _c, α, β, γ = Float64.(cell_parameters(c.cell))
        println(f, "\tGROUP\t\"", join(split(c.cell.spacegroup, ' ')), "\"")
        println(f, "\tCELL\t", _a, ' ', _b, ' ', _c, ' ', α, ' ', β, ' ', γ, ' ')
        println(f, "\tATOM")
        to_revisit = Int[]
        for i in 1:length(c.types)
            push!(to_revisit, i)
            pos = invmat * c.pos[:,i]
            println(f, "\t\t", i, ' ', degree(c.graph, i), ' ', pos[1], ' ',
                    pos[2], ' ', pos[3])
        end
        println(f, "\tEDGE")
        for i in to_revisit
            for e in neighbors(c.graph, i)
                e.v < i && continue
                dest = invmat * c.pos[:,e.v] .+ e.ofs
                println(f, "\t\t", i, '\t', dest[1], ' ', dest[2], ' ', dest[3])
            end
        end
        println(f, "\nEND")
    end
end

function export_cgd(file, g::PeriodicGraph)
    mkpath(splitdir(file)[1])
    open(file, write=true) do f
        println(f, "PERIODIC_GRAPH\n")
        println(f, "ID ", basename(splitext(file)[1]), '\n')
        println(f, "EDGES")
        repr = reverse(split(string(g)))
        n = parse(Int, pop!(repr))
        m = length(repr) ÷ (n+2)
        @assert iszero(length(repr) % (n+2))
        for _ in 1:m
            src = pop!(repr)
            dst = pop!(repr)
            ofs = Vector{String}(undef, n)
            for i in 1:n
                ofs[i] = pop!(repr)
            end
            println(f, '\t', src, ' ', dst, ' ', join(ofs, ' '))
        end
        println(f, "END\n")
    end
end
export_cgd(file, c::CrystalNet) = export_cgd(file, c.graph)

function export_clusters(crystal::Crystal{Clusters}, path=joinpath(tempdir(),tempname()))
    frame = Chemfiles.Frame()
    m = length(crystal.clusters.classes)
    # residues = [Chemfiles.Residue(string(i)) for i in 1:m]
    # for i in 1:m
    #     Chemfiles.set_property!(residues[i], "chainname", string(i))
    #     Chemfiles.set_property!(residues[i], "chainid", string(i))
    #     Chemfiles.add_residue!(frame, residues[i])
    # end
    _a, _b, _c, α, β, γ = cell_parameters(crystal.cell)
    Chemfiles.set_cell!(frame, Chemfiles.UnitCell(_a, _b, _c, α, β, γ))
    recenter::SVector{3,Float64} = minimum(crystal.pos; dims=2)
    for i in 1:length(crystal.types)
        # resid = crystal.clusters.attributions[i]
        atom = Chemfiles.Atom(string(crystal.clusters.classes[crystal.clusters.attributions[i]]))
        Chemfiles.set_type!(atom, string(crystal.types[i]))
        Chemfiles.add_atom!(frame, atom, collect(Float64.(crystal.pos[:,i] .- recenter)))
        # Chemfiles.add_atom!(residues[resid], i)
    end
    for e in edges(crystal.graph)
        iszero(PeriodicGraphs.ofs(e)) || continue
        Chemfiles.add_bond!(frame, src(e)-1, dst(e)-1)
    end
    target = isempty(last(splitext(path))) ? path*".pdb" : path
    output = Chemfiles.Trajectory(target, 'w')
    write(output, frame)
    close(output)
    return target
end

"""
    export_arc(path)

Export the current archive to the specified path.
"""
function export_arc(path, empty=false, arc=CRYSTAL_NETS_ARCHIVE)
    mkpath(splitdir(path)[1])
    open(path, "w") do f
        println(f, "Made by CrystalNets.jl v$CRYSTAL_NETS_VERSION")
        if !empty
            println(f)
            pairs = sort(collect(arc); by=last)
            for (genome, id) in pairs
                println(f, "key\t", genome)
                println(f, "id\t", id)
                println(f)
            end
        end
    end
end


export
    RGBDDataset,
    NYUDataset,
    get_image,
    CDN,
    CDNImage,
    CDNColorant,
    LabelledCDNImage,
    color_image,
    z_image,
    to_array,
    to_rgbdnimage,
    to_image,
    rgb2lab!,
    preprocess,
    colors,
    distances,
    normals,
    z,
    name
import Base.download
using FITSIO
using MAT
using ..RGBDSegmentation
using ColorTypes
using StaticArrays
import Base: show, +,-,*,/,getindex, setindex!,zeros, zero
import ColorTypes: comp1, comp2, comp3
import Images: accum

abstract type RGBDDataset end
import Base.getindex
function getindex(d::RGBDDataset, i::Integer)
    return get_image(d, i)
end

const CDNImage{T} = Array{T, 3} where T<:Number
@enum ImageType Color=1 Depth=2 Normal=3
function CDNImage(RGB::Array{T1, 3}, D::Array{T2, 3}, N::Array{T3, 3}) where T1 where T2 where T3
    if size(RGB,3) == 3
        RGB = permutedims(RGB./255.0, [3 1 2])
    end
    if size(D, 3) == 3
        D = permutedims(D, [3 1 2])
    end
    if size(N, 3) == 3
        N = permutedims(N, [3 1 2])
    end
    return cat(RGB, D,  N, dims=1)
end
function number_of_pixels(img::CDNImage)
    return size(img,2)*size(img,3)
end
function colors(img::CDNImage)
    return colorview(RGB, img[1:3, :, :])
end
function distances(img::CDNImage)
    return img[4:6, :, :]
end
function z(img::CDNImage)
    return img[6, :, :]
end
function normals(img::CDNImage)
    return img[7:9, :, :]
end
struct LabelledCDNImage
    image::CDNImage
    ground_truth::SegmentedImage
end
function LabelledCDNImage(image::CDNImage, labels::AbstractArray{T}) where T<:Integer
    return LabelledCDNImage(image,
                            SegmentedImage(image, labels))
end


@testset "CoreModelCoupled generic interface" begin
    model = load_model(CoreModelCoupled, model_paths["e_coli_core.mat"])

    @test reaction_stoichiometry(model, "EX_ac_e") == Dict("ac_e" => -1)
    @test reaction_stoichiometry(model, 44) == Dict("ac_e" => -1)

end

@testset "Conversion from and to StandardModel" begin
    cm = load_model(CoreModelCoupled, model_paths["e_coli_core.mat"])

    sm = convert(StandardModel, cm)
    cm2 = convert(CoreModelCoupled, sm)

    @test Set(reactions(cm)) == Set(reactions(sm))
    @test Set(reactions(cm)) == Set(reactions(cm2))

    @test reaction_gene_association(sm, reactions(sm)[1]) ==
          reaction_gene_association(cm, reactions(sm)[1])

    @test reaction_gene_association_vec(cm)[1:3] == [
        [["b3916"], ["b1723"]],
        [
            ["b0902", "b0903", "b2579"],
            ["b0902", "b0903"],
            ["b0902", "b3114"],
            ["b3951", "b3952"],
        ],
        [["b4025"]],
    ]
end


module StrandbeestRobot

using RigidBodyDynamics
using RigidBodyDynamics.PDControl
using RigidBodyDynamics.Contact
const rbd = RigidBodyDynamics
using LightXML
using StaticArrays
using Optim
using LinearAlgebra
using RigidBodySim

urdfpath() = joinpath(@__DIR__, "..", "data", "Strandbeest.urdf")

function mechanism(::Type{T} = Float64;
                   floating = true,
                   contact_model = default_contact_model(),
                   add_flat_ground = true) where {T}
    mechanism = parse_urdf(T, urdfpath())
    add_loop_joints!(mechanism, urdfpath())

    if floating
        crossbar = findbody(mechanism, "crossbar")
        basejoint = joint_to_parent(crossbar, mechanism)
        floatingjoint = Joint("floating_base", frame_before(basejoint),
            frame_after(basejoint), QuaternionFloating{Float64}())
        replace_joint!(mechanism, basejoint, floatingjoint)
    end

    add_foot_contacts!(mechanism, contact_model)
    if add_flat_ground
        add_environment_primitive!(mechanism,
            HalfSpace3D(
                Point3D(root_frame(mechanism), 0., 0, 0),
                FreeVector3D(root_frame(mechanism), 0., 0, 1)))
    end

    mechanism
end

"""
Parse the Drake-specific `loop_joint` tags from a URDF and use them
to add non-tree joints to the given mechanism.
"""
function add_loop_joints!(mechanism::Mechanism{T}, urdf::AbstractString) where T
    doc = parse_file(urdf)
    xml_root = LightXML.root(doc)
    xml_loops = get_elements_by_tagname(xml_root, "loop_joint")
    for xml_loop in xml_loops
        name = attribute(xml_loop, "name")
        @assert attribute(xml_loop, "type") == "continuous"
        axis = SVector{3}(rbd.parse_vector(T, find_element(xml_loop, "axis"), "xyz", "1 0 0"))
        joint = Joint(name, Revolute(axis))
        xml_link1 = find_element(xml_loop, "link1")
        body1 = findbody(mechanism, attribute(xml_link1, "link"))
        H1 = Transform3D(frame_before(joint), default_frame(body1),
            rbd.parse_pose(T, xml_link1)...)
        xml_link2 = find_element(xml_loop, "link2")
        body2 = findbody(mechanism, attribute(xml_link2, "link"))
        H2 = Transform3D(frame_after(joint), default_frame(body2),
            rbd.parse_pose(T, xml_link2)...)
        attach!(mechanism, body1, body2, joint,
            joint_pose = H1,
            successor_pose = inv(H2))
    end
end

"""
Return a function which maps a configuration vector
for the given mechanism to the total squared error in
the alignment of the robot's non-tree joints.
"""
function loop_joint_error(mechanism::Mechanism)
    let statecache = StateCache(mechanism), resultcache = DynamicsResultCache(mechanism)
        function(q)
            state = statecache[eltype(q)]
            result = resultcache[eltype(q)]
            set_configuration!(state, q)
            zero_velocity!(state)
            gains = rbd.CustomCollections.ConstDict{JointID}(SE3PDGains(PDGains(1, 0), PDGains(1, 0)))
            rbd.constraint_bias!(result, state; stabilization_gains=gains)
            ret = result.constraintbias ⋅ result.constraintbias
        end
    end
end

default_contact_model() = SoftContactModel(hunt_crossley_hertz(k = 500e3), ViscoelasticCoulombModel(1.0, 20e3, 100.))

function add_foot_contacts!(mechanism::Mechanism, contactmodel=default_contact_model())
    for body in bodies(mechanism)
        if occursin("bars_g_h_i", string(body))
            frame = default_frame(body)
            point = Point3D(frame, 0, 0, 0.49)
            add_contact_point!(body, ContactPoint(
                point, contactmodel))

        end
    end
end

function solve_initial_configuration!(state::MechanismState{T}) where T
    # Find an initial state of the robot which correctly aligns all
    # the loop joints.
    mechanism = state.mechanism

    # We do this before adding the contact points to work around
    # https://github.com/JuliaRobotics/RigidBodyDynamics.jl/issues/483

    # First, initialize the state with a configuration within the
    # joint limits:
    lb = similar(configuration(state))
    ub = similar(configuration(state))
    for joint in tree_joints(mechanism)
        lb_joint = lb[joint]
        ub_joint = ub[joint]
        if isfloating(joint)
            # use identity transform for floating joint
            tf = one(Transform3D{T}, frame_after(joint), frame_before(joint))
            set_configuration!(state, joint, tf)
            # Optim can't handle equal bounds, so give it a bit of wiggle room
            lb_joint .= configuration(state, joint) .- 1e-3
            ub_joint .= configuration(state, joint) .+ 1e-3
        else
            lb_joint .= rbd.lower.(position_bounds(joint))
            ub_joint .= rbd.upper.(position_bounds(joint))
        end
    end
    q = configuration(state)
    for i in eachindex(q)
        if isfinite(lb[i]) && isfinite(ub[i])
            q[i] = (lb[i] + ub[i]) / 2
        else
            q[i] = clamp(0, lb[i], ub[i])
        end
    end
    setdirty!(state)

    # Use Optim's Fminbox to minimize the loop joint error within
    # the joint limits:
    cost = loop_joint_error(mechanism)
    result = Optim.optimize(cost, lb, ub,
        Vector(configuration(state)),
        Fminbox(LBFGS()), autodiff=:forward)

    # Verify that we've actually closed all the loops
    @assert Optim.minimum(result) < 1e-9

    set_configuration!(state, Optim.minimizer(result))
    normalize_configuration!(state)
    state
end

end


const initialObjValue = 1.529659701346959e-5
const objGap = 0.02*initialObjValue

const riskLowerBound = 1E-04

user_initial_guess = Float64[0.0008209427811086344 0.0007636835675518444 0.0005701047968333946 0.0008769116263480982 0.00041324527675481526 0.0005800434950967022 0.0007210045551505841 0.0006276031212332207 0.0007932877009665201 0.0008881171637927818 0.0008766018612865671 0.0006320318164184507 0.00046864233622053727 0.0007972152206997032 0.00044410020082646863 0.00048751053639709734 0.0006271324945815114 0.000833594526671095 0.000637517908021743 0.0005758235103795233 0.0008211858446615176 0.0005506595813249986 0.0004937790049517456 0.00030326690290466917 0.0007858265233236424 0.0004283766060798492 0.0007023563531744242 0.0006383812893485663 0.00041660529637406083 0.0008869414840670435 0.0008486309841526661 0.0007878373981276149 0.0007025895738211321 0.0006221266154370171 0.0007253224649432664 0.0007179847285153958 0.0005743155113023336 0.0004651830659334972 0.0006994539928083114 0.000813757012867089 0.0008426123267428326 0.0008129090919633393 0.0002224035375153371 0.0007633996283991556 0.0008060377452304927 0.0005805474993823768 0.0006758568767157172 0.0008487093317498039 0.0007273095454223649 0.0007672506824693867 0.00034776956635203585 0.000667823526204176 0.0008073333001026818 0.0008118463955352388 0.0008965575169526982 0.0004495653113838658 0.0006647997123165969 0.0006472528325401913 0.0009356519965734018 0.0010672497699516243 0.0011253462942174997 0.0006076383475545551 0.0006686631677018113 0.0008866874124157135 0.0007440327743919704 0.0009304131014551529 0.0005814299397160758 0.0010648434477783509 0.0009183489627890017 0.0004239725944496445 0.0005341670597258343 0.000576451191307805 0.0006589548224670541 0.0003622754702797734 0.0007495302760552053 0.0008448696097240493 0.0007138969914752706 0.0006575246169877966 0.0007249316169717899 0.0005741152648363746 0.0006147471826141572 0.0007997356796553269 0.0007778347652323488 0.0008690003186054852 0.0007256524288646954 0.0009299331492444872 0.0007743840354389615 0.0009591622728523907 0.0009020007349513178 0.000559755222833856 0.0006619681901540276 0.000713979237294059 0.0005889631583622089 0.000863440443912283 0.0006145512583291466 0.0007468192567333447 0.0006992921241140729 0.0007103558563154973 0.0008633419054996466 0.0005733737387722774 0.0007914126339008551 0.0012514110982711668 0.0008622021127955821 0.0005980703775501003 0.000700653564168626 0.0009120405539474175 0.000875585515525431 0.0007266650895836344 0.0004265783693812117 0.0006521518718345092 0.0005124619065816503 0.0004923006408059255 0.0007999620305938937 0.0012081612851983468 0.0008935328351092937 0.0009747528164537189 0.0007526206802701704 0.0009562998353617571 0.0009221615951393487 0.0004526283997919561 0.0007501979401291613 0.0006792407033835312 0.0009418241224837042 0.0009933871804938816 0.0008090888857165224 0.0010207535073878482 0.000824300854913575 0.0004411332775433643 0.0009822455247934403 0.0005584278250656128 0.0008877703972421927 0.0007580793962589777 0.0007663863102773332 0.0008605922280018232 0.0008086846120149574 0.000530847508782214 0.0005727563851087061 0.0008013865698785838 0.0011107323548744997 0.0010099690714047958 0.00037319999111541034 0.0008140689739195794 0.0005375612612210462 0.00063910747383142 0.0005425847990026908 0.0008037415767010002 0.0009097868293051404 0.0006095811690976755 0.0008416701707228973 0.0005423072476081243 0.0006486589625278647 0.0005367637947850263 0.0007693494308252405 0.000711008758454373 0.0008921888683283542 0.0008969566774610786 0.000842119786840999 0.0007265316634713728 0.0009227744525842186 0.0010881828110205858 0.000769784759877319 0.0002852821218126892 0.00078600213965086 0.0010140796668548107 0.0006790073530653253 0.0008440295525002012 0.0010101242557135272 0.0010180177710397583 0.0008462467667346788 0.0012228989677190464 0.0007445113149683374 0.0005706095400447219 0.0010449249119368167 0.0009395015242278339 0.0006968910296485923 0.0007501283795220349 0.000715685305438086 0.00048027337416580015 0.0010704746686966546 0.0008188810546293447 0.0009314103657909904 0.0005720666361682608 0.0012387022323280977 0.0007503426653680849 0.000601405282514845 0.0006481905274730637 0.0005601701838525855 0.0009236596566139973 0.0009677158401210994 0.0005632497639895169 0.0007675964782709339 0.0007563850652900144 0.0007182119902737943 0.0004906045529927696 0.0009007167560036245 0.000510807650533841 0.0007430984749102268 0.0007165927858411207 0.0005430495355291725 0.0007928403533148245 0.000876228226677745 0.0011251550047914488 0.0007529985067706059 0.0008333912348460475 0.000400081733537378 0.0009072571985814809 0.0005241527459227279 0.0010254036212703876 0.0007134183946210163 0.0009456546923780793 0.0006388634133473063 0.0009655762101426236 0.0006483664726507628 0.0006965500023781428 0.0006376498163030177 0.0003968958336100884 0.0004439596257893074 0.001130271687360053 0.000731480515879121 0.0005537891201286045 0.0008802977512820636 0.0009153645368225062 0.0005705148052624951 0.0007219917166105136 0.0004702245397953825 0.0007875775442662136 0.0010030209383180364 0.0007317042652216768 0.0012226228604856567 0.0014028119865632988 0.0005776767988588316 0.0009484531736954914 0.0009494748579465421 0.0007427817773937607 0.0007808385679751186 0.0011736248374371225 0.0011019851590723247 0.0006842860707850102 0.0006734618869398631 0.0007494740493314883 0.0009918981406768924 0.0009600657183434738 0.0004846155959183507 0.000841114666643091 0.0007772963410989715 0.0005115938318186999 0.0007122948496523184 0.0012270643951943446 0.0009600382451763617 0.0005977865737263712 0.0010619091723500583 0.0011519071770915248 0.0005393873336186615 0.0009928481976629834 0.0009084426302312609 0.0009933356693733499 0.0011135369000948644 0.0008355638213707572 0.0009594533651078643 0.00036590315768847995 0.0008008540464168165 0.0007920229835437993 0.0009819292408506208 0.0006139326998714905 0.0009092167669549826 0.0009835473860453111 0.0003476301786949354 0.0006498930670309982 0.0008240009448767369 0.001263893894153865 0.001214130826401816 0.000668180172453217 0.0010539626322570315 0.0005631421322522217 0.0007381333784743991 0.0013186753827719368 0.001192240939472783 0.0009910717530359802 0.0007157048270570837 0.0010651465270070418 0.0012565237924333024 0.000563033935165114 0.0005277377029486118 0.0007589859529664756 0.0004109897417342138 0.0006119105501833137 0.0009465129944708208 0.0001993255884156562 0.0008213984184820397 0.0008562907355875043 0.0010898114992350246 0.0005712057723414628 0.0008103874270986649 0.0010981089108472403 0.001172242375814217 0.0006561275173663624 0.0009301670280869443 0.0005544976110261636 0.0009525322948247709 0.0006271825626415179 0.001125280900550546 0.0010093400073160233 0.0004893138851278724 0.0007102925069228844 0.0004796042507685992 0.0008320952766252863 0.0010268093580907487 0.0009174270220115308 0.0009298890379354706 0.0007664565439431096 0.00096542694894379 0.0005316454657471359 0.001362406051977638 0.00126029833439878 0.0015002729509467304 0.0011691531469095273 0.0008467625790500536 0.0008035687520204203 0.0009614578466938769 0.0008010521678074861 0.00037759130230045174 0.0006589673155935706 0.0005303753656813506 0.0010126210280955714 0.0006258020382088416 0.0009121267248172789 0.0012313691157549722 0.001049065630478657 0.0011048100935120862 0.0006227969652011532 0.0007965724264743089 0.001387224270608414 0.0009522428644574411 0.0010089755272122036 0.0009822089101359638 0.0007677170225958986 0.0006316490723250939 0.0009377560502012278 0.001060987234789023 0.0007861907282201283 0.0012709343870376227 0.0010250970204928515 0.0010837771081950056 0.0009908547264889213 0.0010502152733814748 0.0005120218122460347 0.0007165366252357839 0.0012037113405740831 0.0006707983109259369 0.001389885479633844 0.0005303380245454641 0.0013554719630867475 0.000804813587291104 0.0005007846117488728 0.0008680012852792285 0.0008039899127962578 0.0005295077751178488 0.0010650687700154119 0.0007551182578845476 0.0007263976670925004 0.001212962372405825 0.00104476983342287 0.0008537036247692432 0.0014445821073785932 0.0012449688888018058 0.0015208393742885764 0.0007046003973507587 0.0013934568629435386 0.000914859236698767 0.0007052111127436618 0.0009159172738197191 0.0008778581715316999 0.0010087097539912607 0.0008271830385666752 0.0009360392861293037 0.0008113447613887953 0.00047815824762114945 0.0011067457725130881 0.0008807406021714571 0.0010666916953821813 0.0009577045844013372 0.0006305692589931498 0.0005800668701273395 0.0002048618975592093 0.0006069979497977569 0.000531196898201318 0.0008557851943482072 0.0009750142819443672 0.0009724848942666668 0.0012779415772484239 0.0008074434576826656 0.0008418842441015366 0.0010521237578863631 0.000778198047897382 0.00090376221369533 0.00043035694323763943 0.0008742855799754223 0.001293281901675596 0.0007827360789566337 0.0009870681543964024 0.001004847452866001 0.0012184345961668555 0.0007936755337952622 0.0010459014382046903 0.0009546634995085316 0.001257624167365295 0.0006774271075431505 0.0013639558210267054 0.001100473564927773 0.0012441102707288038 0.0014373857870888854 0.0009150286223557653 0.001209298528005329 0.0007652899916084731 0.0006996104804969957 0.0011260722891758136 0.0011549935018552535 0.0008307287635270463 0.0015058746045194355 0.0010516149892078987 0.0012029447762088293 0.0009149306184527473 0.0009740695125500079 0.0010794086464580963 0.0011349346613733482 0.0008975004855891622 0.001304870649803291 0.001215273303293898 0.0010795893817305833 0.0015884936773312161 0.0010729659479968908 0.0008753081888062908 0.0010974408944288906 0.0009228422478649721 0.0009758794772812533 0.00085299246003311 0.0008867787868433317 0.00124727278366698 0.0012353667203989379 0.0012427596285666108 0.0009222149551261983 0.0011844073632694104 0.0009044989823418179 0.001299040704567225 0.0013880283383311122 0.0010371915040593496 0.0010732780727615372 0.0009135192289410564 0.0012440904907508677 0.0009786975952165416 0.0005696430386582136 0.0011109918483655822 0.0012981570873818445 0.0012858002985194441 0.0011333129761868558 0.0009878432188965528 0.0009053147596283131 0.0009747363835909339 0.0007355339906187744 0.0010330339574952474 0.0009102039828842749 0.001056966511953838 0.0006762026603291838 0.001199470853799652 0.0012545676792113203 0.0007826816289715383 0.0010494180006627143 0.0008986454825122008 0.0008900709684340947 0.0006860818154401112 0.0009857076361018831 0.0003457653958055086 0.0005140022302885193 0.001423235888113857 0.0017553574125951088 0.001197126284837738 0.0007414652441424827 0.0008362348738906361 0.0008132245591619835 0.0011378570794584276 0.0015874003027109434 0.0009255086666005051 0.0009427273245087456 0.0010762379737519164 0.00098788523595807 0.001571019431385693 0.0002909997699635088 0.0010184338942693999 0.0014719939019106052 0.0012457295331779418 0.0011220279251867103 0.001320708704275924 0.0012399075686666584 0.001159364024075537 0.001244669399989565 0.0012067782961340619 0.0013108062922103081 0.0009618976139673397 0.0008071311636264934 0.0012848658844038863 0.0009171019086273589 0.0007996176051974812 0.0008036062149084962 0.0009072798476765115 0.000793089012937557 0.0009600554648462686 0.000992645257032026 0.0009686687940966474 0.0012243837521141365 0.001125806252093119 0.0007654719821346565 0.0013642862324606225 0.0011670301374839115 0.001604578434894504 0.0008374568788123102 0.0014986235006204862 0.0015701356411262878 0.0006657434022405598 0.0013845353661062413 0.001219384596471336 0.0008737768939108806 0.0010494510684420574 0.0009222086839748313 0.0012732347456471667 0.0011858684105957279 0.0019179000784992626 0.0013681221021490028 0.0013281836906744392 0.0012187595843030589 0.0013111683867050915 0.0011719458203592623 0.0013572137824939884 0.0010337098311431982 0.001638871465698339 0.0015604667950890121 0.001323508987236831 0.0020464741506832484 0.0006944878825225153 0.0011046410928042814 0.0014414804611452834 0.0018725218466401593 0.0002404324247984538 0.0010138722934900173 0.0013926979063741147 0.002070058176189564 0.0013019284617489528 0.0013038680434372446 0.0005397204832431734 0.0011923521522380076 0.0008886501333491859 0.0015010097035359769 0.0011247537153967642 0.0007845845804796413 0.0006219419716377376 0.0009462713723758497 0.001649494886433042 0.0014082798737966899 0.001072304075655831 0.0007545879915296843 0.0013142903219559607 0.0012673741886097404 0.0007686825089079977 0.0006515034405852395 0.0004003290511664357 0.0014643531373883926 0.0014171798819934101 0.0016449012501703376 0.002028157630601896 0.0006128944335396143 0.0016631742826592123 0.000723549828572489 0.0013736669987269787 0.0013823399008336121 0.0015810376356261368 0.0011949894572673926 0.001114578942759561 0.0018730136673697695 0.001196997041645648 0.0011722255970589652 0.0008370746042579171 0.0013974757415557363 0.0012688066121312469 0.0010424252641943875 0.001264977755400999 0.0010700943734532474 0.0017819146559548851 0.0005098494196810078 0.0018873505208473866 0.00105391094104933 0.0012466699162386396 0.0004211146137099592 0.001413056211171595 0.001274979235601838 0.0006545247733292137 0.001554598304774834 0.0011463541680346255 0.0014534359957203103 0.0013726032339994264 0.0011308452038855611 0.0013526290665061553 0.0007292296475314194 0.001469852110624965 0.0012256122574814128 0.0004780914210506476 0.0018306464300324543 0.0018831192531126237 0.0010960935296654283 0.0015125460384125887 0.000908354783661358 0.0011002741525626839 0.0018637255066985203 0.0013853153136580672 0.0015969960353949771 0.0009589572055989069 0.0008768182801056807 0.0007037782765468832 0.001902492166644584 0.0020202920907032972 0.0016411461982764198 0.0007455373517780349 0.0013984409733384937 0.0005833112587010457 0.0013941249451607156 0.0016020617420586616 0.0015153846262352937 0.002375857809950477 0.0005340868158132004 0.0008373689872912367 0.001308614247178495 0.0016996779250921489 0.0014141078045271478 0.0015680206621989548 0.0002940026955455891 0.0017030333611348387 0.001195380218694524 0.001806235301558795 0.0010875513922395253 0.0010830733497866878 0.000822317038180907 0.000567722870710013 0.000583568662936589 0.001134118688231472 0.0010025636436904871 0.0005573334444183738 0.0019287841154301936 0.0007634713852630136 0.0012448222272147791 0.0015882791656856592 0.0016850549761983043 0.0008877585912039505 0.00104969676415809 0.0014213751108184158 0.0015169425682421463 0.0004905417410219729 0.0024447412335600665 0.0020835912947756554 0.0018555286072899496 0.0017174195510289902 0.0013937253225056964 0.001467637469255164 0.002532695504216574 0.0015400182099285682 0.0013298598275832457 0.0009850053591270987 0.001571467265404312 0.0013749691850085618 0.002009538941291626 0.0020580068757073606 0.0015042087068884266 0.0024265628730326897 0.0011724249798647946 0.0012647737606120157 0.0014387314832090801 0.0019044943735031447 0.0014317002205580668 0.001062507055779056 0.0013849884348921858 0.0011984775734875351 0.0010121218745125307 0.0014409430939625546 0.0012407462596961648 0.0003437158726665955 0.0016447246383843974 0.0011422292134610546 0.000563097823904911 0.0015765128968950914 0.001505823615162113 0.0017650607218414195 0.0010931119982941712 0.0015125881789526316 0.001361209028585102 0.0006658072320213478 0.00136505909113486 0.001795809851104738 0.0012095923801153348 0.0009776784083181758 0.001627853713850087 0.003023318742709796 0.0016614407782219395 0.0018098704690599993 5.81011264298683e-7 0.0016751546512679446 0.0010344924476244288 0.0008936426090566522 0.0009398686913127734 0.00047546344948500084 0.00084704209531371 0.0005997781107622372 0.0009820247785840939 0.0016008420932675368 0.001361524326786464 0.0011681406116652465 0.0007100794777450526 0.001622593049100077 0.0019312043562599314 0.001955834940734043 0.001300186661130563 0.0008473059639885839 0.0008468296030323799 0.001797805964100868 0.0010032311409894144 0.0008254800640423295 0.0018715669514508166 0.0014605540571477505 0.0018592754796626031 0.0020846651286845486 0.0009005262697593812 0.0009571855578512618 0.0007755804271232539 0.002529844628074969 0.0013351552794298276 0.00129525142595085 0.0022310082213243803 0.0006986957744983211 0.002103264606915408 0.0011158356780163537 0.0011111673909054121 0.0029762941670259916 0.0015918008277034997 0.0014271555841840762 0.0012700077047275167 0.0020324459041304106 0.001442032023090879 0.0020929878441526186 0.002340007075902152 0.0012774323788884406 0.0015995389199063644 0.0021912024186273352 0.0016967336259560508 0.001153798746508344 0.0008731025062233751 0.0013082968275565116 0.001690519897314452 0.0012804431431907733 0.0020546201946592267 0.0013561974403580727 0.001266082254498757 0.001892325369974752 0.0007921558062621175 0.0017086246313910208 6.949962975098176e-6 0.0009307040010744052 0.0005748508210090921 0.0015303236497021776 0.0017467876141135695 0.001302979438345667 0.0018759164807104684 0.00029491744444495616 0.0015776081131344969 0.0025614829713955503 0.0008605819278327193 0.0012447307949637822 0.00043463089799912214 0.0016578433246198525 0.000815010473203796 0.0004851853614444267 0.0019280940726462315 0.001617598884767852 0.0019658389117368016 0.00199361309800802 0.0018354567334096302 0.00254454699600168 0.002217324416727356 0.00020554106276034323 0.0028145695570283478 0.002013541408952003 0.001427609469811477 0.001326367675877524 0.001859069809545089 0.001927764537984837 0.0008969343729452774 0.002078938862771566 0.0021103409071436545 0.0006483807209883566 0.0015428676568223507 0.001257208211624372 0.0008834221452461365 0.0015244953917749255 0.0028480773692240026 0.0018884456043282579 0.0025703151678374485 0.0009733015287772254 0.0012406806699404723 0.000825756285923428 0.001661907000244604 0.0013403525174255958 0.0018343613983112187 0.0014820134576503778 0.002440746528070622 0.0014174054859715324 0.0018759569749941705 0.0016341317196305928 0.001541797791549588 0.000992750475183311 0.0016770105069908708 0.002235592362869389 0.0015017703376287138 0.0017447638527936112 0.0012431075209464906 0.0003018377398732817 0.0015932204974133267 0.001230847072724235 0.0014563964500322218 0.0016790782662605439 0.00026283670948781037 0.0005712391136852477 0.0028085016499335125 0.0018019701057805566 0.0003064629036626991 0.0014811442176186575 0.0018039038548974629 0.00028317267206120215 0.00024862450480226295 0.0018686430495926452 0.0017277721090188012 0.0024070333194674606 0.0001558851165436762 0.0010741548411589336 0.0004448681469482655 0.00019681498154585912 0.001446134649364775 0.0001782602423671078 0.001572887742816197 0.0009669269619178688 0.0025982218580254167 0.00015167827544595842 0.0013517075511813733 0.00205691555242792 0.0006073052864186683 0.0020412513788273775 0.003393092480228378 0.0018897397822988816 0.001149990511353958 0.0010821906311782121 0.002851342859243226 0.0014734653533826877 0.0016072821764510686 0.0019847898616487747 0.0019621598795614608 0.0014846721652844804 0.0007214393476589174 0.0008889890825044406 0.0023808823038167274 0.0011340363384867312 0.0014418176734341326 0.0014944995775858685 0.001990670436878294 0.002006162646462098 2.069100195298758e-8 0.0009544520116727861 0.0017156444671249482 0.0005094174546272584 0.0010241570292064117 0.0010685199587167325 0.0004952150553333657 0.0004213770763419892 0.002216648793494301 0.0026943170022276945 0.001834457520298822 0.0019885956520496236 0.0020503542686098673 0.0006304769498759658 0.0014401053336811508 6.302480035028645e-5 0.001782867032118198 0.0012045258835450886 0.0004823897040521397 0.0005625297908514997 0.001440197815383483 0.0013486029633553785 1.9390007323059745e-7 0.0005823586970054247 0.0024205585903942714 0.0011474378590624731 3.8045824671111825e-7 7.693052885519039e-6 0.0004444258117343846 3.781876683450104e-8 0.0005520312008172473 0.0015434974485103474 0.002117148334089267 0.0001070764197520101 0.0018981206819769396 0.0017519276443864714 0.0014854224702340135 0.000221177581938812 7.633865306674882e-8 1.0009620529554144e-8 0.0018401050947693945 0.0017163227800480927 0.000931208569586141 0.0015806656968959905 0.003321682773356534 0.0016368228694436297 0.0009913897099568239 1.473166501306075e-8 0.0014498251618949955 1.3948583551688246e-8 0.0009316047544549441 5.937529256299154e-9 0.00193065769549726 0.0005305235461135298 0.0006493219161140218 2.8573732186876556e-6 0.0014537525382374566 0.00015266633814638752 0.0002541750223931143 0.0009667350373490413 0.0015993939952722888 2.8402054409721615e-5 9.443498416971869e-5 0.0008573864516430176 0.0007455686150873769 7.320381847940819e-5 0.0007127251076300434 1.944331045862745e-5 0.0012674457981787619 0.0019188022246958497 0.0001402654102448549 9.8748760962498e-9 0.002054071260533966 0.00032206338860479985 5.2255491489495424e-5 0.00028572546709341377 5.8588948699880206e-9 2.930928263810757e-7 0.0008615162795497163 1.6969783960814416e-8 4.677856600305364e-6 6.128424942519684e-6 5.147568343143185e-7 1.5363778323122724e-8 1.719564046229023e-8 4.2253200728789876e-7 0.0030336343505796854 2.6013959898789606e-8 7.408058897525929e-7 2.5964703214561624e-8 1.4781902086346327e-7 0.0016789457944545157 4.0792785976707164e-9 1.184126431594396e-8 0.0014212035553242807 7.918256581302968e-9 1.376527069421637e-8 1.2551192458839895e-6 6.625738746305198e-8 1.1051207905687645e-8 1.0112057334760872e-8 8.763479755492427e-9 3.170223736631503e-8 7.754808371656717e-8 1.0753547828904632e-7 0.0001441470990366897 4.667039213137739e-9 8.086761919908992e-9 3.1025670985930674e-9 2.0358879212685844e-9 2.4875895111223026e-8 7.94646708984332e-9 4.802628614294363e-9 6.479091082103819e-9 3.186251897212768e-8 1.3168881627474778e-8 4.767294979004821e-9 1.314198741983399e-8 1.47370549584402e-8 5.817726412970424e-9 9.867884881361152e-7 3.2943264682696006e-8 5.410852189184902e-9 1.352715654281171e-8 3.635064422711831e-9 3.026385420868993e-9 5.699668300672657e-9 2.567713288161805e-8 3.946371264193405e-9 3.6927865416000654e-9 2.838089517254119e-9 3.962423881746877e-9]
user_initial_guess = user_initial_guess'


function write_pdb(data::Data, name_file::AbstractString)
    io = open(name_file, "w")
    write_pdb_CRYST1(io, data)
    write_pdb_ATOM(io, data)
    write_pdb_END(io)
    close(io)
end

function write_pdb_CRYST1(io, data::Data)
    box_size = data.data_basic.box_size[:, 2] - data.data_basic.box_size[:, 1]
    box_tilt = data.data_basic.box_tilt
    x = sqrt(box_size[1]^2 + box_tilt[1]^2 + box_tilt[2]^2)
    y = sqrt(box_size[2]^2 + box_tilt[3]^2)
    z = box_size[3]

    @printf(io, "CRYST1")
    @printf(io, "%9.3f", x)
    @printf(io, "%9.3f", y)
    @printf(io, "%9.3f", z)
    @printf(io, "%7.2f", 90)
    @printf(io, "%7.2f", 90)
    @printf(io, "%7.2f", 90)
    @printf(io, "\n")
end

function write_pdb_ATOM(io, data::Data_Unit)
    atom_vec = data.vec_atom
    atom_name = data.data_str.atom_name
    res_name = data.data_str.res_name
    if isa(data.data_str, Structure_Wat)
        type_name = "WT1"
        chain_id = "W"
    elseif isa(data.data_str, Structure_Ion)
        type_name = "ION"
        chain_id = "I"
    else
        type_name = "U0"
        chain_id = "U"
    end
    for (id_atom, atom) in enumerate(atom_vec)
        @printf(io, "%-6s", "ATOM")                            #  1 -  6        Record name   "ATOM  "
        @printf(io, "%5d", atom.atom)                   #  7 - 11        Integer       serial       Atom  serial number.
        @printf(io, "%2s%-3s", " ", uppercase(atom_name[atom.typ]))
        #@printf(io, "%5s", atom_name[atom.typ])        # 13 - 16        Atom          name         Atom name.
        #@printf(io, " ")                                # 17             Character     altLoc       Alternate location indicator.
        @printf(io, "%1s%-4s", " ", res_name)                    # 18 - 20        Residue name  resName      Residue name.
        @printf(io, "%1s", chain_id)                   # 22             Character     chainID      Chain identifier.
        @printf(io, "%4d", atom.mol)                    # 23 - 26        Integer       resSeq       Residue sequence number.
        @printf(io, "%4s", " ")                        # 27             AChar         iCode        Code for insertion of residues.
        @printf(io, "%8.3f", atom.coord[1])             # 31 - 38        Real(8.3)     x            Orthogonal coordinates for X in Angstroms.
        @printf(io, "%8.3f", atom.coord[2])             # 39 - 46        Real(8.3)     y            Orthogonal coordinates for Y in Angstroms.
        @printf(io, "%8.3f", atom.coord[3])             # 47 - 54        Real(8.3)     z            Orthogonal coordinates for Z in Angstroms.
        @printf(io, "%6.2f", 1.0)                       # 55 - 60        Real(6.2)     occupancy    Occupancy.
        @printf(io, "%6.2f", 0.0)                       # 61 - 66        Real(6.2)     tempFactor   Temperature  factor.
        @printf(io, "%6s", " ")
        @printf(io, "%-3s%3s", type_name, uppercase(atom_name[atom.typ][1:end-1]))# 73 - 78        LString(2)    element      Element symbol, right-justified.
        # 79 - 80        LString(2)    charge       Charge  on the atom.
        @printf(io, "\n")
    end
end

function write_pdb_ATOM(io, data::Data_Sum)
    atom_vec = data.vec_atom
    atom_name = []
    res_name = []
    type_name = []
    chain_name = []
    chain_id = []
    for (id, str) in enumerate(data.vec_str)
        atom_name = vcat(atom_name, str.atom_name)
        res_name = vcat(res_name, str.res_name) # Residue name
        append!(chain_id, ones(Int64, str.num_atom_types)*id)
        if isa(str, Structure_Wat)
            type_name = vcat(type_name, "WT1")
            chain_name = vcat(chain_name, "W")
        elseif isa(str, Structure_Ion)
            type_name = vcat(type_name, "ION")
            chain_name = vcat(chain_name, "I")
        else
            type_name = vcat(type_name, "U0")
            chain_name = vcat(chain_name, "U")
        end
    end
    for (id_atom, atom) in enumerate(atom_vec)
        @printf(io, "%-6s", "ATOM")                                         #  1 -  6        Record name   "ATOM  "
        @printf(io, "%5d", atom.atom)                                       #  7 - 11        Integer       serial       Atom  serial number.
        @printf(io, "%2s%-3s", " ", uppercase(atom_name[atom.typ]))         # 13 - 16        Atom          name         Atom name.
        @printf(io, "%1s%-4s", " ",res_name[chain_id[atom.typ]])            # 17 - 20        Residue name  resName      Residue name.
        @printf(io, "%1s", chain_name[chain_id[atom.typ]])                  # 22             Character     chainID      Chain identifier.
        @printf(io, "%4d", atom.mol)                                        # 23 - 26        Integer       resSeq       Residue sequence number.
        @printf(io, "%4s", " ")                                             # 27             AChar         iCode        Code for insertion of residues.
        @printf(io, "%8.3f", atom.coord[1])                                 # 31 - 38        Real(8.3)     x            Orthogonal coordinates for X in Angstroms.
        @printf(io, "%8.3f", atom.coord[2])                                 # 39 - 46        Real(8.3)     y            Orthogonal coordinates for Y in Angstroms.
        @printf(io, "%8.3f", atom.coord[3])                                 # 47 - 54        Real(8.3)     z            Orthogonal coordinates for Z in Angstroms.
        @printf(io, "%6.2f", 0.5)                                           # 55 - 60        Real(6.2)     occupancy    Occupancy.
        @printf(io, "%6.2f", 300)                                           # 61 - 66        Real(6.2)     tempFactor   Temperature  factor.
        @printf(io, "%6s", " ")
        @printf(io, "%-3s%3s", type_name[chain_id[atom.typ]], uppercase(atom_name[atom.typ][1:end-1]))# 73 - 78        LString(2)    element      Element symbol, right-justified.
        #@printf(io, "%9s%6s", type_name[chain_id[atom.typ]], atom_name[atom.typ])# 73 - 78        LString(2)    element      Element symbol, right-justified.
        # 79 - 80        LString(2)    charge       Charge  on the atom.
        @printf(io, "\n")
    end
end

function write_pdb_END(io)
    @printf(io, "END")
end


### Indexing utils

Base.@propagate_inbounds rearrange(x::Tuple, val::Val{D}) where {D} = rearrange((), x, val)
Base.@propagate_inbounds rearrange(x::Tuple, y::Tuple, val::Val{D}) where {D} = rearrange((x...,first(y)), Base.tail(y), Val{D-1}())
Base.@propagate_inbounds rearrange(x::NTuple{N,Int}, y::Tuple, val::Val{0}) where {N} = CartesianIndex{N}(x),y...

Base.@propagate_inbounds Base.getindex(lattice::RegularLattice, i::Int, I...) = Base.getindex(lattice::RegularLattice, (i, I...))
Base.@propagate_inbounds Base.getindex(lattice::RegularLattice{D,T,PB,CT,L}, I::NTuple{D′,Int}) where {D,T,PB,CT,L,D′} = Base.getindex(lattice, rearrange(I, Val{D}())...)

Base.@propagate_inbounds Base.getindex(lattice::RegularLattice, i::Tuple{CartesianIndex, Vararg{Int}}) = Base.getindex(lattice, i...)

Base.@propagate_inbounds relative_coordinate(lattice::RegularLattice{D}, i1::NTuple{D′, Int}, i2::NTuple{D′, Int}) where {D, D′} = relative_coordinate(lattice, rearrange(i1, Val{D}()), rearrange(i2, Val{D}()))

### Sorting utils

"""
`takes_precedence(i1, i2)`

Returns `true` if the index `i1` occurs before `i2` while iterating over a node collection.
"""
@inline takes_precedence(i1::T, i2::T) where {T<:Union{Int, CartesianIndex}} = isless(i1, i2)
@inline takes_precedence(i1::T, i2::T) where {T<:NTuple{D,Int}} where D = takes_precedence(CartesianIndex(i1), CartesianIndex(i2))
@inline takes_precedence(i1::T, i2::T) where {T<:Tuple{CartesianIndex, Int, Vararg{Int}}} = takes_precedence(CartesianIndex(first(i1).I..., Base.tail(i1)...), CartesianIndex(first(i2).I..., Base.tail(i2)...))

### Iteration utils

Base.eachindex(cell::AbstractCell) = cell |> length |> Base.OneTo
Base.eachindex(lattice::RegularLattice) = Iterators.product(CartesianIndices(lattice.dims), eachindex(lattice.unit_cell))


using ITensors
using ITensors.NDTensors
using Test

@testset "SymmetryStyle trait" begin
  i = Index(2)
  iqn = Index([QN(0) => 1, QN(1) => 2])

  @test @inferred(ITensors.symmetrystyle(i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i,))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i])) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(i', i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i', i))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i', i])) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(i'', i', i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i'', i', i))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i'', i', i])) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(i''', i'', i', i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i''', i'', i', i))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i''', i'', i', i])) == ITensors.NonQN()

  @test @inferred(ITensors.symmetrystyle(iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((iqn,))) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle([iqn])) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(iqn', iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((iqn', iqn))) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle([iqn', iqn])) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(iqn'', iqn', iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((iqn'', iqn', iqn))) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle([iqn'', iqn', iqn])) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(iqn', i)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(i', i, iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((i', i, iqn))) == ITensors.HasQNs()
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle([i', i, iqn])) ==
    ITensors.HasQNs()

  A = randomITensor(i', dag(i))
  Aqn = randomITensor(iqn', dag(iqn))

  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(A)) == ITensors.NonQN()
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(Aqn)) == ITensors.HasQNs()

  T = Tensor(A)
  Tqn = Tensor(Aqn)

  @test @inferred(ITensors.symmetrystyle(T)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(Tqn)) == ITensors.HasQNs()

  sqn = siteinds("S=1/2", 10; conserve_qns=true)
  s = removeqns(sqn)
  psi = MPS(s)
  psiqn = MPS(sqn)
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(psi)) == ITensors.NonQN()
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(psiqn)) ==
    ITensors.HasQNs()
end


@testset "FileDFG Tests" begin

    if typeof(dfg) <: CloudGraphsDFG
        @warn "TEST: Nuking all data for user '$(dfg.userId)', robot '$(dfg.robotId)'!"
        clearRobot!!(dfg)
    end

    # Same graph as iifInterfaceTests.jl
    numNodes = 10

    #change ready and backendset for x7,x8 for improved tests on x7x8f1
    verts = map(n -> addVariable!(dfg, Symbol("x$n"), ContinuousScalar, labels = [:POSE]), 1:numNodes)
    #TODO fix this to use accessors
    verts[7].ready = 1
    # verts[7].backendset = 0
    verts[8].ready = 0
    verts[8].backendset = 1
    #call update to set it on cloud
    updateVariable!(dfg, verts[7])
    updateVariable!(dfg, verts[8])

    facts = map(n -> addFactor!(dfg, [verts[n], verts[n+1]], LinearConditional(Normal(50.0,2.0))), 1:(numNodes-1))

    # Save and load the graph to test.
    saveFolder = "/tmp/fileDFG"
    saveDFG(dfg, saveFolder)

    copyDfg = DistributedFactorGraphs._getDuplicatedEmptyDFG(dfg)
    retDFG = loadDFG(saveFolder, Main, copyDfg)
    @test symdiff(ls(dfg), ls(retDFG)) == []
    @test symdiff(lsf(dfg), lsf(retDFG)) == []
    for var in ls(dfg)
        @test getVariable(dfg, var) == getVariable(retDFG, var)
    end
    for fact in lsf(dfg)
        @test getFactor(dfg, fact) == getFactor(retDFG, fact)
    end
end


import ThreeJS
import FileIO: load
import Reactive: every

ascent = load("assets/ascent.png")
a = reinterpret(UInt8, ascent.data)

images = cycle((a, a[257:512, 257:512], a[292:375, 355:463]))

main(window) =  begin
    push!(window.assets,("ThreeJS","threejs"))
    push!(window.assets,"nested-props")

    i = start(images)
    map(map(t -> next(images, i), every(1))) do pair
        a, i = pair
        ThreeJS.outerdiv() <<
        (ThreeJS.initscene() <<
            [
                ThreeJS.mesh(0.0, 0.0, 0.0) <<
                [
                    ThreeJS.plane(convert(Float64, size(a, 1)), convert(Float64, size(a, 2))),
                    ThreeJS.shadermaterial(open(readall, "assets/datatexture.vert", "r"), open(readall, "assets/datatexture.frag", "r")) <<
                        ThreeJS.datatexture("data", a; :minfilter => "LinearFilter")
                ],
                ThreeJS.camera(0.0, 0.0, 768.0)
            ]
        )
    end
end


import ParticleSwarmOptimizer
const pso = ParticleSwarmOptimizer
using Test

@testset "ParticleSwarmOptimizer" begin
    include("TestFunctions.jl")
    include("helper.jl")
    include("neighbourhood.jl")
    include("pso.jl")
end


module GradientStarvation

# activate environment
using Pkg
if isfile("Project.toml") && isfile("Manifest.toml")
    Pkg.activate(".")
end

using DataStructures

include("data.jl")
include("experiments.jl")
include("figures.jl")
include("neuralnetwork.jl")
include("plotlib.jl")
include("plotutils.jl")
include("regularization.jl")
include("structs.jl")

# includes & export needed for documentation
export Data, Figures, Experiments, NeuralNetwork, PlotLib, PlotUtils, Regularization

# figure 1
Figures.create_figure_1("medium")

# figures 2-13
# run different experiments on the 2D classification task (moons)
experiments = OrderedDict(
    "GD Δ0.5" => Experiment("GD", 0.5, false),
    "GD Δ1.0" => Experiment("GD", 1.0, false),
    "GD Δ0.5 + SD" => Experiment("GD", 0.5, true),
    "GD Δ1.0 + SD" => Experiment("GD", 1.0, true),

    "SGD Δ0.5" => Experiment("SGD", 0.5, false),
    "SGD Δ1.0" => Experiment("SGD", 1.0, false),
    "SGD Δ0.5 + SD" => Experiment("SGD", 0.5, true),
    "SGD Δ1.0 + SD" => Experiment("SGD", 1.0, true),

    "WD Δ0.5" => Experiment("WD", 0.5, false),
    "WD Δ1.0" => Experiment("WD", 1.0, false),
    "WD Δ0.5 + SD" => Experiment("WD", 0.5, true),
    "WD Δ1.0 + SD" => Experiment("WD", 1.0, true),

    "ADAM Δ1.0" => Experiment("ADAM", 1.0, false),
    "ADAM Δ0.5" => Experiment("ADAM", 0.5, false),
    "ADAM Δ0.5 + SD" => Experiment("ADAM", 0.5, true),
    "ADAM Δ1.0 + SD" => Experiment("ADAM", 1.0, true),
)
Experiments.run_experiments(experiments, show_plots = false, random_moons = false)

# additional experiments with different parameters

# figure 14
Experiments.run_experiments(
    OrderedDict("GD Δ0.5 + SD 10000" => Experiment("GD", 0.5, true)),
    show_plots = false, random_moons = false, epochs = 10000
)
# figure 15, 16
add_exp_adam = OrderedDict("ADAM Δ0.5 + SD" => Experiment("ADAM", 0.5, true))
Experiments.run_experiments(add_exp_adam, show_plots = false, random_moons = false, learning_rate = 1e-4, prefix = "1e-4")
Experiments.run_experiments(add_exp_adam, show_plots = false, random_moons = false, learning_rate = 1e-3, prefix = "1e-3")

end