################################################################################
#
#  Small groups
#
################################################################################

using Pkg.Artifacts

#export number_of_small_groups, small_groups_limit
export small_group
#CF: as of now, nothing works here and the exports cause errors in Oscar

mutable struct SmallGroupDB
  path::String
  max_order::Int
  db::Vector{Vector{NamedTuple{(:id, :name, :gens, :rels, :nontrivrels,
                                :orderdis, :ordersubdis, :isabelian, :iscyclic,
                                :issolvable, :isnilpotent, :autorder,
                                :aut_gens, :nchars, :dims, :schur, :galrep,
                                :fields, :mod),
                              Tuple{Tuple{Int, Int}, String, Vector{Perm{Int}},
                                    Vector{Vector{Int}}, Vector{Vector{Int}},
                                    Vector{Tuple{Int,Int}},
                                    Vector{Tuple{Int,Int}},
                                    Bool, Bool, Bool, Bool, fmpz,
                                    Vector{Vector{Vector{Int}}}, Int,
                                    Vector{Int},Vector{Int}, Vector{Int},
                                    Vector{Vector{fmpq}},
                                    Vector{Vector{Vector{Vector{fmpq}}}}}}}}

end

function SmallGroupDB(path::String)
  db = Vector{NamedTuple{(:id, :name, :gens, :rels, :nontrivrels, :orderdis,
                          :ordersubdis, :isabelian, :iscyclic, :issolvable,
                          :isnilpotent, :autorder, :aut_gens, :nchars, :dims,
                          :schur, :galrep, :fields, :mod),
                         Tuple{Tuple{Int, Int}, String, Vector{Perm{Int}},
                               Vector{Vector{Int}}, Vector{Vector{Int}},
                               Vector{Tuple{Int,Int}},
                               Vector{Tuple{Int,Int}}, Bool, Bool, Bool,
                               Bool, fmpz, Vector{Vector{Vector{Int}}}, Int,
                               Vector{Int},Vector{Int}, Vector{Int},
                               Vector{Vector{fmpq}},
                               Vector{Vector{Vector{Vector{fmpq}}}}}}}[]
  z = eltype(eltype(db))[]
  open(path) do io
    while !eof(io)
    e = _parse_row(io)
    if isempty(z) || e.id[1] == z[1].id[1]
      push!(z, e)
    else
      push!(db, z)
      z = eltype(eltype(db))[]
      push!(z, e)
    end
  end
  push!(db, z)
  end
  max_order = length(db)
  return SmallGroupDB(path, max_order, db)
end

function SmallGroupDBLegacy(path::String)

  f = function(z)
    if z isa BigInt
      return fmpz(z)
    elseif z isa Rational{BigInt}
      return fmpq()
    elseif z isa Vector{Vector{Rational{BigInt}}}
      return Vector{fmpq}[ fmpq.(v) for v in z]
    elseif z isa Vector{Vector{Vector{Vector{Rational{BigInt}}}}}
      zz = Vector{Vector{Vector{Vector{fmpq}}}}(undef, length(z))
      for i in 1:length(z)
        zz[i] = Vector{Vector{Vector{fmpq}}}(undef, length(z[i]))
        for j in 1:length(z[i])
          zz[i][j] = Vector{Vector{fmpq}}(undef, length(z[i][j]))
          for k in 1:length(z[i][j])
            zz[i][j][k] = fmpq.(z[i][j][k])
          end
        end
      end
      return zz
    else
      return z
    end
  end

  dbnew = Vector{NamedTuple{(:id, :name, :gens, :rels, :nontrivrels, :orderdis,
                          :ordersubdis, :isabelian, :iscyclic, :issolvable,
                          :isnilpotent, :autorder, :aut_gens, :nchars, :dims,
                          :schur, :galrep, :fields, :mod),
                         Tuple{Tuple{Int, Int}, String, Vector{Perm{Int}},
                               Vector{Vector{Int}}, Vector{Vector{Int}},
                               Vector{Tuple{Int,Int}},
                               Vector{Tuple{Int,Int}}, Bool, Bool, Bool,
                               Bool, fmpz, Vector{Vector{Vector{Int}}}, Int,
                               Vector{Int},Vector{Int}, Vector{Int},
                               Vector{Vector{fmpq}},
                               Vector{Vector{Vector{Vector{fmpq}}}}}}}[]

  db = Hecke.eval(Meta.parse(Base.read(path, String)))
  for i in 1:length(db)
    push!(dbnew, eltype(dbnew)(undef, length(db[i])))
    for j in 1:length(db[i])
      v = db[i][j]
      vv = map(f, v)
      dbnew[i][j] = (id = (i, j), vv...)
    end
  end

  max_order = length(dbnew)
  return SmallGroupDB(path, max_order, dbnew)
end

function show(io::IO, L::SmallGroupDB)
  print(io, "Database of small groups (order limit = ", L.max_order, ")")
end

const legacy_default_small_group_db = [joinpath(pkgdir, "data/small_groups_extended"), joinpath(pkgdir, "data/small_groups_default")]

function small_group_database()
  st = artifact"SmallGroupDB"
  return SmallGroupDB(joinpath(st, "SmallGroupDB", "data"))
end

const _DefaultSmallGroupDB = Ref{Any}(nothing)

function DefaultSmallGroupDB()
  _DB = _DefaultSmallGroupDB[]
  if _DefaultSmallGroupDB[] === nothing
    DB = small_group_database()
    _DefaultSmallGroupDB[] = DB
    return DB::SmallGroupDB
  else
    return _DB::SmallGroupDB
  end
end

isfrom_db(G::GrpGen) = G.isfromdb

function small_group(i, j; DB = DefaultSmallGroupDB())
  data = DB.db[i][j]
  #i < 1 || i > 63 && error("Group order ($i) must be between 1 and $small_groups_limit")
  #j < 1 || j > number_of_small_groups(i) && error("Index ($j) must be between 1 and $(number_of_small_groups(i))")
  P = SymmetricGroup(i)
  Gens = [one(P)]
  for p in data.gens
    push!(Gens, P(p))
  end
  G, PtoG, GtoP = generic_group(closure(Gens, *, one(P)), *)

  # The small groups have generators, for which we know
  # generators for the relations.
  # This is used in the homomorphism computation.
  # To make this work, the generators must much exactly
  # with the generators in the database.
  G.gens = Vector{Int}(undef, length(data.gens))
  for k in 1:length(data.gens)
    p = data.gens[k]
    G.gens[k] = PtoG[P(p)][]
  end

  G.issolvable = data.issolvable
  G.isnilpotent = data.isnilpotent
  G.isfromdb = true
  G.small_group_id = (i, j)
  return G
end

################################################################################
#
#  Better parsing
#
################################################################################

function _parse_row(io::IO)
  b = Base.read(io, UInt8)
  # id
  b, id = _parse(Tuple{Int, Int}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, name = _parse(String, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, gens = _parse(Vector{Perm{Int}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, rels = _parse(Vector{Vector{Int}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, nontrivrels = _parse(Vector{Vector{Int}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, orderdis = _parse(Vector{Tuple{Int, Int}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, ordersubdis = _parse(Vector{Tuple{Int, Int}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, isabelian = _parse(Bool, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, iscyclic = _parse(Bool, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, issolvable = _parse(Bool, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, isnilpotent = _parse(Bool, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, autorder = _parse(fmpz, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, aut_gens = _parse(Vector{Vector{Vector{Int}}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, nchars = _parse(Int, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, dims = _parse(Vector{Int}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, schur = _parse(Vector{Int}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, galrep = _parse(Vector{Int}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, fields = _parse(Vector{Vector{fmpq}}, io, b)
  @assert b == UInt(',')
  b = Base.read(io, UInt8)
  b, mod = _parse(Vector{Vector{Vector{Vector{fmpq}}}}, io, b)
  return (id = id, name = name, gens = gens, rels = rels, nontrivrels = nontrivrels, orderdis = orderdis,
          ordersubdis = ordersubdis, isabelian = isabelian, iscyclic = iscyclic, issolvable = issolvable,
          isnilpotent = isnilpotent, autorder = autorder, aut_gens = aut_gens, nchars = nchars, dims = dims,
          schur = schur, galrep = galrep, fields = fields, mod = mod)
end


"""
    get_state_series(
        res::SimulationResults,
        ref::Tuple{String, Symbol}
    )
    end

Function to obtain series of states out of DAE Solution.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
- `ref:Tuple{String, Symbol}` : Tuple used to identify the dynamic device, via its name, as a `String`, and the associated state as a `Symbol`.
"""
function get_state_series(res::SimulationResults, ref::Tuple{String, Symbol})
    if allunique(res.solution.t)
        return res.solution.t, post_proc_state_series(res, ref)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], post_proc_state_series(res, ref)[ix]
    end
    return
end

"""
    get_voltage_magnitude_series(
        res::SimulationResults,
        bus_number::Int
    )

Function to obtain the voltage magnitude series out of the DAE Solution.

# Arguments:

- `res::SimulationResults` : Simulation Results object that contains the solution
- `bus_number::Int` : Bus number identifier
"""
function get_voltage_magnitude_series(res::SimulationResults, bus_number::Int)
    n_buses = get_bus_count(res)
    bus_ix = get(get_bus_lookup(res), bus_number, 0)
    V_R, V_I = post_proc_voltage_series(res.solution, bus_ix, n_buses)
    if allunique(res.solution.t)
        return res.solution.t, sqrt.(V_R .^ 2 + V_I .^ 2)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], sqrt.(V_R[ix] .^ 2 + V_I[ix] .^ 2)
    end
end

"""
    get_voltage_angle_series(
        res::SimulationResults,
        bus_number::Int
    )

Function to obtain the voltage angle series out of the DAE Solution.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
- `bus_number::Int` : Bus number identifier
"""
function get_voltage_angle_series(res::SimulationResults, bus_number::Int)
    n_buses = get_bus_count(res)
    bus_ix = get(get_bus_lookup(res), bus_number, 0)
    V_R, V_I = post_proc_voltage_series(res.solution, bus_ix, n_buses)
    if allunique(res.solution.t)
        return res.solution.t, atan.(V_I ./ V_R)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], atan.(V_I[ix] ./ V_R[ix])
    end
end

"""
    get_real_current_series(
            res::SimulationResults,
            name::String,
    )

Function to obtain the real current time series of a Dynamic Injection series out of the DAE Solution.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
- `name::String` : Name to identify the specified device
"""
function get_real_current_series(res::SimulationResults, name::String)
    if allunique(res.solution.t)
        return res.solution.t, post_proc_real_current_series(res, name)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], post_proc_real_current_series(res, name)[ix]
    end
end

"""
    get_imaginary_current_series(
            res::SimulationResults,
            name::String,
    )

Function to obtain the imaginary current time series of a Dynamic Injection series out of the DAE Solution.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
- `name::String` : Name to identify the specified device
"""
function get_imaginary_current_series(res::SimulationResults, name::String)
    if allunique(res.solution.t)
        return res.solution.t, post_proc_imaginary_current_series(res, name)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], post_proc_imaginary_current_series(res, name)[ix]
    end
end

"""
    get_activepower_series(
            res::SimulationResults,
            name::String,
    )

Function to obtain the active power output time series of a Dynamic Injection series out of the DAE Solution.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
- `name::String` : Name to identify the specified device
"""
function get_activepower_series(res::SimulationResults, name::String)
    if allunique(res.solution.t)
        return res.solution.t, post_proc_activepower_series(res, name)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], post_proc_activepower_series(res, name)[ix]
    end
end

"""
    get_reactivepower_series(
            res::SimulationResults,
            name::String,
    )

Function to obtain the reactive power output time series of a Dynamic Injection series out of the DAE Solution.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
- `name::String` : Name to identify the specified device
"""
function get_reactivepower_series(res::SimulationResults, name::String)
    if allunique(res.solution.t)
        return res.solution.t, post_proc_reactivepower_series(res, name)
    else
        @debug "found repeated time steps removing repetitions"
        ix = unique(i -> res.solution.t[i], eachindex(res.solution.t))
        return res.solution.t[ix], post_proc_reactivepower_series(res, name)[ix]
    end
end

"""
    show_states_initial_value(res::SimulationResults)

Function to print initial states.

# Arguments

- `res::SimulationResults` : Simulation Results object that contains the solution
"""
function show_states_initial_value(res::SimulationResults)
    bus_size = get_bus_count(res)
    system = get_system(res)
    x0_init = res.solution.u[1]
    println("Voltage Variables")
    println("====================")
    buses_sorted =
        sort(collect(PSY.get_components(PSY.Bus, system)); by = x -> PSY.get_number(x))
    for bus in buses_sorted
        name = PSY.get_name(bus)
        println(name)
        println("====================")
        bus_n = PSY.get_number(bus)
        bus_ix = get_bus_lookup(res)[bus_n]
        V_R = x0_init[bus_ix]
        V_I = x0_init[bus_ix + bus_size]
        Vm = sqrt(V_R^2 + V_I^2)
        θ = angle(V_R + V_I * 1im)
        print("Vm ", round(Vm, digits = 4), "\n")
        print("θ ", round(θ, digits = 4), "\n")
        println("====================")
    end
    println("====================")
    for device in PSY.get_components(PSY.DynamicInjection, system)
        states = PSY.get_states(device)
        name = PSY.get_name(device)
        println("Differential States")
        println(name)
        println("====================")
        global_index = get_global_index(res)[name]
        for s in states
            print(s, " ", round(x0_init[global_index[s]], digits = 4), "\n")
        end
        println("====================")
    end
    dyn_branches = PSY.get_components(PSY.DynamicBranch, system)
    if !isempty(dyn_branches)
        println("====================")
        println("Line Current States")
        for br in dyn_branches
            states = PSY.get_states(br)
            name = PSY.get_name(br)
            printed_name = "Line " * name
            println("====================")
            println(printed_name)
            global_index = get_global_index(res)[name]
            for (i, s) in enumerate(states)
                print(s, " ", round(x0_init[global_index[s]], digits = 5), "\n")
            end
            println("====================")
        end
    end
    return
end

"""
    show_states_initial_value(sim::Simulation)

Function to print initial states.

# Arguments

- `sim::Simulation` : Simulation object that contains the initial condition
"""
function show_states_initial_value(sim::Simulation)
    inputs = sim.inputs
    bus_size = get_bus_count(inputs)
    system = get_system(sim)
    global_state_map = make_global_state_map(inputs)
    println("Voltage Variables")
    println("====================")
    buses_sorted =
        sort(collect(PSY.get_components(PSY.Bus, system)); by = x -> PSY.get_number(x))
    for bus in buses_sorted
        name = PSY.get_name(bus)
        println(name)
        println("====================")
        bus_n = PSY.get_number(bus)
        bus_ix = PSID.get_lookup(sim.inputs)[bus_n]
        V_R = sim.x0_init[bus_ix]
        V_I = sim.x0_init[bus_ix + bus_size]
        Vm = sqrt(V_R^2 + V_I^2)
        θ = angle(V_R + V_I * 1im)
        print("Vm ", round(Vm, digits = 4), "\n")
        print("θ ", round(θ, digits = 4), "\n")
        println("====================")
    end
    println("====================")
    for device in get_dynamic_injectors(inputs)
        states = PSY.get_states(device)
        name = PSY.get_name(device)
        println("Differential States")
        println(name)
        println("====================")
        global_index = global_state_map[name]
        for s in states
            print(s, " ", round(sim.x0_init[global_index[s]], digits = 4), "\n")
        end
        println("====================")
    end
    dyn_branches = get_dynamic_branches(inputs)
    if !isempty(dyn_branches)
        println("====================")
        println("Line Current States")
        for br in dyn_branches
            states = PSY.get_states(br)
            name = PSY.get_name(br)
            printed_name = "Line " * name
            println("====================")
            println(printed_name)
            global_index = global_state_map[name]
            x0_br = Dict{Symbol, Float64}()
            for (i, s) in enumerate(states)
                print(s, " ", round(sim.x0_init[global_index[s]], digits = 5), "\n")
            end
            println("====================")
        end
    end
    return
end


using CNNVisualize, Metalhead

@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

img = load_image("../images/original/cat_dog.png")
model = VGG19()

@testset "BackPropagation" begin
  @test_nowarn begin
    vbprop = VanillaBackprop(model.layers)
    grads = vbprop(img)
    pgrads, ngrads = positive_negative_saliency(grads[1][1])
  end

  @test_nowarn begin
    gbprop = GuidedBackprop(model.layers)
    grads = gbprop(img)
  end

  @test_nowarn begin
    deconv = Deconvolution(model.layers)
    grads = deconv(img, 2) # Test in case multiple visualizations are to be generated
  end

  save_gradient_images(grads[1][1], "../images/results/save1.png")
  save_grayscale_gradient(grads[2][1], "../images/results/save1.png")
end

@testset "Grad CAM" begin
  @test_nowarn begin
    gcam = GradCAM(model.layers, 12)
    grads = gcam(img)
    save_gradcam(grads[1][1], "../images/original/cat_dog.png", "../images/results/save2.png", "../images/results/save3.png", "../images/results/save4.png")
  end

  @test_warn "target should be less than the model length" GradCAM(model.layers, length(model.layer.layers) + 5)

  @test_nowarn begin
    ggcam = GuidedGradCAM(model.layers, 12)
    grads = ggcam(img)
  end
end


using IGC: year_to_short_string, number_to_char, char_to_number
using IGC: IGCLongFilename, IGCShortFilename, filename

@testset "filename" begin
    @testset "short file name" begin
        @testset "year" begin
            @test year_to_short_string(2019) == "9"
        end

        @testset "number_to_char" begin
            @test number_to_char(0) == '0'
            @test number_to_char(9) == '9'
            @test number_to_char(10) == 'A'
            @test number_to_char(31) == 'V'
            @test number_to_char(35) == 'Z'
        end

        @testset "char_to_number" begin
            @test char_to_number('0') == 0
            @test char_to_number('9') == 9
            @test char_to_number('A') == 10
            @test char_to_number('V') == 31
            @test char_to_number('Z') == 35
        end

        @testset "get short file name" begin
            f = IGCShortFilename(
                Date(2015, 06, 17),
                'C',
                "XXX",
                2
            )
            fname_expected = "56HCXXX2.IGC"
            returned_filename = filename(f)
            @test returned_filename == fname_expected

            parsed_filename = parse(IGCShortFilename, fname_expected, year=2010)
            @test parsed_filename.date == f.date
            @test parsed_filename.manufacturer == f.manufacturer
            @test parsed_filename.serial_id == f.serial_id
            @test parsed_filename.flight_of_day == f.flight_of_day
        end
    end

    @testset "get long file name" begin
        f = IGCLongFilename(
            Date(2015, 06, 17),
            "MMM",
            "XXX",
            2
        )
        fname_expected = "2015-06-17-MMM-XXX-02.IGC"
        returned_filename = filename(f)
        @test returned_filename == fname_expected

        parsed_filename = parse(IGCLongFilename, fname_expected)
        @test parsed_filename.date == f.date
        @test parsed_filename.manufacturer == f.manufacturer
        @test parsed_filename.serial_id == f.serial_id
        @test parsed_filename.flight_of_day == f.flight_of_day
    end
end


abstract type Material end

function emit(material::Material, u::AbstractFloat, v::AbstractFloat, point::AbstractVec)::AbstractVec
    Vec(0.0, 0.0, 0.0)
end


module Wire

using ..Ahorn, Maple

const placements = Ahorn.PlacementDict(
    "Wire" => Ahorn.EntityPlacement(
        Maple.Wire,
        "line"
    )
)

Ahorn.nodeLimits(entity::Maple.Wire) = 1, 1

function Ahorn.selection(entity::Maple.Wire)
    nodes = get(entity.data, "nodes", ())
    x, y = Ahorn.position(entity)

    res = Ahorn.Rectangle[Ahorn.Rectangle(x - 4, y - 4, 8, 8)]

    for node in nodes
        nx, ny = node

        push!(res, Ahorn.Rectangle(nx - 4, ny - 4, 8, 8))
    end

    return res
end

wireColorHex = "595866"
wireColor = (89, 88, 102, 1) ./ (255, 255, 255, 1)
wireColorSelected = (Ahorn.colors.selection_selected_fc)

function renderWire(ctx::Ahorn.Cairo.CairoContext, entity::Maple.Wire, color::Ahorn.colorTupleType=wireColor)
    x, y = Ahorn.position(entity)

    start = (x, y)
    stop = get(entity.data, "nodes", [start])[1]
    control = (start .+ stop) ./ 2 .+ (0, 24)

    curve = Ahorn.SimpleCurve(start, stop, control)
    Ahorn.drawSimpleCurve(ctx, curve, color, thickness=1)
end

Ahorn.renderSelectedAbs(ctx::Ahorn.Cairo.CairoContext, entity::Maple.Wire) = renderWire(ctx, entity, wireColorSelected)

function Ahorn.renderAbs(ctx::Ahorn.Cairo.CairoContext, entity::Maple.Wire, room::Maple.Room)
    # Make sure Alpha is 1
    rawColor = Ahorn.argb32ToRGBATuple(parse(Int, get(entity.data, "color", wireColorHex), base=16))[1:3] ./ 255
    color = (rawColor..., 1.0)

    renderWire(ctx, entity, color)
end

end


# This file is a part of RRRMC.jl. License is MIT: http://github.com/carlobaldassi/RRRMC.jl/LICENCE.md

module Common

export Vec, Vec2, IVec, IVec2, unsafe_bitflip!, flipbits!, discretize,
       LocalFields, AllButOne

import Base: iterate

const Vec  = Vector{Float64}
const Vec2 = Vector{Vec}
const IVec = Vector{Int}
const IVec2 = Vector{IVec}

@inline function unsafe_bitflip!(Bc::Array{UInt64}, i::Int)
    i1, i2 = Base.get_chunks_id(i)
    u = UInt64(1) << i2
    @inbounds begin
        #Bc[i1] ⊻= u
        Bc[i1] = Bc[i1] ⊻ u
    end
end
@inline unsafe_bitflip!(B::BitArray, i::Int) = unsafe_bitflip!(B.chunks, i)

flipbits!(B::BitArray) = B .= .!(B)

mutable struct LocalFields{ET}
    lfields::Vector{ET}
    lfields_last::Vector{ET}
    move_last::Int
    function LocalFields{ET}(N::Int) where {ET}
        lfields = zeros(ET, N)
        lfields_last = zeros(ET, N)
        return new{ET}(lfields, lfields_last, 0)
    end
end

function discretize(x::T, LEV::Tuple{S,Vararg{S}}) where {T<:Real,S<:Real}
    d = LEV[1]
    r = x - d
    for l = 2:length(LEV)
        d1 = LEV[l]
        r1 = x - d1
        if abs(r1) < abs(r)
            d, r = d1, r1
        end
    end
    return d, r
end

function discretize(cvec::Vector{T}, LEV::Tuple{S,Vararg{S}}) where {T<:Real,S<:Real}
    N = length(cvec)
    fields = Array{S}(undef, N)
    rfields = Array{T}(undef, N)
    for (i,x) in enumerate(cvec)
        d, r = discretize(x, LEV)
        fields[i] = d
        rfields[i] = r
    end
    return fields, rfields
end

function discretize(cvec::NTuple{N,T}, LEV::Tuple{S,Vararg{S}}) where {N,T<:Real,S<:Real}
    fields = Array{S}(undef, N)
    rfields = Array{T}(undef, N)
    for (i,x) in enumerate(cvec)
        d, r = discretize(x, LEV)
        fields[i] = d
        rfields[i] = r
    end
    return tuple(fields...), tuple(rfields...)
end

# iterate over 1:N except i, i.e.
#   1, 2, 3, ..., i-1, i+1, ..., N-1, N
# useful for fully-connected models
mutable struct AllButOne
    N::Int
    i::Int
    function AllButOne(N::Integer, i::Integer)
        N ≥ 1 || throw(ArgumentError("N must me larger than 1, given: $N"))
        1 ≤ i ≤ N || throw(ArgumentError("expected 1 ≤ i ≤ N, given i=$i N=$N"))
        return new(N, i)
    end
end

function iterate(n::AllButOne, j = 1 + (n.i == 1))
    j == n.N + 1 && return nothing
    return j, j + 1 + (j == n.i - 1)
end
Base.length(n::AllButOne) = n.N - 1

end # module


### Abstract parameter NStates

abstract type ParameterNState{S<:ValueSupport, F<:VariateForm} <: VariableNState{F} end

const MarkovChain = ParameterNState

DiscreteParameterNState{F<:VariateForm} = ParameterNState{Discrete, F}
ContinuousParameterNState{F<:VariateForm} = ParameterNState{Continuous, F}

UnivariateParameterNState{S<:ValueSupport} = ParameterNState{S, Univariate}
MultivariateParameterNState{S<:ValueSupport} = ParameterNState{S, Multivariate}
MatrixvariateParameterNState{S<:ValueSupport} = ParameterNState{S, Matrixvariate}

diagnostics(nstate::ParameterNState) =
  Dict(zip(nstate.diagnostickeys, Any[nstate.diagnosticvalues[i, :][:] for i = 1:size(nstate.diagnosticvalues, 1)]))

copy!(nstate::ParameterNState{S, F}, state::ParameterState{S, F}, i::Integer) where {S<:ValueSupport, F<:VariateForm} =
  nstate.copy(nstate, state, i)


module Foo

using Hyperspecialize

import Bar

@concretize AlsoNotAType Set{Type}([Int64])
@concretize Float64 Set{Type}([UInt16])

# Make sure that other modules can also change concretizations
function __init__()
  @widen Bar.Float64 Set{Type}([UInt16, Float32])
end

end #module


struct BlocksInfoF90
    nblocks    :: Cint
    nelectrons :: Cint
    ncsfstotal :: Cint
    norbitals  :: Cint
    nvectotal  :: Cint
    nvecsize   :: Cint
end

struct BlockF90
    blockid       :: Cint
    ncsfs         :: Cint
    nevs          :: Cint
    iatjp         :: Cint
    iaspa         :: Cint
    eav           :: Cdouble
    eigenstates   :: Ptr{Cdouble}
    eigenenergies :: Ptr{Cdouble}
end

struct MixingBlock
    eav :: Float64
    energies :: Vector{Float64}
    states :: Matrix{Float64}
    _block_f90 :: BlockF90
end

struct MixingFile
    nelectrons :: Int
    blocks :: Vector{MixingBlock}
    _blocksinfo_f90 :: BlocksInfoF90
end

function read_rmix(filename)
    isfile(filename) || error("Unable to open $(filename)")

    blockinfo = Ref{BlocksInfoF90}()
    blocks = Ref{Ptr{BlockF90}}()
    status = ccall( (:mixread, libgrasp_so),
        Cint, (Cstring, Ref{BlocksInfoF90}, Ref{Ptr{BlockF90}}),
        filename, blockinfo, blocks
    )
    if status != zero(status)
        error("mixread returned with status $(status)")
    end
    blockinfo = blockinfo.x
    blocks = unsafe_wrap(Array{BlockF90}, blocks.x, blockinfo.nblocks, own=true)
    blocks = map(blocks) do block
        energies = unsafe_wrap(Array{Cdouble}, block.eigenenergies, (block.nevs,), own=true)
        states = unsafe_wrap(Array{Cdouble}, block.eigenstates, (block.ncsfs, block.nevs), own=true)
        MixingBlock(block.eav, energies, states, block)
    end
    MixingFile(blockinfo.nelectrons, blocks, blockinfo)
end


import Base

"""
Struct for WilsonFermion
"""
struct WilsonFermion_4D_wing{NC,NDW} <: WilsonFermion_4D{NC} #AbstractFermionfields_4D{NC}

    f::Array{ComplexF64,6}
    NC::Int64
    NX::Int64
    NY::Int64
    NZ::Int64
    NT::Int64
    NG::Int64
    NDW::Int64
    Dirac_operator::String
    #BoundaryCondition::Vector{Int8}

    function WilsonFermion_4D_wing(NC::T,NX::T,NY::T,NZ::T,NT::T) where T<: Integer
        NG = 4
        NDW = 1
        #@assert NDW == 1 "only NDW = 1 is supported. Now NDW = $NDW"
        f = zeros(ComplexF64,NC,NX+2NDW,NY+2NDW,NZ+2NDW,NT+2NDW,NG)
        Dirac_operator = "Wilson"
        return new{NC,NDW}(f,NC,NX,NY,NZ,NT,NG,NDW,Dirac_operator)
    end

    function WilsonFermion_4D_wing{NC}(NX::T,NY::T,NZ::T,NT::T) where {T<: Integer,NC}
        NG = 4
        NDW = 1
        #@assert NDW == 1 "only NDW = 1 is supported. Now NDW = $NDW"
        f = zeros(ComplexF64,NC,NX+2NDW,NY+2NDW,NZ+2NDW,NT+2NDW,NG)
        Dirac_operator = "Wilson"
        return new{NC,NDW}(f,NC,NX,NY,NZ,NT,NG,NDW,Dirac_operator)
    end

end

function Base.length(x::WilsonFermion_4D_wing{NC,NDW}) where {NC,NDW}
    return NC*x.NX*x.NY*x.NZ*x.NT*x.NG
end

function Base.setindex!(x::WilsonFermion_4D_wing{NC,NDW},v,i1,i2,i3,i4,i5,i6)  where {NC,NDW}
    @inbounds x.f[i1,i2 + NDW,i3 + NDW,i4 + NDW,i5 + NDW,i6] = v
end

function Base.getindex(x::WilsonFermion_4D_wing{NC,NDW},i1,i2,i3,i4,i5,i6) where {NC,NDW}
    @inbounds return x.f[i1,i2 .+ NDW,i3 .+ NDW,i4 .+ NDW,i5 .+ NDW,i6]
end

function Base.getindex(x::WilsonFermion_4D_wing{NC,NDW},i1::N,i2::N,i3::N,i4::N,i5::N,i6::N) where {NC,NDW,N<: Integer}
    @inbounds return  x.f[i1,i2 + NDW,i3 + NDW,i4 + NDW,i5 + NDW,i6]
end

function setindex_global!()
end

#=
function Base.getindex(F::Shifted_fermionfields_4D{NC,WilsonFermion_4D_wing{NC,NDW}},i1::N,i2::N,i3::N,i4::N,i5::N,i6::N)  where {NC,NDW,N <: Integer}

    @inbounds begin
    si2 = i2 + NDW  + F.shift[1]
    si3 = i3 + NDW  + F.shift[2]
    si4 = i4 + NDW  + F.shift[3]
    si5 = i5 + NDW  + F.shift[4]
    end
    #v = F.parent.f
    #return  v[i1,si2,si3,si4,si5,i6]

   return  @inbounds F.parent.f[i1,si2,si3,si4,si5,i6]
end
=#

#=

function Base.getindex(x::WilsonFermion_4D_wing{NC,NDW},i1,i2,i3,i4,i5,i6) where {NC,NDW}
    #=
    i2new = i2 .+ x.NDW
    i3new = i3 .+ x.NDW
    i4new = i4 .+ x.NDW
    i5new = i5 .+ x.NDW
    @inbounds return x.f[i1,i2new,i3new,i4new,i5new,i6]
    =#
    #return @inbounds Base.getindex(x.f,i1,i2 .+ NDW,i3 .+ NDW,i4 .+ NDW,i5 .+ NDW,i6)
    @inbounds return x.f[i1,i2 .+ NDW,i3 .+ NDW,i4 .+ NDW,i5 .+ NDW,i6]
end

=#

function Base.similar(x::T) where T <: WilsonFermion_4D_wing
    return WilsonFermion_4D_wing(x.NC,x.NX,x.NY,x.NZ,x.NT)
end

#=
function Base.getindex(x::T,i1,i2,i3,i4,i5,i6) where T <: WilsonFermion_4D_wing{NC} where NC
    i2new = i2 + x.NDW
    i3new = i3 + x.NDW
    i4new = i4 + x.NDW
    i5new = i5 + x.NDW
    f1 = x.f
    #println((i1,i2new,i3new,i4new,i5new,i6))
    @inbounds v = f1[i1,i2new,i3new,i4new,i5new,i6]
    #x.f[i1,i2new,i3new,i4new,i5new,i6]
    return v
    @inbounds return x.f[i1,i2new,i3new,i4new,i5new,i6]
    #@inbounds return x.f[i1,i2 .+ x.NDW,i3 .+ x.NDW,i4 .+ x.NDW,i5 .+ x.NDW,i6]
end
=#

function set_wing_fermion!(a::WilsonFermion_4D_wing{NC,NDW},boundarycondition) where {NC,NDW}
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX

    #!  X-direction
    for ialpha=1:4
        for it=1:NT
            for iz = 1:NZ
                for iy=1:NY
                    @simd for k=1:NC
                        a[k,0,iy,iz,it,ialpha] = boundarycondition[1]*a[k,NX,iy,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for k=1:NC
                        a[k,NX+1,iy,iz,it,ialpha] = boundarycondition[1]*a[k,1,iy,iz,it,ialpha]
                    end
                end
            end
        end
    end

    #Y-direction
    for ialpha = 1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,0,iz,it,ialpha] = boundarycondition[2]*a[k,ix,NY,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,NY+1,iz,it,ialpha] = boundarycondition[2]*a[k,ix,1,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        # Z-direction
        for it=1:NT
            for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,iy,0,it,ialpha] = boundarycondition[3]*a[k,ix,iy,NZ,it,ialpha]
                        a[k,ix,iy,NZ+1,it,ialpha] = boundarycondition[3]*a[k,ix,iy,1,it,ialpha]

                    end
                end
            end
        end

        #T-direction
        for iz=1:NZ
            for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,iy,iz,0,ialpha] = boundarycondition[4]*a[k,ix,iy,iz,NT,ialpha]
                        a[k,ix,iy,iz,NT+1,ialpha] =boundarycondition[4]*a[k,ix,iy,iz,1,ialpha]
                    end
                end
            end
        end

    end

end

function set_wing_fermion!(a::WilsonFermion_4D_wing{NC,NDW},boundarycondition,iseven::Bool) where {NC,NDW}
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX

    #!  X-direction
    @inbounds for ialpha=1:4
        for it=1:NT
            for iz = 1:NZ
                for iy=1:NY
                    ix = NX
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,0,iy,iz,it,ialpha] = boundarycondition[1]*a[k,NX,iy,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end

    @inbounds for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    ix = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,NX+1,iy,iz,it,ialpha] = boundarycondition[1]*a[k,1,iy,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end

    #Y-direction
    @inbounds for ialpha = 1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    iy = NY
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven

                        @simd for k=1:NC
                            a[k,ix,0,iz,it,ialpha] = boundarycondition[2]*a[k,ix,NY,iz,it,ialpha]
                        end
                    end
                end

            end
        end
    end

    @inbounds for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    iy = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven

                        @simd for k=1:NC
                            a[k,ix,NY+1,iz,it,ialpha] = boundarycondition[2]*a[k,ix,1,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end

    @inbounds for ialpha=1:4
        # Z-direction
        for it=1:NT
            for iy=1:NY
                for ix=1:NX
                    iz = NX
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,ix,iy,0,it,ialpha] = boundarycondition[3]*a[k,ix,iy,NZ,it,ialpha]
                        end
                    end

                    iz = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,ix,iy,NZ+1,it,ialpha] = boundarycondition[3]*a[k,ix,iy,1,it,ialpha]
                        end
                    end
                end
            end
        end

        #T-direction
        for iz=1:NZ
            for iy=1:NY
                for ix=1:NX
                    it = NT
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,ix,iy,iz,0,ialpha] = boundarycondition[4]*a[k,ix,iy,iz,NT,ialpha]
                        end
                    end

                    it = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NC
                            a[k,ix,iy,iz,NT+1,ialpha] =boundarycondition[4]*a[k,ix,iy,iz,1,ialpha]
                        end
                    end

                end
            end
        end

    end

end

function add_fermion!(c::WilsonFermion_4D_wing{NC,NDW},α::Number,a::T1,β::Number,b::T2,iseven) where {NC,NDW,T1 <: Abstractfermion,T2 <: Abstractfermion}#c += alpha*a + beta*b
    n1,n2,n3,n4,n5,n6 = size(c.f)

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5 -NDW
            for i4=1:n4
                iz = i4 -NDW
                for i3=1:n3
                    iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse((ix + iy + iz + it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                #println(a.f[i1,i2,i3,i4,i5,i6],"\t",b.f[i1,i2,i3,i4,i5,i6] )
                                c.f[i1,i2,i3,i4,i5,i6] += α*a.f[i1,i2,i3,i4,i5,i6] + β*b.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
    end
    return
end

function add_fermion!(c::WilsonFermion_4D_wing{NC,NDW},α::Number,a::T1,iseven::Bool) where {NC,NDW,T1 <: Abstractfermion,T2 <: Abstractfermion}#c += alpha*a + beta*b
    n1,n2,n3,n4,n5,n6 = size(c.f)

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5 -NDW
            for i4=1:n4
                iz = i4 -NDW
                for i3=1:n3
                    iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse((ix + iy + iz + it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                #println(a.f[i1,i2,i3,i4,i5,i6],"\t",b.f[i1,i2,i3,i4,i5,i6] )
                                c.f[i1,i2,i3,i4,i5,i6] += α*a.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
    end
    return
end

function WWx!(xout::T,U::Array{G,1},x::T,A)  where  {T <: WilsonFermion_4D_wing,G <: AbstractGaugefields} #(1 - K^2 Teo Toe) xe
    iseven = true
    isodd = false
    temp = A._temporary_fermi[7]#temps[4]
    temp2 = A._temporary_fermi[6]#temps[4]

    #println("Wx")
    #@time Wx!(xout,U,x,A)
    clear_fermion!(xout,iseven)

    #Tx!(temp,U,x,A)
    Toex!(temp,U,x,A,iseven) #Toe
    #Tx!(temp2,U,temp,A)
    Toex!(temp2,U,temp,A,isodd) #Teo

    #set_wing_fermion!(temp,A.boundarycondition)
    #add_fermion!(xout,1,x,-1,temp2)
    add_fermion!(xout,1,x,-1,temp2,iseven)
    set_wing_fermion!(xout,A.boundarycondition,iseven)

    #Wx!(xout,U,x,A)
    return

    #Tx!(temp2,U,x,A) #Toe

    #Toex!(temp2,U,x,A,iseven) #Toe

    #Toex!(temp,U,temp2,A,isodd) #Teo
    #Tx!(temp,U,temp2,A) #Toe

    Tx!(temp,U,x,A) #Toe

    add_fermion!(xout,1,x,-1,temp)
    #add_fermion!(xout,1,x,-1,temp,iseven)

    iseven = true
    set_wing_fermion!(xout,A.boundarycondition)

    return
end

function WWdagx!(xout::T,U::Array{G,1},x::T,A)  where  {T <: WilsonFermion_4D_wing,G <: AbstractGaugefields} #(1 - K^2 Teo Toe) xe
    iseven = true
    isodd = false
    temp = A._temporary_fermi[7]#temps[4]
    temp2 = A._temporary_fermi[6]#temps[4]

    clear_fermion!(xout)

    #Tx!(temp,U,x,A)
    Tdagoex!(temp,U,x,A,iseven) #Toe
    #Tx!(temp2,U,temp,A)
    Tdagoex!(temp2,U,temp,A,isodd) #Teo

    #set_wing_fermion!(temp,A.boundarycondition)
    #add_fermion!(xout,1,x,-1,temp2)
    add_fermion!(xout,1,x,-1,temp2,iseven)
    set_wing_fermion!(xout,A.boundarycondition)

    return
end

function clear_fermion!(a::WilsonFermion_4D_wing{NC,NDW} ,iseven) where {NC,NDW}
    n1,n2,n3,n4,n5,n6 = size(a.f)
    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5-NDW
            for i4=1:n4
                iz = i4-NDW
                for i3=1:n3
                    iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                a.f[i1,i2,i3,i4,i5,i6]= 0
                            end
                        end
                    end
                end
            end
        end
    end
end

function LinearAlgebra.mul!(x::WilsonFermion_4D_wing{NC,NDW},A::TA) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            x[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            x[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            x[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            x[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(x::WilsonFermion_4D_wing{NC,NDW},A::TA,iseven::Bool) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            x[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            x[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            x[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            x[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4
                        end
                    end
                end
            end
        end
    end

end

#function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing{NC,NDW},A::TA,x::WilsonFermion_4D_wing{NC}) where {TA <: AbstractMatrix, NC,NDW}
function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing{NC,NDW},A::TA,x::Abstractfermion) where {TA <: AbstractMatrix, NC,NDW}

    NX = xout.NX
    NY = xout.NY
    NZ = xout.NZ
    NT = xout.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            xout[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            xout[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            xout[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing{NC,NDW},A::TA,x::WilsonFermion_4D_wing{NC},iseven::Bool) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            xout[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            xout[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            xout[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4
                        end

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing{NC,NDW},x::WilsonFermion_4D_wing{NC},A::TA) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[2,1]*e2+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix,iy,iz,it,2] = A[1,2]*e1+A[2,2]*e2+A[3,2]*e3+A[4,2]*e4
                            xout[ic,ix,iy,iz,it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            xout[ic,ix,iy,iz,it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing{NC,NDW},x::AbstractFermionfields_4D{NC},A::TA) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[2,1]*e2+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix,iy,iz,it,2] = A[1,2]*e1+A[2,2]*e2+A[3,2]*e3+A[4,2]*e4
                            xout[ic,ix,iy,iz,it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            xout[ic,ix,iy,iz,it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing{NC,NDW},x::WilsonFermion_4D_wing{NC},A::TA,iseven::Bool) where {TA <: AbstractMatrix, NC,NDW}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NC
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[2,1]*e2+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix,iy,iz,it,2] = A[1,2]*e1+A[2,2]*e2+A[3,2]*e3+A[4,2]*e4
                            xout[ic,ix,iy,iz,it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            xout[ic,ix,iy,iz,it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4
                        end
                    end
                end
            end
        end
    end

end

function LinearAlgebra.dot(a::WilsonFermion_4D_wing{NC,NDW},b::WilsonFermion_4D_wing{NC,NDW},iseven::Bool) where {NC,NDW}
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX
    NG = a.NG

    c = 0.0im
    @inbounds for α=1:NG
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for ic=1:NC
                                c+= conj(a[ic,ix,iy,iz,it,α])*b[ic,ix,iy,iz,it,α]
                            end
                        end
                    end
                end
            end
        end
    end
    return c
end

#Overwrite Y with X*a + Y*b, where a and b are scalars. Return Y.
function LinearAlgebra.axpby!(a::Number, X::T, b::Number, Y::WilsonFermion_4D_wing{NC,NDW},iseven::Bool) where {NC,NDW,T <: AbstractFermionfields_4D}
    n1,n2,n3,n4,n5,n6 = size(Y.f)
    #println("axpby")

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5+NDW
            for i4=1:n4
                iz = i4+NDW
                for i3=1:n3
                    iy = i3+NDW
                    for i2=1:n2
                        ix = i2+NDW
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NC
                                Y.f[i1,i2,i3,i4,i5,i6] = a*X.f[i1,i2,i3,i4,i5,i6] + b*Y.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
    end
end

#=
function set_wing_fermion!(a::WilsonFermion_4D_wing{NC},boundarycondition) where NC
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX

    #!  X-direction
    for ialpha=1:4
        for it=1:NT
            for iz = 1:NZ
                for iy=1:NY
                    @simd for k=1:NC
                        a[k,0,iy,iz,it,ialpha] = boundarycondition[1]*a[k,NX,iy,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for k=1:NC
                        a[k,NX+1,iy,iz,it,ialpha] =boundarycondition[1]*a[k,1,iy,iz,it,ialpha]
                    end
                end
            end
        end
    end

    #Y-direction
    for ialpha = 1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,0,iz,it,ialpha] =boundarycondition[2]*a[k,ix,NY,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,NY+1,iz,it,ialpha] = boundarycondition[2]*a[k,ix,1,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        # Z-direction
        for it=1:NT
            for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,iy,0,it,ialpha] = boundarycondition[3]*a[k,ix,iy,NZ,it,ialpha]
                        a[k,ix,iy,NZ+1,it,ialpha] = boundarycondition[3]*a[k,ix,iy,1,it,ialpha]

                    end
                end
            end
        end

        #T-direction
        for iz=1:NZ
            for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NC
                        a[k,ix,iy,iz,0,ialpha] = boundarycondition[4]*a[k,ix,iy,iz,NT,ialpha]
                        a[k,ix,iy,iz,NT+1,ialpha] =boundarycondition[4]*a[k,ix,iy,iz,1,ialpha]
                    end
                end
            end
        end

    end

end

=#

"""
c--------------------------------------------------------------------------c
c     y = gamma_5 * x
c     here
c                  ( -1       )
c        GAMMA5 =  (   -1     )
c                  (     +1   )
c                  (       +1 )
c--------------------------------------------------------------------------c
    """
    function mul_γ5x!(y::WilsonFermion_4D_wing{NC},x::WilsonFermion_4D_wing{NC}) where NC
        n1,n2,n3,n4,n5,n6 = size(x.f)
        @inbounds for i6=1:n6
            for i5=1:n5
                #it = i5+NDW
                for i4=1:n4
                    #iz = i4+NDW
                    for i3=1:n3
                        #iy = i3+NDW
                        for i2=1:n2
                            #ix = i2+NDW
                            @simd for ic=1:NC
                                y.f[i1,i2,i3,i4,i5,i6]=x.f[i1,i2,i3,i4,i5,i6]*ifelse(i6 <= 2,-1,1)
                            end
                        end
                    end
                end
            end
        end

        #=
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        for ig=1:4
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            for ix=1:NX
                                @simd for ic=1:NC
                                    y[ic,ix,iy,iz,it,ig] =x[ic,ix,iy,iz,it,ig]*ifelse(ig <= 2,-1,1)
                                end
                            end
                        end
                    end
                end
            end
        end
        =#

    end

    function apply_γ5!(x::WilsonFermion_4D_wing{NC}) where NC
        n1,n2,n3,n4,n5,n6 = size(x.f)
        #println("axpby")

        @inbounds for i6=1:n6
            for i5=1:n5
                #it = i5+NDW
                for i4=1:n4
                    #iz = i4+NDW
                    for i3=1:n3
                        #iy = i3+NDW
                        for i2=1:n2
                            #ix = i2+NDW
                            @simd for ic=1:NC
                                x.f[i1,i2,i3,i4,i5,i6]=x.f[i1,i2,i3,i4,i5,i6]*ifelse(i6 <= 2,-1,1)
                            end
                        end
                    end
                end
            end
        end

    end

    function mul_1plusγ5x!(y::WilsonFermion_4D_wing{NC},x::WilsonFermion_4D_wing{NC})  where NC#(1+gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,1] = 0#-1*x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,2] = 0#-1*x[ic,ix,iy,iz,it,2]
                            y[ic,ix,iy,iz,it,3] = x[ic,ix,iy,iz,it,3]
                            y[ic,ix,iy,iz,it,4] = x[ic,ix,iy,iz,it,4]
                        end
                    end
                end
            end
        end
    end

    function mul_1plusγ5x_add!(y::WilsonFermion_4D_wing{NC},x::WilsonFermion_4D_wing{NC},factor) where NC#x = x +(1+gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,3] += factor*x[ic,ix,iy,iz,it,3]
                            y[ic,ix,iy,iz,it,4] += factor*x[ic,ix,iy,iz,it,4]
                        end
                    end
                end
            end
        end
    end

    function mul_1minusγ5x!(y::WilsonFermion_4D_wing{NC},x::WilsonFermion_4D_wing{NC}) where NC#(1-gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,1] = x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,2] = x[ic,ix,iy,iz,it,2]
                            y[ic,ix,iy,iz,it,3] = 0#x[ic,ix,iy,iz,it,3]
                            y[ic,ix,iy,iz,it,4] = 0#x[ic,ix,iy,iz,it,4]
                        end
                    end
                end
            end
        end
    end

    function mul_1minusγ5x_add!(y::WilsonFermion_4D_wing{NC},x::WilsonFermion_4D_wing{NC},factor) where NC#+(1-gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,1] += factor*x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,2] += factor*x[ic,ix,iy,iz,it,2]
                        end
                    end
                end
            end
        end
    end

    """
               (       -i )              (       -1 )
     GAMMA1 =  (     -i   )     GAMMA2 = (     +1   )
               (   +i     )              (   +1     )
               ( +i       )              ( -1       )
    """

    function mul_1minusγ1x!(y::WilsonFermion_4D_wing{NC},x) where NC#(1-gamma_5)/2
        NX = y.NX
        NY = y.NY
        NZ = y.NZ
        NT = y.NT
        #NC = x.NC
        @inbounds for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        for ic=1:NC
                            v1 = x[ic,ix,iy,iz,it,1] + im*x[ic,ix,iy,iz,it,4]
                            v2 = x[ic,ix,iy,iz,it,2] + im*x[ic,ix,iy,iz,it,3]
                            v3 = x[ic,ix,iy,iz,it,3] - im*x[ic,ix,iy,iz,it,2]
                            v4 = x[ic,ix,iy,iz,it,4] - im*x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,1] = v1
                            y[ic,ix,iy,iz,it,2] = v2
                            y[ic,ix,iy,iz,it,3] = v3
                            y[ic,ix,iy,iz,it,4] = v4
                        end
                    end
                end
            end

        end
    end

    function mul_1plusγ1x!(y::WilsonFermion_4D_wing{NC},x) where NC#(1-gamma_5)/2
        NX = y.NX
        NY = y.NY
        NZ = y.NZ
        NT = y.NT
        #NC = x.NC
        @inbounds for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            v1 = x[ic,ix,iy,iz,it,1] - im*x[ic,ix,iy,iz,it,4]
                            v2 = x[ic,ix,iy,iz,it,2] - im*x[ic,ix,iy,iz,it,3]
                            v3 = x[ic,ix,iy,iz,it,3] + im*x[ic,ix,iy,iz,it,2]
                            v4 = x[ic,ix,iy,iz,it,4] + im*x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,1] = v1
                            y[ic,ix,iy,iz,it,2] = v2
                            y[ic,ix,iy,iz,it,3] = v3
                            y[ic,ix,iy,iz,it,4] = v4
                        end
                    end
                end
            end
        end
    end


module RobustControl
using LinearAlgebra, PyPlot, Printf

# include all the files
include("StateSpace.jl");
include("Utilities.jl");
include("Response.jl");
include("Plotting.jl");
include("Interconnections.jl");
include("ModelReduction.jl");
include("Realization.jl");
end # module


module ExperimentalSetup

using DataStructures, JSON

export  Factor, Rep,
        delete, replace!,
        saveJSON, loadJSON

import Base: ∈, push!, delete!, empty!

struct Factor
    name::String
    levels::Vector{String}
end

struct Setup
    levels::Vector{Int}
    n::Base.RefValue{Int}
end

struct Metadata
    factors::Vector{Factor}
    setups::Vector{Setup}
    nfactors::Int

    Metadata(factors::Vector{Factor}, setups::Vector{Setup}) = new(factors, setups, length(factors))
end

Metadata(factors::Vector{Factor}) = Metadata(factors, Setup[])

Metadata() = Metadata(Factor[])

mutable struct Rep
    setup::Int
    replicate::Int
    comment::String
end

mutable struct Log
    md::Metadata
    last::UInt64

    reps::SortedDict{UInt64, Rep}
end

Log(md::Metadata) = Log(md, 0, SortedDict{UInt64, Rep}())

Log() = Log(Metadata())

# findfirst

# findfirst(a::Log, x::Rep) = findfirst(a.reps, x)

# in

# repID(md::Metadata, x::Rep) = hash((md.setups[x.setup],x.replicate, x.comment))

∈(k::UInt64, a::Log) = haskey(a.reps, k)
∈(k::T, a::Log) where T <: Integer = UInt64(k) ∈ a

# pushes

# Push a set of levels
function push!(md::Metadata, str_levels::Vector{String})
    @assert length(str_levels) == md.nfactors "Number of levels doesn't match metadata"
    levels = Int[]

    # Ensure that each of the levels supplied exist in the Metatdata
    for (factor, l) in zip(md.factors, str_levels)
        i = findfirst(factor.levels, l)
        @assert i ≠ 0 "Factor $(factor.name) does not have level $l"
        push!(levels, i)
    end

    # ??? Try to find an existing setup with the same levels perhaps?
    i = findfirst(x -> x.levels == levels, md.setups)
    if i == 0
        push!(md.setups, Setup(levels, Ref(0)))
        i = length(md.setups)
    end

    # Increment the level count
    md.setups[i].n.x += 1

    # Return the index and the count
    return (i, md.setups[i].n.x)
end

function push!(a::Log, str_levels::Vector{String}, comment::String)
    i, replicate = push!(a.md, str_levels)
    a.last += 1
    a.reps[a.last] = Rep(i, replicate, comment)
    return a
end
# empty

function empty!(a::Log)
    empty!(a.reps)
    empty!(a.md.setups)
    return a
end

# pop

pop(md::Metadata, x::Rep) = ([f.levels[i] for (i, f) in zip(md.setups[x.setup].levels, md.factors)], x.comment)

pop(a::Log, k::UInt64) = pop(a.md, a.reps[k])

pop(a::Log, k::T) where T <: Integer = pop(a, UInt64(k))

# deletes

function _delete!(md::Metadata, setup::Int)
    step_setup = md.setups[setup].n.x == 1
    if step_setup
        deleteat!(md.setups, setup)
    else
        md.setups[setup].n.x -= 1
    end
    return step_setup
end

function delete!(a::Log, k::UInt64)
    x = a.reps[k]
    step_setup = _delete!(a.md, x.setup)
    for (i,r) in a.reps
        if r.setup == x.setup
            if r.replicate > x.replicate
                r.replicate -= 1
            end
        elseif step_setup && r.setup > x.setup
            r.setup -= 1
        end
    end
    delete!(a.reps, k)
    return a
end

delete!(a::Log, k::T) where T <: Integer = delete!(a, UInt64(k))

function delete(a::Log, x::Factor)
    i = findfirst(y -> y == x, a.md.factors)
    @assert i ≠ 0 "factor not found in log"
    factors = deepcopy(a.md.factors)
    deleteat!(factors, i)
    b = Log(Metadata(factors))
    for r in values(a.reps)
        levels, comment = pop(a.md, r)
        deleteat!(levels, i)
        push!(b, levels, comment)
    end
    return b
end

function delete(a::Log, x::Factor, level::String)
    i = findfirst(y -> y == x, a.md.factors)
    @assert i ≠ 0 "factor not found in log"
    j = findfirst(a.md.factors[i].levels, level)
    @assert j ≠ 0 "level not found in log"
    nlevels = length(a.md.factors[i].levels)
    if nlevels == 1
        return delete(a, x)
    end
    factors = deepcopy(a.md.factors)
    deleteat!(factors[i].levels, j)
    b = Log(Metadata(factors))
    # if the level to be removed is first in the list of levels, then replace it with the second one. If it's the last, replace it with the next to last one. If it's neither, replace it with one level bellow.
    alternative = j == 1 ? a.md.factors[i].levels[2] : j == nlevels ? a.md.factors[i].levels[nlevels - 1] : a.md.factors[i].levels[j - 1]
    for r in values(a.reps)
        levels, comment = pop(a.md, r)
        if levels[i] == level
            levels[i] = alternative
        end
        push!(b, levels, comment)
    end
    return b
end

#=function delete!(a::Log, x::Rep)
    x ∉ a && return a
    step_setup = false
    if a.md.setups[x.setup].n.x == 1
        deleteat!(a.md.setups, x.setup)
        step_setup = true
    else
        a.md.setups[x.setup].n.x -= 1
    end
    ind = 0
    for (i,r) in enumerate(a.reps)
        if r.setup == x.setup
            if r == x
                ind = copy(i)
            else
                if r.replicate > x.replicate
                    r.replicate -= 1
                end
            end
        elseif r.setup > x.setup
            if step_setup
                r.setup -= 1
            end
        end
    end
    deleteat!(a.reps, ind)
    filter!(x -> x.rep ≠ ind, a.asss)
    for s in a.asss
        if s.rep > ind
            s.rep -= 1
        end
    end
    return a
end=#

# replace

function replace!(a::Log, o::UInt64, str_levels::Vector{String}, comment::String)
    @assert haskey(a.reps, o) ≠ 0 "old run not found"
    delete!(a, o)
    push!(a, str_levels, comment)
    a.reps[o] = pop!(a.reps, a.last)
    return a
end

replace!(a::Log, o::T, str_levels::Vector{String}, comment::String) where T <: Integer = replace!(a, UInt64(o), str_levels, comment)

# combine

function combine_factors(logs::Log...)
    factors = Factor[]
    for log in logs, new_factor in log.md.factors
        for old_factor in factors
            if new_factor.name == old_factor.name
                for new_level in new_factor.levels
                    new_level ∉ old_factor.levels && push!(old_factor.levels, new_level)
                end
                break
            end
        end
        push!(factors, new_factor)
    end
    return Metadata(factors)
end

function combine(logs::Log...)
    md = combine_factors(logs...)
    a = Log(md)
    x = OrderedDict{String, String}(f.name => f.levels[1] for f in md.factors)
    for log in logs, r in values(log.reps)
        for f in md.factors
            x[f.name] = f.levels[1] # non existant factors recieve a default first level
        end
        for (i, f) in zip(log.md.setups[r.setup].levels, log.md.factors)
            x[f.name] = f.levels[i]
        end
        push!(a, collect(values(x)), r.comment)
    end
    return a
end

# Converts a JSON representation of a Log to a Log object
function parseJSON(jsonString::AbstractString)
    # create the log objects
    json = JSON.parse(jsonString)
    log = Log(
        Metadata(
            map(x -> Factor(x["name"], x["levels"]), json["md"]["factors"]),
            map(y -> Setup(y["levels"], Base.RefValue{Int64}(y["n"]["x"])), json["md"]["setups"])
        ),
        json["last"],
        SortedDict{UInt64, Rep}()
    )

    # add the reps
    for (k,v) in json["reps"]
       log.reps[parse(UInt64, k)] = Rep(v["setup"], v["replicate"], v["comment"])
    end

    return log
end

# Reads the supplied file and tries to convert the containing JSON to a Log object
function loadJSON(filename::String)
    @assert isfile(filename) "File does not exist."

    json = ""
    # read and parse
    open(filename, "r") do f
        json = readstring(f)
    end

    @assert length(json) > 0 "Unable to read from file, or file is empty."
    return parseJSON(json)
end

# Converts the suppplied Log object to a JSON string and saves it to a file
function saveJSON(filename::String, log::Log, overwrite::Bool)
    @assert overwrite || !isfile(filename) "File alread exists."

    # write the Log out as a JSON string
    open(filename, "w") do f
        write(f, JSON.json(log))
    end
end

# Overloaded saveJSON method--this instance automatically overwrites the file
saveJSON(filename::String, log::Log) = saveJSON(filename, log, true)

end # module


using Logging

const DEFAULT_COMPILE_THRESHOLD = 1.04

abstract type Optimizer end

mutable struct TopNOptimizer{TCostModel} <: Optimizer
    n::Int
    compile_threshold::Float32
    costmodel::TCostModel
    last_fixtypes
    fixtypes_history
    TopNOptimizer(
        n = 10;
        compile_threshold = DEFAULT_COMPILE_THRESHOLD,
        costmodel = basemodel) = new{typeof(costmodel)}(n, compile_threshold, costmodel, EmptyTypeList, [])
end

fixtypes(opt::TopNOptimizer, ::NoProfiler) = opt.last_fixtypes

percent(x) = round((x) * 10_000) / 100

function find_best_historic_cost(opt, _typefreqs, costmodel)
    best_historic_cost = typemax(ClockCycle)
    best_idx = -1
    for idx in 1:length(opt.fixtypes_history) # TODO Do not calculate for all if the list is lengthy
        old = opt.fixtypes_history[idx]
        oldcost = costof(_typefreqs, old, costmodel)
        if oldcost < best_historic_cost
            best_historic_cost = oldcost
            best_idx = idx
        end
    end
    return best_idx, best_historic_cost
end

function fixtypes(opt::TopNOptimizer, prof::FullProfiler)
    tfs = typefreqs(prof)
    @debug "Profiled type freqs: $(tfs)"
    ideal = ideal_fixtypes(opt, prof)
    idealcost = costof(tfs, ideal, opt.costmodel)
    best_idx, best_historic_cost = find_best_historic_cost(opt, tfs, opt.costmodel)
    if best_historic_cost > opt.compile_threshold * idealcost
        @debug begin
            lastcost = costof(typefreqs(prof), opt.last_fixtypes, opt.costmodel)
            speedup_last = percent(1.0 - idealcost / lastcost)
            speedup_best = best_historic_cost == typemax(Int) ? speedup_last : percent(1.0 - idealcost / best_historic_cost)
            "Selected new fixtype list for $(speedup_last)% speedup vs last used, and $(speedup_best)% vs best: $(Catwalk.pretty(ideal))"
        end
        opt.last_fixtypes = ideal
        push!(opt.fixtypes_history, ideal)
    else
        opt.last_fixtypes = opt.fixtypes_history[best_idx]
        @debug "Selected previous fixtype list: $(Catwalk.pretty(opt.last_fixtypes))"
    end
    return opt.last_fixtypes
end

# The N most frequent first types encoded
function ideal_fixtypes(opt::TopNOptimizer, prof::FullProfiler)
    freqs = sort(collect(typefreqs(prof)); order=Base.Order.ReverseOrdering())
    topfreqs = freqs[1:min(opt.n, length(freqs))]
    return encode(map(f -> f.type, topfreqs)...)
end


# run this on the RCI cluster in the `data/experiment_count` dir
using DrWatson
@quickactivate
using PrettyTables
using Statistics
using FileIO, BSON, DataFrames
using GenerativeAD
using CSV

# just count models, score files and training time
# number of models
data_folders = [datadir("experiments/tabular"), datadir("experiments/images"),
	datadir("experiments/images_leave-one-in")]
#allfiles = vcat(map(f->GenerativeAD.Evaluation.collect_files(f), data_folders)...)

# tabular data
df = load(datadir("evaluation/tabular_eval.bson"))[:df];
ddir = data_folders[1]
models = unique(df.modelname)
models = models[models .!= "vae+ocsvm"]
datasets = unique(df.dataset)

mcount_df_tabular = DataFrame(
	:dataset => datasets)
for model in models
	mcount_df_tabular[Symbol(model)] = zeros(Float32,length(datasets))
end
scount_df_tabular = DataFrame(
	:dataset => datasets)
for model in models
	scount_df_tabular[Symbol(model)] = zeros(Float32,length(datasets))
end

for (di,dataset) in enumerate(datasets)
	for (mi,model) in enumerate(models)
		datapath = joinpath(ddir, model, dataset)
		nseeds = length(readdir(datapath))
		files = GenerativeAD.Evaluation.collect_files(datapath)
		sfiles = filter(x->!occursin("model", x), files)
		mfiles = filter(x->occursin("model", x), files)
		mcount_df_tabular[di,mi+1] = length(mfiles)/nseeds
		scount_df_tabular[di,mi+1] = length(sfiles)/nseeds
	end
end

CSV.write("model_count_df_tabular_raw.csv", mcount_df_tabular)
CSV.write("score_count_df_tabular_raw.csv", scount_df_tabular)

mcount_df_tabular = CSV.read("model_count_df_tabular_raw.csv")
scount_df_tabular = CSV.read("score_count_df_tabular_raw.csv")

# some models dont save modelfiles
nosave_models = ["MO_GAAL", "abod", "hbos", "if", "knn", "loda", "lof", "ocsvm",
	"ocsvm_nu", "ocsvm_rbf", "pidforest"]

for model in nosave_models
	mcount_df_tabular[!,Symbol(model)] .= scount_df_tabular[!,Symbol(model)]
end
mcount_df_tabular.knn .= mcount_df_tabular.knn/3
mcount_df_tabular.vae_knn .= mcount_df_tabular.vae_knn/3

# some dont save score files
nosave_scores = ["vae_knn", "vae_ocsvm"]
for model in nosave_scores
	scount_df_tabular[!,Symbol(model)] .= mcount_df_tabular[!,Symbol(model)]
end
scount_df_tabular.vae_knn .= scount_df_tabular.vae_knn*3

# now save the fit and predict times
fit_ts = []
pred_ts = []
for model in models
	subdf = filter(r->r.modelname==model, df)
	push!(fit_ts, sum(subdf.fit_t) + sum(subdf.fs_fit_t))
	push!(pred_ts, sum(subdf.tr_eval_t) + sum(subdf.tst_eval_t) + sum(subdf.val_eval_t)
		+ sum(subdf.fs_eval_t))
end
push!(mcount_df_tabular, vcat(["fit_t"], fit_ts))
push!(mcount_df_tabular, vcat(["pred_t"], pred_ts))
#CSV.write("model_count_df_tabular.csv", mcount_df_tabular)
push!(scount_df_tabular, vcat(["fit_t"], fit_ts))
push!(scount_df_tabular, vcat(["pred_t"], pred_ts))
#CSV.write("score_count_df_tabular.csv", scount_df_tabular)
mcount_df_tabular.pidf[1:end-2] .= mcount_df_tabular.pidf[1:end-2]/3

mcount_df_tabular = CSV.read("model_count_df_tabular.csv")
scount_df_tabular = CSV.read("score_count_df_tabular.csv")

# rename datasets
include("aliases.jl")
mcount_df_tabular.dataset .= map(x->get(DATASET_ALIAS, x, x), mcount_df_tabular.dataset)
scount_df_tabular.dataset .= map(x->get(DATASET_ALIAS, x, x), scount_df_tabular.dataset)

# rename columns
sum_models = Dict(
	"aae" => ["aae", "aae_full", "aae_vamp"],
	"ocsvm" => ["ocsvm", "ocsvm_nu", "ocsvm_rbf"],
	"vae" => ["vae", "vae_full", "vae_simple"],
	"wae" => ["wae", "wae_full", "wae_vamp"]
	)
models_ordered = ["aae", "avae", "gano", "vae", "wae", "abod", "hbos", "if",
	"knn", "loda", "lof", "orbf", "osvm", "pidf", "maf", "rnvp", "sptn", "fmgn",
	"gan", "mgal", "dsvd", "vaek", "vaeo"]
mc_final_tabular = DataFrame(:dataset => copy(mcount_df_tabular.dataset))
for model in models_ordered
	ind = findfirst(values(MODEL_ALIAS) .== model)
	mc_final_tabular[!,Symbol(model)] .=
		mcount_df_tabular[!,Symbol(collect(keys(MODEL_ALIAS))[ind])]
	if model in keys(sum_models)
		for m in sum_models[model][2:end]
			mc_final_tabular[!,Symbol(model)] .+= mcount_df_tabular[!,Symbol(m)]
		end
	end
end
sc_final_tabular = DataFrame(:dataset => copy(scount_df_tabular.dataset))
for model in models_ordered
	ind = findfirst(values(MODEL_ALIAS) .== model)
	sc_final_tabular[!,Symbol(model)] .=
		scount_df_tabular[!,Symbol(collect(keys(MODEL_ALIAS))[ind])]
	if model in keys(sum_models)
		for m in sum_models[model][2:end]
			sc_final_tabular[!,Symbol(model)] .+= scount_df_tabular[!,Symbol(m)]
		end
	end
end

CSV.write("model_count_df_tabular.csv", mc_final_tabular)
CSV.write("score_count_df_tabular.csv", sc_final_tabular)

mc_final_tabular = CSV.read("model_count_df_tabular.csv")
sc_final_tabular = CSV.write("score_count_df_tabular.csv")

# now finalize them
mc_final_tabular[1:end-2, 2:end] .= round.(mc_final_tabular[1:end-2, 2:end])
smc_tabular = pretty_table(mc_final_tabular, backend = :latex)
vals=Array(mc_final_tabular[end-1,2:end])/3600
@printf("%.1e &", vals[1])
for val in vals
	@printf("%.1e & ", val)
end
for val in Array(mc_final_tabular[end,2:end])/3600
	@printf("%.1e & ", val)
end

sc_final_tabular[1:end-2, 2:end] .= round.(sc_final_tabular[1:end-2, 2:end])
ssc_tabular = pretty_table(sc_final_tabular, backend = :latex)

#################
# now the same for image models
df = load(datadir("evaluation/images_eval.bson"))[:df];
ddir = data_folders[2]
models = unique(df.modelname)
datasets = unique(df.dataset)

mcount_df_images = DataFrame(
	:dataset => datasets)
for model in models
	mcount_df_images[Symbol(model)] = zeros(Float32,length(datasets))
end
scount_df_images = DataFrame(
	:dataset => datasets)
for model in models
	scount_df_images[Symbol(model)] = zeros(Float32,length(datasets))
end

for (di,dataset) in enumerate(datasets)
	for (mi,model) in enumerate(models)
		datapath = joinpath(ddir, model, dataset)
		nac = length(readdir(datapath))
		files = GenerativeAD.Evaluation.collect_files(datapath)
		files = filter(x->occursin("seed=1", x), files)
		sfiles = filter(x->!occursin("model", x), files)
		mfiles = filter(x->occursin("model", x), files)
		mcount_df_images[di,mi+1] = length(mfiles)/nac
		scount_df_images[di,mi+1] = length(sfiles)/nac
	end
end

mcount_df_images.knn .= scount_df_images.knn/3
mcount_df_images.ocsvm .= scount_df_images.ocsvm

nosave_scores = ["vae_knn", "vae_ocsvm", "aae_ocsvm"]
for model in nosave_scores
	scount_df_images[!,Symbol(model)] .= mcount_df_images[!,Symbol(model)]
end
scount_df_tabular.vae_knn .= scount_df_tabular.vae_knn*3

CSV.write("model_count_df_images_raw.csv", mcount_df_images)
CSV.write("score_count_df_images_raw.csv", scount_df_images)

mcount_df_images = CSV.read("model_count_df_images_raw.csv")
scount_df_images = CSV.read("score_count_df_images_raw.csv")

models_ordered = ["aae", "gano", "skip", "vae", "wae", "knn", "osvm", "fano",
	"fmgn", "dsvd", "vaek", "vaeo"]
mc_final_images = DataFrame(:dataset => copy(mcount_df_images.dataset))
for model in models_ordered
	ind = findlast(values(MODEL_ALIAS) .== model)
	mc_final_images[!,Symbol(model)] .=
		mcount_df_images[!,Symbol(collect(keys(MODEL_ALIAS))[ind])]
end
sc_final_images = DataFrame(:dataset => copy(scount_df_images.dataset))
for model in models_ordered
	ind = findlast(values(MODEL_ALIAS) .== model)
	sc_final_images[!,Symbol(model)] .=
		scount_df_images[!,Symbol(collect(keys(MODEL_ALIAS))[ind])]
end

CSV.write("model_count_df_images.csv", mc_final_images)
CSV.write("score_count_df_images.csv", sc_final_images)

mc_final_images[1:end, 2:end] .= round.(mc_final_images[1:end, 2:end])
smc_images = pretty_table(mc_final_images, backend = :latex, nosubheader = true)

sc_final_images[1:end, 2:end] .= round.(sc_final_images[1:end, 2:end])
ssc_images = pretty_table(sc_final_images, backend = :latex, nosubheader = true)

fit_t_tabular = sum(df.fit_t)/3600
pred_t_tabular = (sum(df.tr_eval_t) + sum(df.tst_eval_t) + sum(df.val_eval_t))/3600
nexp_images = round(Int,sum(Array(CSV.read("model_count_df_images.csv")[1:end,2:end]))*10)
nscore_images = round(Int,sum(Array(CSV.read("score_count_df_images.csv")[1:end,2:end]))*10)

######## LOI ##############
df = load(datadir("evaluation/images_leave-one-in_eval.bson"))[:df];
ddir = data_folders[3]
models = unique(df.modelname)
datasets = unique(df.dataset)

mcount_df_images = DataFrame(
	:dataset => datasets)
for model in models
	mcount_df_images[Symbol(model)] = zeros(Float32,length(datasets))
end
scount_df_images = DataFrame(
	:dataset => datasets)
for model in models
	scount_df_images[Symbol(model)] = zeros(Float32,length(datasets))
end

for (di,dataset) in enumerate(datasets)
	for (mi,model) in enumerate(models)
		datapath = joinpath(ddir, model, dataset)
		nac = length(readdir(datapath))
		files = GenerativeAD.Evaluation.collect_files(datapath)
		files = filter(x->occursin("seed=1", x), files)
		sfiles = filter(x->!occursin("model", x), files)
		mfiles = filter(x->occursin("model", x), files)
		mcount_df_images[di,mi+1] = length(mfiles)/nac
		scount_df_images[di,mi+1] = length(sfiles)/nac
	end
end

mcount_df_images.knn .= scount_df_images.knn/3
mcount_df_images.ocsvm .= scount_df_images.ocsvm

nosave_scores = ["vae_knn", "vae_ocsvm", "aae_ocsvm"]
for model in nosave_scores
	scount_df_images[!,Symbol(model)] .= mcount_df_images[!,Symbol(model)]
end
scount_df_tabular.vae_knn .= scount_df_tabular.vae_knn*3

CSV.write("model_count_df_images_loi_raw.csv", mcount_df_images)
CSV.write("score_count_df_images_loi_raw.csv", scount_df_images)

mcount_df_images = CSV.read("model_count_df_images_loi_raw.csv")
scount_df_images = CSV.read("score_count_df_images_loi_raw.csv")

models_ordered = ["aae", "gano", "skip", "vae", "wae", "knn", "osvm", "fano",
	"fmgn", "dsvd", "vaek", "vaeo"]
mc_final_images = DataFrame(:dataset => copy(mcount_df_images.dataset))
for model in models_ordered
	ind = findlast(values(MODEL_ALIAS) .== model)
	mc_final_images[!,Symbol(model)] .=
		mcount_df_images[!,Symbol(collect(keys(MODEL_ALIAS))[ind])]
end
sc_final_images = DataFrame(:dataset => copy(scount_df_images.dataset))
for model in models_ordered
	ind = findlast(values(MODEL_ALIAS) .== model)
	sc_final_images[!,Symbol(model)] .=
		scount_df_images[!,Symbol(collect(keys(MODEL_ALIAS))[ind])]
end

CSV.write("model_count_df_images_loi.csv", mc_final_images)
CSV.write("score_count_df_images_loi.csv", sc_final_images)

mc_final_images[1:end, 2:end] .= round.(mc_final_images[1:end, 2:end])
smc_images = pretty_table(mc_final_images, backend = :latex, nosubheader = true)

sc_final_images[1:end, 2:end] .= round.(sc_final_images[1:end, 2:end])
ssc_images = pretty_table(sc_final_images, backend = :latex, nosubheader = true)

fit_t_tabular = sum(df.fit_t)/3600
pred_t_tabular = (sum(df.tr_eval_t) + sum(df.tst_eval_t) + sum(df.val_eval_t))/3600
nexp_images_loi = round(Int,sum(Array(CSV.read("model_count_df_images_loi.csv")[1:end,2:end]))*10)
nscore_images_loi = round(Int,sum(Array(CSV.read("score_count_df_images_loi.csv")[1:end,2:end]))*10)

# total times and counts
# tabular
df = load(datadir("evaluation/tabular_eval.bson"))[:df];
nexp_tabular = round(Int,sum(Array(CSV.read("model_count_df_tabular.csv")[1:end-2,2:end]))*5)
nscore_tabular = round(Int,sum(Array(CSV.read("score_count_df_tabular.csv")[1:end-2,2:end]))*5)

cdf = CSV.read("model_count_df_tabular.csv")
sdf = CSV.read("score_count_df_tabular.csv")
multi_models = ["aae", "wae", "vae", "knn", "pidf"]
scales = []
for m in multi_models
	push!(scales, sum(sdf[1:end-2,Symbol(m)])/sum(cdf[1:end-2,Symbol(m)]))
end
fit_t_tabular = 0
for model in unique(df.modelname)
	ft = sum(filter(r->r.modelname==model, df)[!,:fit_t])
	i = findfirst(map(x->occursin(x,model),multi_models))
	if i != nothing && !(occursin("ocsvm", model))
		global fit_t_tabular += ft/scales[i]
	else
		global fit_t_tabular += ft
	end
end
fit_t_tabular = fit_t_tabular/3600/24
pred_t_tabular = (sum(df.tr_eval_t) + sum(df.tst_eval_t) + sum(df.val_eval_t))/3600/24

# images
df = load(datadir("evaluation/images_eval.bson"))[:df];
nexp_images = round(Int,sum(Array(CSV.read("model_count_df_images.csv")[1:end,2:end]))*10)
nscore_images = round(Int,sum(Array(CSV.read("score_count_df_images.csv")[1:end,2:end]))*10)

cdf = CSV.read("model_count_df_images.csv")
sdf = CSV.read("score_count_df_images.csv")
multi_models = ["aae", "gano", "skip", "wae", "vae", "knn"]
scales = []
for m in multi_models
	push!(scales, sum(sdf[1:end-2,Symbol(m)])/sum(cdf[1:end-2,Symbol(m)]))
end
multi_models2 = ["aae", "GANo", "skip", "wae", "vae", "knn"]
fit_t_images = 0
for model in unique(df.modelname)
	ft = sum(filter(r->r.modelname==model, df)[!,:fit_t])
	i = findfirst(map(x->occursin(x,model),multi_models2))
	if i != nothing && !(occursin("vae_knn", model))
		global fit_t_images += ft/scales[i]
	else
		global fit_t_images += ft
	end
end
fit_t_images = fit_t_images/3600/24
pred_t_images = (sum(df.tr_eval_t) + sum(df.tst_eval_t) + sum(df.val_eval_t))/3600/24

#images LOI
df = load(datadir("evaluation/images_leave-one-in_eval.bson"))[:df];
nexp_images_loi = round(Int,sum(Array(CSV.read("model_count_df_images_loi.csv")[1:end,2:end]))*10)
nscore_images_loi = round(Int,sum(Array(CSV.read("score_count_df_images_loi.csv")[1:end,2:end]))*10)

cdf = CSV.read("model_count_df_images_loi.csv")
sdf = CSV.read("score_count_df_images_loi.csv")
multi_models = ["aae", "gano", "skip", "wae", "vae", "knn"]
scales = []
for m in multi_models
	push!(scales, sum(sdf[1:end-2,Symbol(m)])/sum(cdf[1:end-2,Symbol(m)]))
end
multi_models2 = ["aae", "GANo", "skip", "wae", "vae", "knn"]
fit_t_images_loi = 0
for model in unique(df.modelname)
	ft = sum(filter(r->r.modelname==model, df)[!,:fit_t])
	i = findfirst(map(x->occursin(x,model),multi_models2))
	if i != nothing && !(occursin("vae_knn", model))
		global fit_t_images_loi += ft/scales[i]
	else
		global fit_t_images_loi += ft
	end
end
fit_t_images_loi = fit_t_images_loi/3600/24
pred_t_images_loi = (sum(df.tr_eval_t) + sum(df.tst_eval_t) + sum(df.val_eval_t))/3600/24

total_nmodels = nexp_images + nexp_images_loi + nexp_tabular
total_nscores = nscore_images + nscore_images_loi + nscore_tabular
total_fit_t = fit_t_tabular + fit_t_images + fit_t_images_loi
total_pred_t = pred_t_tabular + pred_t_images + pred_t_images_loi


using NamedTupleTools: namedtuple, merge

# -------------- Implementation of Product Basis

struct Product1pBasis{NB, TB <: Tuple, VALB} <: OneParticleBasis{Any}
   bases::TB
   indices::Vector{NTuple{NB, Int}}
   B_pool::VectorPool{VALB}
end

function Product1pBasis(bases)
   NB = length(bases)
   return Product1pBasis(bases, NTuple{NB, Int}[])
end

function Product1pBasis(bases, indices)
   # TODO: discuss whether to construct an optimal ordering, e.g.
   #       should the discrete bases come first once we implement the
   #       "strongzero" method?
   VT = _valtype(bases)
   Product1pBasis( tuple(bases...), indices, VectorPool{VT}() )
end

import Base.*
*(B1::OneParticleBasis, B2::OneParticleBasis) =
      Product1pBasis((B1, B2))
*(B1::Product1pBasis, B2::OneParticleBasis) =
      Product1pBasis((B1.bases..., B2))
*(B1::OneParticleBasis, B2::Product1pBasis) =
      Product1pBasis((B1, B2.bases...))
*(B1::Product1pBasis, B2::Product1pBasis) =
      Product1pBasis((B1.bases..., B2.bases...))

_numb(b::Product1pBasis{NB}) where {NB} = NB

Base.length(basis::Product1pBasis) = length(basis.indices)

# ------------------------- FIO CODES

==(B1::Product1pBasis, B2::Product1pBasis) =
      ( all(B1.bases .== B2.bases) &&
        B1.indices == B2.indices )

write_dict(B::Product1pBasis) =
      Dict("__id__" => "ACE_Product1pBasis",
            "bases" => write_dict.(B.bases),
          "indices" => B.indices )

function read_dict(::Val{:ACE_Product1pBasis}, D::Dict)
   bases = tuple( read_dict.(D["bases"])... )
   indices = [ tuple(v...) for v in D["indices"] ]
   return Product1pBasis(bases, indices)
end

# # -------- temporary hack for 1.6, should not be needed from 1.7 onwards

# function acquire_B!(basis::Product1pBasis, args...)
#    VT = valtype(basis, args...)
#    return acquire!(basis.B_pool, length(basis), VT)
# end

# function release_B!(basis::Product1pBasis, B)
#    return release!(basis.B_pool, B)
# end

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

valtype(basis::Product1pBasis) = _valtype(basis.bases)

_valtype(bases) = promote_type(valtype.(bases)...)

valtype(basis::Product1pBasis, X::AbstractState) =
      promote_type(valtype.(basis.bases, Ref(X))...)

# valtype(basis::Product1pBasis, cfg::AbstractConfiguration) =
#       promote_type( valtype.(basis.bases, Ref(first(cfg)))... )

function valtype(basis::Product1pBasis, cfg::AbstractConfiguration)
   X = zero(eltype(cfg))
   return promote_type( valtype.(basis.bases, Ref(X))... )
end

gradtype(basis::Product1pBasis, cfg::Union{AbstractConfiguration, AbstractVector}) =
      gradtype(basis, zero(eltype(cfg)))

function gradtype(basis::Product1pBasis, X::AbstractState)
   VALT = valtype(basis, X)
   return dstate_type(VALT, X)
end

import Base.Cartesian: @nexprs

@generated function add_into_A!(A, basis::Product1pBasis{NB}, X) where {NB}
   quote
      @nexprs $NB i -> begin
         bas_i = basis.bases[i]
         B_i = acquire_B!(bas_i, X)
         evaluate!(B_i, bas_i, X)
      end
      for (iA, ϕ) in enumerate(basis.indices)
         t = one(eltype(A))
         @nexprs $NB i -> (t *= B_i[ϕ[i]])
         A[iA] += t
      end
      @nexprs $NB i -> release_B!(basis.bases[i], B_i)
      return nothing
   end
end

# this is a hack to resolve a method ambiguity.
add_into_A_dA!(A, dA, basis::Product1pBasis, X) =
               _add_into_A_dA!(A, dA, basis, X)
add_into_A_dA!(A, dA, basis::Product1pBasis, X, sym::Symbol) =
               _add_into_A_dA!(A, dA, basis, X, sym)

# args... could be a symbol to enable partial derivatives
# at the moment this is just passed into evaluate_ed!(...)
# to not evaluate 1p basis derivatives that aren't needed, but
# in the future a more complex construction could be envisioned that
# might considerably reduce the computational cost here...
@generated function _add_into_A_dA!(A, dA, basis::Product1pBasis{NB}, X,
                                   args...) where {NB}
   quote
      Base.Cartesian.@nexprs($NB, i -> begin   # for i = 1:NB
         bas_i = basis.bases[i]
         if !(bas_i isa Discrete1pBasis)
            # only evaluate basis gradients for a continuous basis
            B_i = acquire_B!(bas_i, X)
            dB_i = acquire_dB!(bas_i, X)
            Bt, dBt = evaluate_ed!(B_i, dB_i, bas_i, X, args...)
         else
            # we still need the basis values for the discrete basis though
            # TODO: maybe the d part should be a no-op and remove this
            # case distinction ...
            B_i = acquire_B!(bas_i, X)
            evaluate!(B_i, bas_i, X)
         end
      end)
      for (iA, ϕ) in enumerate(basis.indices)
         # evaluate A
         t = one(eltype(A))
         @nexprs $NB i -> (t *= B_i[ϕ[i]])
         A[iA] += t

         # evaluate dA
         # TODO: redo this with adjoints!!!!
         #     also reverse order of operations to make fewer multiplications!
         dA[iA] = zero(eltype(dA))
         Base.Cartesian.@nexprs($NB, a -> begin  # for a = 1:NB
            if !(basis.bases[a] isa Discrete1pBasis)
               dt = dB_a[ϕ[a]]
               Base.Cartesian.@nexprs($NB, b -> begin  # for b = 1:NB
                  if b != a
                     dt *= B_b[ϕ[b]]
                  end
               end)
               dA[iA] += dt
            end
         end)
      end
      Base.Cartesian.@nexprs($NB, i -> ( begin   # for i = 1:NB
         release_B!(bas_i, B_i)
         if !(basis.bases[i] isa Discrete1pBasis)
            release_dB!(bas_i, dB_i)
         end
      end))
      return nothing
   end
end

_symbols_prod(bases) = tuple(union( symbols.(bases)... )...)

symbols(basis::Product1pBasis) = _symbols_prod(basis.bases)

function indexrange(basis::Product1pBasis)
   allsyms = tuple(symbols(basis)...)
   rg = Dict{Symbol, Vector{Any}}([ sym => [] for sym in allsyms]...)
   for b in basis.bases
      rgb = indexrange(b)
      for sym in allsyms
         if haskey(rgb, sym)
            rg[sym] = union(rg[sym], rgb[sym])
         end
      end
   end
   # HACK: fix the m range based on the maximal l-range
   #       this needs to be suitably generalised if we have multiple
   #       (l, m) pairs, e.g. (l1, m1), (l2, m2)
   if haskey(rg, :m)
      maxl = maximum(rg[:l])
      rg[:m] = collect(-maxl:maxl)
   end

   # convert the range into a named tuple so that we remember the order!!
   return NamedTuple{allsyms}(ntuple(i -> rg[allsyms[i]], length(allsyms)))
end

isadmissible(b, basis::Product1pBasis) = all(isadmissible.(Ref(b), basis.bases))

function set_spec!(basis::Product1pBasis{NB}, spec) where {NB}
   empty!(basis.indices)
   for b in spec
      inds = ntuple(i -> get_index(basis.bases[i], b), NB)
      push!(basis.indices, inds)
   end
   return basis
end

get_spec(basis::Product1pBasis) = [ get_spec(basis, i) for i = 1:length(basis) ]

function get_spec(basis::Product1pBasis, i::Integer)
   inds = basis.indices[i]
   specs = get_spec.(basis.bases, inds)
   # TODO: here we should check that we are only merging compatible tuples,
   #       e.g. (n = 5, l = 2), (l = 2, m = -1) is ok
   #       but  (n = 5, l = 2), (l = 3, m = -1) is forbidden!
   return merge(specs...)
end

degree(b, basis::Product1pBasis) = sum( degree(b, B) for B in basis.bases )

degree(b::NamedTuple, basis::Product1pBasis, weight::Dict) =
      sum( degree(b, B, weight) for B in basis.bases )

# TODO: this looks like a horrible hack ...
function rand_radial(basis::Product1pBasis)
   for B in basis.bases
      if B isa ScalarACEBasis
         return rand_radial(B)
      end
   end
   return nothing
end

# --------------- AD codes

import ChainRules: rrule, NoTangent, ZeroTangent

_evaluate_bases(basis::Product1pBasis{NB}, X::AbstractState) where {NB} =
      ntuple(i -> evaluate(basis.bases[i], X), NB)

_evaluate_A(basis::Product1pBasis{NB}, BB) where {NB} =
      [ prod(BB[i][ϕ[i]] for i = 1:NB) for ϕ in basis.indices ]

evaluate(basis::Product1pBasis, X::AbstractState) =
      _evaluate_A(basis, _evaluate_bases(basis, X))

function _rrule_evaluate(basis::Product1pBasis{NB}, X::AbstractState,
                         w::AbstractVector{<: Number},
                         BB = _evaluate_bases(basis, X)) where {NB}
   VT = promote_type(valtype(basis, X), eltype(w))

   # dB = evaluate_d(basis, X)
   # return sum( (real(w) * real(db) + imag(w) * imag(db))
   #             for (w, db) in zip(w, dB) )

   # Compute the differentials for the individual sub-bases
   Wsub = ntuple(i -> zeros(VT, length(BB[i])), NB)
   for (ivv, vv) in enumerate(basis.indices)
      for t = 1:NB
         _A = one(VT)
         for s = 1:NB
            if s != t
               _A *= BB[s][vv[s]]
            end
         end
         Wsub[t][vv[t]] += w[ivv] * conj(_A)
      end
   end

   # now these can be propagated into the inner basis
   #  -> type instab to be fixed here
   g = sum( _rrule_evaluate(basis.bases[t], X, Wsub[t] )
            for t = 1:NB )
   return g
end

function rrule(::typeof(evaluate), basis::Product1pBasis, X::AbstractState)
   BB = _evaluate_bases(basis, X)
   A = _evaluate_A(basis, BB)
   return A,
      w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w, BB))
end

#    function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState,
#       w::AbstractVector{<: Number})
# x = _val(X, basis)
# a = _rrule_evaluate(basis.P, x, w)
# TDX = ACE.dstate_type(a, X)
# return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
# end