# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

#=
- read meshes from different formats
- reorder connectivity, create element sets, node sets, ...
- create partitions for parallel runs
- renumber elements / nodes
- maybe precheck for bad elements
- check surface normal direction in boundary elements
- orientation of 2d elements
- etc only topology related stuff
=#

mutable struct Mesh
    nodes :: Dict{Int, Vector{Float64}}
    node_sets :: Dict{Symbol, Set{Int}}
    elements :: Dict{Int, Vector{Int}}
    element_types :: Dict{Int, Symbol}
    element_codes :: Dict{Int, Symbol}
    element_sets :: Dict{Symbol, Set{Int}}
    surface_sets :: Dict{Symbol, Vector{Tuple{Int, Symbol}}}
    surface_types :: Dict{Symbol, Symbol}
    coloring::Union{Nothing, Vector{Vector{Int}}} # Each vector contains a list of elements that do not share nodes
end

function Mesh()
    return Mesh(Dict(), Dict(), Dict(), Dict(), Dict(), Dict(), Dict(), Dict(), nothing)
end

"""
    Mesh(m::Dict)

Create a new `Mesh` using data `m`. It is assumed that `m` is in format what
`abaqus_read_mesh` in `AbaqusReader.jl` is returning.
"""
function Mesh(m::Dict)
    mesh = Mesh()
    mesh.nodes = m["nodes"]
    mesh.elements = m["elements"]
    mesh.element_types = m["element_types"]
    for (k, v) in m["surface_types"]
        mesh.surface_types[Symbol(k)] = v
    end
    for (nset_name, node_ids) in m["node_sets"]
        mesh.node_sets[Symbol(nset_name)] = Set(node_ids)
    end
    for (elset_name, element_ids) in m["element_sets"]
        mesh.element_sets[Symbol(elset_name)] = Set(element_ids)
    end
    for (surfset_name, surfaces) in m["surface_sets"]
        mesh.surface_sets[Symbol(surfset_name)] = surfaces
    end
    return mesh
end

"""
    add_node!(mesh, nid, ncoords)

Add node into the mesh. `nid` is node id and `ncoords` are the node
coordinates.
"""
function add_node!(mesh::Mesh, nid::Int, ncoords::Vector{Float64})
    mesh.nodes[nid] = ncoords
end

"""
    add_nodes!(mesh, nodes)

Add nodes into the mesh.
"""
function add_nodes!(mesh::Mesh, nodes::Dict{Int, Vector{Float64}})
    for (nid, ncoords) in nodes
        add_node!(mesh, nid, ncoords)
    end
end

"""
    add_node_to_node_set!(mesh, nid, ncoords)

Add nodes into a node set. `set_name` is the name of the set and `nids...`
are all the node id:s that wants to be added.
"""
function add_node_to_node_set!(mesh::Mesh, set_name, nids...)
    if !haskey(mesh.node_sets, set_name)
        mesh.node_sets[set_name] = Set{Int}()
    end
    push!(mesh.node_sets[set_name], nids...)
    return
end

"""
    create_node_set_from_element_set!(mesh, set_names...)

Create a new node set from the nodes in an element set. ´set_names...´ are all
the set names to be inserted in the function.
"""
function create_node_set_from_element_set!(mesh::Mesh, set_names::String...)
    for set_name in set_names
        set_name = Symbol(set_name)
        @info("Creating node set $set_name from element set")
        node_ids = Set{Int}()
        for elid in mesh.element_sets[set_name]
            push!(node_ids, mesh.elements[elid]...)
        end
        mesh.node_sets[set_name] = node_ids
    end
    return
end

"""
    create_node_set_from_element_set!(mesh, set_name)

Create a new node set from an element set.
"""
function create_node_set_from_element_set!(mesh::Mesh, set_name::Symbol)
    create_node_set_from_element_set!(mesh, string(set_name))
end

"""
    add_element!(mesh, elid, eltype, connectivity)

Add an element into the mesh. ´elid´ is the element id, ´eltype´ is the type of
the element and ´connectivity´ is the connectivity of the element.
"""
function FEMBase.add_element!(mesh::Mesh, elid, eltype, connectivity)
    mesh.elements[elid] = connectivity
    mesh.element_types[elid] = eltype
    return nothing
end

"""
    add_elements!(mesh, elements)

Add elements into the mesh.
"""
function FEMBase.add_elements!(mesh::Mesh, elements::Dict{Int, Tuple{Symbol, Vector{Int}}})
    for (elid, (eltype, elcon)) in elements
        add_element!(mesh, elid, eltype, elcon)
    end
    return nothing
end

"""
    add_element_to_element_set!(mesh, set_name, elids...)

Add elements into the mesh. ´set_name´ is the name of the element set and
´elids..´ are id:s of all the elements that wants to be added.
"""
function add_element_to_element_set!(mesh::Mesh, set_name, elids...)
    if !haskey(mesh.element_sets, set_name)
        mesh.element_sets[set_name] = Set{Int}()
    end
    push!(mesh.element_sets[set_name], elids...)
end

"""
    copy(mesh)

Return a copy of the mesh.
"""
function Base.copy(mesh::Mesh)
    mesh2 = Mesh()
    mesh2.nodes = copy(mesh.nodes)
    mesh2.node_sets = copy(mesh.node_sets)
    mesh2.elements = copy(mesh.elements)
    mesh2.element_types = copy(mesh.element_types)
    mesh2.element_sets = copy(mesh.element_sets)
    return mesh2
end

"""
    filter_by_element_id(mesh, element_ids)

Filter elements by their id's.
"""
function filter_by_element_id(mesh::Mesh, element_ids::Vector{Int})
    mesh2 = copy(mesh)
    mesh2.elements = Dict()
    for elid in element_ids
        if haskey(mesh.elements, elid)
            mesh2.elements[elid] = mesh.elements[elid]
        end
    end
    return mesh2
end

"""
    filter_by_element_set(mesh, set_name)

Filter elements by an element set.
"""
function filter_by_element_set(mesh::Mesh, set_name)
    filter_by_element_id(mesh::Mesh, collect(mesh.element_sets[set_name]))
end

"""
    create_element(mesh, id)

Create an element from the mesh by it's id.
"""
function create_element(mesh::Mesh, id::Int)
    connectivity = mesh.elements[id]
    element_type = getfield(JuliaFEM, mesh.element_types[id])
    element = Element(element_type, connectivity)
    element.id = id
    update!(element, "geometry", mesh.nodes)
    return element
end

function create_elements(mesh::Mesh; element_type=nothing)
    element_ids = collect(keys(mesh.elements))
    if element_type != nothing
        filter!(id -> mesh.element_types[id] == element_type, element_ids)
    end
    elements = [create_element(mesh, id) for id in element_ids]
    return elements
end

function create_elements(mesh::Mesh, element_sets::Symbol...; element_type=nothing)
    if isempty(element_sets)
        element_ids = collect(keys(mesh.elements))
    else
        element_ids = Set{Int}()
        for set_name in element_sets
            element_ids = union(element_ids, mesh.element_sets[set_name])
        end
    end

    if element_type != nothing
        filter!(id -> mesh.element_types[id] == element_type, element_ids)
    end

    elements = [create_element(mesh, id) for id in element_ids]

    nelements = length(elements)
    content = Dict{Symbol, Int}()
    for elid in element_ids
        eltype = mesh.element_types[elid]
        content[eltype] = get(content, eltype, 0) + 1
    end
    s = join(("$v x $k" for (k, v) in content), ", ")
    v = join(element_sets, ", ")
    @info("Created $nelements elements ($s) from element set: $v.")

    return elements
end

"""
    create_elements(mesh::Mesh, element_set::String)

# Examples

Suppose that there is a `mesh` with element set `Body_1`. Creating elements
based on that element set is done then

```julia
create_elements(mesh, "Body_1")
```
"""
function create_elements(mesh::Mesh, element_sets::String...)
    return create_elements(mesh, map(Symbol, element_sets)...)
end

"""
    find_nearest_nodes(mesh, coords, npts=1; node_set=nothing)

find npts nearest nodes from the mesh and return their id numbers as a list.
"""
function find_nearest_nodes(mesh::Mesh, coords::Vector{Float64}, npts::Int=1; node_set=nothing)
    dist = Dict{Int, Float64}()
    for (nid, c) in mesh.nodes
        if node_set != nothing && !(nid in mesh.node_sets[Symbol(node_set)])
            continue
        end
        dist[nid] = norm(coords-c)
    end
    s = sort(collect(dist), by=x->x[2])
    nd = s[1:npts] # [(id1, dist1), (id2, dist2), ..., (id_npts, dist_npts)]
    node_ids = [n[1] for n in nd]
    return node_ids
end

function find_nearest_node(mesh::Mesh, coords::Vector{Float64}; node_set=nothing)
    return first(find_nearest_nodes(mesh, coords, 1; node_set=node_set))
end

"""
    reorder_element_connectivity!(mesh, mapping)

Apply a new node ordering to elements. JuliaFEM uses the same node ordering as
ABAQUS. If the mesh is parsed from FEM format with some other node ordering,
this function can be used to reorder the nodes.

Parameters
----------
mapping :: Dict{Symbol, Vector{Int}}
    e.g. :Tet10, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

"""
function reorder_element_connectivity!(mesh::Mesh, mapping::Dict{Symbol, Vector{Int}})
    for (elid, eltype) in mesh.element_types
        haskey(mapping, eltype) || continue
        new_order = mapping[eltype]
        element_connectivity = mesh.elements[elid]
        new_element_connectivity = element_connectivity[new_order]
        mesh.elements[elid] = new_element_connectivity
    end
end

function JuliaFEM.Problem(mesh::Mesh, ::Type{P}, name::AbstractString, dimension::Int) where P<:FieldProblem
    problem = Problem(P, name, dimension)
    problem.elements = create_elements(mesh, name)
    return problem
end

function JuliaFEM.Problem(mesh::Mesh, ::Type{P}, name, dimension, parent_field_name) where P<:BoundaryProblem
    problem = Problem(P, name, dimension, parent_field_name)
    problem.elements = create_elements(mesh, name)
    return problem
end

"""
    create_coloring!(mesh::Mesh)

Greedy algorithm for coloring a grid such that no two cells with the same node
have the same color.
This function sets the `coloring` field in `mesh` to a `Vector{Vector{Int}}` where
each vector contains vectors of elements that do not share any nodes.
It is therefore safe to assemble in parallel each element vector by vector.
"""
function create_coloring!(mesh::Mesh)
    # Contains the elements that each node contain
    cell_containing_node = Dict{Int, Set{Int}}()
    for (cellid, nodes) in mesh.elements
        for v in nodes
            if !haskey(cell_containing_node, v)
                cell_containing_node[v] = Set{Int}()
            end
            push!(cell_containing_node[v], cellid)
        end
    end

    I, J, V = Int[], Int[], Bool[]
    for (node, cells) in cell_containing_node
        for cell1 in cells # All these cells have a neighboring node
            for cell2 in cells
                if cell1 != cell2
                    push!(I, cell1)
                    push!(J, cell2)
                    push!(V, true)
                end
            end
        end
    end

    incidence_matrix = sparse(I, J, V)
    # cell -> color of cell
    cell_colors = Dict{Int, Int}()
    # color -> list of cells
    final_colors = Vector{Int}[]
    occupied_colors = Set{Int}()
    # Zero represents no color set yet
    for (cellid, _) in mesh.elements
        cell_colors[cellid] = 0
    end
    total_colors = 0
    for (cellid, _) in mesh.elements
        empty!(occupied_colors)
        # loop over neighbors
        for r in nzrange(incidence_matrix, cellid)
            cell_neighbour = incidence_matrix.rowval[r]
            color = cell_colors[cell_neighbour]
            if color != 0
                push!(occupied_colors, color)
            end
        end

        # occupied colors now contains all the colors we are not allowed to use
        free_color = 0
        for attempt_color in 1:total_colors
            if attempt_color ∉ occupied_colors
                free_color = attempt_color
                break
            end
        end

        if free_color == 0 # no free color found, need to bump max colors
            total_colors += 1
            free_color = total_colors
            push!(final_colors, Int[])
        end

        cell_colors[cellid] = free_color
        push!(final_colors[free_color], cellid)
    end

    mesh.coloring = final_colors
    return mesh
end


struct SomeGroup <: Group end

struct SomeGroupElement <: GroupElement
    elts::Vector{Int} # SomeGroupElement is not isbits anymore
end

@testset "SomeGroup: No interface implemented" begin

    @testset "InterfaceNotImplemented exception" begin
        @test GroupsCore.InterfaceNotImplemented(
            :Castle,
            "Aaaaarghhhhh....",
        ) isa Exception
        ex = GroupsCore.InterfaceNotImplemented(:Castle, "Aaaaarghhhhh....")
        @test ex isa GroupsCore.InterfaceNotImplemented

        @test sprint(showerror, ex) ==
              "Missing method from Castle interface: `Aaaaarghhhhh....`"
    end

    INI = GroupsCore.InterfaceNotImplemented

    @testset "Group Interface" begin

        G = SomeGroup()

        # Iteration
        @test_throws INI eltype(G)
        @test_throws INI iterate(G)
        @test_throws INI iterate(G, 1)

        # Assumption 1: Groups are finite unless claimed otherwise
        @test Base.IteratorSize(G) == Base.HasLength()
        @test Base.isfinite(G)

        Base.IteratorSize(::Type{SomeGroup}) = Base.HasShape{1}()
        @test Base.isfinite(G)

        Base.IteratorSize(::Type{SomeGroup}) = Base.SizeUnknown()
        @test !Base.isfinite(G)

        Base.IteratorSize(::Type{SomeGroup}) = Base.IsInfinite()
        @test !Base.isfinite(G)

        # return to the default:
        Base.IteratorSize(::Type{SomeGroup}) = Base.HasLength()

        # Assumption 2: Groups have generators:
        @test hasgens(G)

        # Group Interface
        @test_throws INI one(G)
        @test_throws INI order(G)
        @test_throws INI gens(G)

        Base.eltype(::Type{SomeGroup}) = SomeGroupElement
        @test_throws INI rand(G, 2)

        @test_throws INI gens(G, 1)
        @test_throws INI ngens(G)

        @test_throws INI GroupsCore.pseudo_rand(G, 2, 2)
    end

    @testset "GroupElem Interface" begin

        g = SomeGroupElement(Int[])

        @test_throws INI parent(g)
        @test_throws INI GroupsCore.parent_type(g)
        @test_throws INI istrulyequal(g, g)

        @test_throws INI hasorder(g)
        @test_throws INI deepcopy(g)

        @test_throws INI inv(g)
        @test_throws INI g * g
    end

end


##############################################################################
#
# RedistributionFunction.jl
#
# Part of cvortex.jl
# Get vortex redistribution methods.
#
# Copyright 2019 HJA Bird
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
##############################################################################

"""
A vortex redistribution function.

The method by which vortex particles are redistibuted in lagrangian vortex
particles onto a grid to ensure long term stability.

# Examples
```
using CVortex
my_redistribution_kernel = singular_regularisation()
my_redistribution_kernel = winckelmans_regularisation()
my_redistribution_kernel = planetary_regularisation()
my_redistribution_kernel = gaussian_regularisation()
```
"""
struct RedistributionFunction
	func:: Ptr{Cvoid}			# Actually float(*g_3D)(float rho)
	critical_dist :: Float32	# Char[32]
end

#= Functions to to get VortexFunc structures =#

"""
Implements 0th order redistribution method.
The use of this redistribution is not advised.
lambda3_redistribution or m4p_redistribution are advised.
returns RedistributionFunction struct.
"""
function lambda0_redistribution()
	ret = ccall(("cvtx_RedistFunc_lambda0", libcvortex),
        RedistributionFunction, ())
	return ret;
end
"""
Implements 1st order redistribution method.
This method is dissipative.
returns RedistributionFunction struct.
"""
function lambda1_redistribution()
	ret = ccall(("cvtx_RedistFunc_lambda1", libcvortex),
        RedistributionFunction, ())
	return ret;
end
"""
Implements 2nd order redistribution method.
This redistribution is discontinuous, and consequently numerically
less useful than the lambda3_redistribution or m4p_redistribution.
returns RedistributionFunction struct.
"""
function lambda2_redistribution()
	ret = ccall(("cvtx_RedistFunc_lambda2", libcvortex),
        RedistributionFunction, ())
	return ret;
end
"""
Implements 3th order redistribution method.
returns RedistributionFunction struct.
"""
function lambda3_redistribution()
	ret = ccall(("cvtx_RedistFunc_lambda3", libcvortex),
        RedistributionFunction, ())
	return ret;
end
"""
Implements M_4' method of vorticity redistribution.
returns RedistributionFunction struct.
"""
function m4p_redistribution()
	ret = ccall(("cvtx_RedistFunc_m4p", libcvortex),
        RedistributionFunction, ())
	return ret;
end


##############################################################################
###                            Faraday Grid                                ###
###                       Wholesale market model                           ###
###                            Control panel                               ###
###                                v0.1                                    ###
##############################################################################

##############################################################################
###                               options                                  ###
##############################################################################
input_data1 = "/Users/harry/Dropbox/Faraday Grid/Wholesale market model/InputDataUK3.xlsx" #input data file
input_data2 = "/Users/harry/Dropbox/Faraday Grid/Wholesale market model/InputDataUK3.xlsx" #input data file
input_demandfunctions="/Users/harry/Dropbox/Faraday Grid/Wholesale market model/DemandFunctions.xlsx"
input_marketoutcomes="/Users/harry/Dropbox/Faraday Grid/Wholesale market model/MarketOutcome.xlsx"

output_marketoutcomes= "/Users/harry/Dropbox/Faraday Grid/Wholesale market model/MarketOutcome.xlsx" #output data file
output_redispatch= "/Users/harry/Dropbox/Faraday Grid/Wholesale market model/RedispatchOutcome.xlsx"
output_demandfunctions="/Users/harry/Dropbox/Faraday Grid/Wholesale market model/DemandFunctions.xlsx" #output data file

epsilon = -0.2 #demand elasticity
hh=24 #length of horizon in market model in hours, must be between 1 and 8761 but higher numbers may lead to very long solution times
rhh=1 #length of horizon in redispatch model in hours, must be between 1 and 8761 but higher numbers may lead to very long solution times
elastic="false"

##############################################################################
###                               run models                               ###
##############################################################################

include("ModelFunctions.jl")

# Data input
ProcessData(input_data1) #process input data - necessary before all model runs
LoadDemandFunctions(input_demandfunctions) #load previously saved demand functions - only necessary if elastic demand model is used
LoadMarketOutcomes(input_marketoutcomes) #load previously saved market outcomes - only necessary if redispatch model is used without market model

# Data output
WriteMarketData(output_marketoutcomes) #write market outcomes to specified file
WriteRedispatchData(output_redispatch) #write redispatch decisions to specified file
WriteDemandFunctions(output_demandfunctions,epsilon) #generate and write estimated demand functions to specified file

# Plot outputs
GeneratePricePlot() #generate plot of prices over time
GenerateRedispatchCostPlot() #generate plot of redispatch costs over time

# Models
RollingHorizonModel(hh) #run rolling horizon model with specified horizon length
RollingHorizonModelElasticDemand(hh) #run rolling horizon model with elastic demand with specified horizon length
ImperfectCompetitionModel(hh) #run Nash-Cournot imperfectly competitive market model
RedispatchModel(rhh,elastic) #run redispatch model with specified horizon length


# This file is a part of SolidStateDetectors.jl, licensed under the MIT License (MIT).
using SolidStateDetectors
using Test

T = Float32

@testset "Package SolidStateDetectors" begin

    @testset "Simulate example detector: Inverted Coax" begin
        sim = Simulation(SSD_examples[:InvertedCoax])
        simulate!(sim, max_refinements = 1, verbose = true)
        evt = Event(CartesianPoint.([CylindricalPoint{T}(20e-3, deg2rad(10), 40e-3 )]))
        simulate!(evt, sim)
        signalsum = T(0)
        for i in 1:length(evt.signals)
            signalsum += abs(evt.signals[i][end])
        end
        @test isapprox( signalsum, T(2), atol = 5e-3 )
    end
    @testset "Simulate example detector: Coax" begin
        sim = Simulation(SSD_examples[:Coax])
        simulate!(sim, max_refinements = 0, verbose = true)
        evt = Event(CartesianPoint.([CylindricalPoint{T}(20e-3, deg2rad(30), 12e-3 )]))
        simulate!(evt, sim)
        signalsum = T(0)
        for i in 1:length(evt.signals)
            signalsum += abs(evt.signals[i][end])
        end
        @test isapprox( signalsum, T(2), atol = 5e-2 )
    end
    @testset "Simulate example detector: BEGe" begin
        sim = Simulation(SSD_examples[:BEGe])
        simulate!(sim, max_refinements = 1, verbose = true)
        evt = Event(CartesianPoint.([CylindricalPoint{T}(20e-3, deg2rad(10), 20e-3 )]))
        simulate!(evt, sim)
        signalsum = T(0)
        for i in 1:length(evt.signals)
            signalsum += abs(evt.signals[i][end])
        end
        @test isapprox( signalsum, T(2), atol = 1e-2 )
    end
    @testset "Simulate example detector: CGD" begin
        sim = Simulation(SSD_examples[:CGD])
        simulate!(sim, max_refinements = 1, verbose = true)
        evt = Event([CartesianPoint{T}(5e-3, 5e-3, 5e-3)])
        simulate!(evt, sim)
        signalsum = T(0)
        for i in 1:length(evt.signals)
            signalsum += abs(evt.signals[i][end])
        end
        @test isapprox( signalsum, T(2), atol = 1e-2 )
    end
    @testset "Simulate example detector: Spherical" begin
        sim = Simulation(SSD_examples[:Spherical])
        simulate!(sim, max_refinements = 2, verbose = true)
        evt = Event([CartesianPoint{T}(0,0,0)])
        simulate!(evt, sim)
        signalsum = T(0)
        for i in 1:length(evt.signals)
            signalsum += abs(evt.signals[i][end])
        end
        @test isapprox( signalsum, T(2), atol = 1e-2 )
    end
    @testset "Simulate example detector: SigGen PPC" begin
        sim = Simulation(SSD_examples[:SigGen])
        simulate!(sim, max_refinements = 1, verbose = true)
        evt = Event(CartesianPoint.([CylindricalPoint{T}(20e-3, deg2rad(10), 40e-3 )]))
        simulate!(evt, sim)
        signalsum = T(0)
        for i in 1:length(evt.signals)
            signalsum += abs(evt.signals[i][end])
        end
        @test isapprox( signalsum, T(2), atol = 5e-3 )
    end

end


### FlatS2 types
# (spin-2 fields are just special FieldTuples)
const FlatQUMap{P,T,M}     = FieldTuple{QUMap,     NamedTuple{(:Q,:U),NTuple{2,FlatMap{P,T,M}}}, T}
const FlatQUFourier{P,T,M} = FieldTuple{QUFourier, NamedTuple{(:Q,:U),NTuple{2,FlatFourier{P,T,M}}}, Complex{T}}
const FlatEBMap{P,T,M}     = FieldTuple{EBMap,     NamedTuple{(:E,:B),NTuple{2,FlatMap{P,T,M}}}, T}
const FlatEBFourier{P,T,M} = FieldTuple{EBFourier, NamedTuple{(:E,:B),NTuple{2,FlatFourier{P,T,M}}}, Complex{T}}
# some handy Unions
const FlatS2{P,T,M}=Union{FlatEBMap{P,T,M},FlatEBFourier{P,T,M},FlatQUMap{P,T,M},FlatQUFourier{P,T,M}}
const FlatQU{P,T,M}=Union{FlatQUMap{P,T,M},FlatQUFourier{P,T,M}}
const FlatEB{P,T,M}=Union{FlatEBMap{P,T,M},FlatEBFourier{P,T,M}}
const FlatS2Map{P,T,M}=Union{FlatQUMap{P,T,M},FlatEBMap{P,T,M}}
const FlatS2Fourier{P,T,M}=Union{FlatQUFourier{P,T,M},FlatEBFourier{P,T,M}}

spin(::Type{<:FlatS2}) = S2

### convenience constructors
for (F,F0,(X,Y),T) in [
        (:FlatQUMap,     :FlatMap,     (:Qx,:Ux), :T),
        (:FlatQUFourier, :FlatFourier, (:Ql,:Ul), :(Complex{T})),
        (:FlatEBMap,     :FlatMap,     (:Ex,:Bx), :T),
        (:FlatEBFourier, :FlatFourier, (:El,:Bl), :(Complex{T}))
    ]
    doc = """
        # main constructor:
        $F(
            $X::AbstractArray, $Y::AbstractArray; $(F0==:FlatFourier ? "\n        Nside, # required, size of the map in pixels" : "")
            θpix,  # optional, resolution in arcmin (default: 1)
            ∂mode, # optional, fourier∂ or map∂ (default: fourier∂)
        )

        # more low-level:
        $F{P}($X::AbstractArray, $Y::AbstractArray) # specify pixelization P explicilty
        $F{P,T}($X::AbstractArray, $Y::AbstractArray) # additionally, convert elements to type $T
        $F{P,T,M<:AbstractArray{$T}}($X::M, $Y::M) # specify everything explicilty

    Construct a `$F` object. The top form of the constructor is most convenient
    for interactive work, while the others may be more useful for low-level code.
    """
    @eval begin
        @doc $doc $F
        $F($X::AbstractRank2or3Array, $Y::AbstractRank2or3Array; kwargs...) =
            $F(($F0($X; kwargs...), $F0($Y; kwargs...)))
        $F{P}($X::AbstractRank2or3Array, $Y::AbstractRank2or3Array) where {P} =
            $F{P}(($F0{P}($X), $F0{P}($Y)))
        $F{P,T}($X::AbstractRank2or3Array, $Y::AbstractRank2or3Array) where {P,T} =
            $F{P,T}($F0{P,T}($X), $F0{P,T}($Y))
        $F{P,T,M}($X::AbstractRank2or3Array, $Y::AbstractRank2or3Array) where {P,T,M} =
            $F{P,T,M}($F0{P,T,M}($X), $F0{P,T,M}($Y))
    end
end

### properties
function propertynames(f::FlatS2)
    (:fs, propertynames(f.fs)...,
     (Symbol(string(k,(f isa FlatMap ? "x" : "l"))) for (k,f) in pairs(f.fs))...)
end
getproperty(f::FlatQUMap,     ::Val{:Qx}) = getfield(f,:fs).Q.Ix
getproperty(f::FlatQUMap,     ::Val{:Ux}) = getfield(f,:fs).U.Ix
getproperty(f::FlatQUFourier, ::Val{:Ql}) = getfield(f,:fs).Q.Il
getproperty(f::FlatQUFourier, ::Val{:Ul}) = getfield(f,:fs).U.Il
getproperty(f::FlatEBMap,     ::Val{:Ex}) = getfield(f,:fs).E.Ix
getproperty(f::FlatEBMap,     ::Val{:Bx}) = getfield(f,:fs).B.Ix
getproperty(f::FlatEBFourier, ::Val{:El}) = getfield(f,:fs).E.Il
getproperty(f::FlatEBFourier, ::Val{:Bl}) = getfield(f,:fs).B.Il
getproperty(f::FlatS2,        ::Val{:P})  = f
function getindex(f::FlatS2, k::Symbol; full_plane=false)
    maybe_unfold = (full_plane && k in [:El,:Bl,:Ql,:Ul]) ? x->unfold(x,fieldinfo(f).Ny) : identity
    B = @match k begin
        (:P)         => identity
        (:E  || :B)  => @match f begin
            _ :: FlatQUMap     => EBMap
            _ :: FlatQUFourier => EBFourier
            _                  => identity
        end
        (:Q  || :U)  => @match f begin
            _ :: FlatEBMap     => QUMap
            _ :: FlatEBFourier => QUFourier
            _                  => identity
        end
        (:Ex || :Bx) => EBMap
        (:El || :Bl) => EBFourier
        (:Qx || :Ux) => QUMap
        (:Ql || :Ul) => QUFourier
        _ => throw(ArgumentError("Invalid FlatS2 index: $k"))
    end
    maybe_unfold(getproperty(B(f),k))
end

function getindex(D::DiagOp{<:FlatEBFourier}, k::Symbol)
    @unpack El, Bl = diag(D)
    @unpack sin2ϕ, cos2ϕ, P = fieldinfo(diag(D))
    f = @match k begin
        (:QQ)        => FlatFourier{P}(@. Bl*sin2ϕ^2 + El*cos2ϕ^2)
        (:QU || :UQ) => FlatFourier{P}(@. (El - Bl) * sin2ϕ * cos2ϕ)
        (:UU)        => FlatFourier{P}(@. Bl*cos2ϕ^2 + El*sin2ϕ^2)
        _            => getproperty(D.diag, k)
    end
    Diagonal(f)
end

### basis conversion

QUFourier(f::FlatQUMap) = FlatQUFourier(map(Fourier,f.fs))
QUFourier(f::FlatEBMap) = f |> EBFourier |> QUFourier
QUFourier(f::FlatEBFourier{P,T}) where {P,T} = begin
    @unpack sin2ϕ, cos2ϕ = fieldinfo(f)
    Ql = @. - f.El * cos2ϕ + f.Bl * sin2ϕ
    Ul = @. - f.El * sin2ϕ - f.Bl * cos2ϕ
    FlatQUFourier(Q=FlatFourier{P}(Ql), U=FlatFourier{P}(Ul))
end

QUMap(f::FlatQUFourier)  = FlatQUMap(map(Map,f.fs))
QUMap(f::FlatEBMap)      = f |> EBFourier |> QUFourier |> QUMap
QUMap(f::FlatEBFourier)  = f |> QUFourier |> QUMap

EBFourier(f::FlatEBMap) = FlatEBFourier(map(Fourier,f.fs))
EBFourier(f::FlatQUMap) = f |> QUFourier |> EBFourier
EBFourier(f::FlatQUFourier{P,T}) where {P,T} = begin
    @unpack sin2ϕ, cos2ϕ = fieldinfo(f)
    El = @. - f.Ql * cos2ϕ - f.Ul * sin2ϕ
    Bl = @.   f.Ql * sin2ϕ - f.Ul * cos2ϕ
    FlatEBFourier(E=FlatFourier{P}(El), B=FlatFourier{P}(Bl))
end

EBMap(f::FlatEBFourier) = FlatEBMap(map(Map,f.fs))
EBMap(f::FlatQUMap)     = f |> QUFourier |> EBFourier |> EBMap
EBMap(f::FlatQUFourier) = f |> EBFourier |> EBMap

QUFourier(f′::FlatQUFourier, f::FlatQUMap) = (map(Fourier,f′.fs,f.fs); f′)
QUMap(f′::FlatQUMap, f::FlatQUFourier) = (map(Map,f′.fs,f.fs); f′)

Map(f::FlatQUFourier) = QUMap(f)
Map(f::FlatEBFourier) = EBMap(f)
Map(f::FlatS2Map) = f
Fourier(f::FlatQUMap) = QUFourier(f)
Fourier(f::FlatEBMap) = EBFourier(f)
Fourier(f::FlatS2Fourier) = f

### simulation and power spectra
function white_noise(rng::AbstractRNG, ::Type{F2}) where {F2<:FlatS2}
    F = (((::Type{<:FieldTuple{B,NamedTuple{Names,NTuple{2,F}}}}) where {B,Names,F}) -> F)(F2)
    FlatEBMap(E=white_noise(rng,F), B=white_noise(rng,F))
end
function Cℓ_to_Cov(::Type{P}, ::Type{T}, ::Type{S2}, CℓEE::InterpolatedCℓs, CℓBB::InterpolatedCℓs; kwargs...) where {P,T,M}
    Diagonal(FlatEBFourier(E=diag(Cℓ_to_Cov(P,T,S0,CℓEE;kwargs...)), B=diag(Cℓ_to_Cov(P,T,S0,CℓBB;kwargs...))))
end
function Cℓ_to_Cov(::Type{P}, ::Type{T}, ::Type{S2}, (CℓEE, ℓedges, θname)::Tuple, CℓBB::InterpolatedCℓs; units=fieldinfo(P).Ωpix) where {P,T}
    C₀ = Cℓ_to_Cov(P, T, S2, CℓEE, CℓBB, units=units)
    Cbins = Diagonal.(FlatEBFourier.(MidPasses(ℓedges) .* [diag(C₀).E], [zero(diag(C₀).B)]))
    BinRescaledOp(C₀,Cbins,θname)
end

function get_Cℓ(f1::FlatS2, f2::FlatS2=f1; which=(:EE,:BB), kwargs...)
    Cℓ = (;[Symbol(x1*x2) => get_Cℓ(getindex(f1,Symbol(x1)),getindex(f2,Symbol(x2)); kwargs...) for (x1,x2) in split.(string.(ensure1d(which)),"")]...)
    which isa Symbol ? Cℓ[1] : Cℓ
end

function ud_grade(f::FlatS2{P}, args...; kwargs...) where {P}
    f′ = FieldTuple(map(f->ud_grade(f, args...; kwargs...), f.fs))
    B′ = (f′[1] isa FlatMap) ? (f isa FlatQU ? QUMap : EBMap) : (f isa FlatQU ? QUFourier : EBFourier)
    FieldTuple{B′}(f′)
end


export MADDPGManager

"""
    MADDPGManager(; agents::Dict{<:Any, <:Agent}, args...)
Multi-agent Deep Deterministic Policy Gradient(MADDPG) implemented in Julia. Here only works for simultaneous games whose action space is discrete.
See the paper https://arxiv.org/abs/1706.02275 for more details.

# Keyword arguments
- `agents::Dict{<:Any, <:NamedPolicy{<:Agent{<:DDPGPolicy, <:AbstractTrajectory}, <:Any}}`, here each agent collects its own information. While updating the policy, each `critic` will assemble all agents' trajectory to update its own network.
- `batch_size::Int`
- `update_freq::Int`
- `update_step::Int`, count the step.
- `rng::AbstractRNG`.
"""
mutable struct MADDPGManager{P<:DDPGPolicy, T<:AbstractTrajectory, N<:Any} <: AbstractPolicy
    agents::Dict{<:N, <:Agent{<:NamedPolicy{<:P, <:N}, <:T}}
    batch_size::Int
    update_freq::Int
    update_step::Int
    rng::AbstractRNG
end

# for simultaneous game with a discrete action space.
function (π::MADDPGManager)(env::AbstractEnv)
    while current_player(env) == chance_player(env)
        env |> legal_action_space |> rand |> env
    end
    Dict((player, ceil(agent.policy(env))) for (player, agent) in π.agents)
end

function (π::MADDPGManager)(stage::Union{PreEpisodeStage, PostActStage}, env::AbstractEnv)
    # only need to update trajectory.
    for (_, agent) in π.agents
        update!(agent.trajectory, agent.policy, env, stage)
    end
end

function (π::MADDPGManager)(stage::PreActStage, env::AbstractEnv, actions)
    # update each agent's trajectory.
    for (player, agent) in π.agents
        update!(agent.trajectory, agent.policy, env, stage, actions[player])
    end

    # update policy
    update!(π)
end

function (π::MADDPGManager)(stage::PostEpisodeStage, env::AbstractEnv)
    # collect state and a dummy action to each agent's trajectory here.
    for (_, agent) in π.agents
        update!(agent.trajectory, agent.policy, env, stage)
    end

    # update policy
    update!(π)
end

# update policy
function RLBase.update!(π::MADDPGManager)
    π.update_step += 1
    π.update_step % π.update_freq == 0 || return

    for (_, agent) in π.agents
        length(agent.trajectory) > agent.policy.policy.update_after || return
        length(agent.trajectory) > π.batch_size || return
    end

    # get training data
    temp_player = collect(keys(π.agents))[1]
    t = π.agents[temp_player].trajectory
    inds = rand(π.rng, 1:length(t), π.batch_size)
    batches = Dict((player, RLCore.fetch!(BatchSampler{SARTS}(π.batch_size), agent.trajectory, inds))
                for (player, agent) in π.agents)

    # get s, a, s′ for critic
    s = Flux.stack((batches[player][:state] for (player, _) in π.agents), 1)
    a = Flux.stack((batches[player][:action] for (player, _) in π.agents), 1)
    s′ = Flux.stack((batches[player][:next_state] for (player, _) in π.agents), 1)

    # for training behavior_actor
    mu_actions = Flux.stack(
        ((
            batches[player][:state] |> # get personal state information
            x -> send_to_device(device(agent.policy.policy.behavior_actor), x) |>
            agent.policy.policy.behavior_actor |> send_to_host
        ) for (player, agent) in π.agents), 1
    )
    # for training behavior_critic
    new_actions = Flux.stack(
        ((
            batches[player][:next_state] |> # get personal next_state information
            x -> send_to_device(device(agent.policy.policy.target_actor), x) |>
            agent.policy.policy.target_actor |> send_to_host
        ) for (player, agent) in π.agents), 1
    )

    for (player, agent) in π.agents
        p = agent.policy.policy # get DDPGPolicy struct
        A = p.behavior_actor
        C = p.behavior_critic
        Aₜ = p.target_actor
        Cₜ = p.target_critic

        γ = p.γ
        ρ = p.ρ

        _device(x) = send_to_device(device(A), x)

        # Note that here default A, C, Aₜ, Cₜ on the same device.
        s, a, s′ = _device((s, a, s′))
        mu_actions = _device(mu_actions)
        new_actions = _device(new_actions)
        r = _device(batches[player][:reward])
        t = _device(batches[player][:terminal])

        qₜ = Cₜ(vcat(s′, new_actions)) |> vec
        y = r .+ γ .* (1 .- t) .* qₜ

        gs1 = gradient(Flux.params(C)) do
            q = C(vcat(s, a)) |> vec
            loss = mean((y .- q) .^ 2)
            ignore() do
                p.critic_loss = loss
            end
            loss
        end

        update!(C, gs1)

        gs2 = gradient(Flux.params(A)) do
            loss = -mean(C(vcat(s, mu_actions)))
            ignore() do
                p.actor_loss = loss
            end
            loss
        end

        update!(A, gs2)

        # polyak averaging
        for (dest, src) in zip(Flux.params([Aₜ, Cₜ]), Flux.params([A, C]))
            dest .= ρ .* dest .+ (1 - ρ) .* src
        end
    end
end


#Testbed for synthetic spike recovery from example background
using BioBackgroundModels, BioMotifInference, Random, Distributed, Distributions, Serialization
Random.seed!(786)
#CONSTANTS
no_obs=500
obsl=100:200

const folder_path="/bench/PhD/NGS_binaries/BBM/refined_folders"

report_folders=deserialize(folder_path)

bhmm_vec=Vector{BHMM}()
push!(bhmm_vec,report_folders["intergenic"].partition_report.best_model[2])
push!(bhmm_vec,report_folders["periexonic"].partition_report.best_model[2])
push!(bhmm_vec,report_folders["exon"].partition_report.best_model[2])

bhmm_dist=Categorical([.6,.2,.2])

struc_sig_1=[.1 .7 .1 .1
           .1 .1 .1 .7
           .1 .7 .1 .1]
struc_sig_2=[.1 .1 .7 .1
           .1 .1 .7 .1
           .7 .1 .1 .1]
periodicity=8
struc_frac_obs=.75

tata_box=[.05 .05 .05 .85
          .85 .05 .05 .05
          .05 .05 .05 .85
          .85 .05 .05 .05
          .425 .075 .075 .425
          .85 .05 .05 .05
          .425 .075 .075 .425]
caat_box=[.15 .15 .55 .15
          .15 .15 .55 .15
          .05 .85 .05 .05
          .05 .85 .05 .05
          .85 .05 .05 .05
          .85 .05 .05 .05
          .05 .05 .05 .85
          .15 .55 .15 .15
          .15 .15 .15 .55]
motif_frac_obs=.7
motif_recur_range=1:4

@info "Constructing synthetic sample set 1..."
obs1, bg_scores1, hmm_truth1, spike_truth1 = synthetic_sample(no_obs,obsl,bhmm_vec,bhmm_dist,[struc_sig_1,tata_box],[(true,(struc_frac_obs,periodicity)),(false,(motif_frac_obs,motif_recur_range))])

@info "Constructing synthetic sample set 2..."
obs2, bg_scores2, hmm_truth2, spike_truth2 = synthetic_sample(no_obs,obsl,bhmm_vec,bhmm_dist,[struc_sig_2,caat_box],[(true,(struc_frac_obs,periodicity)),(false,(motif_frac_obs,motif_recur_range))])

@info "Assembling combined sample set..."
obs3=hcat(obs1,obs2)
bg_scores3=hcat(bg_scores1, bg_scores2)

@info "Assembling worker pool..."

#DISTRIBUTED CLUSTERs CONSTANTS
remote_machine = "10.0.0.3"
no_local_processes = 2
no_remote_processes = 6

@info "Spawning local cluster workers..."
worker_pool=addprocs(no_local_processes, topology=:master_worker)
remote_pool=addprocs([(remote_machine,no_remote_processes)], tunnel=true, topology=:master_worker)

worker_pool=vcat(worker_pool, remote_pool)

@info "Loading worker libraries everywhere..."
@everywhere using BioMotifInference, Random
@everywhere Random.seed!(myid())

e1 = "/bench/PhD/NGS_binaries/BMI/e1"
e2 = "/bench/PhD/NGS_binaries/BMI/e2"
e3 = "/bench/PhD/NGS_binaries/BMI/e3"

#JOB CONSTANTS
const ensemble_size = 250
const no_sources = 2
const source_min_bases = 3
const source_max_bases = 12
const source_length_range= source_min_bases:source_max_bases
const mixing_prior = .5
const models_to_permute = ensemble_size * 3
funcvec=full_perm_funcvec
push!(funcvec, BioMotifInference.permute_source)
push!(funcvec, BioMotifInference.permute_source)
args=[Vector{Tuple{Symbol,Any}}() for i in 1:length(funcvec)]
args[end-1]=[(:weight_shift_freq,0.),(:length_change_freq,1.),(:length_perm_range,1:1)]
args[end]=[(:weight_shift_freq,.1),(:length_change_freq,0.),(:weight_shift_dist,Uniform(.00001,.01))]

instruct = Permute_Instruct(funcvec, ones(length(funcvec))./length(funcvec),models_to_permute,100, .02; args=args)

@info "Assembling source priors..."
prior_array= Vector{Matrix{Float64}}()
source_priors = BioMotifInference.assemble_source_priors(no_sources, prior_array)

@info "Assembling ensemble 1..."
isfile(e1*"/ens") ? (ens1 = deserialize(e1*"/ens")) :
    (ens1 = IPM_Ensemble(worker_pool, e1, ensemble_size, source_priors, (falses(0,0), mixing_prior), bg_scores1, obs1, source_length_range))

@info "Converging ensemble 1..."
logZ1 = BioMotifInference.converge_ensemble!(ens1, instruct, worker_pool, .001, backup=(true,250), wk_disp=false, tuning_disp=true, ens_disp=false, conv_plot=true, src_disp=true, lh_disp=false, liwi_disp=false)

@info "Assembling ensemble 2..."
isfile(e2*"/ens") ? (ens1 = deserialize(e2*"/ens")) :
    (ens2 = IPM_Ensemble(worker_pool, e2, ensemble_size, source_priors, (falses(0,0), mixing_prior), bg_scores2, obs2, source_length_range))

@info "Converging ensemble 2..."
logZ2 = BioMotifInference.converge_ensemble!(ens2, instruct, worker_pool, .001, backup=(true,250), wk_disp=false, tuning_disp=true, ens_disp=false, conv_plot=true, src_disp=true, lh_disp=false, liwi_disp=false)

@info "Assembling ensemble 3..."
isfile(e3*"/ens") ? (ens1 = deserialize(e3*"/ens")) :
    (ens3 = IPM_Ensemble(worker_pool, e3, ensemble_size, source_priors, (falses(0,0), mixing_prior), bg_scores3, obs3, source_length_range))

@info "Converging ensemble 3..."

logZ3 = BioMotifInference.converge_ensemble!(ens3, instruct, worker_pool, .001, backup=(true,250), wk_disp=false, tuning_disp=true, ens_disp=false, conv_plot=true, src_disp=true, lh_disp=false, liwi_disp=false)

@info "Evidence ratio of joint 1&2 / 3:"
ratio=exp((logZ1+logZ2)-logZ3)
println(ratio)

rmprocs(worker_pool)


using Random
using Statistics
using StatsBase
using Plots
import Dates
import Base.+
import Base.-

include("Point_Struct.jl")
include("fun.jl")
include("constants.jl")
include("integrate.jl")

#################### Integration #####################

calls = Int64(1E3)
it    = 5
dimensions = 6.

Δₐ = 0.4
ss = 0.04 #0.02 rad  / sec

fₛ = ss / Δₐ

Φₛ = π/2         # start point
vₛ = 2*π / 60.   # speed rad/sec
dt = 0.3         # Time interval

Start_Time = 0   #parse(Float64, ARGS[1])
Stop_Time  = 35 #parse(Float64, ARGS[2])
file_name  = "ciao.txt" #ARGS[3]

Step_start = Start_Time / dt
Step_stop  = Stop_Time / dt

time_s = Array{Float64, 1}(undef, 0)
Coo  = Array{Float64, 1}(undef, 0)
std_err = Array{Float64,1}(undef, 0)
el_arr = Array{Float64,1}(undef, 0)
az_arr = Array{Float64,1}(undef, 0)

el = π/(4.5)

# Raster scan
# xu = Point{Float64}(4000.0,  π/(4.5) + θᵦ, Φₛ + θᵦ, undef, undef, undef)
# xl = Point{Float64}(0.,      π/(4.5) - θᵦ,  Φₛ - θᵦ, undef, undef, undef )

xu = Point{Float64}(2000.0,  el + θᵦ, Φₛ + θᵦ, undef, undef, undef)
xl = Point{Float64}(00.,     el - θᵦ, Φₛ - θᵦ, undef, undef, undef)

start = Dates.value(Dates.now())

#hit=Array{Float64}(undef,0)

for t = Step_start:Step_stop
    tempo = Float64(t*dt)
    # xu = Point{Float64}(2000.0,  el + θᵦ, Φₛ + θᵦ, undef, undef, undef)
    # xl = Point{Float64}(00.,     el - θᵦ, Φₛ - θᵦ, undef, undef, undef)
    # Rivedere gli estremi di integrazione.
    # xu_p = Point{Float64}(5035.,   2.8902652413026098,  (Φₛ + vₛ*t*dt) + θᵦ, undef, undef, undef)
    # xl_p = Point{Float64}(0.,     0.12566370614359174,  (Φₛ + vₛ*t*dt) - θᵦ , undef, undef ,undef)

    # Raster scan
    # xu_p = Point{Float64}(4000.,   π/(4.5) + θᵦ,  (Φₛ + (Δₐ/2)*sin(2*π*fₛ*t*dt)) + θᵦ, undef, undef, undef)
    # xl_p = Point{Float64}(0.,      π/(4.5) - θᵦ,  (Φₛ + (Δₐ/2)*sin(2*π*fₛ*t*dt)) - θᵦ , undef, undef ,undef)

    xu_p = Point{Float64}(2000.,   el + θᵦ,  (Φₛ + (Δₐ/2)*sin(2*π*fₛ*t*dt)) + θᵦ, undef, undef, undef)
    xl_p = Point{Float64}(0.,      el - θᵦ,  (Φₛ + (Δₐ/2)*sin(2*π*fₛ*t*dt)) - θᵦ, undef, undef ,undef)

    result, chi_square = Integrate(corr, xl, xu, xl_p, xu_p, calls, it, dimensions, tempo )
    error = (chi_square/(it+1))^0.5
    println(t*dt," ", result, " ± ", error)

    append!(time_s, t*dt)
    append!(Coo, result)
    append!(std_err,error)

    if t*dt % 10 == 0
        global el -= 2*θᵦ
    end
    println(el, " ", Φₛ + (Δₐ/2)*sin(2*π*fₛ*t*dt))
    append!(el_arr, el)
    append!(az_arr, Φₛ + (Δₐ/2)*sin(2*π*fₛ*t*dt))
end

ord = sortperm(time_s)
time_s = time_s[ord]
Coo  = Coo[ord]

#
# # Print to file
# f = open(file_name, "w")
# for j = 1:min(length(Coo), length(time_s))
#     println(f, time_s[j], " ", Coo[j])
# end
# close(f)
x = Array{Float64}(undef, 0)
y = Array{Float64}(undef, 0)

x_w = Array{Float64}(undef, 0)
y_w = Array{Float64}(undef, 0)

time_arr = Array{Float64}(undef, 0)
Coo_arr  = Array{Float64}(undef, 0)

phi_cir = ((0:1000) ./ 1000) .* π
punti_cir = [Point{Float64}(2500, π/4., i, undef, undef, undef)() for i in phi_cir]

x_c = [i.x/2000 for i in punti_cir ]
y_c = [i.y/2000 for i in punti_cir ]

@gif for i = 1:length(az_arr)

    pun = Point{Float64}(2500., π/4., az_arr[i], undef, undef, undef )()
    append!(x, pun.x/2000)
    append!(y, pun.y/2000)

    append!(x_w, -1-(45/2000)*i*cos((π/2)+(π/7)) )
    append!(y_w, 0+(45/2000)*i*sin((π/2)+(π/7)) )

    p1=plot(x, y, xlims=(-1000/2000,1000/2000), ylims=(-1,1), arrow=arrow(), legend=false, ylabel="Y", xlabel="X")
    p1=plot!(x_c,y_c, xlims=(-1000/2000,1000/2000), ylims=(-1,1))

    for j = -5:5
        for k = 0:8
            p1=plot!(x_w .+ k*(500/2000) , y_w .- j*(500/2000), arrow=arrow(), legend = false)
        end
    end
    append!(time_arr, time_s[i])
    append!(Coo_arr, Coo[i]/findmax(Coo)[1])

    p2 = plot(time_arr, Coo_arr, xlims=(0,35), ylims=(0,1.2), arrow=arrow(),  ylabel="Autocorrelation C_00^0t [normalized unit]", xlabel="time [s]")

    plot(p1,p2, layout =(1,2))

end every 1

#plot(time_s, Coo/findmax(Coo)[1], seriestype=:scatter, ylims=(0,1.1), title="Test", xlabel = "Time [seconds]", ylabel="Autocorrelation Level [Normalized Units]", m=(1, :circle, 1), bg=RGB(1,1,1))#, yerr=(std_err ./ findmax(Coo)[1]))
#plot(hit)
#plot!(hit, seriestype=:scatter)
#
# using DSP
# period = periodogram(Coo, fs=1)
# Freq   = freq(period)
# A      = power(period)

#stop = Dates.value(Dates.now())
#println("Performance: ", ( stop - start ) / (1000 * (Step_stop - Step_start)), " sec/step" )


module Ls

import LsoBase
import Obj
import Obj.Objective

"""
    LineSearch(name, strategy)

    Creates an LineSearch object for the given strategy method.
    The strategy method should have the following signature:

    (w, s, obj, b, fw, gw) -> (α, iter)

    Where:
       w: The current point
       s: The current step direction
     obj: The objective function
       b: Index batch to evaluate objective function on. May be ignored by
          the implementation.
      fw: f(w) (present for performance reasons)
      gw: g(w) (present for performance reasons)
       α: The resulting step size
    iter: The number of LS iterations required
"""
type LineSearch
        name::ASCIIString
    strategy::Function
end

"""
    ls(linesearch, w, s [, obj, b, fw, gw])

    Performs a Line Search with the given strategy for current point w and step
    direction s. When the LS strategy depends on an objective function value, that
    function should be specified using the obj parameter.

    Optionally, b can specify a batch of indices to evaluate the objective function
    on. For performance reasons, f(w) and g(w) can be provided as fw and gw.

    This helper function should ease the access to objects of type Ls.LineSearch.
"""
function ls(linesearch::LineSearch, w::Array{Float64,1}, s::Array{Float64,1},
            obj::Union{Objective,Void}=nothing, b::Array{Int,1}=Int[],
            fw::Float64=NaN, gw::Array{Float64,1}=Float64[])

    # sanitize arguments
    if obj != nothing
      if isnan(fw)
          if length(b) == 0
              fw = Obj.f(obj, w)
          else
              fw = Obj.f(obj, w, b)
          end
      end
      if length(gw) == 0
          if length(b) == 0
              gw = Obj.g(obj, w)
          else
              gw = Obj.g(obj, w, b)
          end
      end
    end

    # run ls
    return linesearch.strategy(w, s, obj, b, fw, gw)

end

# include ./opt/*
LsoBase.includedir(dirname(@__FILE__)*"/ls")

end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    abstract type DifferentiationAlgorithm

*Experimental feature, not part of stable public API.*

Abstract type for integrated autocorrelation length estimation algorithms.
"""
abstract type DifferentiationAlgorithm end
export DifferentiationAlgorithm

MaybeDiffAlg = Union{DifferentiationAlgorithm,Missing}

function vjp_algorithm end
vjp_algorithm(::Any) = missing

function promote_vjp_algorithm end
promote_vjp_algorithm() = missing
promote_vjp_algorithm(a::MaybeDiffAlg) = a
promote_vjp_algorithm(a::DifferentiationAlgorithm, b::DifferentiationAlgorithm) = a
promote_vjp_algorithm(a::DifferentiationAlgorithm, b::Missing) = a
promote_vjp_algorithm(a::Missing, b::DifferentiationAlgorithm) = b
promote_vjp_algorithm(a::Missing, b::Missing) = b
function promote_vjp_algorithm(a::MaybeDiffAlg, b::MaybeDiffAlg, cs::Vararg{MaybeDiffAlg,N}) where {N}
    promote_vjp_algorithm(promote_vjp_algorithm(a, b), cs...)
end

generic_vjp_algorithm(obj::T) where {T} = promote_vjp_algorithm(map(vjp_algorithm, object_contents(obj))...)
vjp_algorithm(f::F) where {F<:Function} = generic_vjp_algorithm(f)
vjp_algorithm(d::AbstractDensity) = generic_vjp_algorithm(d)

# ToDo: Switch default from ForwardDiff to Zygote:
default_vjp_algorithm() = ForwardDiffAD()
#default_vjp_algorithm() = ZygoteAD()

@inline function some_vjp_algorithm(x)
    alg = vjp_algorithm(x)
    ismissing(alg) ? default_vjp_algorithm() : alg
end

function jvp_algorithm end
jvp_algorithm(::Any) = missing

function promote_jvp_algorithm end
promote_jvp_algorithm() = missing
promote_jvp_algorithm(a::MaybeDiffAlg) = a
promote_jvp_algorithm(a::DifferentiationAlgorithm, b::DifferentiationAlgorithm) = a
promote_jvp_algorithm(a::DifferentiationAlgorithm, b::Missing) = a
promote_jvp_algorithm(a::Missing, b::DifferentiationAlgorithm) = b
promote_jvp_algorithm(a::Missing, b::Missing) = b
function promote_jvp_algorithm(a::MaybeDiffAlg, b::MaybeDiffAlg, cs::Vararg{MaybeDiffAlg,N}) where {N}
    promote_jvp_algorithm(promote_jvp_algorithm(a, b), cs...)
end

generic_jvp_algorithm(obj::T) where {T} = promote_jvp_algorithm(map(jvp_algorithm, object_contents(obj))...)
jvp_algorithm(f::F) where {F<:Function} = generic_jvp_algorithm(f)

default_jvp_algorithm() = ForwardDiffAD()

function some_jvp_algorithm(x)
    alg = jvp_algorithm(x)
    ismissing(alg) ? default_jvp_algorithm() : alg
end

"""
    valgradof(f::Function, algorithm::DifferentiationAlgorithm  = vjp_algorithm(f))

*Experimental feature, not part of stable public API.*

Generates a function that calculates both value and gradient of `f` at given
points. The differentiation algorithm used depends on `f`.

The function `f` must support `ValueShapes.varshape(f)` and
`ValueShapes.unshaped(f)`.

Returns function `valgrad_f` with

```julia
f_at_x, grad_of_f_at_x = valgrad_f(x)

grad_of_f_at_x = zero(x)
f_at_x = valgrad_f(!, grad_of_f_at_x, x)
```
"""
function valgradof end
export valgradof

abstract type GradientFunction <: Function end


function copy_self(obj::SubHyperplane)
    return jcall(obj, "copySelf", SubHyperplane, ())
end

function get_hyperplane(obj::SubHyperplane)
    return jcall(obj, "getHyperplane", Hyperplane, ())
end

function get_size(obj::SubHyperplane)
    return jcall(obj, "getSize", jdouble, ())
end

function is_empty(obj::SubHyperplane)
    return jcall(obj, "isEmpty", jboolean, ())
end

function reunite(obj::SubHyperplane, arg0::SubHyperplane)
    return jcall(obj, "reunite", SubHyperplane, (SubHyperplane,), arg0)
end

function split(obj::SubHyperplane, arg0::Hyperplane)
    return jcall(obj, "split", SubHyperplane_SplitSubHyperplane, (Hyperplane,), arg0)
end


type CtrDrbg  <: AbstractRNG
    data::Ptr{Void}
    entropy::Nullable{Entropy}

    function CtrDrbg()
        ctx = new()
        ctx.data = Libc.malloc(1000)  # 344
        ccall((:mbedtls_ctr_drbg_init, MBED_CRYPTO), Void, (Ptr{Void},), ctx.data)
        finalizer(ctx, ctx->begin
            ccall((:mbedtls_ctr_drbg_free, MBED_CRYPTO), Void, (Ptr{Void},), ctx.data)
            Libc.free(ctx.data)
        end
        )
        ctx
    end

    CrtDrbg(data) = new(data)
end

function f_rng(c_ctx, c_buf, sz)
    jl_ctx = unsafe_pointer_to_objref(c_ctx)
    jl_buf = unsafe_wrap(Array, c_buf, sz, false)
    rand!(jl_ctx, jl_buf)
    MBED_SUCCESS
end

function seed!(rng::CtrDrbg, entropy, pdata)
    rng.entropy = Nullable(entropy)
    entropy_func = cglobal((:mbedtls_entropy_func, MBED_CRYPTO))
    @err_check ccall((:mbedtls_ctr_drbg_seed, MBED_CRYPTO), Cint,
        (Ptr{Void}, Ptr{Void}, Ptr{Void}, Ptr{Void}, Csize_t),
        rng.data, entropy_func, entropy.data, pdata, sizeof(pdata))
    rng
end

seed!(rng::CtrDrbg, entropy) = seed!(rng, entropy, UInt8[])

function Base.rand!(rng::CtrDrbg, buf::Array)
    @err_check ccall((:mbedtls_ctr_drbg_random, MBED_CRYPTO), Cint,
            (Ptr{Void}, Ptr{Void}, Csize_t),
        rng.data, buf, sizeof(buf))
    buf
end

function Base.rand(rng::CtrDrbg, size::Integer)
    buf = Vector{UInt8}(size)
    rand!(rng, buf)
end

function __ctr_drbg__init__()
    const global c_rng=cfunction(f_rng, Cint, (Ptr{Void}, Ptr{UInt8}, Csize_t))
end


struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
    groupindex::I
    ndrange::NDRange
    iterspace::Iterspace

    # CPU variant
    function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB}
        if ndrange !== nothing
            ndrange = CartesianIndices(ndrange)
        end
        return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace)
    end

    # GPU variante: index is given implicit
    function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB}
        if ndrange !== nothing
            ndrange = CartesianIndices(ndrange)
        end
        return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace)
    end
end

@inline __iterspace(cm::CompilerMetadata)  = cm.iterspace
@inline __groupindex(cm::CompilerMetadata) = cm.groupindex
@inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm)))
@inline __dynamic_checkbounds(::CompilerMetadata{NDRange, CB}) where {NDRange, CB} = CB <: DynamicCheck
@inline __ndrange(::CompilerMetadata{NDRange}) where {NDRange<:StaticSize}  = CartesianIndices(get(NDRange))
@inline __ndrange(cm::CompilerMetadata{NDRange}) where {NDRange<:DynamicSize} = cm.ndrange
@inline groupsize(ctx) = __groupsize(ctx)
@inline __workitems_iterspace(ctx) = workitems(__iterspace(ctx))

include("compiler/contract.jl")


using LinearAlgebra

export flattened_tup_chain, flattened_named_tuple, flattened_tuple
export FlattenType, FlattenArr, RetainArr

abstract type FlattenType end

"""
    FlattenArr

Flatten arrays in `flattened_tup_chain`
and `flattened_named_tuple`.
"""
struct FlattenArr <: FlattenType end

"""
    RetainArr

Do _not_ flatten arrays in `flattened_tup_chain`
and `flattened_named_tuple`.
"""
struct RetainArr <: FlattenType end

# The Vars instance has many empty entries.
# Keeping all of the keys results in many
# duplicated values. So, it's best we
# "prune" the tree by removing the keys:
flattened_tup_chain(
    ::Type{NamedTuple{(), Tuple{}}},
    ::FlattenType = FlattenArr();
    prefix = (Symbol(),),
) = ()

flattened_tup_chain(
    ::Type{T},
    ::FlattenType;
    prefix = (Symbol(),),
) where {T <: Real} = (prefix,)

flattened_tup_chain(
    ::Type{T},
    ::RetainArr;
    prefix = (Symbol(),),
) where {T <: SArray} = (prefix,)
flattened_tup_chain(
    ::Type{T},
    ::FlattenArr;
    prefix = (Symbol(),),
) where {T <: SArray} = ntuple(i -> (prefix..., i), length(T))

flattened_tup_chain(
    ::Type{T},
    ::RetainArr;
    prefix = (Symbol(),),
) where {T <: SHermitianCompact} = (prefix,)
flattened_tup_chain(
    ::Type{T},
    ::FlattenType;
    prefix = (Symbol(),),
) where {T <: SHermitianCompact} =
    ntuple(i -> (prefix..., i), length(StaticArrays.lowertriangletype(T)))

flattened_tup_chain(
    ::Type{T},
    ::RetainArr;
    prefix = (Symbol(),),
) where {N, TA, T <: Diagonal{N, TA}} = (prefix,)
flattened_tup_chain(
    ::Type{T},
    ::FlattenArr;
    prefix = (Symbol(),),
) where {N, TA, T <: Diagonal{N, TA}} = ntuple(i -> (prefix..., i), length(TA))

flattened_tup_chain(::Type{T}, ::FlattenType; prefix = (Symbol(),)) where {T} =
    (prefix,)

"""
    flattened_tup_chain(::Type{T}) where {T <: Union{NamedTuple,NTuple}}

An array of tuples, containing symbols
and integers for every combination of
each field in the `Vars` array.
"""
function flattened_tup_chain(
    ::Type{T},
    ft::FlattenType = FlattenArr();
    prefix = (Symbol(),),
) where {T <: Union{NamedTuple, NTuple}}
    map(1:fieldcount(T)) do i
        Ti = fieldtype(T, i)
        name = fieldname(T, i)
        sname = name isa Int ? name : Symbol(name)
        flattened_tup_chain(
            Ti,
            ft;
            prefix = prefix == (Symbol(),) ? (sname,) : (prefix..., sname),
        )
    end |>
    Iterators.flatten |>
    collect
end
flattened_tup_chain(
    ::AbstractVars{S},
    ft::FlattenType = FlattenArr(),
) where {S} = flattened_tup_chain(S, ft)

"""
    flattened_named_tuple

A flattened NamedTuple, given a
`Vars` or nested `NamedTuple` instance.

# Example:

```julia
using Test
using ClimateMachine.VariableTemplates
nt = (x = 1, a = (y = 2, z = 3, b = ((a = 1,), (a = 2,), (a = 3,))));
fnt = flattened_named_tuple(nt);
@test keys(fnt) == (:x, :a_y, :a_z, :a_b_1_a, :a_b_2_a, :a_b_3_a)
@test length(fnt) == 6
@test fnt.x == 1
@test fnt.a_y == 2
@test fnt.a_z == 3
@test fnt.a_b_1_a == 1
@test fnt.a_b_2_a == 2
@test fnt.a_b_3_a == 3
```
"""
function flattened_named_tuple end

function flattened_named_tuple(v::AbstractVars, ft::FlattenType = FlattenArr())
    ftc = flattened_tup_chain(v, ft)
    keys_ = Symbol.(join.(ftc, :_))
    vals = map(x -> getproperty(v, wrap_val.(x)), ftc)
    length(keys_) == length(vals) || error("key-value mismatch")
    return (; zip(keys_, vals)...)
end

function flattened_named_tuple(nt::NamedTuple, ft::FlattenType = FlattenArr())
    ftc = flattened_tup_chain(typeof(nt), ft)
    keys_ = Symbol.(join.(ftc, :_))
    vals = flattened_tuple(ft, nt)
    length(keys_) == length(vals) || error("key-value mismatch")
    return (; zip(keys_, vals)...)
end

flattened_tuple(::FlattenArr, a::AbstractArray) = tuple(a...)
flattened_tuple(::RetainArr, a::AbstractArray) = tuple(a)

flattened_tuple(::FlattenArr, a::Diagonal) = tuple(a.diag...)
flattened_tuple(::RetainArr, a::Diagonal) = tuple(a.diag)

flattened_tuple(::FlattenArr, a::SHermitianCompact) = tuple(a.lowertriangle...)
flattened_tuple(::RetainArr, a::SHermitianCompact) = tuple(a.lowertriangle)

# when we splat an empty tuple `b` into `flattened_tuple(ft, b...)`
flattened_tuple(::FlattenType) = ()

# for structs
flattened_tuple(::FlattenType, a) = (a,)

# Divide and conquer:
flattened_tuple(ft::FlattenType, a, b...) =
    tuple(flattened_tuple(ft, a)..., flattened_tuple(ft, b...)...)

flattened_tuple(ft::FlattenType, a::Tuple) = flattened_tuple(ft, a...)

flattened_tuple(ft::FlattenType, a::NamedTuple) = flattened_tuple(ft, Tuple(a))


const pi = 3.14;
println(pi);


using InformationDistances
using CodecXz: XzCompressor
using TranscodingStreams: Noop

using Test

struct MockCompressor <: InformationDistances.AbstractCompressor end

function InformationDistances.compressed_length(::MockCompressor, data::InformationDistances.ByteData)

    return length(data)
end

s0 = ""
s1 = "x"
s2 = "xy"

@testset "InformationDistances.jl" begin

    @testset "AbstractCompressor interface" begin
        @test compressed_lengths(MockCompressor(), [s0, s1, s2]) ==
            map(s -> compressed_length(MockCompressor(), s), [s0, s1, s2])

        @test compressed_length(MockCompressor(), s2) ==
            compressed_length(MockCompressor(), codeunits(s2))
    end

    @testset "NormalizedCompressionDistance" begin

        @test NormalizedCompressionDistance() isa NormalizedCompressionDistance{CodecCompressor{XzCompressor}}

        @test NormalizedCompressionDistance(MockCompressor()) isa
                NormalizedCompressionDistance{MockCompressor}

        len_s1 = compressed_length(MockCompressor(), s1)
        len_s2 = compressed_length(MockCompressor(), s2)
        len_s1s2 = compressed_length(MockCompressor(), s1 * s2)

        @test NormalizedCompressionDistance(MockCompressor())(s1, s2) ==
            (len_s1s2 - min(len_s1, len_s2)) / max(len_s1, len_s2)

        @test NormalizedCompressionDistance()("", "") == 0.0
        @test NormalizedCompressionDistance()(s2, s2) == 0.0
        # NCD is not symmetric in general, but it should be for the empty string
        @test NormalizedCompressionDistance()("", s2) ==
            NormalizedCompressionDistance()(s2, "")
        @test 0.0 <= NormalizedCompressionDistance()(s1, s2) <= 1.0
    end

    @testset "LibDeflateCompressor" begin

        @test LibDeflateCompressor().compresslevel == 12
        @test LibDeflateCompressor(;compresslevel=5).compresslevel == 5

        compressor = LibDeflateCompressor()

        data_CodeUnits = codeunits.([s0, s1, s2])
        data_Vector = Vector(data_CodeUnits)

        @test compressed_lengths(compressor, data_CodeUnits) ==
              compressed_lengths(compressor, data_Vector) ==
              map(s -> compressed_length(compressor, s), data_CodeUnits) ==
              map(s -> compressed_length(compressor, s), data_Vector)

    end

    @testset "CodecCompressor" begin

        @test CodecCompressor{Noop}() isa CodecCompressor{Noop}
        @test collect(CodecCompressor{Noop}(arg1="1", arg2=2).kwargs) ==
            [:arg1 => "1", :arg2 => 2]

        compressor = CodecCompressor{Noop}()

        data_CodeUnits = codeunits.([s0, s1, s2])
        data_Vector = Vector(data_CodeUnits)

        @test compressed_lengths(compressor, data_CodeUnits) ==
              compressed_lengths(compressor, data_Vector) ==
              map(s -> compressed_length(compressor, s), data_CodeUnits) ==
              map(s -> compressed_length(compressor, s), data_Vector)

    end

end


using Knet:MB

import Knet.nll
import Knet.accuracy

"""
`nll(model::KnetModule, data::Knet.MB; o...)`
wrapper over `Knet.nll`.
"""
nll(model::KnetModule, data::MB; o...) =
    nll(model, data, (m, x)->m(actx(), x); o...)

"""
`accuracy(model::KnetModule, data::Knet.MB; o...)`
wrapper over `Knet.accuracy`.
"""
accuracy(model::KnetModule, data::MB; o...) =
    accuracy(model, data, (m, x)->m(actx(), x); o...)


module Tree

#using StatsBase
include("../utils.jl")

export DecisionTree, MV, RandomForest

include("utils.jl")
include("DT.jl")
include("RF.jl")
include("MV.jl")

end


using Base.Threads: @threads

function (nw::Network)(du, u::T, p, t) where {T}
    # for now just one layer
    layer = nw.nl
    dupt = (du, u, p, t)

    # get cache array for accumulator of appropriate type
    _acc = accumulator_cache(layer, T)
    fill!(_acc, zero(eltype(T)))

    # go through all colobatches separatly to avoid writing conflicts to accumulator
    unroll_colorbatches!(nw, layer, layer.colorbatches, dupt, _acc)

    # can be run parallel
    process_batches!(nw, layer, nw.vertexbatches, dupt, _acc)
    return nothing
end

@unroll function unroll_colorbatches!(nw, layer, colorbatches, dupt, _acc)
    @unroll for cbatch in colorbatches
        process_batches!(nw, layer, cbatch.edgebatches, dupt, _acc)
    end
end

@unroll function process_batches!(nw, layer, batches, dupt, acc)
    @unroll for batch in batches
        process_batch!(nw, layer, batch, dupt, acc)
    end
end

function process_batch!(nw::Network{:seq}, layer, batch, dupt, acc)
    for idx in 1:length(batch)
        element_kernel!(layer, batch, dupt, acc, idx)
    end
end

function process_batch!(nw::Network{:threaded}, layer, batch, dupt, acc)
    @threads for idx in 1:length(batch)
        element_kernel!(layer, batch, dupt, acc, idx)
    end
end

@inline function element_kernel!(layer, batch::VertexBatch{F}, dupt, acc, i) where {F}
    @inbounds begin
    du, u, p, t = dupt
    (v_r, acc_r) = vertex_ranges(layer, batch, i)
    pv_r = parameter_range(batch, i)

    vdu = view(du, v_r)
    vu  = view(u, v_r)
    vacc = view(acc, acc_r)
    pv = p===nothing ? nothing : view(p, pv_r)

    batch.fun.f(vdu, vu, vacc, pv, t)
    end
end

@inline function element_kernel!(layer, batch::EdgeBatch{F}, dupt::T, acc, i) where {F<:StaticEdge, T}
    @inbounds begin
    # a cache of size dim per thread
    _, u, p, t = dupt

    # collect all the ranges to index into the data arrays
    (src_r, dst_r, src_acc_r, dst_acc_r, _) = src_dst_ranges(layer, batch, i)
    pe_r = parameter_range(batch, i)

    # create the views into the data & parameters
    vs = view(u, src_r)
    vd = view(u, dst_r)
    pe = p===nothing ? nothing : view(p, pe_r)

    # apply the edge function
    e = batch.fun.f(vs, vd, pe, t)

    apply_accumulation!(coupling(F), layer.accumulator, acc, src_acc_r, dst_acc_r, e)
    end
end

@inline function element_kernel!(layer, batch::EdgeBatch{F}, dupt, acc, i) where {F<:ODEEdge}
    @inbounds begin
    du, u, p, t = dupt
    # for i in 1:length(batch)
    # collect all the ranges to index into the data arrays
    (src_r, dst_r, src_acc_r, dst_acc_r, e_r) = src_dst_ranges(layer, batch, i)
    pe_r = parameter_range(batch, i)

    # create the views into the data & parameters
    de = view(du, e_r)
    e  = view(u, e_r)
    vs = view(u, src_r)
    vd = view(u, dst_r)
    pe = p===nothing ? nothing : view(p, pe_r)

    # apply the edge function
    batch.fun.f(de, e, vs, vd, pe, t)

    apply_accumulation!(coupling(F), layer.accumulator, acc, src_acc_r, dst_acc_r, e)
    end
end

Base.@propagate_inbounds function apply_accumulation!(::AntiSymmetric, f, acc, src_acc_r, dst_acc_r, edge)
    _acc_dst = view(acc, dst_acc_r)
    _acc_src = view(acc, src_acc_r)

    for i in 1:length(dst_acc_r)
        _acc_dst[i] = f(_acc_dst[i],  edge[i])
        _acc_src[i] = f(_acc_src[i], -edge[i])
    end
end

Base.@propagate_inbounds function apply_accumulation!(::Symmetric, f, _acc, src_acc_r, dst_acc_r, edge)
    _acc_dst = view(_acc, dst_acc_r)
    _acc_src = view(_acc, src_acc_r)

    for i in 1:length(dst_acc_r)
        _acc_dst[i] = f(_acc_dst[i], edge[i])
        _acc_src[i] = f(_acc_src[i], edge[i])
    end
end

Base.@propagate_inbounds function apply_accumulation!(::Directed, f, _acc, _, dst_acc_r, edge)
    _acc_dst = view(_acc, dst_acc_r)

    for i in 1:length(dst_acc_r)
        _acc_dst[i] = f(_acc_dst[i], edge[i])
    end
end

Base.@propagate_inbounds function apply_accumulation!(::Fiducial, f, _acc, src_acc_r, dst_acc_r, edge)
    _acc_dst = view(_acc, dst_acc_r)
    _acc_src = view(_acc, src_acc_r)
    offset = Int(length(edge) / 2)

    for i in 1:length(dst_acc_r)
        _acc_dst[i] = f(_acc_dst[i], edge[i])
        _acc_src[i] = f(_acc_src[i], edge[offset + i])
    end
end


struct MockDofBasis{P} <: AbstractVector{PointValue}
  nodes::Vector{P}
end

function return_cache(b::MockDofBasis,field)
  return_cache(field,b.x)
end

@inline function evaluate!(cache,b::MockDofBasis,field)
  evaluate!(cache,field,b.x)
end

@inline Base.size(a::MockDofBasis) = (length(a.nodes),)
@inline Base.axes(a::MockDofBasis) = (axes(a.nodes,1),)
# @santiagobadia : Not sure we want to create the moment dofs
@inline Base.getindex(a::MockDofBasis,i::Integer) = PointValue(a.x[i])
@inline Base.IndexStyle(::MockDofBasis) = IndexLinear()


import PDESolver: evalJacobian, evalJacobianStrong

"""
  Euler implementation of `evalJacobian`.
"""
function evalJacobian(mesh::AbstractMesh, sbp::AbstractOperator, eqn::EulerData,
                      opts::Dict, assembler::AssembleElementData, t=0.0;
                      start_comm=false)

  #println("\nEntered evalJacobian")
  time = eqn.params.time
  eqn.params.t = t  # record t to params
  myrank = mesh.myrank

  @assert opts["parallel_type"] == 2
#  println("entered evalResidual")
  time.t_send += @elapsed if start_comm
    setParallelData(eqn.shared_data, opts["parallel_data"])
    startSolutionExchange(mesh, sbp, eqn, opts)
  end

  time.t_dataprep_diff += @elapsed dataPrep_diff(mesh, sbp, eqn, opts)
  #println("dataPrep @time printed above")

  time.t_volume_diff += @elapsed if opts["addVolumeIntegrals"]
    evalVolumeIntegrals_diff(mesh, sbp, eqn, opts, assembler)
    #println("volume integral diff @time printed above")
  end

  if opts["use_GLS"]
    error("use_GLS not supported by evalJacobian")
  end

  time.t_bndry_diff += @elapsed if opts["addBoundaryIntegrals"]
    evalBoundaryIntegrals_diff(mesh, sbp, eqn, opts, assembler)
    #println("boundary integral diff @time printed above")
  end

  time.t_stab_diff += @elapsed if opts["addStabilization"]
    addStabilization_diff(mesh, sbp, eqn, opts, assembler)
#    println("stabilizing @time printed above")
  end

  time.t_face_diff += @elapsed if mesh.isDG && opts["addFaceIntegrals"]
    evalFaceIntegrals_diff(mesh, sbp, eqn, opts, assembler)
    #println("face integral diff @time printed above")
  end

  time.t_shock_diff += @elapsed if opts["addShockCapturing"]
    evalShockCapturing_diff(mesh, sbp, eqn, opts, assembler)
  end

  time.t_sharedface_diff += @elapsed if mesh.commsize > 1
    evalSharedFaceIntegrals_diff(mesh, sbp, eqn, opts, assembler)
#    println("evalSharedFaceIntegrals @time printed above")
  end

  time.t_source_diff += @elapsed evalSourceTerm_diff(mesh, sbp, eqn, opts, assembler)
  #println("source integral diff @time printed above")

  # apply inverse mass matrix to eqn.res, necessary for CN
  #TODO: this will have to be done at the element level
  if opts["use_Minv"]
#    error("use_Minv must be applied at the element level")
#    applyMassMatrixInverse3D(mesh, sbp, eqn, opts, eqn.res)
  end

  return nothing
end

function evalJacobianStrong(mesh::AbstractMesh, sbp::AbstractOperator,
                      eqn::EulerData,
                      opts::Dict, assembler::AssembleElementData, t=0.0;
                      start_comm=false)
# currently this function neglects the SAT terms (including boundary conditions)

  time = eqn.params.time
  eqn.params.t = t  # record t to params
  myrank = mesh.myrank

  time.t_volume_diff += @elapsed if opts["addVolumeIntegrals"]
    calcVolumeIntegralsStrong_nopre_diff(mesh, sbp, eqn, opts, assembler)
#    println("volume integral @time printed above")
  end

  return nothing
end

"""
  Differentiated version of [`dataPrep`](@ref).  The Jacobian computation
  does use any of the arrays of precomputed values (options keys
  `precompute_*`).  Those keys can be specified by the user and will affect
  [`evalResidual`](@ref), but not [`evalJacobian`](@ref).

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
"""
function dataPrep_diff(mesh::AbstractMesh{Tmsh},
                      sbp::AbstractOperator,
                      eqn::AbstractEulerData{Tsol, Tres}, opts) where {Tmsh, Tsol, Tres}
# gather up all the data needed to do vectorized operatinos on the mesh
# calculates all mesh wide quantities in eqn
# this is almost the exact list of everything we *shouldn't* be storing, but
# rather recalculating on the fly

#println("Entered dataPrep()")

#  println("typeof(eqn) = ", typeof(eqn))
#  println("typeof(eqn.params) = ", typeof(eqn.params))

  # apply filtering to input
  if eqn.params.use_filter
    applyFilter(mesh, sbp, eqn, eqn.q, opts)
  end

  # zero out res
#  fill!(eqn.res, 0.0)
#  fill!(eqn.res_edge, 0.0)

  getAuxVars(mesh, eqn)
#  println("  getAuxVars @time printed above")

  if opts["check_density"]
    checkDensity(eqn, mesh)
#    println("  checkDensity @time printed above")
  end

  if opts["check_pressure"]
#    throw(ErrorException("I'm done"))
    checkPressure(eqn, mesh)
#    println("  checkPressure @time printed above")
  end

  # calculate fluxes

  if opts["use_staggered_grid"]
    error("staggered grid not supported by calcJacobian()")
  end

  # the jacobian computation doesn't precompute anything

  if eqn.params.use_edgestab
    error("Use_edgestab not supported by calcJacobian()")
    stabscale(mesh, sbp, eqn)
  end
#  println("  stabscale @time printed above")

  return nothing
end # end function dataPrep

"""
  Differentiated version of [`evalVolumeIntegrals`](@ref).  Throws an error
  for unsupported schemes.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler
"""
function evalVolumeIntegrals_diff(mesh::AbstractMesh{Tmsh},
                       sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim},
                       opts, assembler::AssembleElementData
                       ) where {Tmsh,  Tsol, Tres, Tdim}

  integral_type = opts["volume_integral_type"]

  if opts["Q_transpose"] == false
    error("Q_transpose == false not supported by evalJacobian")
  end  # end if Q_transpose

  if integral_type == 1  # regular volume integrals
    calcVolumeIntegrals_nopre_diff(mesh, sbp, eqn, opts, assembler)
  elseif integral_type == 2  # entropy stable formulation
    calcVolumeIntegralsSplitForm_diff(mesh, sbp, eqn, opts,
                                      eqn.volume_flux_func_diff, assembler)
  else
    throw(ErrorException("Unsupported volume integral type = $integral_type"))
  end

end  # end evalVolumeIntegrals

"""
  Differentiated version of [`evalBoundaryIntegrals`](@ref).

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler.
"""
function evalBoundaryIntegrals_diff(mesh::AbstractMesh{Tmsh},
                               sbp::AbstractOperator,
                               eqn::EulerData{Tsol, Tres, Tdim}, opts,
                               assembler::AssembleElementData) where {Tmsh, Tsol, Tres, Tdim}

  if mesh.isDG
    # when precompute_boundary_flux == false, this fuunction does the
    # integration too, updating res
    getBCFluxes_diff(mesh, sbp, eqn, opts, assembler)
  else
    error("CG meshes not supported by evalJacobian()")
  end

  return nothing

end  # end evalBoundaryIntegrals

"""
  Differentiated version of [`addStabilization`](@ref).

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler
"""
function addStabilization_diff(mesh::AbstractMesh{Tmsh},
                          sbp::AbstractOperator, eqn::EulerData{Tsol}, opts,
                          assembler::AssembleElementData) where {Tmsh,  Tsol}

#  println("==== start of addStabilization ====")

  # ----- Edge Stabilization -----#
  if eqn.params.use_edgestab  #TODO: use options dictionary instead
#    println("applying edge stabilization")
    error("edge stabilization not supported by evalJacobian()")
  end

  # ----- Filtering -----
  if eqn.params.use_res_filter
    error("use_res_filter not supported by evalJacobian()")
  end

  # ----- Artificial Dissipation -----
  if eqn.params.use_dissipation
    error("use_dissipation not supported by evalJacobian()")
  end

  if opts["use_GLS2"]
    error("use_GLS2 not supported by evalJacobian()")
     applyGLS3(mesh, sbp, eqn, opts)
#    test_GLS(mesh, sbp, eqn, opts)
  end

  if opts["use_lps"]
    applyLPStab_diff(mesh, sbp, eqn, opts, assembler)
  end
#  println("==== end of addStabilization ====")

  return nothing
end

"""
  Differentiated version of [`evalFaceIntegrals_diff`](@ref).  Throws an
  error for unsupported schemes.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler
"""
function evalFaceIntegrals_diff(mesh::AbstractDGMesh{Tmsh},
                    sbp::AbstractOperator,
                    eqn::EulerData{Tsol}, opts,
                    assembler::AssembleElementData) where {Tmsh, Tsol}

  face_integral_type = opts["face_integral_type"]
  if face_integral_type == 1
      calcFaceIntegral_nopre_diff(mesh, sbp, eqn, opts, eqn.flux_func_diff, mesh.interfaces, assembler)

  elseif face_integral_type == 2

    if opts["use_staggered_grid"]
      error("explicit jacobian calculation does not support staggered grids")
    else
      getFaceElementIntegral_diff(mesh, sbp, eqn,
                  eqn.face_element_integral_func,
                  eqn.flux_func_diff, mesh.sbpface, mesh.interfaces, assembler)
    end

  else
    throw(ErrorException("Unsupported face integral type = $face_integral_type"))
  end

  return nothing
end

"""
  Differentiated version of [`evalSharedFaceIntegrals`](@ref).
  Currently only supports `opts[parallel_data] == element` because it is
  used for Jacobian calculation.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler
"""
function evalSharedFaceIntegrals_diff(mesh::AbstractDGMesh, sbp, eqn, opts,
                                      assembler::AssembleElementData)

  face_integral_type = opts["face_integral_type"]
  eqn.assembler = assembler  # stash this here to get retrieved later
  if face_integral_type == 1

    finishExchangeData(mesh, sbp, eqn, opts, eqn.shared_data, calcSharedFaceIntegrals_element_diff)

  elseif face_integral_type == 2
    finishExchangeData(mesh, sbp, eqn, opts, eqn.shared_data, calcSharedFaceElementIntegrals_element_diff)
  else
    throw(ErrorException("unsupported face integral type = $face_integral_type"))
  end

  # we don't want to keep a reference to the matrix inside the assembler
  # that would prevent it from getting freed
  eqn.assembler = NullAssembleElementData
  return nothing
end

"""
  Differentiated version of [`evalSourceTerm`](@ref).  Currently, source
  terms are not supported.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler
"""
function evalSourceTerm_diff(mesh::AbstractMesh{Tmsh},
                     sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim},
                     opts, assembler::AssembleElementData) where {Tmsh, Tsol, Tres, Tdim}

  # source terms are not a function of q, so they don't affect the Jacobian

  #=
  # placeholder for multiple source term functionality (similar to how
  # boundary conditions are done)
  if opts["use_src_term"]
    error("source terms not supported by evalJacobian()")
  end
  =#
  return nothing
end  # end function

"""
  Differentiated version of [`evalShockCapturing`](@ref).

 **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * assembler
"""
function evalShockCapturing_diff(mesh::AbstractMesh{Tmsh},
                       sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim},
                       opts, assembler::AssembleElementData
                       ) where {Tmsh,  Tsol, Tres, Tdim}

  capture = eqn.shock_capturing
  applyShockCapturing_diff(mesh, sbp, eqn, opts, capture, assembler)

  return nothing
end


module FusionDirect

using OpenGene
using OpenGene.Algorithm
using OpenGene.Reference

# package code goes here

# make it compatible for different version of Julia
include("compat.jl")

export detect

import Base: -,
    abs

include("index/index.jl")
include("detect/detect.jl")

end # module


__precompile__()

module LoopThrottle

export
    @throttle

function loop_throttle_params(; max_rate = 1., min_sleep_time = 0.01)
    min_sleep_time >= 0.001 || error("min_sleep_time must be at least 0.001")
    Float64(max_rate), Float64(min_sleep_time)
end

"""
    @throttle

Throttle a `for` loop or `while` loop by calling `sleep` at the beginning of each
loop iteration (if necessary), so that a designated variable increases
at a rate of at most `max_rate` (compared to wall time).

The first argument that `@throttle` takes is the variable to be rate-limited. This
variable must be in scope inside the loop body. The second argument is a for loop
or while loop. Optionally, the following keyword arguments may be passed in after
the loop argument:
* `max_rate`: specifies the maximum rate at which the designated variable should
increase. Defaults to `1.`.
* `min_sleep_time`: specifies the minimum time to sleep (in seconds). Defaults to
`0.01`, and must be at least 0.001 (due to the accuracy of the `sleep` function).

# Examples
```julia
x = 0
@throttle t for t = 1 : 0.01 : 2
    x += 1
end max_rate = 2.
```
will finish in approximately 0.5 seconds.

```julia
x = 0.
@throttle x for i = 0 : 10
    x += 0.1
end
```
will use the default `max_rate` value of `1.` and thus finish in approximately 1 second.
```julia
i = 0
@throttle i while i <= 10
    println(i)
    i += 1
end min_sleep_time=1.5 max_rate=1
```
will print the numbers from 0 to 10 at an average rate of one per second, while never
sleeping for less than 1.5 second.
"""
macro throttle(t::Symbol, loopexpr::Expr, params::Expr...)
    foreach(params) do expr
        @assert expr.head == :(=)
        @assert expr.args[1] isa Symbol
        expr.head = :kw
        expr.args[2] = esc(expr.args[2])
    end

    setup = quote
        max_rate, min_sleep_time = loop_throttle_params($(params...))
        firstloop = true
        local t0
        local walltime0
    end

    @assert loopexpr.head ∈ (:while, :for)
    loopcondition = loopexpr.args[1]
    loopbody = loopexpr.args[2]
    newloopbody = quote
        if firstloop
            t0 = $(esc(t))
            walltime0 = time()
            firstloop = false
        else
            if !isinf(max_rate)
                Δwalltime = time() - walltime0
                Δt = $(esc(t)) - t0
                sleeptime = Δt / max_rate - Δwalltime
                if sleeptime > min_sleep_time
                    sleep(sleeptime)
                end
            end
        end

        $(esc(loopbody))
    end
    loop = Expr(loopexpr.head, :($(esc(loopcondition))), newloopbody)

    quote
        $setup
        $loop
    end
end

end # module


#=

This file contains the entry point for running the RT simulation, rt_run.

There are two implementations: one that accepts the raw parameters, and one that accepts
the model. The latter should generally be used by users.

=#

"""
    $(FUNCTIONNAME)(pol_type, obs_geom::ObsGeometry, τ_rayl, τ_aer, quad_points::QuadPoints, max_m, aerosol_optics, greek_rayleigh, τ_abs, brdf, architecture::AbstractArchitecture)

Perform Radiative Transfer calculations using given parameters

"""
function rt_run(pol_type::AbstractPolarizationType,   # Polarization type (IQUV)
                obs_geom::ObsGeometry,                # Solar Zenith, Viewing Zenith, Viewing Azimuthal
                τ_rayl,                               # Rayleigh optical depth
                τ_aer,                                # Aerosol optical depth and single-scattering albedo
                quad_points::QuadPoints,              # Quadrature points and weights
                max_m,                                # Max Fourier terms
                aerosol_optics,                       # AerosolOptics (greek_coefs, ω̃, k, fᵗ)
                greek_rayleigh::GreekCoefs,           # Greek coefficients of Rayleigh Phase Function
                τ_abs,                                # nSpec x Nz matrix of absorption
                brdf,                                 # BRDF surface type
                architecture::AbstractArchitecture)   # Whether to use CPU / GPU

    @unpack obs_alt, sza, vza, vaz = obs_geom   # Observational geometry properties
    @unpack qp_μ, wt_μ, qp_μN, wt_μN, iμ₀Nstart, μ₀, iμ₀,Nquad = quad_points # All quadrature points
    FT = eltype(sza)                    # Get the float-type to use
    Nz = length(τ_rayl)                 # Number of vertical slices
    nSpec = size(τ_abs, 1)              # Number of spectral points
    arr_type = array_type(architecture) # Type of array to use
    SFI = true                          # SFI flag
    NquadN = Nquad * pol_type.n         # Nquad (multiplied by Stokes n)
    dims = (NquadN,NquadN)              # nxn dims
    nAer  = length(aerosol_optics)      # Number of aerosols

    # Need to check this a bit better in the future!
    FT_dual = length(τ_aer) > 0 ? typeof(τ_aer[1]) : Float64

    # Output variables: Reflected and transmitted solar irradiation at TOA and BOA respectively # Might need Dual later!!
    R = zeros(FT_dual, length(vza), pol_type.n, nSpec)
    T = zeros(FT_dual, length(vza), pol_type.n, nSpec)
    R_SFI = zeros(FT_dual, length(vza), pol_type.n, nSpec)
    T_SFI = zeros(FT_dual, length(vza), pol_type.n, nSpec)

    # Notify user of processing parameters
    msg =
    """
    Processing on: $(architecture)
    With FT: $(FT)
    Source Function Integration: $(SFI)
    Dimensions: $((NquadN, NquadN, nSpec))
    """
    @info msg

    # Create arrays
    @timeit "Creating layers" added_layer         = make_added_layer(FT_dual, arr_type, dims, nSpec)
    @timeit "Creating layers" added_layer_surface = make_added_layer(FT_dual, arr_type, dims, nSpec)
    @timeit "Creating layers" composite_layer     = make_composite_layer(FT_dual, arr_type, dims, nSpec)
    @timeit "Creating arrays" Aer𝐙⁺⁺ = arr_type(zeros(FT_dual, (dims[1], dims[2], nAer)))
    @timeit "Creating arrays" Aer𝐙⁻⁺ = similar(Aer𝐙⁺⁺)
    @timeit "Creating arrays" I_static = Diagonal(arr_type(Diagonal{FT}(ones(dims[1]))));

    println("Finished initializing arrays")

    # Loop over fourier moments
    for m = 0:max_m - 1

        println("Fourier Moment: ", m, "/", max_m-1)

        # Azimuthal weighting
        weight = m == 0 ? FT(0.5) : FT(1.0)

        # Compute Z-moments of the Rayleigh phase matrix
        # For m>=3, Rayleigh matrices will be 0, can catch with if statement if wanted
        @timeit "Z moments" Rayl𝐙⁺⁺, Rayl𝐙⁻⁺ = Scattering.compute_Z_moments(pol_type, Array(qp_μ), greek_rayleigh, m, arr_type = arr_type);

        # Need to make sure arrays are 0:
        # TBD here

        # Compute aerosol Z-matrices for all aerosols
        for i = 1:nAer
            @timeit "Z moments"  Aer𝐙⁺⁺[:,:,i], Aer𝐙⁻⁺[:,:,i] = Scattering.compute_Z_moments(pol_type, Array(qp_μ), aerosol_optics[i].greek_coefs, m, arr_type = arr_type)
        end

        @timeit "Creating arrays" τ_sum_old = arr_type(zeros(FT, nSpec)) # Suniti: declaring τ_sum to be of length nSpec

        # Loop over all layers and pre-compute all properties before performing core RT
        @timeit "Computing Layer Properties" computed_atmosphere_properties = construct_all_atm_layers(FT, nSpec, Nz, NquadN, τ_rayl, τ_aer, aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type, qp_μ, μ₀, m)

        # Loop over vertical layers:
        @showprogress 1 "Looping over layers ..." for iz = 1:Nz  # Count from TOA to BOA

            # Construct the atmospheric layer
            # From Rayleigh and aerosol τ, ϖ, compute overall layer τ, ϖ
            computed_layer_properties = get_layer_properties(computed_atmosphere_properties, iz, arr_type)

            # Perform Core RT (doubling/elemental/interaction)
            rt_kernel!(pol_type, SFI, added_layer, composite_layer, computed_layer_properties, m, quad_points, I_static, architecture, qp_μN, iz)
        end

        # Create surface matrices:
        create_surface_layer!(brdf, added_layer, SFI, m, pol_type, quad_points, arr_type(computed_atmosphere_properties.τ_sum_all[:,end]), architecture);

        # One last interaction with surface:
        @timeit "interaction" interaction!(computed_atmosphere_properties.scattering_interfaces_all[end], SFI, composite_layer, added_layer, I_static)

        # Postprocess and weight according to vza
        postprocessing_vza!(iμ₀, pol_type, composite_layer, vza, qp_μ, m, vaz, μ₀, weight, nSpec, SFI, R, R_SFI, T, T_SFI)
    end

    # Show timing statistics
    print_timer()
    reset_timer!()

    # Return R_SFI or R, depending on the flag
    return SFI ? R_SFI : R
end

"""
    $(FUNCTIONNAME)(model::vSmartMOM_Model, i_band::Integer = -1)

Perform Radiative Transfer calculations using parameters passed in through the
vSmartMOM_Model struct

"""
function rt_run(model::vSmartMOM_Model; i_band::Integer = -1)

    # Number of bands total
    n_bands = length(model.params.spec_bands)

    # Check that i_band is valid
    @assert (i_band == -1 || i_band in collect(1:n_bands)) "i_band is $(i_band) but there are only $(n_bands) bands"

    # User wants a specific band
    if i_band != -1
        return rt_run(model.params.polarization_type,
                      model.obs_geom::ObsGeometry,
                      model.τ_rayl[i_band],
                      model.τ_aer[i_band],
                      model.quad_points,
                      model.params.max_m,
                      model.aerosol_optics[i_band],
                      model.greek_rayleigh,
                      model.τ_abs[i_band],
                      model.params.brdf[i_band],
                      model.params.architecture)

    # User doesn't specify band, but there's only one
    elseif n_bands == 1

        return rt_run(model.params.polarization_type,
                      model.obs_geom::ObsGeometry,
                      model.τ_rayl[1],
                      model.τ_aer[1],
                      model.quad_points,
                      model.params.max_m,
                      model.aerosol_optics[1],
                      model.greek_rayleigh,
                      model.τ_abs[1],
                      model.params.brdf[1],
                      model.params.architecture)

    # User wants all bands
    else

        Rs = []

        for i in 1:n_bands

            println("------------------------------")
            println("Computing R for band #$(i)")
            println("------------------------------")

            R = rt_run(model.params.polarization_type,
                       model.obs_geom::ObsGeometry,
                       model.τ_rayl[i],
                       model.τ_aer[i],
                       model.quad_points,
                       model.params.max_m,
                       model.aerosol_optics[i],
                       model.greek_rayleigh,
                       model.τ_abs[i],
                       model.params.brdf[i],
                       model.params.architecture)
            push!(Rs, R);
        end

        return Rs
    end

end