using Test
using Qecsim.Model
using Qecsim:QecsimError
using Qecsim.BasicModels:BasicCode
using Qecsim.PauliTools:to_bsf

# test stub code using duck-typing
struct _DuckCode end
Model.label(::_DuckCode) = "duck-code"
Model.stabilizers(::_DuckCode) = to_bsf(["XZZXI", "IXZZX", "XIXZZ", "ZXIXZ"])
Model.logical_xs(::_DuckCode) = to_bsf(["XXXXX"])
Model.logical_zs(::_DuckCode) = to_bsf(["ZZZZZ"])
Model.nkd(::_DuckCode) = (5, 1, 3)

@testset "StabilizerCode" begin
    # 5-qubit code
    code = BasicCode(["XZZXI", "IXZZX", "XIXZZ", "ZXIXZ"], ["XXXXX"], ["ZZZZZ"])
    @test logicals(code) == to_bsf(["XXXXX", "ZZZZZ"])
    @test validate(code) === nothing
    # Non-commuting stabilizers
    code = BasicCode(["XXXXI", "IXZZX", "XIXZZ", "ZXIXZ"], ["XXXXX"], ["ZZZZZ"])
    @test_throws QecsimError validate(code)
    # Non-commuting stabilizers with logicals
    code = BasicCode(["XZZXI", "IXZZX", "XIXZZ", "ZXIXZ"], ["XXXII"], ["IIZZZ"])
    @test_throws QecsimError validate(code)
    # Commuting logicals
    code = BasicCode(["XZZXI", "IXZZX", "XIXZZ", "ZXIXZ"], ["IIIII"], ["IIIII"])
    @test_throws QecsimError validate(code)
end

@testset "StabilizerCode-duck-typing" begin
    # duck code
    code = _DuckCode()
    @test logicals(code) == to_bsf(["XXXXX", "ZZZZZ"])
    @test validate(code) === nothing
end

@testset "DecodeResult" begin
    # typical valid kwargs/args constructors no exception
    # recovery only (typical)
    DecodeResult(recovery=BitVector([0, 0, 0, 0]))
    DecodeResult(nothing, BitVector([0, 0, 0, 0]), nothing, nothing)
    # success only (override)
    DecodeResult(success=true)
    DecodeResult(true, nothing, nothing, nothing)
    # success + logical_commutations (override)
    DecodeResult(success=true, logical_commutations=BitVector([0, 0]))
    DecodeResult(true, nothing, BitVector([0, 0]), nothing)
    # unusual but valid kwargs/args constructors no exception, type inferred
    # success + recovery (unusual)
    DecodeResult(success=true, recovery=[0, 0, 0, 0])
    DecodeResult(true, [0, 0, 0, 0], nothing, nothing)
    # recovery + logical_commutations (unusual)
    DecodeResult(recovery=[0, 0, 0, 0], logical_commutations=[0, 0])
    DecodeResult(nothing, [0, 0, 0, 0], [0, 0], nothing)
    # success + recovery + logical_commutations (unusual)
    DecodeResult(success=true, recovery=[0, 0, 0, 0],
        logical_commutations=[0, 0])
    DecodeResult(true, [0, 0, 0, 0], [0, 0], nothing)
    # custom_values valid kwargs/args constructors test type inferred
    # unspecified custom_values
    @test isa(DecodeResult(success=true), DecodeResult{Nothing})
    @test isa(DecodeResult(true, nothing, nothing, nothing), DecodeResult{Nothing})
    # custom_values Int
    @test isa(DecodeResult(success=true, custom_values=[1, 2]), DecodeResult{Vector{Int}})
    @test isa(DecodeResult(true, nothing, nothing, [1, 2]), DecodeResult{Vector{Int}})
    # custom_values Rational
    @test isa(DecodeResult(success=true, custom_values=[1//2, 1//3]),
        DecodeResult{Vector{Rational{Int}}})
    @test isa(DecodeResult(true, nothing, nothing, [1//2, 1//3]),
        DecodeResult{Vector{Rational{Int}}})
    # custom_values Float64
    @test isa(DecodeResult(success=true, custom_values=[1., 2.]),
        DecodeResult{Vector{Float64}})
    @test isa(DecodeResult(true, nothing, nothing, [1., 2.]),
        DecodeResult{Vector{Float64}})
    # custom_values Int + Float64 promoted
    @test isa(DecodeResult(success=true, custom_values=[1, 2.]),
        DecodeResult{Vector{Float64}})
    @test isa(DecodeResult(true, nothing, nothing, [1, 2.]),
        DecodeResult{Vector{Float64}})
    # invalid kwargs/args constructors throw exception
    @test_throws QecsimError DecodeResult()
    @test_throws QecsimError DecodeResult(nothing, nothing, nothing, nothing)
    @test_throws QecsimError DecodeResult(logical_commutations=[1, 0])
    @test_throws QecsimError DecodeResult(nothing, nothing, [1, 0], nothing)
end


module Master

include("./constraints/geometric/variables.jl")
include("./constraints/geometric/sentinel.jl")
include("./constraints/geometric/rotations_currence.jl")
include("./constraints/geometric/non_overlapping.jl")
include("./constraints/geometric/inside_pallet.jl")
include("./constraints/static_stability/partial_base/0_partial_base.jl")
include("./constraints/static_stability/partial_base_2/0_partial_base.jl")
include("./constraints/static_stability/on_top/0_on_top.jl")
include("./constraints/static_stability/ss_mult_boxes/0_ss_mult_boxes.jl")
include("./constraints/dynamic_stability/guillotine/0_guillotine.jl")
include("./constraints/dynamic_stability/guillotine/callback.jl")
include("./constraints/dynamic_stability/middle/0_middle.jl")
include("./constraints/objective/constraints/objective.jl")
include("./constraints/geometric/constants.jl")
include("./constraints/dynamic_stability/mult_boxes/0_mult_boxes.jl")
include("./constraints/load_balance/0_load_balancing.jl")
include("./constraints/robot/robot_pattern/0_robot_pattern.jl")
include("./constraints/robot/simple/0_robot_simple.jl")
include("./constraints/robot/simple/0_robot_simple_fixed.jl")
include("./constraints/load_bearing/partial_base/0_partial_base_lb.jl")
include("./constraints/load_bearing/simple/0_pressure_lb.jl")
include("./constraints/load_bearing/simple/0_number_lb.jl")
include("./results/write_to_file.jl")
include("./results/write_to_json.jl")
using JuMP, Gurobi
using .Variables
using .Sentinel
using .RotationsCurrence
using .NonOverlapping
using .InsidePallet
using .PartialBase
using .PartialBase2
using .OnTop
using .CMultipleBoxes
using .Guillotine
using .GuillotineCallback
using .Middle
using .Objective
using .Constants
using .MultBoxes
using .LoadBalancing
using .RobotPattern
using .RobotSimple
using .RobotSimpleFixed
using .LBPartialBase
using .LBSimple
using .LBNumber
using .WriteToFile
using .WriteToJSON

function init_model()
    return direct_model(Gurobi.Optimizer())
end

export main

function main(mode::Int64)
    m = init_model()
    m = init_variables(m)
    m = init_sentinel(m)
    m = cons_rotations_currence(m)
    m = cons_non_overlapping(m)
    m = cons_inside_pallet(m)

    if (mode == 1)
        # LB

        m = cons_lb_load_balancing(m)
    elseif (mode == 2)
        # LBP

        m = cons_lb_simple_load_bearing(m)
    elseif (mode == 3)
        # Robot

        m = cons_rp_robot_simple_fixed(m)
    elseif (mode == 4)
        #PBS1

        m = cons_ss_partial_base(m)
    elseif (mode == 5)
        #PBS2

        m = cons_ss_partial_base_two(m)
    elseif (mode == 6)
        # LBP OnTop

        m = cons_ss_on_top(m)
        m = cons_lb_simple_load_bearing(m)
    elseif (mode == 7)
        # OnTop

        m = cons_ss_on_top(m)
    elseif (mode == 8)
        # Guillotine

        m = cons_ds_guillotine(m)
    elseif (mode == 9)
        # Guillotine OnTop

        m = cons_ss_on_top(m)
        m = cons_ds_guillotine(m)
    elseif (mode == 10)
        # Guillotine OnTop LB LBP

        m = cons_ss_on_top(m)
        m = cons_ds_guillotine(m)
        m = cons_lb_load_balancing(m)
        m = cons_lb_simple_load_bearing(m)
    elseif (mode == 11)
        # PBS2 LB LBP

        m = cons_ss_partial_base_two(m)
        m = cons_lb_load_balancing(m)
        m = cons_lb_simple_load_bearing(m)
    elseif (mode == 12)
        # PBS2 Guillotine LB LBP

        m = cons_ss_partial_base_two(m)
        m = cons_ds_guillotine(m)
        m = cons_lb_load_balancing(m)
        m = cons_lb_simple_load_bearing(m)
    elseif (mode == 13)
        #PBS1 LB LBP MIDDLE MB

        m = cons_ss_partial_base(m)
        m = cons_lb_load_balancing(m)
        m = cons_lb_simple_load_bearing(m)
        m = cons_ds_mult_boxes(m)
        m = cons_ds_middle(m)
    elseif (mode == 14)
        # MB OnTop

        m = cons_ss_on_top(m)
        m = cons_ds_mult_boxes(m)
    elseif (mode == 15)
        # MIDDLE OnTop

        m = cons_ss_on_top(m)
        m = cons_ds_mult_boxes(m)
    elseif (mode == 16)
        #MBSS ONTOP

        m = cons_ss_partial_base_two(m)
        m = cons_ss_mult_boxes(m)
    end

    if (mode == 8 || mode == 9 || mode == 10 || mode == 12)
        # Guillotine

        m = objective(m, 1)
    elseif (mode == 13)
        # MIDDLE MB

        m = objective(m, 2)
    elseif (mode == 14)
        # MB

        m = objective(m, 3)
    elseif (mode == 13)
        # MIDDLE

        m = objective(m, 4)
    else
        # NONE

        m = objective(m, 0)
    end

    function callback_function(cb_data)
        #guillotine_callback(cb_data, m)
        return
    end

    MOI.set(m, MOI.RawParameter("TimeLimit"), 2400)

    try
        optimize!(m)
    catch err
        print("\n ERROR \n")
    end

    try
        print("\n MaxHEIGHT \n")
        print(value(m[:maxHeight]))
        print('\n')
    catch err
        print("\n ERROR \n")
    end

    WriteToFile.append_to_file(m, mode)
    """
    if (has_values(m))
        WriteToJSON.write_to_json_simple(m)
    end
    """

end

#SPECIFY MODE
for i in [5, 11, 12]
    main(i)
end

end


module spiavalonfpga

#  Derived from spiavagif4 July 22, 2019.
#  By Gregory Raven, Plantation Florida
#  KF5N

#include("command_line.jl")
include("SPIManager.jl")

export send

# For opening and working with images
using Images
using Printf

# For progress information on the command line
using ProgressMeter

# Address of the max_frame_count register.
# System counts to this value before moving to the next frame.
framecount_reg_address() = 0x0400_0000

# Address of the max_image_count register.
# Number of images in the GIF.
imagecount_reg_address() = 0x0400_0010

dram_base_address() = zero(UInt)

#  Modify this function to work with SPI.
function send(image::String)  # image is a string path.
    @info "Loading Image"
    img = load(image)

    # Configure frame count register.
    #  This writes the number of frames in the GIF -1 (zero based).
    gif_frames = zeros(UInt8,1)  #  view will not take array size of 1.
    gif_frames[1] = size(img, 3) - 1
    Printf.@printf("The number of frames in the GIF is %d\n", gif_frames[1] + 1)
    SPIManager.write(imagecount_reg_address(), gif_frames)

    @info "Packing Image"
    # By default, use 4 bits per color and 16 bits per pixel.
    data = pack(flatten(img), 4, 16)

    @info "Sending Image"
    #  write function for SPI:
    #  write(address, data)
    SPIManager.write(dram_base_address(), data)
end

# Orient images correctly so standard iteration ends up in the correct format
# for the display unit.
#
# Display unit expects images to be row major.
flatten(img::AbstractArray{T,2}) where T = reshape(img', :)
flatten(img::AbstractArray{T,3}) where T = permutedims(img, (2,1,3))

function pack(img::Array, bits_per_color, bits_per_pixel)
    # Make sure numbers match.
    num_colors = 3
    if num_colors * bits_per_color > bits_per_pixel
        throw(Error())
    end
    # Final return type, just a collection of bytes.
    packed_data = UInt8[]

    # Basic Idea:
    #
    # Go through each item in the flattened image. Pull out the red, green, and
    # blue channels, taking the requested number of bits from the top of each
    # value. As these are decoded, they will be appended to a vector.
    #
    # Each time the vector exceeds 8 bits, we will convert the leading 8 bits
    # to a UInt8 and remove them from the vector.
    #
    # Pad zeros to the desired "bits_per_pixel".
    #
    # Note that this is not super efficient.
    bit_buffer = UInt8[]
    colors = (red, green, blue)

    @showprogress 1 for pixel in img
        newbits = upperbits.(bits.(colors, pixel), bits_per_color)
        for vec in newbits
            append!(bit_buffer, vec)
        end

        # Zero pad if necessary
        pad = zeros(UInt8, bits_per_pixel - 3 * bits_per_color)
        append!(bit_buffer, pad)

        # Pull out groups of 8 bits and put them in the return data.
        groupbits!(packed_data, bit_buffer)
    end

    return packed_data
end

bits(f, x) = f(x).i
upperbits(x::Unsigned, n) = digits(UInt8, x; base = 2, pad = 8 * sizeof(UInt8))[end+1-n:end]

function groupbits!(target, buffer)
    while length(buffer) >= 8
        result = zero(UInt8)
        for i in 0:7
            b = popfirst!(buffer)
            result |= UInt8(b * 2^(i))
        end
        push!(target, result)
    end
end

end # module


module GenML

abstract MLAlgorithm{F} #general type for any machine learning algorithm.
abstract Layer{F} <: MLAlgorithm{F} #a generalized "layer" type

# interface definition for MLAlgorithm.

# basics.  input and output vector size count.
inputs{T <: MLAlgorithm}(::Type{T})     = throw(ArgumentError("inputs for $T not properly implemented"))
outputs{T <: MLAlgorithm}(::Type{T})    = throw(ArgumentError("outputs for $T not properly implemented"))
parameters{T <: MLAlgorithm}(::Type{T}) = throw(ArgumentError("parameters for $T not properly implemented"))

inputs{T <: MLAlgorithm}(::T)     = inputs(T)
outputs{T <: MLAlgorithm}(::T)    = outputs(T)
parameters{T <: MLAlgorithm}(::T) = parameters(T)

@generated function flatten!{F}(::Vector{F}, mla::MLAlgorithm{F}, offset::Integer = 0)
  T = mla
  :(throw(ArgumentError("flatten! for $T not propertly implemented")))
end
@generated function unflatten!{F}(mla::MLAlgorithm{F}, ::Vector{F}, offset::Integer = 0)
  T = mla
  :(throw(ArgumentError("unflatten! for $T not propertly implemented")))
end

################################################################################
## nonparametric data storage.

doc"""
  `GenML.Storage(::Type{MLAlgorithm}, N = :v)`
  `GenML.Storage(::MLAlgorithm, N = :v)`

  takes an ML algorithm and allocates all non-parametric memory required to
  evaluate the algorithm.
"""
abstract Storage{T <: MLAlgorithm, N}

doc"""
  `GenML.BackpropStorage(::Type{MLAlgorithm})`
  `GenML.BackpropStorage(::MLAlgorithm)`

  takes an ML algorithm and allocates all non-parametric memory required to
  evaluate AND backpropagate the algorithm.  Implementations of backprop storage
  should be written in such a way that it be passed as a parameter to the
  forward evaluation which will persistently populate the layer data for the
  backwards propagation.
"""
abstract BackpropStorage{T, N} <: Storage{T, N}

#storage and backpropstorage for layer types is generally not allowed.
(::Type{Storage}){T <: Layer}(::Type{T})         = throw(ArgumentError("storage of layer types not allowed"))
(::Type{Storage}){T <: Layer}(::T)               = throw(ArgumentError("storage of layer types not allowed"))
(::Type{BackpropStorage}){T <: Layer}(::Type{T}) = throw(ArgumentError("storage of layer types not allowed"))
(::Type{BackpropStorage}){T <: Layer}(::T)       = throw(ArgumentError("storage of layer types not allowed"))

################################################################################
## evaluation

# MLAlgorithms should implement evaluate() methods.
@generated function evaluate!{F}(::AbstractArray{F}, mla::MLAlgorithm{F}, ::AbstractArray, ::Storage)
  if mla <: Layer
    :(throw(ArgumentError("Layer types don't implement storage evaluation.")))
  else
    :(throw(MethodError(evaluate!, (AbstractArray{F}, MLAlgorithm{F}, Array, Storage))))
  end
end

@generated function evaluate!{F}(output::AbstractArray{F}, mla::MLAlgorithm{F}, input::AbstractArray)
  #for the situation where the memory has not been preallocated, do the following.
  if input <: AbstractVector
    #catch mismatched input/output type parameters at the compile stage.
    (output <: AbstractVector) || throw(ArgumentError("vector/matrix status of input and output must match."))
    quote
      #generate the temporary storage
      sbuf = Storage(mla, :v)
      #go!
      evaluate!(output, mla, input, sbuf)
    end
  elseif input <: AbstractMatrix
    #catch mismatched input/output type parameters at the compile stage.
    (output <: AbstractMatrix) || throw(ArgumentError("vector/matrix status of input and output must match."))
    quote
      #do a second check to make sure the size of the batches are the same.
      size(output, 2) == size(input, 2) || throw(ArgumentError("input and output matrices must have matching batch sizes"))
      #generate the temporary storage with the correct batch size.
      sbuf = Storage(mla, size(input, 2))
      #go!
      evaluate!(output, mla, input, sbuf)
    end
  end
end

@generated function ml_call{F}(mla::MLAlgorithm{F}, input::Array)
  if input <: Vector
    quote
      output = Array{F,1}(outputs(mla))
      evaluate!(output, mla, input)
      output
    end
  elseif input <: Matrix
    quote
      output = Array{F,2}(outputs(mla), size(input, 2))
      evaluate!(output, mla, input)
      output
    end
  end
end

#optionally, an MLAlgorithm can implement backpropagation.
hasbackpropagation{T <: MLAlgorithm}(::Type{T}) = false  #default does not implement backpropagation

#to give backpropagation algorithms the option of not returning the next deltas for the last layer.
typealias VoidableDeltas{F} Union{AbstractArray{F}, Void}

#compound ml algorithms with backpropagation should implement this function.
@generated function backpropagate!{F}(mla::MLAlgorithm{F},                   #algorithm to be backpropagated.
                                      inp::AbstractArray,                    #input values to be trained
                                      oup::AbstractArray{F},                 #expected output values for training
                                      ouδ::AbstractArray{F},                 #deltas for these output values
                                      sto::BackpropStorage,                  #backprop storage data.
                                                                             #input deltas for passing to the next layer.
                                      inδ::VoidableDeltas{F} = nothing)
  T = mla
  :(throw(ArgumentError("backpropagate! for $T not properly implemented.")))
end

#layers should implement backpropagation without storage function.
function backpropagate!{F}(mla::Layer{F},
                           inp::AbstractArray,
                           oup::AbstractArray{F},
                           ouδ::AbstractArray{F},
                           inδ::VoidableDeltas{F} = nothing)
  T = mla
  :(throw(ArgumentError("backpropagate! for $T not properly implemented.")))
end

macro import_interface()
  quote
    #key functions in the interface
    import ..inputs; import ..outputs; import ..parameters; import ..flatten!; import ..unflatten!
    import ..evaluate!; import ..hasbackpropagation; import ..backpropagate!
    #key type definitions in the interface
    import ..MLAlgorithm; import ..Layer; import ..VoidableDeltas
    import ..Storage; import ..BackpropStorage
    #calling convenience function
    import ..ml_call
    #possibly useful modules.
    import ..TF; import ..CF
    #differentiation interface
    import ..d; import ..dxasy
  end
end

#mathematical utilities
#tools for differentiation
include("./differentiation.jl")
#transfer functions.
include("./transferfunctions.jl")
#include general cost functions
include("./costfunctions.jl")

#subtypes
include("./FCL/FCL.jl") #fully connected layer
include("./MLP/MLP.jl") #multilayer perceptrons

#include general optimizers
include("optimizers/optimizers.jl")

end # module


"""
Converts a dictionary parsed by PowerModels to a System.
Currently Supports MATPOWER and PSSE data files parsed by PowerModels.
Supports kwargs to supply formatters for different device types,
such as `bus_name_formatter` or `gen_name_formatter`.

# Examples
```julia
sys = PSY.pm2ps_dict(pm_data, configpath = "ACTIVSg25k_validation.json",
                    bus_name_formatter = x->string(x["name"]*"-"*string(x["index"])),
                    load_name_formatter = x->strip(join(x["source_id"], "_")))
```
"""
function pm2ps_dict(data::Dict{String, Any}; kwargs...)
    if length(data["bus"]) < 1
        throw(DataFormatError("There are no buses in this file."))
    end

    @info "Constructing System from Power Models" data["name"] data["source_type"]

    sys = System(data["baseMVA"]; kwargs...)

    bus_number_to_bus = read_bus!(sys, data; kwargs...)
    read_loads!(sys, data, bus_number_to_bus; kwargs...)
    read_loadzones!(sys, data, bus_number_to_bus; kwargs...)
    read_gen!(sys, data, bus_number_to_bus; kwargs...)
    read_branch!(sys, data, bus_number_to_bus; kwargs...)
    read_shunt!(sys, data, bus_number_to_bus; kwargs...)
    read_dcline!(sys, data, bus_number_to_bus; kwargs...)

    check!(sys)
    return sys
end

"""
Internal component name retreval from pm2ps_dict
"""
function _get_pm_dict_name(device_dict)
    return get(device_dict, "name", string(device_dict["index"]))
end

"""
Creates a PowerSystems.Bus from a PowerSystems bus dictionary
"""
function make_bus(bus_dict::Dict{String, Any})
    bus = Bus(bus_dict["number"],
                     bus_dict["name"],
                     bus_dict["bustype"],
                     bus_dict["angle"],
                     bus_dict["voltage"],
                     bus_dict["voltagelimits"],
                     bus_dict["basevoltage"]
                     )
     return bus
 end

function make_bus(bus_name, bus_number, d, bus_types)
    bus = make_bus(Dict{String, Any}("name" => bus_name ,
                            "number" => bus_number,
                            "bustype" => bus_types[d["bus_type"]],
                            "angle" => d["va"],
                            "voltage" => d["vm"],
                            "voltagelimits" => (min=d["vmin"], max=d["vmax"]),
                            "basevoltage" => d["base_kv"]
                            ))
    return bus
end

# "From http://www.pserc.cornell.edu/matpower/MATPOWER-manual.pdf Table B-1"
@enum MatpowerBusType begin
    MATPOWER_PQ = 1
    MATPOWER_PV = 2
    MATPOWER_REF = 3
    MATPOWER_ISOLATED = 4
end

function Base.convert(::Type{BusType}, x::MatpowerBusType)
    map = Dict(MATPOWER_ISOLATED => ISOLATED,
               MATPOWER_PQ => PQ,
               MATPOWER_PV => PV,
               MATPOWER_REF => REF)
    return map[x]
end

function read_bus!(sys::System, data; kwargs...)
    @info "Reading bus data"
    bus_number_to_bus = Dict{Int, Bus}()

    bus_types = instances(MatpowerBusType)
    data = sort(collect(data["bus"]), by = x->parse(Int64, x[1]))

    if length(data) == 0
        @error "No bus data found" # TODO : need for a model without a bus
    end

    _get_name = get(kwargs, :bus_name_formatter, _get_pm_dict_name)

    for (i, (d_key, d)) in enumerate(data)
        # d id the data dict for each bus
        # d_key is bus key
        d["name"] = get(d, "name", string(d["bus_i"]))
        bus_name = _get_name(d)
        bus_number = Int(d["bus_i"])
        bus = make_bus(bus_name, bus_number, d, bus_types)
        bus_number_to_bus[bus.number] = bus
        add_component!(sys, bus; skip_validation=SKIP_PM_VALIDATION)
    end

    return bus_number_to_bus
end

function make_load(d, bus; kwargs...)
    _get_name = get(kwargs, :load_name_formatter, x->strip(join(x["source_id"])))
    return PowerLoad(;
        name= _get_name(d),
        available=true,
        model = ConstantPower::LoadModel,
        bus=bus,
        activepower=d["pd"],
        reactivepower=d["qd"],
        maxactivepower=d["pd"],
        maxreactivepower=d["qd"],
    )
end

function read_loads!(sys::System, data, bus_number_to_bus::Dict{Int, Bus}; kwargs...)
    if !haskey(data, "load")
        @error "There are no loads in this file"
        return
    end

    for d_key in keys(data["load"])
        d = data["load"][d_key]
        if d["pd"] != 0.0
            bus = bus_number_to_bus[d["load_bus"]]
            load = make_load(d, bus; kwargs...)

            add_component!(sys, load; skip_validation=SKIP_PM_VALIDATION)
        end
    end
end

function make_loadzones(d, bus_l, activepower, reactivepower; kwargs...)
    _get_name = get(kwargs, :loadzone_name_formatter, _get_pm_dict_name)
    return LoadZones(;
        number=d["index"],
        name=_get_name(d),
        buses=bus_l,
        maxactivepower=sum(activepower),
        maxreactivepower=sum(reactivepower),
    )
end

function read_loadzones!(sys::System, data, bus_number_to_bus::Dict{Int, Bus}; kwargs...)
    if !haskey(data, "areas")
        @info "There are no Load Zones data in this file"
        return
    end

    for (d_key, d) in data["areas"]
        buses = [bus_number_to_bus[b["bus_i"]] for (b_key, b) in data["bus"]
                 if b["area"] == d["index"]]
        bus_names = Set{String}()
        for bus in buses
            push!(bus_names, get_name(bus))
        end

        active_power = Vector{Float64}()
        reactive_power = Vector{Float64}()

        for (key, load) in data["load"]
            load_bus = bus_number_to_bus[load["load_bus"]]
            if get_name(load_bus) in bus_names
                push!(active_power, load["pd"])
                push!(reactive_power, load["qd"])
            end
        end

        d["name"] = get(d, "name", d_key)

        load_zones = make_loadzones(d, buses, active_power, reactive_power; kwargs...)
        add_component!(sys, load_zones; skip_validation=SKIP_PM_VALIDATION)
    end
end

function make_hydro_gen(gen_name, d, bus)
    ramp_agc = get(d, "ramp_agc", get(d, "ramp_10", get(d, "ramp_30", d["pmax"])))
    tech = TechHydro(;
        rating=calculate_rating(d["pmax"], d["qmax"]),
        primemover=convert(PrimeMovers, d["type"]),
        activepowerlimits=(min=d["pmin"], max=d["pmax"]),
        reactivepowerlimits=(min=d["qmin"], max=d["qmax"]),
        ramplimits=(up=ramp_agc / d["mbase"], down=ramp_agc / d["mbase"]),
        timelimits=nothing,
    )

    curtailcost = TwoPartCost(0.0, 0.0)

    return HydroDispatch(name = gen_name,
                         available = Bool(d["gen_status"]),
                         bus = bus,
                         activepower = d["pg"],
                         reactivepower = d["qg"],
                         tech = tech,
                         op_cost = curtailcost)
end

function make_tech_renewable(d)
    tech = TechRenewable(;
        rating=float(d["pmax"]),
        primemover=convert(PrimeMovers, d["type"]),
        reactivepowerlimits=(min=d["qmin"], max=d["qmax"]),
        powerfactor=1.0,
    )

    return tech
end

function make_renewable_dispatch(gen_name, d, bus)
    tech = make_tech_renewable(d)
    cost = TwoPartCost(0.0, 0.0)
    generator = RenewableDispatch(;
        name=gen_name,
        available=Bool(d["gen_status"]),
        bus=bus,
        activepower = d["pg"],
        reactivepower = d["qg"],
        tech=tech,
        op_cost=cost,
    )

    return generator
end

function make_renewable_fix(gen_name, d, bus)
    tech = make_tech_renewable(d)
    generator = RenewableFix(;
        name=gen_name,
        available=Bool(d["gen_status"]),
        bus=bus,
        activepower = d["pg"],
        reactivepower = d["qg"],
        tech=tech,
    )

    return generator
end

function make_generic_battery(gen_name, d, bus)

    # TODO: placeholder
    #battery=GenericBattery(;
    #    name=gen_name,
    #    available=Bool(d["gen_status"]),
    #    bus=bus,
    #    energy=,
    #    capacity=,
    #    rating=,
    #    activepower=,
    #    inputactivepowerlimits=,
    #    outputactivepowerlimits=,
    #    efficiency=,
    #    reactivepower=,
    #    reactivepowerlimits=,
    #)
    #return battery
end

"""
The polynomial term follows the convention that for an n-degree polynomial, at least n + 1 components are needed.
    c(p) = c_n*p^n+...+c_1p+c_0
    c_o is stored in the  field in of the Econ Struct
"""
function make_thermal_gen(gen_name::AbstractString, d::Dict, bus::Bus)
    if haskey(d, "model")
        model = GeneratorCostModel(d["model"])
        if model == PIECEWISE_LINEAR::GeneratorCostModel
            cost_component = d["cost"]
            power_p = [i for (ix, i) in enumerate(cost_component) if isodd(ix)]
            cost_p =  [i for (ix, i) in enumerate(cost_component) if iseven(ix)]
            cost = [(p, c) for (p, c) in zip(cost_p, power_p)]
            fixed = max(0.0, cost[1][1] - (cost[2][1] - cost[1][1])/(cost[2][2] - cost[1][2])*cost[1][2])
            cost = [(c[1] - fixed, c[2]) for c in cost]
        elseif model == POLYNOMIAL::GeneratorCostModel
            if d["ncost"] == 0
                cost = (0.0, 0.0)
                fixed = 0.0
            elseif d["ncost"] == 1
                cost = (0.0, 0.0)
                fixed = d["cost"][1]
            elseif d["ncost"] == 2
                cost = (0.0, d["cost"][1])
                fixed = d["cost"][2]
            elseif d["ncost"] == 3
                cost = (d["cost"][1], d["cost"][2])
                fixed = d["cost"][3]
            else
                throw(DataFormatError("invalid value for ncost: $(d["ncost"]). PowerSystems only supports polynomials up to second degree"))
            end
        end
        startup = d["startup"]
        shutdn = d["shutdown"]
    else
        @warn "Generator cost data not included for Generator: $gen_name"
        tmpcost = ThreePartCost(nothing)
        cost = tmpcost.variable
        fixed = tmpcost.fixed
        startup = tmpcost.startup
        shutdn = tmpcost.shutdn
    end

    # TODO GitHub #148: ramp_agc isn't always present. This value may not be correct.
    ramp_agc = get(d, "ramp_agc", get(d, "ramp_10", get(d, "ramp_30", d["pmax"])))

    tech = TechThermal(;
        rating = sqrt(d["pmax"]^2 + d["qmax"]^2),
        primemover = convert(PrimeMovers, d["type"]),
        fuel = convert(ThermalFuels, d["fuel"]),
        activepowerlimits = (min = d["pmin"], max = d["pmax"]),
        reactivepowerlimits = (min = d["qmin"], max = d["qmax"]),
        ramplimits = (up = ramp_agc / d["mbase"], down = ramp_agc / d["mbase"]),
        timelimits = nothing,
    )
    op_cost = ThreePartCost(;
        variable=cost,
        fixed=fixed,
        startup=startup,
        shutdn=shutdn,
    )

    thermal_gen = ThermalStandard(
        name=gen_name,
        available=Bool(d["gen_status"]),
        bus=bus,
        activepower = d["pg"],
        reactivepower = d["qg"],
        tech=tech,
        op_cost=op_cost,
    )

    return thermal_gen
end

"""
Transfer generators to ps_dict according to their classification
"""
function read_gen!(sys::System, data, bus_number_to_bus::Dict{Int, Bus}; kwargs...)
    @info "Reading generator data"

    if !haskey(data, "gen")
        @error "There are no Generators in this file"
        return nothing
    end

    genmap_file = get(kwargs, :genmap_file, nothing)
    genmap = get_generator_mapping(genmap_file)

    for (name, pm_gen) in data["gen"]
        if haskey(kwargs, :gen_name_formatter)
            _get_name = kwargs[:gen_name_formatter]
        elseif haskey(pm_gen, "name")
            _get_name = _get_pm_dict_name
        elseif haskey(pm_gen, "source_id")
            _get_name = d -> strip(string(d["source_id"][1]) * "-" * string(d["source_id"][2]) * "-" * string(d["index"]))
        end

        gen_name = _get_name(pm_gen)

        bus = bus_number_to_bus[pm_gen["gen_bus"]]
        pm_gen["fuel"] = get(pm_gen, "fuel", "OTHER")
        pm_gen["type"] = get(pm_gen, "type", "OT")
        @debug "Found generator" gen_name bus pm_gen["fuel"] pm_gen["type"]

        gen_type = get_generator_type(pm_gen["fuel"], pm_gen["type"], genmap)
        if gen_type == ThermalStandard
            generator = make_thermal_gen(gen_name, pm_gen, bus)
        elseif gen_type == HydroDispatch
            generator = make_hydro_gen(gen_name, pm_gen, bus)
        elseif gen_type == RenewableDispatch
            generator = make_renewable_dispatch(gen_name, pm_gen, bus)
        elseif gen_type == RenewableFix
            generator = make_renewable_fix(gen_name, pm_gen, bus)
        elseif gen_type == GenericBattery
            @warn "Skipping GenericBattery"
            continue
            # TODO
            #generator = make_generic_battery(gen_name, pm_gen, bus)
        else
            @error "Skipping unsupported generator" gen_type
            continue
        end

        add_component!(sys, generator; skip_validation=SKIP_PM_VALIDATION)
    end
end

function make_branch(name, d, bus_f, bus_t)
    primary_shunt = d["b_fr"]
    alpha = d["shift"]
    branch_type = get_branch_type(d["tap"], alpha)

    if d["transformer"]
        if branch_type == Line
            throw(DataFormatError("Data is mismatched; this cannot be a line. $d"))
        elseif branch_type == Transformer2W
            value = make_transformer_2w(name, d, bus_f, bus_t)
        elseif branch_type == TapTransformer
            value = make_tap_transformer(name, d, bus_f, bus_t)
        elseif branch_type == PhaseShiftingTransformer
            value = make_phase_shifting_transformer(name, d, bus_f, bus_t, alpha)
        else
            error("Unsupported branch type $branch_type")
        end
    else
        # The get_branch_type() logic doesn't work for this data.
        # tap can be 1.0 for this data.
        value = make_line(name, d, bus_f, bus_t)
    end

    return value
end

function make_line(name, d, bus_f, bus_t)
    pf = get(d, "pf", 0.0)
    qf = get(d, "qf", 0.0)

    return Line(;
        name=name,
        available=Bool(d["br_status"]),
        activepower_flow = pf,
        reactivepower_flow = qf,
        arc=Arc(bus_f, bus_t),
        r=d["br_r"],
        x=d["br_x"],
        b=(from=d["b_fr"], to=d["b_to"]),
        rate=d["rate_a"],
        anglelimits=(min=d["angmin"], max=d["angmax"]),
    )
end

function make_transformer_2w(name, d, bus_f, bus_t)
    pf = get(d, "pf", 0.0)
    qf = get(d, "qf", 0.0)
    return Transformer2W(;
        name=name,
        available=Bool(d["br_status"]),
        activepower_flow = pf,
        reactivepower_flow = qf,
        arc=Arc(bus_f, bus_t),
        r=d["br_r"],
        x=d["br_x"],
        primaryshunt=d["b_fr"],  # TODO: which b ??
        rate=d["rate_a"],
    )
end

function make_tap_transformer(name, d, bus_f, bus_t)
    pf = get(d, "pf", 0.0)
    qf = get(d, "qf", 0.0)
    return TapTransformer(;
        name=name,
        available=Bool(d["br_status"]),
        activepower_flow = pf,
        reactivepower_flow = qf,
        arc=Arc(bus_f, bus_t),
        r=d["br_r"],
        x=d["br_x"],
        tap=d["tap"],
        primaryshunt=d["b_fr"],  # TODO: which b ??
        rate=d["rate_a"],
    )
end

function make_phase_shifting_transformer(name, d, bus_f, bus_t, alpha)
    pf = get(d, "pf", 0.0)
    qf = get(d, "qf", 0.0)
    return PhaseShiftingTransformer(;
        name=name,
        available=Bool(d["br_status"]),
        activepower_flow = pf,
        reactivepower_flow = qf,
        arc=Arc(bus_f, bus_t),
        r=d["br_r"],
        x=d["br_x"],
        tap=d["tap"],
        primaryshunt=d["b_fr"],  # TODO: which b ??
        α=alpha,
        rate=d["rate_a"],
    )
end

function read_branch!(sys::System, data, bus_number_to_bus::Dict{Int, Bus}; kwargs...)
    @info "Reading branch data"
    if !haskey(data, "branch")
        @info "There is no Branch data in this file"
        return
    end

    _get_name = get(kwargs, :branch_name_formatter, _get_pm_dict_name)

    for (d_key, d) in data["branch"]
        d["name"] = get(d, "name", d_key)
        name = _get_name(d)
        bus_f = bus_number_to_bus[d["f_bus"]]
        bus_t = bus_number_to_bus[d["t_bus"]]
        value = make_branch(name, d, bus_f, bus_t)

        add_component!(sys, value; skip_validation=SKIP_PM_VALIDATION)
    end
end

function make_dcline(name, d, bus_f, bus_t)
    return HVDCLine(;
        name=name,
        available=Bool(d["br_status"]),
        activepower_flow = get(d, "pf", 0.0),
        arc=Arc(bus_f, bus_t),
        activepowerlimits_from=(min=d["pminf"] , max=d["pmaxf"]),
        activepowerlimits_to=(min=d["pmint"], max=d["pmaxt"]),
        reactivepowerlimits_from=(min=d["qminf"], max=d["qmaxf"]),
        reactivepowerlimits_to=(min=d["qmint"], max=d["qmaxt"]),
        loss=(l0=d["loss0"], l1 =d["loss1"]),
    )
end

function read_dcline!(sys::System, data, bus_number_to_bus::Dict{Int, Bus}; kwargs...)
    @info "Reading DC Line data"
    if !haskey(data, "dcline")
        @info "There is no DClines data in this file"
        return
    end

    _get_name = get(kwargs, :branch_name_formatter, _get_pm_dict_name)

    for (d_key, d) in data["dcline"]
        d["name"] = get(d, "name", d_key)
        name = _get_name(d)
        bus_f = bus_number_to_bus[d["f_bus"]]
        bus_t = bus_number_to_bus[d["t_bus"]]

        dcline = make_dcline(name, d, bus_f, bus_t)
        add_component!(sys, dcline, skip_validation=SKIP_PM_VALIDATION)
    end
end

function make_shunt(name, d, bus)
    return FixedAdmittance(;
        name=name,
        available=Bool(d["status"]),
        bus=bus,
        Y=(-d["gs"] + d["bs"]im),
    )
end

function read_shunt!(sys::System, data, bus_number_to_bus::Dict{Int, Bus}; kwargs...)
    @info "Reading branch data"
    if !haskey(data, "shunt")
        @info "There is no shunt data in this file"
        return
    end

    _get_name = get(kwargs, :shunt_name_formatter, _get_pm_dict_name)

    for (d_key, d) in data["shunt"]
        d["name"] = get(d, "name", d_key)
        name = _get_name(d)
        bus = bus_number_to_bus[d["shunt_bus"]]
        shunt = make_shunt(name, d, bus)

        add_component!(sys, shunt; skip_validation=SKIP_PM_VALIDATION)
    end
end


module CP

using Compat: eachrow, isnothing
using LinearAlgebra: det
using Parameters: @with_kw
using Setfield: get, set, @lens, @set
using StaticArrays: SVector, SMatrix, FieldVector
using Unitful
using UnitfulAtomic

using ..Inputs:
    Namelist,
    QuantumESPRESSOInput,
    Card,
    optionof,
    optionpool,
    allnamelists,
    allcards,
    required_namelists,
    optional_namelists,
    required_cards,
    optional_cards

import AbInitioSoftwareBase.Inputs: asstring, groupname
import Crystallography: Bravais, Lattice
# import Pseudopotentials: pseudoformat
import ..Inputs:
    optionpool,
    allnamelists,
    allcards,
    optionof,
    required_namelists,
    optional_namelists,
    required_cards,
    optional_cards
using ..Formatter: delimiter, newline, indent, floatfmt, intfmt

export ControlNamelist,
    SystemNamelist,
    ElectronsNamelist,
    IonsNamelist,
    CellNamelist,
    PressAiNamelist,
    WannierNamelist,
    AtomicSpecies,
    AtomicSpeciesCard,
    AtomicPosition,
    AtomicPositionsCard,
    CellParametersCard,
    RefCellParametersCard,
    AtomicVelocity,
    AtomicVelocitiesCard,
    AtomicForce,
    AtomicForcesCard,
    CPInput,
    optconvert

include("namelists.jl")
include("cards.jl")

"""
    CPInput <: QuantumESPRESSOInput
    CPInput(control, system, electrons, ions, cell, press_ai, wannier, atomic_species, atomic_positions, atomic_velocities, cell_parameters, ref_cell_parameters, constraints, occupations, atomic_forces, plot_wannier, autopilot)

Construct a `CPInput` which represents the input of program `cp.x`.

# Arguments
- `control::ControlNamelist=ControlNamelist()`: the `CONTROL` namelist of the input. Optional.
- `system::SystemNamelist=SystemNamelist()`: the `SYSTEM` namelist of the input. Optional.
- `electrons::ElectronsNamelist=ElectronsNamelist()`: the `ELECTRONS` namelist of the input. Optional.
- `ions::IonsNamelist=IonsNamelist()`: the `IONS` namelist of the input. Optional.
- `cell::CellNamelist=CellNamelist()`: the `CELL` namelist of the input. Optional.
- `press_ai::PressAiNamelist=PressAiNamelist()`: the `PRESS_AI` namelist of the input. Optional.
- `wannier::WannierNamelist=WannierNamelist()`: the `WANNIER` namelist of the input. Optional.
- `atomic_species::AtomicSpeciesCard`: the `ATOMIC_SPECIES` card of the input. Must be provided explicitly.
- `atomic_positions::AtomicPositionsCard`: the `ATOMIC_POSITIONS` card of the input. Must be provided explicitly.
- `atomic_velocities::AtomicVelocitiesCard`: the `ATOMIC_VELOCITIES` card of the input. Must be provided explicitly.
- `cell_parameters::Union{Nothing,CellParametersCard}`: the `CELL_PARAMETERS` card of the input. Must be either `nothing` or a `CellParametersCard`.
- `ref_cell_parameters::Union{Nothing,RefCellParametersCard}`: the `REF_CELL_PARAMETERS` card of the input. Must be either `nothing` or a `CellParametersCard`.
"""
@with_kw struct CPInput <: QuantumESPRESSOInput
    control::ControlNamelist = ControlNamelist()
    system::SystemNamelist = SystemNamelist()
    electrons::ElectronsNamelist = ElectronsNamelist()
    ions::IonsNamelist = IonsNamelist()
    cell::CellNamelist = CellNamelist()
    press_ai::PressAiNamelist = PressAiNamelist()
    wannier::WannierNamelist = WannierNamelist()
    atomic_species::AtomicSpeciesCard
    atomic_positions::AtomicPositionsCard
    atomic_velocities::AtomicVelocitiesCard
    cell_parameters::Union{Nothing,CellParametersCard} = nothing
    ref_cell_parameters::Union{Nothing,RefCellParametersCard} = nothing
    constraints::Union{Nothing,Float64} = nothing
    occupations::Union{Nothing,Float64} = nothing
    atomic_forces::Union{Nothing,AtomicForcesCard} = nothing
    plot_wannier::Union{Nothing,Float64} = nothing
    autopilot::Union{Nothing,Float64} = nothing
    @assert(
        !(isnothing(cell_parameters) && system.ibrav == 0),
        "Cannot specify `ibrav = 0` with an empty `cell_parameters`!"
    )
end # struct CPInput

"""
    optconvert(new_option::AbstractString, card::AbstractCellParametersCard)

Convert the option of an `AbstractCellParametersCard` from "bohr" to "angstrom", or its reverse.

!!! warning
    It does not support conversion between `"alat"` and the others.
"""
function optconvert(new_option::AbstractString, card::AbstractCellParametersCard)
    old_option = optionof(card)
    if new_option == old_option
        return card  # No conversion is needed
    else
        return typeof(card)(
            if (old_option => new_option) == ("bohr" => "angstrom")
                @. ustrip(u"angstrom", card.data * u"bohr")
            elseif (old_option => new_option) == ("angstrom" => "bohr")
                @. ustrip(u"bohr", card.data * u"angstrom")
            else
                error("unknown conversion rule $(old_option => new_option)!")
            end,
        )
    end
end # function optconvert

groupname(::Type{ControlNamelist}) = "CONTROL"
groupname(::Type{SystemNamelist}) = "SYSTEM"
groupname(::Type{ElectronsNamelist}) = "ELECTRONS"
groupname(::Type{IonsNamelist}) = "IONS"
groupname(::Type{CellNamelist}) = "CELL"
groupname(::Type{AtomicSpeciesCard}) = "ATOMIC_SPECIES"
groupname(::Type{AtomicPositionsCard}) = "ATOMIC_POSITIONS"
groupname(::Type{<:CellParametersCard}) = "CELL_PARAMETERS"

"""
    Bravais(nml::SystemNamelist)

Return a `Bravais` from a `SystemNamelist`.
"""
Bravais(nml::SystemNamelist) = Bravais(nml.ibrav)

"""
    Lattice(nml::SystemNamelist)

Return a `Lattice` from a `SystemNamelist`.
"""
Lattice(nml::SystemNamelist) = Lattice(Bravais(nml), nml.celldm)

"""
    cellvolume(nml::SystemNamelist)

Return the volume of the cell based on the information given in a `SystemNamelist`, in atomic unit.
"""
cellvolume(nml::SystemNamelist) = cellvolume(Lattice(nml))
"""
    cellvolume(card)

Return the cell volume of a `CellParametersCard` or `RefCellParametersCard`, in atomic unit.

!!! warning
    It will throw an error if the option is `"alat"`.
"""
function cellvolume(card::AbstractCellParametersCard)
    option = optionof(card)
    if option == "bohr"
        abs(det(card.data))
    elseif option == "angstrom"
        ustrip(u"bohr^3", abs(det(card.data)) * u"angstrom^3")
    else  # option == "alat"
        error("information not enough! Parameter `celldm[1]` needed!")
    end
end # function cellvolume

function asstring(data::AtomicSpecies)
    return join(
        (sprintf1("%3s", data.atom), sprintf1(floatfmt(data), data.mass), data.pseudopot),
        delimiter(data),
    )
end
function asstring(card::AtomicSpeciesCard)
    # Using generator expressions in `join` is faster than using `Vector`s.
    return "ATOMIC_SPECIES" *
           newline(card) *
           join((indent(card) * asstring(x) for x in unique(card.data)), newline(card))
end
function asstring(data::AtomicPosition)
    f(x) = x ? "" : "0"
    return join(
        [
            sprintf1("%3s", data.atom)
            map(x -> sprintf1(floatfmt(data), x), data.pos)
            map(f, data.if_pos)
        ],
        delimiter(data),
    )
end
function asstring(card::AtomicPositionsCard)
    return "ATOMIC_POSITIONS { $(optionof(card)) }" *
           newline(card) *
           join((indent(card) * asstring(x) for x in card.data), newline(card))
end
function asstring(card::CellParametersCard)
    it = (
        indent * join((sprintf1(floatfmt(card), x) for x in row), delimiter(card)) for
        row in eachrow(card.data)
    )
    return "CELL_PARAMETERS { $(optionof(card)) }" * newline(card) * join(it, newline)
end

optionof(::AtomicVelocitiesCard) = "a.u"

optionpool(::Type{<:AtomicPositionsCard}) =
    ("alat", "bohr", "angstrom", "crystal", "crystal_sg")
optionpool(::Type{<:CellParametersCard}) = ("alat", "bohr", "angstrom")
optionpool(::Type{<:AtomicVelocity}) = ("a.u",)
optionpool(::Type{<:RefCellParametersCard}) = ("bohr", "angstrom")

allnamelists(::Type{CPInput}) =
    (:control, :system, :electrons, :ions, :cell, :press_ai, :wannier)
allnamelists(x::CPInput) = (getfield(x, f) for f in allnamelists(typeof(x)))

allcards(::Type{CPInput}) = (
    :atomic_species,
    :atomic_positions,
    :atomic_velocities,
    :cell_parameters,
    :ref_cell_parameters,
    :constraints,
    :occupations,
    :atomic_forces,
    :plot_wannier,
    :autopilot,
)
allcards(x::CPInput) = (getfield(x, f) for f in allcards(typeof(x)))

required_namelists(::Type{CPInput}) = (:control, :system, :electrons)
required_namelists(x::CPInput) = (getfield(x, f) for f in required_namelists(typeof(x)))

optional_namelists(::Type{CPInput}) = (:ions, :cell, :press_ai, :wannier)
optional_namelists(x::CPInput) = (getfield(x, f) for f in optional_namelists(typeof(x)))

required_cards(::Type{CPInput}) = (:atomic_species, :atomic_positions)
required_cards(x::CPInput) = (getfield(x, f) for f in required_cards(typeof(x)))

optional_cards(::Type{CPInput}) = (
    :atomic_velocities,
    :cell_parameters,
    :ref_cell_parameters,
    :constraints,
    :occupations,
    :atomic_forces,
    :plot_wannier,
    :autopilot,
)
optional_cards(x::CPInput) = (getfield(x, f) for f in optional_cards(typeof(x)))

end


"""
`module  JAC.DoubleAutoIonization`
    ... a submodel of JAC that contains all methods for computing double Auger and autoionization amplitudes and rates.
"""
module DoubleAutoIonization

    using Printf, JAC, ..AngularMomentum, ..AtomicState, ..AutoIonization, ..Basics, ..ManyElectron, ..TableStrings

    """
    `struct  DoubleAutoIonization.Settings`  ... defines a type for the settings in estimating double-Auger and autoionization rates.

        + green                   ::Array{GreenChannel,1}   ... Precalculated and user-specified Green function of the ion.
        + NoEnergySharings        ::Int64                   ... Number of energy sharings that are used in the computations for each line.
        + printBefore             ::Bool                    ... True, if all energies and lines are printed before their evaluation.
        + minAugerEnergy          ::Float64                 ... Minimum energy of free (Auger) electrons to be included.
        + maxAugerEnergy          ::Float64                 ... Maximum energy of free (Auger) electrons to be included.
        + maxKappa                ::Int64                   ... Maximum kappa value of partial waves to be included.
        + operator                ::AbstractEeInteraction   ... Auger operator that is to be used for evaluating the Auger amplitudes:
                                                                allowed values are: CoulombInteraction(), BreitInteraction(), ...
        + lineSelection           ::LineSelection           ... Specifies the selected levels, if any.
    """
    struct Settings
        green                     ::Array{GreenChannel,1}
        NoEnergySharings          ::Int64
        printBefore               ::Bool
        minAugerEnergy            ::Float64
        maxAugerEnergy            ::Float64
        maxKappa                  ::Int64
        operator                  ::AbstractEeInteraction
        lineSelection             ::LineSelection
    end

    """
    `DoubleAutoIonization.Settings()`  ... constructor for the default values of DoubleAutoIonization line computations
    """
    function Settings()
        Settings(GreenChannel[], 0, false, 0., 10e5, 100, CoulombInteraction(), LineSelection())
    end

    # `Base.show(io::IO, settings::DoubleAutoIonization.Settings)`  ... prepares a proper printout of the variable settings::DoubleAutoIonization.Settings.
    function Base.show(io::IO, settings::DoubleAutoIonization.Settings)
        println(io, "green:                         (settings.green)  ")
        println(io, "NoEnergySharings:             $(settings.NoEnergySharings)  ")
        println(io, "printBefore:                  $(settings.printBefore)  ")
        println(io, "minAugerEnergy:               $(settings.minAugerEnergy)  ")
        println(io, "maxAugerEnergy:               $(settings.maxAugerEnergy)  ")
        println(io, "maxKappa:                     $(settings.maxKappa)  ")
        println(io, "operator:                     $(settings.operator)  ")
        println(io, "lineSelection:                $(settings.lineSelection)  ")
    end

    """
    `struct  DoubleAutoIonization.ReducedChannel`
        ... defines a type for a DoubleAutoIonization (reduced) channel to help characterize a scattering (continuum) state of many
            electron-states with two free electrons.

        + energy1        ::Float64              ... energy of the partial wave_1
        + kappa1         ::Int64                ... partial-wave of the free electron_1
        + phase1         ::Float64              ... phase of the partial wave_1
        + xSymmetry      ::LevelSymmetry        ... angular momentum and parity after coupling electron1.
        + energy2        ::Float64              ... energy of the partial wave_2
        + kappa2         ::Int64                ... partial-wave of the free electron_2
        + phase2         ::Float64              ... phase of the partial wave_2
        + amplitude      ::Complex{Float64}     ... Auger amplitude associated with the given channel.
    """
    struct  ReducedChannel
        energy1          ::Float64
        kappa1           ::Int64
        phase1           ::Float64
        xSymmetry        ::LevelSymmetry
        energy2          ::Float64
        kappa2           ::Int64
        phase2           ::Float64
        amplitude        ::Complex{Float64}
    end

    # `Base.show(io::IO, channel::DoubleAutoIonization.ReducedChannel)`
    #       ... prepares a proper printout of the variable cannel::DoubleAutoIonization.ReducedChannel.
    function Base.show(io::IO, channel::DoubleAutoIonization.ReducedChannel)
        ## sa = "reduced DA channel for J^P=(channel.tSymmetry)  with " *
        ##      "kappa1=$(channel.kappa1), energy1=$(channel.energy1), phase1=$(channel.phase1), " *
        ##      "kappa2=$(channel.kappa2), energy2=$(channel.energy2), phase2=$(channel.phase2), amp=$(channel.amplitude) "
        ## println(io, sa)
        println(io, "multipole:                    $(channel.multipole)  ")
        println(io, "gauge:                        $(channel.gauge)  ")
        println(io, "omega:                        $(channel.omega)  ")
        println(io, "energy1:                      $(channel.energy1)  ")
        println(io, "kappa1:                       $(channel.kappa1)  ")
        println(io, "phase1:                       $(channel.phase1)  ")
        println(io, "xSymmetry:                    $(channel.xSymmetry)  ")
        println(io, "energy2:                      $(channel.energy2)  ")
        println(io, "kappa2:                       $(channel.kappa2)  ")
        println(io, "phase2:                       $(channel.phase2)  ")
        println(io, "amplitude:                    $(channel.amplitude)  ")
    end

    """
    `struct  DoubleAutoIonization.Sharing`
        ... defines a type for a DoubleAutoIonization sharing to help characterize energy sharing between the two emitted electrons.

        + epsilon1         ::Float64         ... Energy of (free) electron 1.
        + epsilon2         ::Float64         ... Energy of (free) electron 2.
        + weight           ::Float64         ... Gauss-Lengendre weight of this sharing for energy-integrated quantities.
        + differentialRate ::EmProperty      ... differential cross section of this energy sharing.
        + channels         ::Array{DoubleAutoIonization.ReducedChannel,1}  ... List of DoubleAutoIonization (reduced) channels of this line.
    """
    struct  Sharing
        epsilon1           ::Float64
        epsilon2           ::Float64
        weight             ::Float64
        differentialRate   ::EmProperty
        channels           ::Array{DoubleAutoIonization.ReducedChannel,1}
    end

    # `Base.show(io::IO, sharing::DoubleAutoIonization.Sharing)`  ... prepares a proper printout of the variable sharing::DoubleAutoIonization.Sharing.
    function Base.show(io::IO, sharing::DoubleAutoIonization.Sharing)
        println(io, "epsilon1:               $(sharing.epsilon1)  ")
        println(io, "epsilon2:               $(sharing.epsilon2)  ")
        println(io, "differentialCs:         $(sharing.differentialCs)  ")
        println(io, "channels:               $(sharing.channels)  ")
    end

    """
    `struct  DoubleAutoIonization.Line`  ... defines a type for a double Auger line that includes sharings and their reduced amplitudes.

        + initialLevel   ::Level            ... initial-(state) level
        + finalLevel     ::Level            ... final-(state) level
        + totalRate      ::EmProperty       ... Total rate of this line.
        + sharings       ::Array{DoubleAutoIonization.Sharing,1}  ... List of DoubleAutoIonization sharings of this line.
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        totalRate        ::EmProperty
        sharings         ::Array{DoubleAutoIonization.Sharing,1}
    end

    # `Base.show(io::IO, line::DoubleAutoIonization.Line)`  ... prepares a proper printout of the variable line::DoubleAutoIonization.Line.
     function Base.show(io::IO, line::DoubleAutoIonization.Line)
        println(io, "initialLevel:           $(line.initialLevel)  ")
        println(io, "finalLevel:             $(line.finalLevel)  ")
        println(io, "totalRate:              $(line.totalRate)  ")
        println(io, "sharings:               $(line.sharings)  ")
    end

    """
    `DoubleAutoIonization.computeAmplitudesProperties(line::DoubleAutoIonization.Line, grid::Radial.Grid, settings::DoubleAutoIonization.Settings)`
        ... to compute all amplitudes and properties of the given line; a line::DoubleAutoIonization.Line is returned for which the amplitudes
            and properties have now been evaluated.
    """
    function  computeAmplitudesProperties(line::DoubleAutoIonization.Line, grid::Radial.Grid, settings::DoubleAutoIonization.Settings)
        newSharings = DoubleAutoIonization.Sharing[]
        for sharing in line.sharings
            newChannels = DoubleAutoIonization.ReducedChannel[]
            for ch in sharing.channels
                # Generate a continuum orbital
                phase     = 1.3
                amplitude = 1.0im
                push!( newChannels, ReducedChannel(ch.energy1, ch.kappa1, phase, ch.xSymmetry, ch.energy2, ch.kappa2, phase, amplitude) )
            end
            # Calculate the differential rate
            diffRate = EmProperty(-1., -1.)
            push!( newSharings, DoubleAutoIonization.Sharing(sharing.epsilon1, sharing.epsilon2, sharing.weight, diffRate, newChannels) )
        end
        # Calculate the total rate
        totalRate = EmProperty(-1., -1.)
        line = DoubleAutoIonization.Line(line.initialLevel, line.finalLevel, totalRate, newSharings)
        return( line )
    end

    """
    `DoubleAutoIonization.computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                       settings::DoubleAutoIonization.Settings; output=true)`
        ... to compute the double-Auger transition amplitudes and all properties as requested by the given settings. A list of
            lines::Array{DoubleAutoIonization.Lines} is returned.
    """
    function  computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                           settings::DoubleAutoIonization.Settings; output=true)
        println("")
        printstyled("DoubleAutoIonization.computeLines(): The computation of double-Auger rates starts now ... \n", color=:light_green)
        printstyled("----------------------------------------------------------------------------------------- \n", color=:light_green)
        println("")
        #
        lines = DoubleAutoIonization.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        if  settings.printBefore    DoubleAutoIonization.displayLines(lines)    end
        # Calculate all amplitudes and requested properties
        newLines = DoubleAutoIonization.Line[]
        for  line in lines
            newLine = DoubleAutoIonization.computeAmplitudesProperties(line, grid, settings)
            push!( newLines, newLine)
        end
        # Print all results to screen
        DoubleAutoIonization.displayTotalRates(stdout, lines, settings)
        DoubleAutoIonization.displayDifferentialRates(stdout, lines, settings)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   DoubleAutoIonization.displayTotalRates(iostream, lines, settings)
                           DoubleAutoIonization.displayDifferentialRates(iostream, lines, settings)     end
        #
        if    output    return( lines )
        else            return( nothing )
        end
    end

    """
    `DoubleAutoIonization.determineLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, settings::DoubleAutoIonization.Settings)`
        ... to determine a list of DoubleAutoIonization.Line's for transitions between levels from the initial- and final-state multiplets,
            and  by taking into account the particular selections and settings for this computation; an Array{DoubleAutoIonization.Line,1} is
            returned. Apart from the level specification and sharing, all physical properties are set to zero during the initialization process.
    """
    function  determineLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, settings::DoubleAutoIonization.Settings)
        lines = DoubleAutoIonization.Line[]
        for  iLevel  in  initialMultiplet.levels
            for  fLevel  in  finalMultiplet.levels
                if  Basics.selectLevelPair(iLevel, fLevel, settings.lineSelection)
                    energy     = iLevel.energy - fLevel.energy
                    if  energy < settings.minAugerEnergy  ||  energy > settings.maxAugerEnergy    continue   end
                    sharings = DoubleAutoIonization.determineSharingsAndChannels(fLevel, iLevel, energy, settings)
                    push!( lines, DoubleAutoIonization.Line(iLevel, fLevel, EmProperty(0., 0.,), sharings) )
                end
            end
        end
        return( lines )
    end

    """
    `DoubleAutoIonization.determineSharingsAndChannels(finalLevel::Level, initialLevel::Level, energy::Float64, settings::DoubleAutoIonization.Settings)`
        ... to determine a list of DoubleAutoIonization Sharing's and Channel's for a transitions from the initial to final level and by taking into
            account the particular settings of for this computation; an Array{DoubleAutoIonization.Sharing,1} is returned.
    """
    function determineSharingsAndChannels(finalLevel::Level, initialLevel::Level, energy::Float64, settings::DoubleAutoIonization.Settings)
        sharings = DoubleAutoIonization.Sharing[];    eSharings = Basics.determineEnergySharings(energy, settings.NoEnergySharings)
        for  es in eSharings
            epsilon1  = es[1];    epsilon2 = es[2];    weight = es[3]
            channels  = DoubleAutoIonization.ReducedChannel[];
            symi      = LevelSymmetry(initialLevel.J, initialLevel.parity);    symf = LevelSymmetry(finalLevel.J, finalLevel.parity)
            couplings = AngularMomentum.allowedDoubleKappaCouplingSequence(symf, symi, settings.maxKappa)
            for  coupling  in  couplings
                push!(channels, ReducedChannel(epsilon1, coupling[1], 0., coupling[2], epsilon2, coupling[3], 0., Complex(0.)) )
            end
            push!(sharings, DoubleAutoIonization.Sharing(epsilon1, epsilon2, weight, EmProperty(0., 0.), channels) )
        end
        return( sharings )
    end

    """
    `DoubleAutoIonization.displayDifferentialRates(stream::IO, lines::Array{DoubleAutoIonization.Line,1}, settings::DoubleAutoIonization.Settings)`
        ... to display all differential rates, etc. of the selected lines. A neat table is printed but nothing is returned otherwise.
    """
    function  displayDifferentialRates(stream::IO, lines::Array{DoubleAutoIonization.Line,1}, settings::DoubleAutoIonization.Settings)
        #
        # First, print lines and sharings
        nx = 128
        println(stream, " ")
        println(stream, "  Energy-differential rates of selected double-Auger lines:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.flushleft(38, "Energies (all in " * TableStrings.inUnits("energy") * ")"; na=3);
        sb = sb * TableStrings.flushleft(38, "  i -- f        epsilon_1     epsilon_2"; na=3)
        sa = sa * TableStrings.center(14, "Weight"; na=0);                            sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(34, "Cou --  diff. rate  -- Bab"; na=3)
        sb = sb * TableStrings.center(34, TableStrings.inUnits("rate") * "        " *
                                          TableStrings.inUnits("rate"); na=3)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";      isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                           fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            energy = line.initialLevel.energy - line.finalLevel.energy
            sa = sa * @sprintf("%.5e", Defaults.convertUnits("energy: from atomic", energy)) * "    "
            #
            for  (is, sharing)  in  enumerate(line.sharings)
                if  is == 1     sb = sa     else    sb = TableStrings.hBlank( length(sa) )    end
                sb = sb * @sprintf("%.4e", Defaults.convertUnits("energy: from atomic", sharing.epsilon1))                 * "    "
                sb = sb * @sprintf("%.4e", Defaults.convertUnits("energy: from atomic", sharing.epsilon2))                 * "    "
                sb = sb * @sprintf("%.4e",                                              sharing.weight)                    * "       "
                sb = sb * @sprintf("%.5e", Defaults.convertUnits("rate: from atomic", sharing.differentialRate.Coulomb))   * "   "
                sb = sb * @sprintf("%.5e", Defaults.convertUnits("rate: from atomic", sharing.differentialRate.Babushkin)) * "   "
                println(stream,  sb )
            end
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

    """
    `DoubleAutoIonization.displayLines(lines::Array{DoubleAutoIonization.Line,1})`
        ... to display a list of lines, sharings and channels that have been selected due to the prior settings. A neat table
            of all selected transitions and energies is printed but nothing is returned otherwise.
    """
    function  displayLines(lines::Array{DoubleAutoIonization.Line,1})
        #
        # First, print lines and sharings
        nx = 82
        println(" ")
        println("  Selected double-Auger lines & sharings:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.flushleft(54, "Energies (all in )" * TableStrings.inUnits("energy") * ")"; na=5);
        sb = sb * TableStrings.flushleft(54, "   i -- f       epsilon_1     epsilon_2"; na=5)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";      isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                           fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            energy = line.initialLevel.energy - line.finalLevel.energy
            sa = sa * @sprintf("%.5e", Defaults.convertUnits("energy: from atomic", energy)) * "    "
            #
            for  (is, sharing)  in  enumerate(line.sharings)
                if  is == 1     sb = sa     else    sb = TableStrings.hBlank( length(sa) )    end
                sb = sb * @sprintf("%.4e", Defaults.convertUnits("energy: from atomic", sharing.epsilon1)) * "    "
                sb = sb * @sprintf("%.4e", Defaults.convertUnits("energy: from atomic", sharing.epsilon2)) * "    "
                println( sb )
            end
        end
        println("  ", TableStrings.hLine(nx))
        #
        #
        # Second, print lines and channles
        nx = 123
        println(" ")
        println("  Selected photo-double ionization lines & channels:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.flushleft(77, "Channels (all energies in " * TableStrings.inUnits("energy") * ")" ; na=5);
        sb = sb * TableStrings.flushleft(77, "(epsilon, kappa)_1   -->    J^P_x  -->   (epsilon, kappa)_2   -->  J^P_t = J^P_i"; na=5)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";      isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                           fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            for (ic, ch) in enumerate(line.sharings[1].channels)
                if  ic == 1     sb = sa     else    sb = TableStrings.hBlank( length(sa) )    end
                skappa1 = "  " * string(ch.kappa1);    skappa2 = "  " * string(ch.kappa2);    sxsym = "      " * string(ch.xSymmetry)
                sb = sb * " (" * @sprintf("%.3e", Defaults.convertUnits("energy: from atomic", ch.energy1)) * ", " * skappa1[end-2:end] * ")    -->"
                sb = sb * sxsym[end-8:end] * "  -->   "
                sb = sb * " (" * @sprintf("%.3e", Defaults.convertUnits("energy: from atomic", ch.energy2)) * ", " * skappa2[end-2:end] * ")    -->     "
                sb = sb * string(isym)
                println( sb )
            end
        end
        println("  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

    """
    `DoubleAutoIonization.displayTotalRates(stream::IO, lines::Array{DoubleAutoIonization.Line,1}, settings::DoubleAutoIonization.Settings)`
        ... to display all total rates, etc. of the selected lines. A neat table is printed but nothing is returned otherwise.
    """
    function  displayTotalRates(stream::IO, lines::Array{DoubleAutoIonization.Line,1}, settings::DoubleAutoIonization.Settings)
        #
        # First, print lines and sharings
        nx = 88
        println(stream, " ")
        println(stream, "  Total rates of selected double-Auger lines:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=2);                         sb = sb * TableStrings.hBlank(21)
        sa = sa * TableStrings.center(12, "i--Energy--f"; na=3)
        sb = sb * TableStrings.center(12,TableStrings.inUnits("energy"); na=3)
        sa = sa * TableStrings.center(30, "Cou --   total rate   -- Bab"; na=4)
        sb = sb * TableStrings.center(30, TableStrings.inUnits("rate") * "          " *
                                          TableStrings.inUnits("rate"); na=3)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";      isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                           fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            energy = line.initialLevel.energy - line.finalLevel.energy
            sa = sa * @sprintf("%.4e", Defaults.convertUnits("energy: from atomic", energy))                 * "      "
            #
            sb = sa * @sprintf("%.5e", Defaults.convertUnits("rate: from atomic", line.totalRate.Coulomb))   * "     "
            sb = sb * @sprintf("%.5e", Defaults.convertUnits("rate: from atomic", line.totalRate.Babushkin)) * "   "
            println(stream,  sb )
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

end # module


module PlotFunctions
	import  ..MasslistFunctions
	using PyPlot

	export massDefectPlot, bananaPlot, addClickToggle

	function massDefectPlot(masses, compositions, concentrations, colors, plotTitle, colorCodeTitle; dotSize = 10, maxMass = 450, maxDefect = 0.25, minConc = 0.02, sumformulas = false)
	  figure()

	  h2o = MasslistFunctions.createCompound(H=2,O=1, Hplus=0)
	  o = MasslistFunctions.createCompound(O=1, Hplus=0)

	  for i=1:length(masses)
	    m = masses[i]
	    adduct = h2o
	    if (inCompositions(compositions[:,i] + adduct[4], compositions))
	      m1 = m + adduct[1]
	      plot([m, m1], [m-round(m), m1 - round(m1)], color="lightblue", zorder=1)
	    end
	    adduct = o
	    if (inCompositions(compositions[:,i] + adduct[4], compositions))
	      m1 = m + adduct[1]
	      plot([m, m1], [m-round(m), m1 - round(m1)], color="red", zorder=1)
	    end
	    if sumformulas text(masses[i],masses[i]-round(masses[i]),sumFormulaStringFromCompositionArray(compositions[:,i]), color="grey", clip_on=true, verticalalignment="center", size=10, zorder=100) end
	  end

	  scatter(masses, masses-round.(masses),dotSize*log.(concentrations./minConc), colors, zorder=10, linewidths=0.5)
	  xlim(0,maxMass)
	  ylim(0,maxDefect)
	  cb=colorbar()
	  cb["ax"]["set_ylabel"](colorCodeTitle)
	  xlabel("Mass [amu]")
	  ylabel("Kendrick Mass Defect")
	  title(plotTitle)
	  grid("on")
	end

	function bananaPlot(xbins,ybins,meshdataXY;subplotAx=0)
	    println("Not implemented yet!!!")
	    #if subplotAx == 0
	    #    figure()
	    #    subplotAx = subplot(1,1,1)
	    #end
	    pcolormesh([DateTime(2018,05,05),DateTime(2018,05,06),DateTime(2018,05,07)],[1,7,9],[1 2 3; 4 5 6; 7 8 9][1:end-1,1:end-1])
	    #yscale("log")
	end

	function addClickToggle(plotAxis)
	    f = plotAxis.get_figure()
	    plotLines = plotAxis.get_lines()
	    legLines = plotAxis.get_legend().get_lines()
	    for ll in legLines
	        ll.set_picker(5)
	    end
	    linesDict = Dict(zip(legLines,plotLines))
	    function onLegendPick(event)
	        legLine = event.artist
	        plotLine = linesDict[legLine]
	        isVisible = plotLine.get_visible()
	        plotLine.set_visible(! isVisible)
	        if ! isVisible
	            legLine.set_alpha(1.0)
	        else
	            legLine.set_alpha(0.1)
	        end
	        f.canvas.draw()
	    end
	    f.canvas.mpl_connect("pick_event",onLegendPick)
	end
	function getVisibleTraces(axis)
		return [ axis.get_lines()[i].get_visible() for i in 1:length(axis.get_lines())]
	end

end

module CalibrationFunctions
	import ..InterpolationFunctions
	import PyPlot

	export generateCalibFactorTrace, generateBgTraces, interpolateBgTraces
	function getMeanOfQuantile(samples,quant)
	    q = quantile(samples,quant)
	    return mean(samples[samples.<q])
	end
	function generateBgTraces(times, traces; slices=10, quant=0.05)
	    dt = false
	    if typeof(times[1]) == DateTime
		times = Dates.datetime2unix.(times)
		dt = true
	    end
	    l=size(traces,1)
	    w=size(traces,2)
	    m=slices
	    bgtraces = Array{Float64}(size(traces,1), size(traces,2))
	    bgtimes = [mean(times[Int(floor((n-1)*l/m+1)):Int(floor(n*l/m)),1]) for n = 1:m]
	    for i=1:w
		bgtraces[:,i] = InterpolationFunctions.interpolate(times, bgtimes, [getMeanOfQuantile(traces[Int(floor((n-1)*l/m+1)):Int(floor(n*l/m)),i],0.08) for n = 1:m])
	    end
	    if dt
		bgtimes = Dates.unix2datetime.(bgtimes)
	    end
	    return bgtraces
	end

	function interpolateBgTraces(times, bgtimes, bgvalues)
	    if typeof(times[1]) == DateTime
		times = Dates.datetime2unix.(times)
	    end
	    if typeof(bgtimes[1]) == DateTime
		bgtimes = Dates.datetime2unix.(bgtimes)
	    end

	    nMasses = size(bgvalues,2)
	    bgtraces = Array{Float64}(size(times,1), nMasses)

	    for i=1:nMasses
		bgtraces[:,i] = InterpolationFunctions.interpolate(times, bgtimes, bgvalues[:,i])
	    end

	    return bgtraces
	end

	function generateCalibFactorTrace(traceTimes, calibTimes, calibValues, transitionTimes)
	    if typeof(traceTimes[1]) == DateTime
		println("Convertung traceTimes to unixtime")
		traceTimes = Dates.datetime2unix(traceTimes)
	    end
	    if typeof(calibTimes[1]) == DateTime
		println("Convertung calibTimes to unixtime")
		calibTimes = Dates.datetime2unix(calibTimes)
	    end
	    if typeof(transitionTimes[1]) == DateTime
		println("Convertung transitionTimes to unixtime")
		transitionTimes = Dates.datetime2unix(transitionTimes)
	    end
	  # Stuff might be unsorted
	  traceTimesSorted = sort(traceTimes)
	  sortedCalibIndices = sortperm(calibTimes)
	  calibTimesSorted = calibTimes[sortedCalibIndices]
	  calibValuesSorted = calibValues[sortedCalibIndices]
	  transitionTimesSorted = sort(transitionTimes)

	  piecewiseTimes = Array{Float64}(0)
	  piecewiseValues = Array{Float64}(0)
	  currentValue = 0
	  currentCalibIndex = 1
	  currentTransitionIndex = 1
	  currentPiecewiseIndex = 1

	  # Drop transitions before start of trace
	  while (transitionTimesSorted[currentTransitionIndex] < traceTimes[1]) && (currentTransitionIndex < length(transitionTimesSorted))
	    currentTransitionIndex += 1
	  end
	  # Drop transitions before first calibration
	  while (transitionTimesSorted[currentTransitionIndex] < calibTimes[1]) && (currentTransitionIndex < length(transitionTimesSorted))
	    currentTransitionIndex += 1
	  end
	  println("First transition within data range: $currentTransitionIndex - $(transitionTimesSorted[currentTransitionIndex])")

	  # If no calib was done before start of trace, insert copy of first calib
	  println("No Calib before start of data range found, copying first successive calib")
	  push!(piecewiseTimes, minimum([traceTimes[1] calibTimes[1]]))
	  push!(piecewiseValues, calibValuesSorted[1])
	  currentPiecewiseIndex += 1

	  # Iterate over calibs and transitions
	  while (currentCalibIndex <= length(calibTimesSorted))
	    if currentTransitionIndex < length(transitionTimesSorted)
	      println("CurrCalib: $currentCalibIndex ($(calibTimesSorted[currentCalibIndex])), CurrTransition: $currentTransitionIndex ($(transitionTimesSorted[currentTransitionIndex]))")
	    end
	    # define piecewise values depending on what comes next, calib or transition
	    if ((currentTransitionIndex > length(transitionTimesSorted)) || (calibTimesSorted[currentCalibIndex] < transitionTimesSorted[currentTransitionIndex])) && (currentCalibIndex <= length(calibTimesSorted))
	      # add one point for a calib
	      println("Adding calib point  $(calibValuesSorted[currentCalibIndex]) at $(calibTimesSorted[currentCalibIndex])")
	      push!(piecewiseTimes, calibTimesSorted[currentCalibIndex])
	      push!(piecewiseValues, calibValuesSorted[currentCalibIndex])
	      currentPiecewiseIndex += 1
	      currentCalibIndex += 1
	    else
	      # skip all following transitions that come before a calibration, since there are no calibs in between
	      while (currentTransitionIndex < length(transitionTimesSorted)) && (transitionTimesSorted[currentTransitionIndex+1] < calibTimesSorted[currentCalibIndex])
		println("Skipping transition $currentTransitionIndex, there are no calibs in between!")
		currentTransitionIndex +=1
	      end
	      # add one point with last value before transition
	      println("Adding pre-transition point $(piecewiseValues[end]) at $(transitionTimesSorted[currentTransitionIndex])")
	      push!(piecewiseTimes, transitionTimesSorted[currentTransitionIndex] - 1e-99)
	      push!(piecewiseValues, piecewiseValues[end])
	      currentPiecewiseIndex += 1

	      # add one point with next value after transition
	      println("Adding post-transition point $(calibValuesSorted[currentCalibIndex]) at $(transitionTimesSorted[currentTransitionIndex])")
	      push!(piecewiseTimes, transitionTimesSorted[currentTransitionIndex] + 1e-99)
	      push!(piecewiseValues, calibValuesSorted[currentCalibIndex])
	      currentPiecewiseIndex += 1

	      currentTransitionIndex += 1
	    end
	  end
	  PyPlot.figure()
	  PyPlot.plot(piecewiseTimes, piecewiseValues, "x-")

	  return InterpolationFunctions.interpolate(traceTimes, piecewiseTimes, piecewiseValues)
	end
end

module ExportFunctions
	using ..InterpolationFunctions
	using Dates
	import ..MasslistFunctions

	export exportTracesCSV, exportTracesCSVLossCorr, toMatlabTime, fromMatlabTime

	function exportTracesCSV(saveFolderPath, elementNames, compositions, times, traces; average=0)
	  sumformulas = MasslistFunctions.sumFormulaStringListFromCompositionArrayList(compositions)
	  f = open("$saveFolderPath/ptr3compositions.txt", "w")
	  writedlm(f, hcat(["Mass" "SumFormula"],reshape(elementNames,(1,length(elementNames)))))
	  writedlm(f, hcat(MasslistFunctions.massFromCompositionArrayList(compositions),sumformulas , compositions'))
	  close(f)
	  f = open("$saveFolderPath/ptr3traces.csv", "w")
	  writedlm(f, hcat(["Time"], reshape(sumformulas,(1,length(sumformulas)))))
	  if (average==0)
	    writedlm(f, hcat(times ,traces))
	  else
	    writedlm(f, hcat(averageSamples(times,average) ,averageSamples(traces,average)))
	  end
	  close(f)
	end

	function exportTracesCSVLossCorr(saveFolderPath, elementNames, compositions, times, traces, lossfactor, lossfactorerr, corrfactor, corrfactorerr, corrnotes; average=0)
	  sumformulas = MasslistFunctions.sumFormulaStringListFromCompositionArrayList(compositions)
	  f = open("$saveFolderPath/ptr3compositions.txt", "w")
	  writedlm(f, hcat(["Mass"	"SumFormula"],reshape(elementNames,(1,length(elementNames))),["LossFactor"	"LossFactorError"	"CorrFactor"	"CorrFactorErr"	"CorrNotes"]))
	  writedlm(f, hcat(MasslistFunctions.massFromCompositionArrayList(compositions),sumformulas , compositions', lossfactor, lossfactorerr, corrfactor, corrfactorerr, corrnotes))
	  close(f)
	  f = open("$saveFolderPath/ptr3tracesInletLossCorr.csv", "w")
	  writedlm(f, hcat(["Time"], reshape(sumformulas,(1,length(sumformulas)))))
	  if (average==0)
	    writedlm(f, hcat(times , (corrfactor.*traces' )' ))
	  else
	    writedlm(f, hcat(averageSamples(times,average) ,(corrfactor.*(averageSamples(traces,average))' )' ))
	  end
	  close(f)
	end

	################ EXAMPLE from Matlab: 737551.673515479 should be 06-May-2019 16:09:51 #############
	function toMatlabTime(t::Dates.DateTime)
	    timespan = ((t+Dates.Day(1)) - Dates.DateTime(0,1,1,0,0,0))
	    millis::Float64 = Dates.value(timespan)
	    return millis/24/3600/1000
	end

	function fromMatlabTime(timestamp::Number)
	    days=Int(floor(timestamp))
	    millisecondsRemainder = Int(round((timestamp-days)*24*3600*1000))
	    return Dates.DateTime(0,1,1,0,0,0)+Dates.Day(days)+Dates.Millisecond(millisecondsRemainder)-Dates.Day(1)
	end

end


"""
This module contains functions and types necessary for performing
nonadiabatic molecular dynamics.

Dynamics is performed using [`DifferentialEquations.jl`](https://diffeq.sciml.ai/stable/).
As such, this module is centered around the implementation of the functions
necessary to integrate the dynamics.

For deterministic Hamiltonian methods, the central function is [`DynamicsMethods.motion!`](@ref),
which is the inplace form of the function
to be integrated by [`DifferentialEquations.jl`](https://diffeq.sciml.ai/stable/).

Further, methods that have discontinuities, such as surface hopping, use the
[callback interface](https://diffeq.sciml.ai/stable/features/callback_functions/#callbacks)
provided by `DifferentialEquations.jl`.
"""
module DynamicsMethods

using DiffEqBase: DiffEqBase
using TypedTables: TypedTables
using UnitfulAtomic: austrip
using Reexport: @reexport
using OrdinaryDiffEq: OrdinaryDiffEq
using ComponentArrays: ComponentVector

using NQCDynamics: AbstractSimulation, Simulation, RingPolymerSimulation,
    DynamicsOutputs
using NQCBase: austrip_kwargs

export DynamicsVariables
export run_trajectory

"""
Each type of dynamics subtypes `Method` which is passed to
the `AbstractSimulation` as a parameter to determine the type of
dynamics desired.
"""
abstract type Method end

"""
    motion!(du, u, sim, t)

As per `DifferentialEquations.jl`, this function is implemented for
each method and defines the time derivatives of the `DynamicalVariables`.

We require that each implementation ensures `du` and `u` are subtypes
of `DynamicalVariables` and `sim` subtypes `AbstractSimulation`.
"""
function motion! end

"Provides the DEProblem for each type of simulation."
create_problem(u0, tspan, sim) =
    OrdinaryDiffEq.ODEProblem(motion!, u0, tspan, sim; callback=get_callbacks(sim))

"Choose a default algorithm for solving the differential equation."
select_algorithm(::AbstractSimulation) = OrdinaryDiffEq.VCABM5()

"Select the default callbacks for this simulation type."
get_callbacks(::AbstractSimulation) = nothing

"""
    DynamicsVariables(::AbstractSimulation, args...)

For each dynamics method this function is implemented to provide the variables for the
dynamics in the appropriate format.

By default, `DynamicsVariables` is set up for the classical case and takes
`sim`, `v`, `r` as arguments and returns a `ComponentVector(v=v, r=r)`
which is used as a container for the velocities and positions during
classical dynamics.
"""
DynamicsVariables(::AbstractSimulation, v, r) = ComponentVector(v=v, r=r)

"""
    run_trajectory(u0, tspan::Tuple, sim::AbstractSimulation;
                   output=(:u,), saveat=[], algorithm=select_algorithm(sim), kwargs...)

Solve a single trajectory starting from `u0` with a timespan `tspan`
for the simulation `sim`.

# Keyword arguments

`output` specifies the quantities that should be saved during the dynamics simulation.
The options for this keyword are any of the functions found in
`src/DynamicsMethods/output.jl`.

The rest of the keywords are the usual arguments for the `solve` function from
`DifferentialEquations.jl`.
It is possible to use `Unitful` quantities for any of the arguments since these are
automatically converted to atomic units internally.

# Output

The function returns a `Table` from `TypedTables.jl` with columns for time `t` and
every quantity specified in the `output` tuple.
"""
function run_trajectory(u0, tspan::Tuple, sim::AbstractSimulation;
        output=(:u,), saveat=[], callback=nothing, algorithm=select_algorithm(sim),
        kwargs...)

    if !(output isa Tuple)
        output = (output,)
    end

    problem = create_problem(u0, austrip.(tspan), sim)

    saving_callback, vals = DynamicsOutputs.create_saving_callback(output; saveat=austrip.(saveat))
    callback_set = DiffEqBase.CallbackSet(callback, saving_callback, get_callbacks(sim))
    problem = DiffEqBase.remake(problem, callback=callback_set)

    stripped_kwargs = austrip_kwargs(;kwargs...)
    @info "Performing a single trajectory."
    stats = @timed DiffEqBase.solve(problem, algorithm; stripped_kwargs...)
    @info "Finished after $(stats.time) seconds."
    out = [(;zip(output, val)...) for val in vals.saveval]
    TypedTables.Table(t=vals.t, out)
end

include("ClassicalMethods/ClassicalMethods.jl")
@reexport using .ClassicalMethods

include("SurfaceHoppingMethods/SurfaceHoppingMethods.jl")
@reexport using .SurfaceHoppingMethods

include("MappingVariableMethods/MappingVariableMethods.jl")
@reexport using .MappingVariableMethods

include("EhrenfestMethods/EhrenfestMethods.jl")
@reexport using .EhrenfestMethods

include("IntegrationAlgorithms/IntegrationAlgorithms.jl")

end # module