###################################################
## selection.jl
##      Build graph for scenario-based network flow method to solve full problem
## Authors: Arthur Delarue, Sébastien Martin, 2018
###################################################

"""
    Represents a node in scenario graph
    Required attributes:
        - id::Int           : unique reference of node
"""
abstract type ScenarioNode end

"""
    Node in scenario graph, represents buses arriving at the beginning of a route for a school
"""
struct ArrivalNode <: ScenarioNode
    id::Int

    "School associated with this node"
    school::Int
    "Route corresponding to the node"
    route::Route
    "Total service time of route"
    serviceTime::Float64
    "Scenario number"
    scenario::Int
end

"""
    Node in scenario graph, represents buses departing from a school
"""
struct DepartureNode <: ScenarioNode
    id::Int

    "School associated with this node"
    school::Int
    "Capacity, i.e. number of buses that can leave"
    capacity::Int
    "Scenario number"
    scenario::Int
end

"""
    Node in scenario graph, represents a bus yard
"""
struct YardNode <: ScenarioNode
    id::Int

    "Yard associated with this node"
    yard::Int
end

"""
    represents graph of scenarios for schools (master problem)
"""
mutable struct ScenarioGraph
    "graph object"
    graph::DiGraph
    "list of nodes"
    nodes::Vector{ScenarioNode}
    "dictionary mapping node type to node ids"
    nodeTypes::Dict{Symbol,Vector{Int}}
    "dictionary mapping pairs of nodes to costs"
    costs::Dict{Tuple{Int,Int}, Float64}
    "number of scenarios per school"
    numScenarios::Vector{Int}
end

"""
    Assumes the scenarios have already been computed in the data
    Builds ScenarioGraph using this input
"""
function buildScenarioGraph(data::SchoolBusData)
    if !data.withRoutingScenarios
        error("Cannot compute scenario graph without scenarios")
    end
    nodes = ScenarioNode[]
    nodeTypes = Dict(elt => Int[] for elt in [:YardNode, :ArrivalNode, :DepartureNode])
    # get the yard nodes
    yardNodes, currentNodeId = getYardNodes!(data, nodeTypes)
    append!(nodes, yardNodes)
    numScenarios = [length(data.scenarios[sID]) for sID=eachindex(data.schools)]
    for schoolID = eachindex(data.schools)
        for scenarioNum = 1:numScenarios[schoolID]
            schoolNodes, currentNodeId = getSchoolScenarioNodes!(data, schoolID,
                                                                 currentNodeId, scenarioNum,
                                                                 nodeTypes)
            append!(nodes, schoolNodes)
        end
    end
    graph = DiGraph(length(nodes))
    costs = Dict{Tuple{Int,Int}, Float64}()
    possibleEdgeTypes = [(:DepartureNode, :ArrivalNode), (:DepartureNode, :YardNode),
                         (:ArrivalNode, :DepartureNode), (:YardNode, :ArrivalNode)]
    for (type1, type2) in possibleEdgeTypes
        for idx1 in nodeTypes[type1], idx2 in nodeTypes[type2]
            node1 = nodes[idx1]
            node2 = nodes[idx2]
            edgeWasAdded, cost = createEdge!(data, graph, node1, node2)
            if edgeWasAdded
                costs[node1.id, node2.id] = cost
            end
        end
    end
    return ScenarioGraph(graph, nodes, nodeTypes, costs, numScenarios)
end

"""
    Get nodes for one scenario for a given school
    Args:
        data::SchoolBusData    : the data
        schoolId::Int          : the id of the school
        currentNodeId          : id of node, incremented for each new node
        scenarioNum            : id of scenario
        nodeTypes              : the attributes of the ScenarioGraph that we are updating
    Returns:
        the nodes for that scenario
        the id of the current node (to be passed along to the next call)
"""
function getSchoolScenarioNodes!(data::SchoolBusData, schoolID::Int,
                                 currentNodeId::Int, scenarioNum::Int,
                                 nodeTypes::Dict{Symbol, Vector{Int}})
    schoolNodes = ScenarioNode[]
    scenario = data.scenarios[schoolID][scenarioNum]
    for routeID in scenario.routeIDs
        route = data.routes[schoolID][routeID]
        # create arrival nodes
        push!(schoolNodes, ArrivalNode(currentNodeId, schoolID, route,
                                       serviceTime(data, schoolID, route),
                                       scenario.id))
        push!(nodeTypes[:ArrivalNode], currentNodeId)
        currentNodeId += 1
    end
    push!(schoolNodes, DepartureNode(currentNodeId,schoolID,length(scenario.routeIDs),scenarioNum))
    push!(nodeTypes[:DepartureNode], currentNodeId)
    currentNodeId += 1
    return schoolNodes, currentNodeId
end

"""
    Get nodes for yards
"""
function getYardNodes!(data::SchoolBusData, nodeTypes::Dict{Symbol,Vector{Int}})
    yardNodes = ScenarioNode[]
    nodeId = 1
    for yard in data.yards
        push!(yardNodes, YardNode(nodeId, yard.id))
        push!(nodeTypes[:YardNode], nodeId)
        nodeId += 1
    end
    return yardNodes, nodeId
end

"""
    Decide whether a bus from one school can serve another school (time feasible)
    By convention, school 1 is the school from which the bus would depart in the morning, and
        the route belings to school 2
"""
function isFeasibleInTime(data::SchoolBusData, school1::Int, school2::Int,
                          route::Route, routeTime::Float64)
    return (data.schools[school1].starttime +
            traveltime(data, data.schools[school1], data.stops[school2][route.stops[1]]) +
            routeTime + data.schools[school2].dwelltime <= data.schools[school2].starttime)
end

"""
    Add edge between two nodes and compute cost of the edge
"""
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::ScenarioNode, node2::ScenarioNode)
    edgeWasAdded = false
    cost = nothing
    return edgeWasAdded, cost
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::DepartureNode, node2::ArrivalNode)
    edgeWasAdded = false
    if node1.school != node2.school
        # can only go to another school and route must be feasible
        if isFeasibleInTime(data, node1.school, node2.school,
                            node2.route, node2.serviceTime)
            edgeWasAdded = add_edge!(graph, node1.id, node2.id)
        end
    end
    return edgeWasAdded, 0.
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::DepartureNode, node2::YardNode)
    return add_edge!(graph, node1.id, node2.id), 0.
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::ArrivalNode, node2::DepartureNode)
    edgeWasAdded = false
    if node1.school == node2.school && node1.scenario == node2.scenario
        edgeWasAdded = add_edge!(graph, node1.id, node2.id)
    end
    return edgeWasAdded, 0.
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::YardNode, node2::ArrivalNode)
    return add_edge!(graph, node1.id, node2.id), 1.
end

"""
    Wrapper function replacing LightGraphs.out_edges since it was deprecated
    Returns a generator with the desired edges
"""
function edges_out(graph::DiGraph, node::Int)
    return (Edge(node, j) for j in outneighbors(graph, node))
end

"""
    Wrapper function replacing LightGraphs.in_edges since it was deprecated
    Returns a generator with the desired edges
"""
function edges_in(graph::DiGraph, node::Int)
    return (Edge(j, node) for j in inneighbors(graph, node))
end

"""
    Flow method for scenario selection
    Args:
        - the data
        - the scenario graph we've constructed
    Returns:
        - a vector of length the number of schools, containing the index of
                the scenario that is used for each school
"""
function selectScenario(data::SchoolBusData, sg::ScenarioGraph; args...)
    edgeCost(edge::Edge) = sg.costs[src(edge), dst(edge)]

    model = Model(solver = GurobiSolver(Threads=getthreads(), MIPGap=0.01; args...))

    @variable(model, useScenario[i=eachindex(data.schools), j=1:sg.numScenarios[i]], Bin)
    @variable(model, busFlow[edge=edges(sg.graph)] >= 0, Int)

    @constraint(model, capacityDeparture[i=sg.nodeTypes[:DepartureNode]],
                sum(busFlow[edge] for edge=edges_out(sg.graph, i)) <= sg.nodes[i].capacity)
    @constraint(model, capacityArrival[i=sg.nodeTypes[:ArrivalNode]],
                sum(busFlow[edge] for edge=edges_out(sg.graph, i)) ==
                useScenario[sg.nodes[i].school, sg.nodes[i].scenario])

    @constraint(model, oneScenario[i=eachindex(data.schools)],
                sum(useScenario[i,j] for j=1:sg.numScenarios[i]) == 1)

    @constraint(model, flowConservation[i=vertices(sg.graph)],
                sum(busFlow[edge] for edge=edges_in(sg.graph, i)) ==
                sum(busFlow[edge] for edge=edges_out(sg.graph, i)))

    @objective(model, Min, sum(busFlow[edge] * edgeCost(edge) for edge=edges(sg.graph)))

    status = solve(model)
    scenarioUsed = [indmax(getvalue(useScenario[i,j])
                           for j=1:sg.numScenarios[i]) for i=eachindex(sg.numScenarios)]
    return scenarioUsed
end

"""
    Represents node in FullRoutingGraph
    Mandatory attributes:
    - id::Int        : unique identifier
    - yardId::Int    : yard number
"""
abstract type FullRoutingNode end

"""
    In full routing flow graph, represents a bus with a certain capacity
"""
struct BusNode <: FullRoutingNode
    id::Int
    yardId::Int
    "Route associated with node"
    route::Route
    "School associated with node"
    school::Int
    "Service time of route"
    serviceTime::Float64
end

"""
    In full routing flow graph, represents a yard
"""
struct FullYardNode <: FullRoutingNode
    id::Int
    yardId::Int
end

"""
    Represents the situation after the scenario has been picked: one bus per route per school
    Similar to ScenarioGraph
"""
mutable struct FullRoutingGraph
    "The graph itself"
    graph::DiGraph
    "list of node objects"
    nodes::Vector{FullRoutingNode}
    "dictionary mapping pairs of nodes to costs"
    costs::Dict{Tuple{Int,Int},Float64}
end

"""
    Given routes object (such as that output by flow master problem), set up full routing graph
"""
function buildFullRoutingGraph(data::SchoolBusData, usedScenario::Vector{Int})
    nodes = FullRoutingNode[]
    costs = Dict{Tuple{Int,Int},Float64}()

    currentId = 1
    # yard nodes
    for yard in data.yards
        push!(nodes, FullYardNode(currentId, yard.id))
        currentId += 1
    end
    # school bus nodes
    for schoolID in eachindex(data.schools)
        scenario = data.scenarios[schoolID][usedScenario[schoolID]]
        for routeID in scenario.routeIDs
            route = data.routes[schoolID][routeID]
            for yard in data.yards
                push!(nodes, BusNode(currentId, yard.id, route, schoolID,
                                     serviceTime(data, schoolID, route)))
                currentId += 1
            end
        end
    end
    graph = DiGraph(length(nodes))
    # edges
    for node1 in nodes, node2 in nodes
        edgeWasAdded, cost = createEdge!(data, graph, node1, node2)
        if edgeWasAdded
            costs[node1.id, node2.id] = cost
        end
    end
    return FullRoutingGraph(graph, nodes, costs)
end

"""
    Method creates an edge between two nodes, with one function per type of node pair
"""
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::FullRoutingNode, node2::FullRoutingNode)
    return false, 0.
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::FullYardNode, node2::BusNode)
    edgeWasAdded, cost = false, 0.
    if node1.yardId == node2.yardId
        edgeWasAdded = add_edge!(graph, node1.id, node2.id)
        cost = traveltime(data, data.yards[node1.yardId],
                          data.stops[node2.school][node2.route.stops[1]])
    end
    return edgeWasAdded, cost
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::BusNode, node2::FullYardNode)
    edgeWasAdded, cost = false, 0.
    if node1.yardId == node2.yardId
        edgeWasAdded = add_edge!(graph, node1.id, node2.id)
        cost = traveltime(data, data.schools[node1.school], data.yards[node2.yardId])
    end
    return edgeWasAdded, cost
end
function createEdge!(data::SchoolBusData, graph::DiGraph,
                     node1::BusNode, node2::BusNode)
    edgeWasAdded, cost = false, 0.
    if node1.yardId == node2.yardId && node1.school != node2.school
        if isFeasibleInTime(data, node1.school,
                            node2.school, node2.route, node2.serviceTime)
            edgeWasAdded = add_edge!(graph, node1.id, node2.id)
            cost = traveltime(data, data.schools[node1.school],
                              data.stops[node2.school][node2.route.stops[1]])
        end
    end
    return edgeWasAdded, cost
end

"""
    Given a full routing graph, constructs a YardId => node dictionary
"""
function yardToNodeDict(frg::FullRoutingGraph)
    d = Dict{Int,Int}()
    for node in frg.nodes
        if typeof(node) == FullYardNode
            d[node.yardId] = node.id
        end
    end
    return d
end

"""
    Bus scheduling problem
"""
function solveFullRouting(data::SchoolBusData, frg::FullRoutingGraph; args...)
    model = Model(solver = GurobiSolver(Threads=getthreads(), MIPGap=0.01; args...))
    # variables - morning
    @variable(model, busFlow[edge=edges(frg.graph)], Bin)
    length([node for node in frg.nodes if typeof(node) == FullYardNode]) > 1 &&
        error("Too many yards")
    # variables - yard capacity
    @variable(model, yardCapacity[y=eachindex(frg.nodes);
                                  typeof(frg.nodes[y]) == FullYardNode] >= 0, Int)
    # a route can be served by at most one bus from a given yard - morning
    @constraint(model, oneBusPerRoute[nodeId=eachindex(frg.nodes);
                                      typeof(frg.nodes[nodeId]) == BusNode],
                sum(busFlow[edge] for edge = edges_in(frg.graph, nodeId)) == 1)
    # flow conservation constraints - morning
    @constraint(model, flowConservation[nodeId=vertices(frg.graph)],
                sum(busFlow[edge] for edge=edges_in(frg.graph, nodeId)) ==
                sum(busFlow[edge] for edge=edges_out(frg.graph, nodeId)))
    @constraint(model, flowBound[y=eachindex(frg.nodes);
                                 typeof(frg.nodes[y]) == FullYardNode],
                sum(busFlow[edge] for edge=edges_out(frg.graph, y)) <= yardCapacity[y])
    # minimize the total cost of buses (First Order = total number of bus)
    @objective(model, Min, sum(yardCapacity[y]
                               for (y,node)=enumerate(frg.nodes) if typeof(node) == FullYardNode))
    status = solve(model)
    flows = Dict(edge => getvalue(busFlow[edge]) for edge=edges(frg.graph))
    return interpretFlows!(data, frg, flows)
end

"""
    Given flows, get final buses
"""
function interpretFlows!(data::SchoolBusData, frg::FullRoutingGraph, flows::Dict)
    finalBuses = Bus[]
    currentBusId = 1
    dict = yardToNodeDict(frg)
    for yard in data.yards
        yardNode = dict[yard.id]
        yardIsEmpty = false
        # keep following bus routes until yard is empty
        while !yardIsEmpty
            schools, routes = followBusAlongRoute!(frg, flows, yardNode)
            if isempty(schools)
                yardIsEmpty = true
            else
                push!(finalBuses, Bus(currentBusId, yard.id, schools, routes))
                currentBusId += 1
            end
        end
    end
    return finalBuses
end

"""
    Extract a bus route from a flows object, removing flows as they are found
"""
function followBusAlongRoute!(frg::FullRoutingGraph, flows::Dict, yardNode::Int)
    schools, routes = Int[], Int[]
    if yardNode == -1 && length(flows) == 0
        return schools, routes
    end
    currentNode = yardNode
    backToYard = false
    while !backToYard
        for edge in edges_out(frg.graph, currentNode)
            if flows[edge] > 0.5
                if typeof(frg.nodes[dst(edge)]) == BusNode
                    push!(schools, frg.nodes[dst(edge)].school)
                    push!(routes, frg.nodes[dst(edge)].route.id)
                    currentNode = dst(edge)
                else
                    backToYard = true
                end
                flows[edge] = 0.
                break
            end
        end
        if isempty(schools) # couldn't find anything leading out of the yard
            backToYard = true
        end
    end
    return schools, routes
end

function routeBuses!(data::SchoolBusData; args...)
    sg = buildScenarioGraph(data)
    data.usedScenario = selectScenario(data, sg; args...)
    frg = buildFullRoutingGraph(data, data.usedScenario)
    data.buses = solveFullRouting(data, frg; args...)
    data.withFinalBuses = true
    return data
end


using Pkg;Pkg.activate(@__DIR__)
using Comrade
using PyCall
using Plots

using GalacticOptim, Optim
using Zygote
using Parameters
using Distributions
using SpeedMapping
using AdvancedHMC
using Dynesty

load_ehtim()
@pyimport eht_dmc as ed

# To download the data visit https://doi.org/10.25739/g85n-f134
obs = ehtim.obsdata.load_uvfits(joinpath(@__DIR__, "SR1_M87_2017_096_hi_hops_netcal_StokesI.uvfits"))
obs.add_scans()

obsavg = obs.avg_coherent(0.0, scan_avg=true)

img = ehtim.image.make_empty(512, μas2rad(250.0), obs.ra, obs.dec, obs.rf, obs.source)
img = img.add_gauss(1.0, [μas2rad(40.0), μas2rad(30.0), π/3, 0.0, 0.0])

obslist = obsavg.split_obs()

gainoff = 0.0
gainp = Dict("AA"=> 0.01,
             "AP" => 0.1,
             "AZ" => 0.1,
             "JC" => 0.1,
             "LM" => 0.2,
             "PV" => 0.1,
             "SM" => 0.1,
             "SR" => 0.1)
obsim = img.observe_same(ehtim.obsdata.merge_obs(obslist[10:15]), ttype="fast",  gainp=gainp, gain_offset=gainoff, ampcal=false, phasecal=false, stabilize_scan_phase=true, stabilize_scan_amp=true)
# now get cal table
ctable = ehtim.calibrating.self_cal.self_cal(obsim, img, caltable=true)
ehtim.caltable.save_caltable(ctable, obsim, datadir=joinpath(@__DIR__ ,"CaltableTest"))

dvis = extract_vis(obsim)
dlcamp = extract_lcamp(obsim; count="min")
dcphase = extract_cphase(obsim, count="min-cut0bl")
damp = extract_amp(obsim)
st = scantable(dvis)

gcache = Comrade.GainCache(st)

struct Model{G}
    gcache::G
end

function Model(st::Comrade.ScanTable)
    gcache = Comrade.GainCache(st)
    return Model{typeof(gcache)}(gcache)
end

function (mod::Model)(θ)
    @unpack f, σ, τ, ξ,  gamp, gphase = θ
    g = gamp.*cis.(gphase)
    m = f*rotated(stretched(Gaussian(), σ*τ, σ), ξ)
    Comrade.GainModel(mod.gcache, g, m)
end

# define the priors
distamp = (AA = LogNormal(0.0, 0.1),
           AP = LogNormal(0.0, 0.1),
           LM = LogNormal(0.0, 0.2),
           AZ = LogNormal(0.0, 0.1),
           JC = LogNormal(0.0, 0.1),
           PV = LogNormal(0.0, 0.1),
           #SM = LogNormal(0.0, 0.1)
         )
distphase = (AA = Normal(0.0, 1e-4),
             AP = Normal(0.0, π),
             LM = Normal(0.0, π),
             AZ = Normal(0.0, 1.0*π),
             JC = Normal(0.0, 1.0*π),
             PV = Normal(0.0, π),
             #SM = Normal(0.0, 1.0*π),
             )
prior = (
          f = Uniform(0.9, 1.1),
          σ = Uniform(μas2rad(10.0), μas2rad(30.0)),
          τ = Uniform(0.1, 1.0),
          ξ = Uniform(-π/2, π/2),
          gamp = Comrade.GainPrior(distamp, st),
          gphase = Comrade.GainPrior(distphase, st)
        )
# Now form the posterior

mms  = Model(st)

cllklhd = RadioLikelihood(dlcamp, dcphase)
clpost = Posterior(cllklhd, prior, mms)

lklhd = RadioLikelihood(dvis)#damp, dcphase)
post = Posterior(lklhd, prior, mms)
tpost = asflat(post)
# We will use HMC to sample the posterior.
# First to reduce burn in we use pathfinder
ndim = dimension(tpost)
f, tr = GalacticOptim.OptimizationFunction(post, GalacticOptim.AutoForwardDiff{15}())
#x0 = inverse(tr, rand(post.prior))

x0 = rand(post.prior, 50)
q, ϕ, inds = multipathfinder(clpost, 100; init_params=x0, ndraws_per_run=50)

res = map(1:150) do i
    #x0 = xx
    #x0.gphase .= randn(length(xx.gphase))
    #x0.gamp .= 1 .+ 0.2*randn(length(xx.gamp))
    prob = GalacticOptim.OptimizationProblem(f, randn(ndim), nothing)
    #sol, xopt = solve(prob, BBO_adaptive_de_rand_1_bin_radiuslimited(), tr; maxiters=200_000)

    #xopt = rand(post.prior)
    #prob = GalacticOptim.OptimizationProblem(f, sol.u, nothing)
    sol, xopt = solve(prob, LBFGS(), tr; iterations=2_000)
    @info "$i/50 $(sol.minimum)"
    xopt, sol.minimum
end

metric = DenseEuclideanMetric(ndim)
chain, stats = sample(post, HMC(;metric), 14_000; nadapts=10_000, init_params=res[ind[1]][1])


# User-level methods to compute lightcurves

struct Lightcurve{T<:Real}
    dt::T             # Exposure time
    tobs::Vector{T}   # Observed times
    fobs::Vector{T}   # Observed flux
    eobs::Vector{T}   # Measurement errors
    nobs::Int64       # number of flux measurements
    flux::Vector{T}   # Computed model flux
    dfdu::Matrix{T}   # Derivative of flux wrt limbdark coefficients
    dfdk::Matrix{T}   # Derivatives wrt the radius ratios
    dfdq0::Matrix{T}  # Derivatives wrt initial Nbody Cartesian coordinates and masses
    dfdr::Vector{T}   # Derivatives wrt stellar radius

    # Transit parameters
    u_n::Vector{T}   # Limbdark coefficients
    k::Vector{T}     # radius ratios
    rstar::Vector{T} # Stellar radius

    # Interal arrays/values
    dtinv::T          # Inverse of the exposure time
    dbdq0::Vector{T}  # derivative of the impact parameter wrt the Nbody initial conditions
    n_params::Int64   # Number of nbody model parameters
    do_grad::Bool     # Compute gradients or not

    function Lightcurve(dt::T, tobs::Vector{T}, fobs::Vector{T}, eobs::Vector{T}, u_n::Vector{T}, k::Vector{T}, rstar::T, n_params::Int64=0) where T<:Real
        @assert (length(tobs) == length(fobs)) && (length(tobs) == length(eobs)) "Data arrays are different sizes"
        @assert n_params >= 0 "Number of model parameters must be a positive integer"
        n_params > 0 ? do_grad = true : do_grad = false
        nobs = length(tobs)
        flux = zeros(T,nobs)
        dtinv = dt == 0.0 ? 0.0 : inv(dt)
        dfdu = zeros(T, nobs, length(u_n))
        dfdk = zeros(T, nobs, length(k))
        dfdq0 = zeros(T, nobs, n_params)
        dbdq0 = zeros(T, n_params)
        dfdr = zeros(T, nobs)
        return new{T}(dt,tobs,fobs,eobs,nobs,flux,dfdu,dfdk,dfdq0,dfdr,u_n,k,[rstar],dtinv,dbdq0,n_params,do_grad)
    end
end

"""Zero out the model arrays"""
function zero_out!(lc::Lightcurve{T}) where T<:Real
    lc.dfdu .= 0.0
    lc.dfdk .= 0.0
    lc.dfdq0 .= 0.0
    lc.dfdr .= 0.0
    lc.flux .= 0.0
end

# Normalize the exposure integration by the exposure time.
function normalize!(lc::Lightcurve{T}, ia::IntegralArrays{T}) where T<:Real
    lc.flux .*= lc.dtinv # Divide by exposure time to get average flux
    if lc.do_grad
        # Do the same for derivatives
        lc.dfdk .*= lc.dtinv
        lc.dfdu .*= lc.dtinv
        lc.dfdq0 .*= lc.dtinv
        lc.dfdr .*= lc.dtinv
    end
end

# If non-integrated exposure, do nothing
@inline normalize!(lc::Lightcurve{T}, ia::T) where T<:Real = nothing

function find_transit_time(t0::T, h::T, points::AbstractMatrix{T}) where T<:Real
    tt = find_zero(t->compute_impact_parameter(t, t0, h, points), t0)
    return tt
end

function points_of_contact_4(tt::T,t0::T,h::T,points::AbstractMatrix{T},k::T) where T<:Real
    t1 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt-h)
    t2 = find_zero(t -> (1.0-k-compute_impact_parameter(t,t0,h,points)), tt-h)
    t3 = find_zero(t -> (1.0-k-compute_impact_parameter(t,t0,h,points)), tt+h)
    t4 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt+h)
    return SVector(t1,t2,t3,t4)
end

function points_of_contact_2(t0::T,tt::T,h::T,points::AbstractMatrix{T},k::T) where T<:Real
    t1 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt-h)
    t4 = find_zero(t -> (1.0+k-compute_impact_parameter(t,t0,h,points)), tt+h)
    return SVector(t1,t4)
end

function compute_lightcurve!(lc::Lightcurve{T}, ts::TransitSeries{T, ProvidedTimes}; tol::T=1e-6, maxdepth::Int64=6) where T<:Real

    zero_out!(lc) # Zero out model arrays

    # Check if we're doing an integrated lightcurve
    ia = lc.dt == 0.0 ? 0.0 : IntegralArrays(lc.do_grad ? (lc.n_params + length(lc.u_n) + length(lc.k) + 2) : 1, maxdepth, tol) # Plus 2 for flux and rstar

    # Make transit structure (will be updated with proper r and b later)
    trans = transit_init(lc.k[1], 0.0, lc.u_n, lc.do_grad)

    # Iterate over each transit time and sum Lightcurve
    rstar = lc.rstar[1]
    for it in eachindex(ts.times)
        # check for transit
        t0 = ts.times[it]   # Get "data" transit time (expansion point)
        ib = ts.bodies[it]  # Get transiting body

        tt = try
            # Compute the transit time by finding where the x component of
            # the impact parameter is equal to zero
            # This assumes the body crosses the mid plane during the transit
            find_transit_time(t0, ts.h, @views(ts.points[ib,it,:,:]./rstar))
        catch e
            # If no convergence, just use the supplied transit time
            t0
        end

        # Get the impact parameter at the transit midpoint (computed above)
        b0 = compute_impact_parameter(t0, tt, ts.h, @views(ts.points[ib,it,:,:]./rstar))
        if b0 > 1.0+lc.k[ib-1]; continue; end

        # Compute points of contact
        # If grazing transit, only two points
        if (b0 + lc.k[ib-1]) >= 1.0
            tc = points_of_contact_2(t0, tt, ts.h, @views(ts.points[ib,it,:,:]./rstar), lc.k[ib-1])
        else
            tc = points_of_contact_4(t0, tt, ts.h, @views(ts.points[ib,it,:,:]./rstar), lc.k[ib-1])
        end

        # Integrate lightcurve
        trans.r = lc.k[ib-1]
        integrate_transit!(ib,it,t0,tc,trans,lc,ts,ia)
    end

    # Normalize by exposure time, if needed.
    normalize!(lc, ia)
    return
end

function integrate_transit!(ib::Int64,it::Int64,t0::T,tc::SVector{N,T},trans::Transit_Struct{T},lc::Lightcurve{T},ts::TransitSeries{T},ia::IntegralArrays{T}) where {N, T<:Real}
    nc = length(tc) # Number of points of contact
    dt = lc.dt
    inv_rstar = inv(lc.rstar[1])

    # Integrate over each Exposure
    for i in 1:lc.nobs
        tstart = lc.tobs[i] - 0.5*dt
        tend = lc.tobs[i] + 0.5*dt

        # Check if remaining exposures are inside transit
        if tstart > tc[end]; break; end # Don't need to continue loop if passed transit
        if tend < tc[1]; continue; end

        # Check if points of contact are within exposure
        tlim = [tstart]
        for j in 1:nc
            if tstart < tc[j] && tc[j] < tend
                push!(tlim, tc[j])
            end
        end
        push!(tlim, tend)

        # Get series expansion components
        xc = components(@views(ts.points[ib,it,:,1].*inv_rstar), ts.h)
        yc = components(@views(ts.points[ib,it,:,2].*inv_rstar), ts.h)

        if lc.do_grad
            n_bodies = length(ts.count)
            dxc = [components(ts.dpoints[ib,it,:,1,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]
            dyc = [components(ts.dpoints[ib,it,:,2,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]

            # integrate over exposure
            for j in 1:length(tlim)-1
                integrate_timestep!(t0, tlim[j], tlim[j+1], xc, yc, dxc, dyc, trans, ia, lc.dbdq0, ib-1)
                lc.flux[i] += ia.I_of_f[1]
                lc.dfdq0[i,:] .+= ia.I_of_f[2:1+n_bodies*7]
                lc.dfdk[i,:] .+= ia.I_of_f[2+n_bodies*7:n_bodies*8]
                lc.dfdu[i,:] .+= trans.dgdu' * ia.I_of_f[end-trans.n-1:end-1]
                lc.dfdr[i] += ia.I_of_f[end] * inv_rstar
            end
        else
            # Integrate over exposure
            for j in 1:length(tlim)-1
                integrate_timestep!(t0, tlim[j], tlim[j+1], xc, yc, trans, ia)
                lc.flux[i] += ia.I_of_f[1]
            end
        end
    end
    return
end

function integrate_transit!(ib::Int64,it::Int64,t0::T,tc::SVector{N,T},trans::Transit_Struct{T},lc::Lightcurve{T},ts::TransitSeries{T},ia::T) where {N, T<:Real}
    inv_rstar = inv(lc.rstar[1])

    # Compute the flux at each point
    for i in 1:lc.nobs
        # Check if observation is outside of a transit
        if lc.tobs[i] > tc[end]; break; end
        if lc.tobs[i] < tc[1]; continue; end

        # Get series expansion components
        xc = components(@views(ts.points[ib,it,:,1].*inv_rstar), ts.h)
        yc = components(@views(ts.points[ib,it,:,2].*inv_rstar), ts.h)

        if lc.do_grad
            n_bodies = length(ts.count)
            dxc = [components(ts.dpoints[ib,it,:,1,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]
            dyc = [components(ts.dpoints[ib,it,:,2,k,i].*inv_rstar, ts.h) for i in 1:7, k in 1:n_bodies][:]

            compute_flux!(lc.tobs[i], t0, xc, yc, dxc, dyc, lc, trans, i, ki, inv_rstar)
        else
            lc.flux[i] += compute_flux(lc.tobs[i], t0, xc, yc, trans)
        end
    end
    return
end

function integrate_timestep!(t0::T, a::T, b::T, xc::SVector{N,T}, yc::SVector{N,T}, trans::Transit_Struct{T}, ia::IntegralArrays{T}) where {N, T<:Real}
    # Computes the flux as function of time.
    # Closure to be passed to integrator function
    transit_flux! = let trans=trans, t0=t0, xc=xc, yc=yc
        (time::T, flux::Vector{T}) -> begin
            trans.b = compute_impact_parameter(time, t0, xc, yc)
            flux[1] = transit_poly_g(trans)-1
        end
    end

    # Integrate transit_flux! over interval [a,b]
    integrate_simpson!(a,b,transit_flux!,ia)
end

function compute_flux(tc::T, t0::T, xc::SVector{N,T}, yc::SVector{N,T}, trans::Transit_Struct{T}) where {N, T<:Real}
    # Compute flux at a particular time
    trans.b = compute_impact_parameter(tc, t0, xc, yc)
    flux = transit_poly_g(trans) - 1
    return flux
end

function integrate_timestep!(t0::T, a::T, b::T, xc::SVector{N,T}, yc::SVector{N,T}, dxc, dyc, trans::Transit_Struct{T}, ia, dbdq0, ki) where {N,T<:Real}
    n_coords = length(dbdq0)

    # Map radius ratio index to the gradient array index.
    k_ind = ki + 1 + n_coords

    # Compute the flux and derivatives
    transit_flux_grad! = let trans=trans, t0=t0, xc=xc, yc=yc, dbdq0=dbdq0, n_coords = n_coords, k_ind=k_ind
        (time::T, dflux::Vector{T}) -> begin
            trans.b = compute_impact_parameter!(time, t0, xc, yc, dxc, dyc, dbdq0)
            dflux .= 0.0
            dflux[1] = transit_poly_g!(trans)-1            # Flux at time
            dflux[2:1+n_coords] .= trans.dfdrb[2] .* dbdq0 # Gradient wrt the initial conditions
            dflux[k_ind] = trans.dfdrb[1]                  # radius ratio (derivative of others are zero)
            dflux[end-trans.n-1:end-1] .= trans.dfdg       # dfdg (greens basis)
            dflux[end] = -trans.dfdrb[2]*trans.b           # dfdrstar * rstar (divide by rstar outside)
        end
    end

    # Integrate flux and derivatives over exposure interval [a,b]
    integrate_simpson!(a,b,transit_flux_grad!,ia)
end

function compute_flux!(tc::T, t0::T, xc, yc, dxc, dyc, lc, trans, i, ki, inv_rstar) where T<:Real
    # Compute flux and derivatives
    trans.b = compute_impact_parameter!(tc, t0, xc, yc, dxc, dyc, lc.dbdq0)
    lc.flux[i] += transit_poly_g!(trans) - 1            # Flux
    lc.dfdq0[i,:] .+= trans.dfdrb[2] .* lc.dbdq0        # Derivative wrt initial conditions
    lc.dfdk[i,ki] += trans.dfdrb[1]                     # Radius ratio
    lc.dfdu[i,:] .+= trans.dgdu' * trans.dfdg           # Limbdarkening params
    lc.dfdr[i] += -trans.dfdrb[2]*trans.b*inv_rstar     # Stellar radius
    return
end


"Values for units for strings found in option line."
const units = Dict{ String, Float64 }(
  "HZ"    =>      1e0,
  "KHZ"   =>      1e3,
  "MHZ"   =>      1e6,
  "GHZ"   =>      1e9
)
"The default unit, if not found in option line."
const default_units = "GHZ"

"Symbols for different parameter formats for strings found in option line"
const parameters = Dict{ String, Symbol }(
  "S" => :ScatteringParameters,
  "Y" => :AdmittanceParameters,
  "Z" => :ImpedanceParameters,
  "H" => :HybridHParameters,
  "G" => :HybridGParameters,
)
"The default parameter format, if not found in option line."
const default_parameter = "S"

"Symbols for data format for strings found in option line."
const formats = Dict{ String, Symbol }(
  "DB" => :DecibelAngle,
  "MA" => :MagnitudeAngle,
  "RI" => :RealImaginary
)
"The default data format, if not found in option line."
const default_format = "MA"
"The default characteristic impedance, if not found in option line."
const default_resistance = 50

"""
    is_commentline( line )

Checks, if a line is a comment line.
"""
is_comment_line( line ) = length( line ) > 0 && line[ 1 ] == '!'
"""
    parse_comment_line( line )

Returns the comment of a comment line.
"""
parse_comment_line( line ) = line[ 2:end ]

"""
    is_option_line( line )

Checks, if a line is a option line.
"""
is_option_line( line ) = length( line ) > 0 && line[ 1 ] == '#'

"""
    is_frequency_unit_option( option )

Checks, if a string is a valid unit in an option line.
"""
is_frequency_unit_option( option ) = haskey( units,  uppercase( option ) )

"""
    parse_frequency_unit_option( option )

Returns the multiplier for a unit string in an option line.
"""
parse_frequency_unit_option( option ) = units[ uppercase( option ) ]

"""
    is_parameter_option( option )

Checks, if a string is a valid parameter format in an option line.
"""
is_parameter_option( option ) = haskey( parameters, uppercase( option ) )

"""
    parse_parameter_option( option )

Returns the symbol for a valid parameter format string in an option line.
"""
parse_parameter_option( option ) = parameters[ uppercase( option ) ]

"""
    is_format_option( option )

Checks, if a string is a valid data format in an option line.
"""
is_format_option( option ) = haskey( formats, uppercase( option ) )

"""
    parse_format_option( option )

Returns the symbol for a valid data format string in an option line.
"""
parse_format_option( option ) = formats[ uppercase( option ) ]

"""
    is_resistance_option( option )

Checks, if a string is a valid resistance in an option line.
"""
is_resistance_option( option ) = uppercase( option ) == "R"

"""
    parse_resistance_option( option )

Returns the value of the characteristic impedance for a valid resistance string in an option line.
"""
parse_resistance_option( option ) = parse( Float64, option )

"""
    is_empy_line( line )

Checks, if a line is empty.
"""
is_empy_line( line ) = length( line ) == 0

"""
    parse_option_line( line )

Returns the option structure for a valid option line.
"""
function parse_option_line( line )
  line = uppercase( line )

  default = Options()
  unit        = default.unit
  parameter   = default.parameter
  format      = default.format
  resistance  = default.resistance

  entries = split( line[ 2:end ] )

  next_is_resistance = false
  for entry in entries
    if next_is_resistance
      resistance = parse_resistance_option( entry )
      next_is_resistance = false
    else
      if is_frequency_unit_option( entry )
        unit = parse_frequency_unit_option( entry )
      elseif is_parameter_option( entry )
        parameter = parse_parameter_option( entry )
      elseif is_format_option( entry )
        format = parse_format_option( entry )
      else
        next_is_resistance =  is_resistance_option( entry )
      end
    end
  end
  Options( unit, parameter, format, resistance )
end

"""
    ma2comp( m, a )

Converts a pair of values in magnitude / angle format to a complex number.
"""
ma2comp( m, a ) =  m * cis( deg2rad( a ) )

dB2Mag( dB ) = 10 ^ ( dB / 20 )

"""
    da2comp( m, a )

Converts a pair of values in dB / angle format to a complex number.
"""
da2comp( d, a ) = ma2comp( dB2Mag( d ), a )

"Holds conversion functions for parameter format symbols when parsing."
const ParseConversions = Dict{ Symbol, Function }(
  :RealImaginary  => complex,
  :MagnitudeAngle => ma2comp,
  :DecibelAngle   => da2comp,
)

"""
    parse_data( line, N, [ options ] )

Returns a data point structure for a valid line of a one-port Touchstone file.

For 3 or more ports, the line should be the concatenation of 3 or more lines from the Touchstone file.
"""
function parse_data( line::String, N, options = Options(), version = 1 )
  vals = map( s -> parse( Float64, s ), split( line ) )
  parse_data( vals, N, options, version )
end

"""
    parse_data( vals, N, [ options ] )

Returns a data point structure for a vector of numbers from a one-port Touchstone file.

For 3 or more ports, the array should contain exactly 2N² + 1 values.
"""
function parse_data( vals::Array{Float64}, N, options = Options(), version = 1, twoPortDataFlipped = true )
  if length( vals ) != 2N * N + 1
    error( "Wrong number of values!" )
  end
  freq = vals[ 1 ]  * options.unit
  pairs = zip( vals[ 2:2:end - 1 ], vals[ 3:2:end ] )
  converted = map( pair -> ParseConversions[ options.format ]( pair... ), pairs )
  # renormalize Z, Y, G and H parameters for version 1.0
  if version == 1
    if options.parameter == :ImpedanceParameters
      converted = map( x -> x * options.resistance, converted )
    end
    if options.parameter == :AdmittanceParameters
      converted = map( x -> x / options.resistance, converted )
    end
    if N == 2
      if options.parameter == :HybridGParameters
        converted[ 1 ] /= options.resistance
        converted[ 4 ] *= options.resistance
      end
      if options.parameter == :HybridHParameters
        converted[ 1 ] *= options.resistance
        converted[ 4 ] /= options.resistance
      end
    end
  end
  res = reshape( converted, N, N )
  if N > 2 || ( N == 2 && !twoPortDataFlipped )
    res = permutedims( res, [ 2, 1 ] )
  end
  return DataPoint( freq, res )
end

function parseNoiseData( vals::Array{Float64}, options = Options(), version = 1 )
  if length( vals ) != 5
    error( "Wrong number of noise data values!" )
  end
  frequency = vals[ 1 ]  * options.unit
  minNoiseFigure = dB2Mag( vals[ 2 ] )
  reflCoeff = ma2comp( vals[ 3 ], vals[ 4 ] )
  effNoiseRes = vals[ 5 ]
  if version == 1
    effNoiseRes *= options.resistance
  end
  return NoiseDataPoint( frequency, minNoiseFigure, reflCoeff, effNoiseRes )
end

function stripCommentFromLine( line )
  comment = ""
  strings = split( line, "!" )
  if length( strings ) > 1
    comment = join( strings[ 2:end ], "!" )
    line = strings[ 1 ]
  end
  return( line, comment)
end

function getValsAndPushComments!( comments, line )
  line, comment = stripCommentFromLine( line )
  if comment != ""
    push!( comments, comment )
  end
  vals = map( s -> parse( Float64, s ), split( line ) )
  return vals
end

"""
    parse_touchstone_stream( stream, [ ports ] )

Returns a TouchstoneData structure for a valid Touchstone data stream.
"""
function parse_touchstone_stream( stream::IO, ports::Integer = 1 )
  version = 1
  options = Options()
  option_line_found = false
  comments = Vector{ String }()
  data = Vector{ DataPoint }()
  noiseData = Vector{ NoiseDataPoint }()
  neededValues = 1 + 2ports * ports
  vals = Vector{ Float64 }()
  keywordparams = Dict{ Symbol, Any }()
  networkdata = false
  noiseDataExpected = false
  twoPortDataFlipped = true
  referenceNextLine = false
  endfound = false
  lastParamFrequency = -1
  for line in eachline( stream )
    if referenceNextLine
      keywordparams[ :Reference ] = parse_reference( split( line ) )
      referenceNextLine = false
    elseif is_empy_line( line )
      ;# nothing
    elseif is_keyword_line( line )
      version = 2
      keyword, params = parse_keyword_line( line )
      keywordparams[ keyword ] = params
      if option_line_found && keyword == :Version
        error( "V2.0: [Version] 2.0 before option line expected." )
      end
      if keyword in unOrderedKeywordSymbols
        keywordString = keywordSymbolStrings[ keyword ]
        if !haskey( keywordparams, :Version ) || keywordparams[ :Version ] != 2
          error( "V2.0: [Version] keyword before [$( keywordString )] keyword expected." )
        end
        if !option_line_found
          error( "V2.0: Option line before [$( keywordString )] keyword expected." )
        end
        if !haskey( keywordparams, :NumberOfPorts )
          error( "V2.0: [Number of Ports] keyword before [$( keywordString )] keyword expected." )
        end
        if haskey( keywordparams, :NetworkData )
          error( "V2.0: [Network Data] keyword after [$( keywordString )] keyword expected." )
        end
      end
      if keyword == :Reference
        if length( keywordparams[ :Reference ] ) == 0
          referenceNextLine = true
        end
      elseif keyword == :NetworkData
        networkdata = true
        if haskey( keywordparams, :NoiseData )
          error( "V2.0: [Noise Data] keyword after [Network Data] keyword expected." )
        end
        if !haskey( keywordparams, :NumberOfFrequencies )
          error( "V2.0: [Number of Frequencies] keyword before [Network Data] keyword expected." )
        end
      elseif keyword == :NumberOfPorts
        ports = keywordparams[ :NumberOfPorts ]
        neededValues = 1 + 2ports * ports
      elseif keyword == :TwoPortDataOrder
        dataOrder = keywordparams[ :TwoPortDataOrder ]
        if dataOrder == "12_21"
          twoPortDataFlipped = true
        elseif dataOrder == "21_12"
          twoPortDataFlipped = false
        end
      elseif keyword == :NoiseData
        noiseDataExpected   = true
        networkdata = false
      elseif keyword == :End
        networkdata = false
        noiseDataExpected   = false
        endfound = true
      end
    elseif is_comment_line( line )
      push!( comments, parse_comment_line( line ) )
    elseif !option_line_found && is_option_line( line )
      if haskey( keywordparams, :Version ) && keywordparams[ :Version ] == 2 && haskey( keywordparams, :NumberOfPorts )
        error( "V2.0: Option line before [Number of Ports] keyword expected." )
      end
      options = parse_option_line( line )
      option_line_found = true
    elseif noiseDataExpected
      vals = getValsAndPushComments!( comments, line )
      push!( noiseData, parseNoiseData( vals, options, version ) )
    elseif !haskey( keywordparams, :Version ) || ( keywordparams[ :Version ] == 2 && networkdata )
      newVals = getValsAndPushComments!( comments, line )
      if version == 1
        # V1.0 Noise Data starts, when frequency decreases.
        if length( vals ) == 0
          if newVals[ 1 ] <= lastParamFrequency
            push!( noiseData, parseNoiseData( newVals, options, 1 ) )
            vals = Vector{Float64}()
            noiseDataExpected = true
          end
          lastParamFrequency = newVals[ 1 ]
        end
        # V1.0: not more then 4 value pairs per line. (#7)
        # The first  line contains the frequency.
        maxVals = length( vals ) == 0 ? 9 : 8
        if length( newVals ) > maxVals
          error( "V1.0: more than four complex values in one parameter line." )
        end
      end
      vals = [ vals; newVals ]
      if length( vals ) >= neededValues
        push!( data, parse_data( vals, ports, options, version, twoPortDataFlipped  ) )
        vals = Vector{Float64}()
      end
    elseif endfound
      error( "V2.0: Non empty or comment line found after [End] keyword." )
    end
  end
  ret = TouchstoneData( data, noiseData, options, comments, keywordparams )
  return ret
end

"""
    parse_touchstone_string( in, [ ports ] )

Returns a TouchstoneData structure for a valid Touchstone data string.
"""
parse_touchstone_string( in::String, ports::Integer = 1 ) = parse_touchstone_stream( IOBuffer( in ), ports )

"""
    parse_touchstone_file( filename, [ ports ] )

Returns a TouchstoneData structure for a valid Touchstone data file.

When called without ports, tries to interpret the file extension, such as ".s2p", or ".S4P".
"""
function parse_touchstone_file( filename::String, ports::Integer = 0 )
  if ports == 0
    extorig =  splitext( filename )[ 2 ][ 2: end ]
    ext = uppercase( extorig )
    if ext[ 1 ] == 'S' && ext[ end ] == 'P'
      ports = parse( Int64, ext[ 2:end - 1 ] )
    else
      error( "File extension ( $( extorig ) ) not recognized, unknown number of ports." )
    end
  end
  open( filename ) do io
    parse_touchstone_stream( open( filename ), ports )
  end
end

# V2.0
"Symbols for different keywords"
const keywordStringSymbols = Dict{ String, Symbol }(
  "VERSION"         => :Version,
  "NUMBER OF PORTS" => :NumberOfPorts,
  "TWO-PORT DATA ORDER"         => :TwoPortDataOrder,
  "NUMBER OF FREQUENCIES"       => :NumberOfFrequencies,
  "NUMBER OF NOISE FREQUENCIES" => :NumberOfNoiseFrequencies,
  "REFERENCE"         => :Reference,
  "MATRIX FORMAT"     => :MatrixFormat,
  "MIXED-MODE ORDER"  => :MixedModeOrder,
  "BEGIN INFORMATION" => :BeginInformation,
  "END INFORMATION"   => :EndInformation,
  "NETWORK DATA"      => :NetworkData,
  "NOISE DATA"        => :NoiseData,
  "END"   => :End,
)

const keywordSymbolStrings = Dict{ Symbol, String }(
  :Version        => "Version",
  :NumberOfPorts  => "Number of Ports",
  :TwoPortDataOrder         => "Two-Port Data Order",
  :NumberOfFrequencies      => "Number of Frequencies",
  :NumberOfNoiseFrequencies => "Number of Noise Frequencies",
  :Reference        => "Reference",
  :MatrixFormat     => "Matrix Format",
  :MixedModeOrder   => "Mixed-Mode Order",
  :BeginInformation => "Begin Information",
  :EndInformation   => "End Information",
  :NetworkData      => "Network Data",
  :NoiseData  => "NoiseData",
  :End        => "End",
)

const unOrderedKeywordSymbols = [
  :TwoPortDataOrder,
  :NumberOfFrequencies,
  :NumberOfNoiseFrequencies,
  :Reference,
  :MatrixFormat,
  :MixedModeOrder,
  :BeginInformation,
  :EndInformation,
]

"""
    is_keyword_line( line )

Checks, if a line is a valid keyword line.
"""
is_keyword_line( line::String ) = length( line ) > 0 && line[ 1 ] == '['

function parse_N_params( params, N::Integer, tp::DataType = String )
  if length( params ) != N
    error( "Exactly $( N ) parameter expected." )
  else
    if tp != String
      params = map( x -> parse( tp, x ), params )
    end
    if N == 1
      params = params[ 1 ]
    end
    return params
  end
end

function parse_N_params( params, tp::DataType )
  N = length( params )
  if tp != String
    params = map( x -> parse( tp, x ), params )
  end
end

function parse_version_params( params )
  versionstring = parse_N_params( params, 1 )
  if versionstring == "2.0"
    return 2
  else
    error( "Unknown version $( versionstring )" )
  end
end

parse_number_of_ports( params ) = parse_N_params( params, 1, Int64 )
parse_two_port_data_order( params )  = parse_N_params( params, 1 )
parse_number_of_frequencies( params )       = parse_N_params( params, 1, Int64 )
parse_number_of_noise_frequencies( params ) = parse_N_params( params, 1, Int64 )
parse_reference( params )         = parse_N_params( params, Float64 )
parse_matrix_format( params )     = []
parse_mixed_mode_order( params )  = []
parse_begin_information( params ) = []
parse_end_information( params )   = []
parse_network_data( params )  = []
parse_noise_data( params )  = []
parse_end( params )  = []

const version_param_parsers = Dict{ Symbol, Function }(
  :Version        => parse_version_params,
  :NumberOfPorts  => parse_number_of_ports,
  :TwoPortDataOrder          => parse_two_port_data_order,
  :NumberOfFrequencies       => parse_number_of_frequencies,
  :NumberOfNoiseFrequencies  => parse_number_of_noise_frequencies,
  :Reference      => parse_reference,
  :MatrixFormat   => parse_matrix_format,
  :MixedModeOrder => parse_mixed_mode_order,
  :BeginInformation => parse_begin_information,
  :EndInformation   => parse_end_information,
  :NetworkData      => parse_network_data,
  :NoiseData  => parse_noise_data,
  :End        => parse_end,
)

"""
    parse_keyword_line( line )

Returns the keyword parameters for a keyword line.
"""
function parse_keyword_line( line )
  keyword, paramstring = split( line[ 2:end ], "]" )
  keyword = uppercase( keyword )
  params = split( paramstring )
  if haskey( keywordStringSymbols, keyword )
    kwSymbol = keywordStringSymbols[ keyword ]
    params = version_param_parsers[ kwSymbol ]( params )
  else
    error( "Not a valid keyword." )
  end
  return kwSymbol, params
end


module Optimisers

using Functors: functor, fmap, isleaf
using LinearAlgebra

include("interface.jl")

include("destructure.jl")
export destructure, total, total2

include("rules.jl")
export Descent, ADAM, Momentum, Nesterov, RMSProp,
       ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW, RADAM, OADAM, AdaBelief,
       WeightDecay, ClipGrad, ClipNorm, OptimiserChain

"""
    Optimisers.apply!(rule::RuleType, state, parameters, gradient) -> (state, gradient)

This defines the action of any optimisation rule. It should return the modified gradient
which will be subtracted from the parameters, and the updated state (if any) for use at
the next iteration, as a tuple `(state, gradient)`.

For efficiency it is free to mutate the old state, but only what is returned will be used.
Ideally this should check `iswriteable(x)`, which the built-in rules do via [`@..`](@ref).

The initial state is `init(rule::RuleType, parameters)`.

# Example
```jldoctest
julia> Optimisers.init(Descent(0.1), Float32[1,2,3]) === nothing
true

julia> Optimisers.apply!(Descent(0.1), nothing, Float32[1,2,3], [4,5,6])
(nothing, Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}}(*, ([4, 5, 6], 0.1f0)))
```
"""
apply!

"""
    Optimisers.init(rule::RuleType, parameters) -> state

Sets up the initial state for a given optimisation rule, and an array of parameters.
This and [`apply!`](@ref) are the two functions which any new optimisation rule must define.

# Examples
```jldoctest
julia> Optimisers.init(Descent(), Float32[1,2,3])  # is `nothing`

julia> Optimisers.init(Momentum(), [1.0, 2.0])
2-element Vector{Float64}:
 0.0
 0.0
```
"""
init

"""
    Optimisers.setup(rule, model) -> tree

Initialises the given optimiser for every trainable parameter within the model.
Returns a tree of the relevant states, which must be passed to [`update`](@ref)
or [`update!`](@ref).

# Example
```jldoctest
julia> Optimisers.setup(Descent(0.1f0), (x = rand(3), y = (true, false), z = tanh))
(x = Leaf(Descent{Float32}(0.1), nothing), y = (nothing, nothing), z = nothing)
```
"""
setup

"""
    Optimisers.update(tree, model, gradient) -> (tree, model)

Uses the optimiser and the gradient to change the trainable parameters in the model.
Returns the improved model, and the optimiser states needed for the next update.
The initial tree of states comes from [`setup`](@ref).

See also [`update!`](@ref), which will be faster for models of ordinary `Array`s or `CuArray`s.

# Example
```jldoctest
julia> m = (x = Float32[1,2,3], y = tanh);

julia> t = Optimisers.setup(Descent(0.1f0), m)
(x = Leaf(Descent{Float32}(0.1), nothing), y = nothing)

julia> g = (x = [1,1,1], y = nothing);  # fake gradient

julia> Optimisers.update(t, m, g)
((x = Leaf(Descent{Float32}(0.1), nothing), y = nothing), (x = Float32[0.9, 1.9, 2.9], y = tanh))
```
"""
update

"""
    Optimisers.update!(tree, model, gradient) -> (tree, model)

Uses the optimiser and the gradient to change the trainable parameters in the model.
Returns the improved model, and the optimiser states needed for the next update.
The initial tree of states comes from [`setup`](@ref).

This is used in exactly the same manner as [`update`](@ref), but because it may mutate
arrays within the old model (and the old state), it will be faster for models of ordinary
`Array`s or `CuArray`s. However, you should not rely on the old model being fully updated.
"""
update!

end # module


"""
    minimalsetup1d(N)

Creates a minimal setup of all important variables.

Such it generates the moments `height`, `velocity` as well as a dummy distribution with its speeds `f0`, `f1` and `f2`.
Further a set of forcing is supplied, they can be used with `foces.slip`, `.h∇p`, `.bathymetry` and `.thermal`.

# Example
```jldoctest
julia> using JuSwalbe

julia> constants, mom, f, distribution = minimalsetup1d(10)
(JuSwalbe.Inputconstants
  lx: Int64 10
  ly: Int64 512
  maxruntime: Int64 100000
  dumping: Int64 1000
  τ: Float64 1.0
  gravity: Float64 0.0
  γ: Float64 0.01
  δ: Float64 1.0
  μ: Float64 0.16666666666666666
  kbt: Float64 0.0
, JuSwalbe.Macroquant{Array{Float64,1},Array{Float64,1}}
  height: Array{Float64}((10,)) [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
  velocity: Array{Float64}((10,)) [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
  pressure: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  energy: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
, JuSwalbe.Forces{Array{Float64,1}}
  slip: Array{Float64}((10,)) [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
  h∇p: Array{Float64}((10,)) [-0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1]
  bathymetry: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  thermal: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
, JuSwalbe.DistributionD1Q3{Array{Float64,1}}
  f0: Array{Float64}((10,)) [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
  f1: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  f2: Array{Float64}((10,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
)

julia> mom.height
10-element Array{Float64,1}:
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
```
See also: [`minimalsetup2d`](@ref)
"""
function minimalsetup1d(N::Int)
    input = Inputconstants(lx=N)
    mom = Macroquant(height = ones(N), velocity=fill(0.1, N), pressure=zeros(N), energy=zeros(N))
    forces = Forces(slip=fill(0.1, N), thermal=zeros(N), h∇p=fill(-0.1, N), bathymetry=zeros(N))
    dist = DistributionD1Q3(f0 = ones(N), f1 = zeros(N), f2 = zeros(N))
    return input, mom, forces, dist
end

"""
    minimalsetup2d(N,M)

Creates a minimal working example set of macroscopic moments, forces and a distribution.

# Example
```jldoctest
julia> using JuSwalbe

julia> input, mom, f, dist = minimalsetup2d(5,5)
(JuSwalbe.Inputconstants
  lx: Int64 5
  ly: Int64 5
  maxruntime: Int64 100000
  dumping: Int64 1000
  τ: Float64 1.0
  gravity: Float64 0.0
  γ: Float64 0.01
  δ: Float64 1.0
  μ: Float64 0.16666666666666666
  kbt: Float64 0.0
, JuSwalbe.Macroquant{Array{Float64,2},JuSwalbe.Twovector{Array{Float64,2}}}
  height: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  velocity: JuSwalbe.Twovector{Array{Float64,2}}
  pressure: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  energy: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
, JuSwalbe.Forces{JuSwalbe.Twovector{Array{Float64,2}}}
  slip: JuSwalbe.Twovector{Array{Float64,2}}
  h∇p: JuSwalbe.Twovector{Array{Float64,2}}
  bathymetry: JuSwalbe.Twovector{Array{Float64,2}}
  thermal: JuSwalbe.Twovector{Array{Float64,2}}
, JuSwalbe.DistributionD2Q9{Array{Float64,2}}
  f0: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  f1: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f2: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f3: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f4: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f5: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f6: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f7: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  f8: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
)

```

See also: [`minimalsetup1d`](@ref)
"""
function minimalsetup2d(N::Int, M::Int)
    input = Inputconstants(lx=N, ly=M)
    vel = Twovector(x=fill(0.02, (N,M)), y=fill(-0.02, (N,M)))
    zerovec = Twovector(x=zeros(N,M), y=zeros(N,M))
    mom = Macroquant(height = ones(N,M), velocity=vel, pressure=zeros(N,M), energy=zeros(N,M))
    forces = Forces(slip=vel, thermal=zerovec, h∇p=vel, bathymetry=zerovec)
    dist = DistributionD2Q9(f0 = ones(N,M), f1 = zeros(N,M), f2 = zeros(N,M),
                            f3 = zeros(N,M), f4 = zeros(N,M), f5 = zeros(N,M),
                            f6 = zeros(N,M), f7 = zeros(N,M), f8 = zeros(N,M))
    return input, mom, forces, dist
end

"""
  simplemoment2d(n,m,type)

Creates moments of type JuSwalbe.Macroquant with dimensions (n,m).

# Example
```jldoctest
julia> using JuSwalbe

julia> mom = simplemoment2d(5,5)
JuSwalbe.Macroquant{Array{Float64,2},JuSwalbe.Twovector{Array{Float64,2}}}
  height: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  velocity: JuSwalbe.Twovector{Array{Float64,2}}
  pressure: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
  energy: Array{Float64}((5, 5)) [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]

julia> mom.height
5×5 Array{Float64,2}:
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0

```
"""
function simplemoment2d(n::Int, m::Int; T=Float64)
  mom = JuSwalbe.Macroquant(height=ones(T, (n,m)),
                            velocity=JuSwalbe.Twovector(x=fill(T(0.1),(n,m)), y=fill(T(-0.1),(n,m))),
                            pressure=zeros(T, (n,m)),
                            energy=zeros(T,(n,m)))
  return mom
end

"""
  simplemoment1d(n,type)

Creates moments of type JuSwalbe.Macroquant with length n.

# Example
```jldoctest
julia> using JuSwalbe

julia> mom = simplemoment1d(5)
JuSwalbe.Macroquant{Array{Float64,1},Array{Float64,1}}
  height: Array{Float64}((5,)) [1.0, 1.0, 1.0, 1.0, 1.0]
  velocity: Array{Float64}((5,)) [0.1, 0.1, 0.1, 0.1, 0.1]
  pressure: Array{Float64}((5,)) [0.0, 0.0, 0.0, 0.0, 0.0]
  energy: Array{Float64}((5,)) [0.0, 0.0, 0.0, 0.0, 0.0]

julia> mom.height
5-element Array{Float64,1}:
 1.0
 1.0
 1.0
 1.0
 1.0

```
"""
function simplemoment1d(n::Int; T=Float64)
  mom = JuSwalbe.Macroquant(height=ones(T, n), velocity=fill(T(0.1),n), pressure=zeros(T, n), energy=zeros(T, n))
  return mom
end

"""
  simpleTwovector(n,m,T)

Generates a JuSwalbe.Twovector of type T and dimension (n,m).

# Example
```jldoctest
julia> using JuSwalbe

julia> twovec = simpleTwovector(5,5)
JuSwalbe.Twovector{Array{Float64,2}}
  x: Array{Float64}((5, 5)) [0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1; … ; 0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1]
  y: Array{Float64}((5, 5)) [-0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1; … ; -0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1]

julia> twovec.x
5×5 Array{Float64,2}:
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
 0.1  0.1  0.1  0.1  0.1
```
"""
function simpleTwovector(n::Int,m::Int; T=Float64)
  xy = JuSwalbe.Twovector(x=fill(T(0.1),(n,m)), y=fill(T(-0.1),(n,m)))
  return xy
end

function zeroTwovector(n::Int,m::Int; T=Float64)
  xy = JuSwalbe.Twovector(x=zeros(T ,(n,m)), y=zeros(T, (n,m)))
  return xy
end

function zeroTwovectorCU(n::Int,m::Int; T=Float32)
  jaja = zeros(n,m)
  xy = JuSwalbe.Twovector(x=CuArray(jaja), y=CuArray(jaja))
  return xy
end

"""
  simpledistD2Q9(n, m, T)

Generates a `D2Q9` distribution function of dimension `n`,`m` with defined fill statements.

# Example
```jldoctest
julia> using JuSwalbe

julia> dist = simpledistD2Q9(5,5)
JuSwalbe.DistributionD2Q9{Array{Float64,2}}
  f0: Array{Float64}((5, 5)) [1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0; … ; 1.0 1.0 … 1.0 1.0; 1.0 1.0 … 1.0 1.0]
  f1: Array{Float64}((5, 5)) [0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1; … ; 0.1 0.1 … 0.1 0.1; 0.1 0.1 … 0.1 0.1]
  f2: Array{Float64}((5, 5)) [-0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1; … ; -0.1 -0.1 … -0.1 -0.1; -0.1 -0.1 … -0.1 -0.1]
  f3: Array{Float64}((5, 5)) [0.01 0.01 … 0.01 0.01; 0.01 0.01 … 0.01 0.01; … ; 0.01 0.01 … 0.01 0.01; 0.01 0.01 … 0.01 0.01]
  f4: Array{Float64}((5, 5)) [-0.01 -0.01 … -0.01 -0.01; -0.01 -0.01 … -0.01 -0.01; … ; -0.01 -0.01 … -0.01 -0.01; -0.01 -0.01 … -0.01 -0.01]
  f5: Array{Float64}((5, 5)) [0.2 0.2 … 0.2 0.2; 0.2 0.2 … 0.2 0.2; … ; 0.2 0.2 … 0.2 0.2; 0.2 0.2 … 0.2 0.2]
  f6: Array{Float64}((5, 5)) [0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02; … ; 0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02]
  f7: Array{Float64}((5, 5)) [0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02; … ; 0.02 0.02 … 0.02 0.02; 0.02 0.02 … 0.02 0.02]
  f8: Array{Float64}((5, 5)) [0.5 0.5 … 0.5 0.5; 0.5 0.5 … 0.5 0.5; … ; 0.5 0.5 … 0.5 0.5; 0.5 0.5 … 0.5 0.5]

julia> dist.f0
5×5 Array{Float64,2}:
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0
 1.0  1.0  1.0  1.0  1.0

julia> typeof(dist)
JuSwalbe.DistributionD2Q9{Array{Float64,2}}
```

See also: [`simpleTwovector`](@ref), [`simplemoment2d`](@ref)

"""
function simpledistD2Q9(n::Int, m::Int; T=Float64)
  dist = JuSwalbe.DistributionD2Q9(f0 = ones(T, (n,m)), f1 = fill(T(0.1), (n,m)), f2 = fill(T(-0.1), (n,m)),
                                   f3 = fill(T(0.01), (n,m)), f4 = fill(T(-0.01), (n,m)), f5 = fill(T(0.2), (n,m)),
                                   f6 = fill(T(0.02), (n,m)), f7 = fill(T(0.02), (n,m)), f8 = fill(T(0.5), (n,m)))

  return dist
end


"""
Generic Fermion module
"""
module AbstractFermion
    #export FermionFields,Wx!,Wdagx!,clear!,substitute_fermion!,Dx!

    #export FermionFields,
    #        Wx!,Wdagx!,clear!,substitute_fermion!,Dx!,fermion_shift!,
    #        fermion_shiftB!,add!,set_wing_fermi!,WdagWx!,apply_periodicity

    import ..Actions:FermiActionParam,FermiActionParam_Wilson,
                FermiActionParam_WilsonClover,FermiActionParam_Staggered
    import ..Gaugefields:GaugeFields,GaugeFields_1d,SU3GaugeFields,SU2GaugeFields,SU3GaugeFields_1d,SU2GaugeFields_1d,
                staggered_phase,SUn,SU2,SU3,SUNGaugeFields,SUNGaugeFields_1d,SU
    using Random
    abstract type FermionFields end

    function clear!(a::FermionFields)
        n1,n2,n3,n4,n5,n6 = size(a.f)
        for i6=1:n6
            for i5=1:n5
                for i4=1:n4
                    for i3=1:n3
                        for i2=1:n2
                            @simd for i1=1:n1
                                a.f[i1,i2,i3,i4,i5,i6]= 0
                            end
                        end
                    end
                end
            end
        end
    end

    function add!(c::FermionFields,alpha::Number,a::FermionFields,beta::Number,b::FermionFields) #c = c + alpha*a + beta*b
        n1,n2,n3,n4,n5,n6 = size(a.f)

        for i6=1:n6
            for i5=1:n5
                for i4=1:n4
                    for i3=1:n3
                        for i2=1:n2
                            @simd for i1=1:n1
                                #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] += alpha*a.f[i1,i2,i3,i4,i5,i6] +beta*b.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
        return
    end

    function add!(c::FermionFields,alpha::Number,a::FermionFields) #c = c + alpha*a
        n1,n2,n3,n4,n5,n6 = size(a.f)

        for i6=1:n6
            for i5=1:n5
                for i4=1:n4
                    for i3=1:n3
                        for i2=1:n2
                            @simd for i1=1:n1
                                #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] += alpha*a.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
        return
    end

    function add!(coeff::Number,c::FermionFields,alpha::Number,a::FermionFields) #c = coeff*c + alpha*a
        n1,n2,n3,n4,n5,n6 = size(a.f)

        for i6=1:n6
            for i5=1:n5
                for i4=1:n4
                    for i3=1:n3
                        for i2=1:n2
                            @simd for i1=1:n1
                                #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] = coeff*c.f[i1,i2,i3,i4,i5,i6] + alpha*a.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
        return
    end

    function Wx!(xout::FermionFields,U::Array{GaugeFields,1},
        x::FermionFields,temps::Array{FermionFields,1},fparam::T)  where T <: FermiActionParam
        error("This FermionFields type is not supported yet in function Wx!. The type is $(typeof(x))")
    end

    function Dx!(xout::T,U::Array{G,1},
        x::T,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        error("This FermionFields type is not supported yet in function Dx!. The type is $(typeof(x))")

        return
    end

    function Ddagx!(xout::T,U::Array{G,1},
        x::T,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        error("This FermionFields type is not supported yet in function Ddagx!. The type is $(typeof(x))")

        return
    end

    function Dxplus!(xout::T,ν::Int64,U::Array{G,1},
        x::T,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        #temp = temps[4]
        error("This FermionFields type is not supported yet in function Dx!. The type is $(typeof(x))")

        return
    end

    function fermion_shiftB!(b::T,evensite,u::Array{GaugeFields{SU{NC}},1},μ,a::T)  where {NC,T <: FermionFields}
        error("This FermionFields type is not supported yet in function fermion_shiftB!. The type is $(typeof(x))")
    end

    function Wdagx!(xout::FermionFields,U::Array{GaugeFields,1},
        x::FermionFields,temps::Array{FermionFields,1},fparam::T)  where T <: FermiActionParam
        error("This FermionFields type is not supported yet in function Wdagx!. The type is $(typeof(x))")
    end

    function WdagWx!(xout::T,U::Array{G,1},x::T,temps::Array{T,1},fparam) where {T <: FermionFields,G <: GaugeFields}
        temp = temps[5]
        Wx!(temp,U,x,temps,fparam)

        Wdagx!(xout,U,temp,temps,fparam)

        return
    end

    function WdagWx!(xout::T,U::Array{G,1},x::T,temps::Array{T,1},fparam,indices) where {T <: FermionFields,G <: GaugeFields}
        WdagWx!(xout,U,x,temps,fparam)
        return
    end

    function substitute_fermion!(H,j,x::FermionFields)
        error("substitute_fermion! is not implemented for $(typeof(x))")
    end

    function set_wing_fermi!(a::FermionFields)
        NT = a.NT
        NZ = a.NZ
        NY = a.NY
        NX = a.NX
        NC = a.NC

        #!  X-direction
        for ialpha=1:4
            for it=1:NT
                for iz = 1:NZ
                    for iy=1:NY
                        for k=1:NC
                            a[k,0,iy,iz,it,ialpha] = a.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
                        for k=1:NC
                            a[k,NX+1,iy,iz,it,ialpha] = a.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
                        for k=1:NC
                            a[k,ix,0,iz,it,ialpha] = a.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
                        for k=1:NC
                            a[k,ix,NY+1,iz,it,ialpha] = a.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
                        for k=1:NC
                            a[k,ix,iy,0,it,ialpha] = a.BoundaryCondition[3]*a[k,ix,iy,NZ,it,ialpha]
                            a[k,ix,iy,NZ+1,it,ialpha] = a.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
                        for k=1:NC
                            a[k,ix,iy,iz,0,ialpha] = a.BoundaryCondition[4]*a[k,ix,iy,iz,NT,ialpha]
                            a[k,ix,iy,iz,NT+1,ialpha] = a.BoundaryCondition[4]*a[k,ix,iy,iz,1,ialpha]
                        end
                    end
                end
            end

        end

    end

    function fermion_shift!(b::F,u::Array{GaugeFields{SU{NC}},1},μ,a::F) where {NC,F <: FermionFields}
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        #NC = a.NC

        if μ > 0
            idel = zeros(Int64,4)
            idel[μ] = 1

            n6 = size(a.f)[6]
            for ialpha=1:n6
                for it=1:NT
                    it1 = it + idel[4]
                    for iz=1:NZ
                        iz1 = iz + idel[3]
                        for iy=1:NY
                            iy1 = iy + idel[2]
                            for ix=1:NX
                                ix1 = ix + idel[1]
                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    @simd for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += u[μ][k1,k2,ix,iy,iz,it]*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                            end
                        end
                    end
                end
            end

        elseif μ < 0
            idel = zeros(Int64,4)
            idel[-μ] = 1
            n6 = size(b.f)[6]
            for ialpha =1:n6
                for it=1:NT
                    it1 = it - idel[4]
                    for iz=1:NZ
                        iz1 = iz -idel[3]
                        for iy=1:NY
                            iy1 = iy -idel[2]
                            for ix=1:NX
                                ix1 = ix -idel[1]

                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    @simd for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += conj(u[-μ][k2,k1,ix1,iy1,iz1,it1])*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shift!(b::F,evensite::Bool,u::Array{GaugeFields{SU{NC}},1},μ,a::F) where {NC,F <: FermionFields}
        if μ == 0
            substitute!(b,a)
            return
        end
        ibush = ifelse(evensite,0,1)

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        #NC = a.NC

        if μ > 0
            idel = zeros(Int64,4)
            idel[μ] = 1

            n6 = size(a.f)[6]
            for ialpha=1:n6
                for it=1:NT
                    it1 = it + idel[4]
                    for iz=1:NZ
                        iz1 = iz + idel[3]
                        for iy=1:NY
                            iy1 = iy + idel[2]
                            xran =1+(1+ibush+iy+iz+it)%2:2:NX
                            for ix in xran
                                ix1 = ix + idel[1]
                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    @simd for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += u[μ][k1,k2,ix,iy,iz,it]*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                            end
                        end
                    end
                end
            end

        elseif μ < 0
            idel = zeros(Int64,4)
            idel[-μ] = 1
            n6 = size(b.f)[6]
            for ialpha =1:n6
                for it=1:NT
                    it1 = it - idel[4]
                    for iz=1:NZ
                        iz1 = iz -idel[3]
                        for iy=1:NY
                            iy1 = iy -idel[2]
                            xran =1+(1+ibush+iy+iz+it)%2:2:NX
                            for ix in xran
                                ix1 = ix -idel[1]

                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    @simd for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += conj(u[-μ][k2,k1,ix1,iy1,iz1,it1])*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end

    end

    function apply_periodicity(ix,iy,iz,it,NX,NY,NZ,NT,BC)

        ixflag1 = (ix < 1)
        ixflag2 = (ix > NX)
        ix += ifelse(ixflag1,NX,0)
        ix += ifelse(ixflag2,-NX,0)

        sign =  ifelse(ixflag1 || ixflag2,BC[1],1)
        #sign =  ifelse(ixflag2,BC[1],1)

        iyflag1 = (iy < 1)
        iyflag2 = (iy > NY)

        iy += ifelse(iyflag1,NY,0)
        iy += ifelse(iyflag2,-NY,0)

        sign = ifelse(iyflag1 || iyflag2,BC[2],1)
        #sign = ifelse(iyflag2,BC[2],1)

        izflag1 = (iz < 1)
        izflag2 = (iz > NZ)

        iz += ifelse(izflag1,NZ,0)
        iz += ifelse(izflag2,-NZ,0)

        sign = ifelse(izflag1 || izflag2,BC[3],1)
        #sign = ifelse(izflag2,BC[3],1)

        itflag1 = (it < 1)
        itflag2 = (it > NT)

        it += ifelse(itflag1,NT,0)
        it += ifelse(itflag2,-NT,0)

        sign = ifelse(itflag1 || itflag2,BC[4],1)
        #sign = ifelse(itflag2,BC[4],1)

        return ix,iy,iz,it,sign
    end

    """
c-------------------------------------------------c
c     Random number function for Gaussian  Noise
    with σ^2 = 1/2
c-------------------------------------------------c
    """
    function gauss_distribution_fermi!(x::FermionFields)
        NC = x.NC
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        n6 = size(x.f)[6]
        σ = sqrt(1/2)

        for ialpha = 1:n6
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        for ix=1:NX
                            for ic=1:NC

                                x[ic,ix,iy,iz,it,ialpha] = σ*randn()+im*σ*randn()
                            end
                        end
                    end
                end
            end
        end

        set_wing_fermi!(x)

        return
    end

    """
c-------------------------------------------------c
c     Random number function for Gaussian  Noise
    with σ^2 = 1/2
c-------------------------------------------------c
    """
    function gauss_distribution_fermi!(x::FermionFields,randomfunc,σ)
        NC = x.NC
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        n6 = size(x.f)[6]
        #σ = sqrt(1/2)

        for mu = 1:n6
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            for ix=1:NX
                                v1 = sqrt(-log(randomfunc()+1e-10))
                                v2 = 2pi*randomfunc()

                                xr = v1*cos(v2)
                                xi = v1 * sin(v2)

                                x[ic,ix,iy,iz,it,mu] = σ*xr + σ*im*xi
                            end
                        end
                    end
                end
            end
        end

        set_wing_fermi!(x)

        return
    end

    function gauss_distribution_fermi!(x::FermionFields,randomfunc)
        σ = 1
        gauss_distribution_fermi!(x,randomfunc,σ)
    end

    function gauss_distribution_fermi_Z2!(x::FermionFields)
        NC = x.NC
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        n6 = size(x.f)[6]
        #σ = sqrt(1/2)

        for mu = 1:n6
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        for ix=1:NX
                            for ic=1:NC
                                x[ic,ix,iy,iz,it,mu] = rand([-1,1])
                            end
                        end
                    end
                end
            end
        end

        set_wing_fermi!(x)

        return
    end

    """
c-------------------------------------------------c
c     Random number function Z4  Noise
c     https://arxiv.org/pdf/1611.01193.pdf
c-------------------------------------------------c
    """
    function Z4_distribution_fermi!(x::FermionFields)
        NC = x.NC
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        n6 = size(x.f)[6]
        θ = 0.0
        N::Int32 = 4
        Ninv = Float64(1/N)
        for ialpha = 1:n6
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        for ix=1:NX
                            for ic=1:NC
                                θ = Float64(rand(0:N-1))*π*Ninv # r \in [0,π/4,2π/4,3π/4]
                                x[ic,ix,iy,iz,it,ialpha] = cos(θ)+im*sin(θ)
                            end
                        end
                    end
                end
            end
        end

        set_wing_fermi!(x)

        return
    end

end