#InspectDR: Base functionnality and types
#-------------------------------------------------------------------------------

#AnnotationArea
#TODO: Themes: Define StyleInfo/style?
#TODO: Themes: Vector/Dict of StyleInfo/Layout?

#==Abstracts
===============================================================================#
abstract Plot
abstract PlotCanvas

#==Constants
===============================================================================#
const DInf = convert(DReal, Inf)
const DNaN = convert(DReal, NaN)

const COLOR_TRANSPARENT = ARGB32(0, 0, 0, 0)

const COLOR_BLACK = RGB24(0, 0, 0)
const COLOR_WHITE = RGB24(1, 1, 1)

const COLOR_RED = RGB24(1, 0, 0)
const COLOR_GREEN = RGB24(0, 1, 0)
const COLOR_BLUE = RGB24(0, 0, 1)

#==Plot-level structures
===============================================================================#
#=Input structures are high-level constructs describing the plot.  They will
be broken down to lower-level structures later on.
=#

type LineAttributes <: AttributeList
	style
	width #[0, 10]
	color
end
line(;style=:solid, width=1, color=COLOR_BLACK) =
	LineAttributes(style, width, color)

#"line" constructor:
#TODO: Inadequate
#eval(genexpr_attriblistbuilder(:line, LineAttributes, reqfieldcnt=0))

type GlyphAttributes <: AttributeList #Don't use "Symbol" - name used by Julia
#==IMPORTANT:
Edge width & color taken from LineAttributes
==#
	shape #because "type" is reserved
	size #of glyph.  edge width taken from LineAttributes
	color #glyph linecolor. = nothing to match line.color
	fillcolor
end
glyph(;shape=:none, size=3, color=nothing, fillcolor=nothing) =
	GlyphAttributes(shape, size, color, fillcolor)

#"glyph" constructor:
#TODO: Inadequate
#eval(genexpr_attriblistbuilder(:glyph, GlyphAttributes, reqfieldcnt=0))

type GridAttributes <: AttributeList
	#Bool values
	vmajor
	vminor
	hmajor
	hminor
end
grid(;vmajor=false, vminor=false, hmajor=false, hminor=false) =
	GridAttributes(vmajor, vminor, hmajor, hminor)

#Dispatchable scale:
immutable AxisScale{T}
	call{T}(t::Type{AxisScale{T}}) = error("$t not supported")
	call(::Type{AxisScale{:lin}}) = new{:lin}()
	call(::Type{AxisScale{:log2}}) = new{:log2}()
	call(::Type{AxisScale{:log10}}) = new{:log10}()
	call(::Type{AxisScale{:dB10}}) = new{:dB10}()
	call(::Type{AxisScale{:dB20}}) = new{:dB20}()

	#Aliases:
	call(::Type{AxisScale{:log}}) = new{:log10}()
end
AxisScale(t::Symbol) = AxisScale{t}()

#Specifies configuration of axes:
abstract Axes

#Rectilinear axis (ex: normal cartesian +logarithmic, ...):
immutable AxesRect <: Axes
	xscale::AxisScale
	yscale::AxisScale
end
AxesRect(xscale::Symbol, yscale::Symbol) = AxesRect(AxisScale{xscale}(), AxisScale{yscale}())
AxesRect() = AxesRect(:lin, :lin)

#Curvilinear axes (ex: polar plots):
immutable AxesCurv <: Axes
	rscale::AxisScale #Radial scale could be
end
AxesCurv(rscale::Symbol=:lin) = AxesCurv(AxisScale{rscale}())

immutable AxesSmith{T} <: Axes
	ref::Float64 #Y/Zref
	call{T}(t::Type{AxesSmith{T}}, ref::Real) = error("$t not supported")
	call(::Type{AxesSmith{:Z}}, ref::Real) = new{:Z}(ref)
	call(::Type{AxesSmith{:Y}}, ref::Real) = new{:Y}(ref)
end
AxesSmith(t::Symbol, ref::Real=1) = AxesSmith{t}(ref)
AxesSmith(t::Symbol, ref::Void) = AxesSmith(t)

type Waveform{T}
	id::DisplayString
	ds::T
	line::LineAttributes
	glyph::GlyphAttributes
	ext::PExtents2D
end

#Input waveform:
typealias IWaveform Waveform{IDataset}

#Display waveform (concrete (x, y) pair):
typealias DWaveform Waveform{Vector{Point2D}}

#=TODO:
typealias DWaveformCplx Waveform{PointComplex}
to track x-values of complex data??
=#

type Annotation
	title::DisplayString
	xlabel::DisplayString
	ylabel::DisplayString
	timestamp::DisplayString
end
Annotation(;title="") = Annotation(title, "", "", Libc.strftime(time()))

type Font
	name::AbstractString
	_size::Float64
	bold::Bool
end
Font(name::AbstractString, _size::Real; bold::Bool=false) = Font(name, _size, bold)
Font(_size::Real; bold::Bool=false) = Font(defaults.fontname, _size, bold)

#Legend layout/style
type LegendLStyle
	enabled::Bool
	#autosize::Bool Auto-compute width from text.
	font::Font
	#position: left/right/...
	width::Float64
	#border::Bool
	#wborder
	vgap::Float64 #Vertical spacing between (% of text height).
	linegap::Float64 #Spacing between line and label (% of M character).
	linelength::Float64 #length of line segment to display.
end

#Plot layout
#TODO: Split Layout into "StyleInfo" - which includes Layout??
type Layout
	htitle::Float64 #Title allocation
	waxlabel::Float64 #Vertical axis label allocation (width)
	haxlabel::Float64 #Horizontal axis label allocation (height)
	wnolabels::Float64 #Width to use where no labels are displayed

	wticklabel::Float64 #y-axis values allocation (width)
	hticklabel::Float64 #x-axis values allocation (height)
	hlabeloffset::Float64
	vlabeloffset::Float64

	tframe::Float64 #Frame thickness

	wdata::Float64 #Suggested width of data (graph) area
	hdata::Float64 #Suggested height of data (graph) area

	fnttitle::Font
	fntaxlabel::Font
	fntticklabel::Font
	fnttime::Font

	grid::GridAttributes
	legend::LegendLStyle
	xlabelformat::TickLabelStyle
	ylabelformat::TickLabelStyle
	showtimestamp::Bool
end
Layout(;fontname::AbstractString=defaults.fontname, fontscale=defaults.fontscale) =
	Layout(
	20*fontscale, 20*fontscale, 20*fontscale, 45, #Title/main labels
	60*fontscale, 15*fontscale, 3, 7, 2, #Ticks/frame
	defaults.wdata, defaults.hdata,
	Font(fontname, fontscale*14, bold=true), #Title
	Font(fontname, fontscale*14), Font(fontname, fontscale*12), #Axis/Tick labels
	Font(fontname, fontscale*8), #Time stamp
	GridAttributes(true, false, true, false),
	LegendLStyle(false, Font(fontname, fontscale*12), 100.0, .25, 0.5, 20),
	TickLabelStyle(), TickLabelStyle(), defaults.showtimestamp
)

#Tag data as being part of a given coordinate system:
immutable CoordSystem{ID}; end
typealias DeviceCoord CoordSystem{:dev}
typealias NormCoord CoordSystem{:norm}
typealias DataCoord CoordSystem{:data}

immutable TypedCoord{CT<:CoordSystem}
	v::DReal
end
#Annotation coordinates can match data or be normalized to plot bounds (0 -> 1):
typealias AnnotationCoord Union{TypedCoord{NormCoord}, TypedCoord{DataCoord}}
coord(s::Symbol, v::DReal) = TypedCoord{CoordSystem{s}}(v)

type TextAnnotation
	text::DisplayString
	pt::Point2D #Data coordinates - set NaN to use offsets only
	xoffset::DReal #Normalized to [0,1] plot bounds
	yoffset::DReal #Normalized to [0,1] plot bounds
	font::Font
	color::Colorant
	angle::DReal #Degrees
	align::Symbol #tl, tc, tr, cl, cc, cr, bl, bc, br
end
#Don't use "text"... high probability of collisions when exported...
atext(text::AbstractString; x::Real=DNaN, y::Real=DNaN, xoffset=0, yoffset=0,
	font=Font(10), color=COLOR_BLACK, angle=0, align=:bl) =
	TextAnnotation(text, Point2D(x,y), xoffset, yoffset, font, color, angle, align)

type HVMarker
	vmarker::Bool #false: hmarker
	pos::DReal
	line::LineAttributes
end
vmarker(pos, line=InspectDR.line()) = HVMarker(true, pos, line)
hmarker(pos, line=InspectDR.line()) = HVMarker(false, pos, line)

#2D plot.
type Plot2D <: Plot
	layout::Layout
	axes::Axes
	annotation::Annotation

	#Plot extents (access using getextents):
	ext_data::PExtents2D #Maximum extents of data
	ext_full::PExtents2D #Defines "full" zoom when combined with ext_data
	ext::PExtents2D #Current/active extents

	data::Vector{IWaveform}
	markers::Vector{HVMarker}
	atext::Vector{TextAnnotation}

	#Display data cache:
	invalid_ddata::Bool #Is cache of display data invalid?
	display_data::Vector{DWaveform} #Clipped to current extents

	#Maximum # of x-pts in display:
	#TODO: move to layout?
	xres::Int
end

Plot2D(;title="") = Plot2D(Layout(), AxesRect(), Annotation(title=title),
	PExtents2D(), PExtents2D(), PExtents2D(), [], [], [], true, [], 1000
)

type Multiplot
	title::DisplayString

	subplots::Vector{Plot}
	ncolumns::Int

	#Default width/height of plots
	wplot::Float64
	hplot::Float64

	htitle::Float64 #Title height allocation
	fnttitle::Font
end
Multiplot(;title="", ncolumns=1, titlefont=Font(defaults.fontname, 20),
		wplot=defaults.wplot, hplot=defaults.hplot) =
	Multiplot(title, [], ncolumns, wplot, hplot, 30, titlefont)

#==Constructor-like functions
===============================================================================#

#axes function:
#Interface is a bit irregular, but should be easy to use...
#-------------------------------------------------------------------------------
function axes(a1::Symbol)
	if :smith == a1
		return AxesSmith(:Z)
	elseif :polar == a1
		return AxesCurv()
	else
		throw(MethodError(axes, (a1,)))
	end
end

function axes(a1::Symbol, a2::Symbol; ref = nothing)
	if :smith == a1
		return AxesSmith(a2, ref)
	elseif ref != nothing
		error("cannot set ref for axes(:$a1, :$a2)")
	end

	if :polar == a1
		return AxesCurv(a2)
	else
		return AxesRect(a1, a2)
	end
end

function axes(a1::Symbol, a2::Symbol, a3::Symbol)
	if :polar == a1
		return AxesCurv(a2, a3)
	else
		throw(MethodError(axes, (a1,a2,a3)))
	end
end

#==Accessors
===============================================================================#
_width(style::LegendLStyle) = style.width

getextents(d::IWaveform) = getextents(d.ds)
function getextents(dlist::Vector{IWaveform})
	result = PExtents2D(DNaN, DNaN, DNaN, DNaN)
	for d in dlist
		result = union(result, d.ext)
	end
	return result
end

#==Mutators
===============================================================================#
settitle(mplot::Plot2D, title::DisplayStringArg) =
	(mplot.annotation.title = DisplayString(title))

settitle(mplot::Multiplot, title::DisplayStringArg) =
	(mplot.title = DisplayString(title))

#=="add" interface
===============================================================================#

function _add(mp::Multiplot, plot::Plot2D)
	push!(mp.subplots, plot)
	return plot
end
_add{T<:Plot}(mp::Multiplot, ::Type{T}) = _add(mp, T())

#Set dataf1=false to overwrite optimizations for functions of 1 argument.
function _add(plot::Plot2D, x::Vector, y::Vector; id::AbstractString="", dataf1=true)
	if dataf1
		dataf1 = isincreasing(x) #Can we use optimizations?
	end
	ext = PExtents2D() #Don't care at the moment
	ds = IWaveform(id, IDataset{dataf1}(x, y), line(), glyph(), ext)
	push!(plot.data, ds)
	return ds
end

function _add(plot::Plot2D, marker::HVMarker)
	push!(plot.markers, marker)
	return marker
end

function _add(plot::Plot2D, a::TextAnnotation)
	push!(plot.atext, a)
	return a
end

#==Mapping/interpolation functions
===============================================================================#

#Mapping functions, depending on axis type:
#-------------------------------------------------------------------------------
datamap{T<:Number}(::Type{T}, ::AxisScale) = (x::T)->DReal(x)
datamap{T<:Number}(::Type{T}, ::AxisScale{:dB10}) = (x::T)->(5*log10(DReal(abs2(x))))
datamap{T<:Number}(::Type{T}, ::AxisScale{:dB20}) = (x::T)->(10*log10(DReal(abs2(x))))
datamap{T<:Number}(::Type{T}, ::AxisScale{:log2}) =
	(x::T)->(x<0? DNaN: log2(DReal(x)))
datamap{T<:Number}(::Type{T}, ::AxisScale{:log10}) =
	(x::T)->(x<0? DNaN: log10(DReal(x)))
#TODO: find a way to show negative values for log10?

datamap_rev{T<:Number}(::Type{T}, ::AxisScale) = (x::T)->DReal(x)
#NOTE: Values dBs remain in dBs for readability
datamap_rev{T<:Number}(::Type{T}, ::AxisScale{:log2}) = (x::T)->(2.0^x)
datamap_rev{T<:Number}(::Type{T}, ::AxisScale{:log10}) = (x::T)->(10.0^x)

datamap_rev{T<:Number}(v::T, s::AxisScale) = datamap_rev(T,s)(v) #One-off conversion

#Extents mapping functions, depending on axis type:
#-------------------------------------------------------------------------------
extentsmap(::AxisScale) = (x::DReal)->x #Most axis types don't need to re-map extents.
#NOTE: Extents in dBs remain extents in dBs
extentsmap(t::AxisScale{:log2}) = datamap(DReal, t)
extentsmap(t::AxisScale{:log10}) = datamap(DReal, t)

extentsmap_rev(::AxisScale) = (x::DReal)->x #Most axis types don't need to re-map extents.
#NOTE: Extents in dBs remain extents in dBs
extentsmap_rev(t::AxisScale{:log2}) = datamap_rev(DReal, t)
extentsmap_rev(t::AxisScale{:log10}) = datamap_rev(DReal, t)

#==Plot extents
===============================================================================#

function invalidate_extents(plot::Plot2D)
	#If extents are no longer valid, neither is display cache:
	plot.invalid_ddata = true
end

function invalidate_datalist(plot::Plot2D)
	plot.invalid_ddata = true
end

rescale(ext::PExtents2D, axes::Axes) = ext #Default
function rescale(ext::PExtents2D, axes::AxesRect)
	xmap = extentsmap(axes.xscale)
	ymap = extentsmap(axes.yscale)
	return PExtents2D(
		xmap(ext.xmin), xmap(ext.xmax),
		ymap(ext.ymin), ymap(ext.ymax)
	)
end
rescale_rev(ext::PExtents2D, axes::Axes) = ext #Default
function rescale_rev(ext::PExtents2D, axes::AxesRect)
	xmap = extentsmap_rev(axes.xscale)
	ymap = extentsmap_rev(axes.yscale)
	return PExtents2D(
		xmap(ext.xmin), xmap(ext.xmax),
		ymap(ext.ymin), ymap(ext.ymax)
	)
end

#Accessor:
getextents(plot::Plot2D) = plot.ext
getextents_xfrm(plot::Plot2D) = rescale(plot.ext, plot.axes)

#Full extents are always merged (ext_full expected to be incomplete):
getextents_full(plot::Plot2D) = merge(plot.ext_data, plot.ext_full)

#Set active plot extents using data coordinates:
function setextents(plot::Plot2D, ext::PExtents2D, hallowed::Bool=true, vallowed::Bool=true)
	xmin = ext.xmin; xmax = ext.xmax; ymin = ext.ymin; ymax = ext.ymax
	if !hallowed
		xmin = plot.ext.xmin
		xmax = plot.ext.xmax
	end
	if !vallowed
		ymin = plot.ext.ymin
		ymax = plot.ext.ymax
	end
	#Automatically fill-in any NaN fields, if possible:
	plot.ext = merge(getextents_full(plot), PExtents2D(xmin, xmax, ymin, ymax))
	invalidate_extents(plot)
end

#Set active plot extents using xfrm display coordinates:
setextents_xfrm(plot::Plot2D, ext::PExtents2D, hallowed::Bool=true, vallowed::Bool=true) =
	setextents(plot, rescale_rev(ext, plot.axes), hallowed, vallowed)

function setextents_full(plot::Plot2D, ext::PExtents2D)
	plot.ext_full = ext
end

#==Plot/graph bounding boxes
===============================================================================#

aspect_square(::Axes) = false
aspect_square(::AxesSmith) = true

#Returns a centered bounding box with square aspect ratio.
function squarebounds(bb::BoundingBox)
	w = width(bb); h = height(bb)
	xmin = bb.xmin; xmax = bb.xmax
	ymin = bb.ymin; ymax = bb.ymax
	if w < h
		hdim = w / 2
		c = (ymin + ymax) / 2
		ymin = c - hdim; ymax = c + hdim
	else
		hdim = h / 2
		c = (xmin + xmax) / 2
		xmin = c - hdim; xmax = c + hdim
	end
	return BoundingBox(xmin, xmax, ymin, ymax)
end

#Get bounding box of graph (plot data area):
function graphbounds(plotb::BoundingBox, lyt::Layout)
	xmin = plotb.xmin + lyt.waxlabel + lyt.wticklabel
	xmax = plotb.xmax
	xmax -= lyt.legend.enabled? _width(lyt.legend): lyt.wnolabels
	ymin = plotb.ymin + lyt.htitle
	ymax = plotb.ymax - lyt.haxlabel - lyt.hticklabel

	#Avoid division by zero, inversions, ...
	if xmin >= xmax
		c = (xmin + xmax)/2
		xmin = c - 0.5
		xmax = c + 0.5
	end
	if ymin >= ymax
		c = (ymin + ymax)/2
		ymin = c - 0.5
		ymax = c + 0.5
	end

	return BoundingBox(xmin, xmax, ymin, ymax)
end

function graphbounds(plotb::BoundingBox, lyt::Layout, axes::Axes)
	graphbb = graphbounds(plotb, lyt)
	if aspect_square(axes)
		graphbb = squarebounds(graphbb)
	end
	return graphbb
end

#Get bounding box of entire plot:
function plotbounds(lyt::Layout, graphw::Float64, graphh::Float64)
	xmax = graphw + lyt.waxlabel + lyt.wticklabel
	xmax += lyt.legend.enabled? _width(lyt.legend): lyt.wnolabels
	ymax = graphh + lyt.htitle + lyt.haxlabel + lyt.hticklabel
	return BoundingBox(0, xmax, 0, ymax)
end

#Get suggested plot bounds:
function plotbounds(lyt::Layout, axes::Axes)
	wdata = lyt.wdata; hdata = lyt.hdata
	if aspect_square(axes)
		wdata = hdata = min(wdata, hdata)
	end
	return plotbounds(lyt, wdata, hdata)
end

#==Pre-processing display data
===============================================================================#

function _reduce(input::IWaveform, ext::PExtents2D, xres_max::Integer)
	return DWaveform(input.id, _reduce(input.ds, ext, xres_max), input.line, input.glyph, input.ext)
end

_reduce(inputlist::Vector{IWaveform}, ext::PExtents2D, xres_max::Integer) =
	map((input)->_reduce(input, ext, xres_max::Integer), inputlist)

#Rescale input dataset:
#-------------------------------------------------------------------------------
function _rescale{T<:Number}(d::Vector{T}, scale::AxisScale)
	#Apparently, passing functions as arguments is not efficient in Julia.
	#-->Specializing on AxisScale, hoping to improve efficiency on dmap:
	dmap = datamap(T, scale)

	result = Array(DReal, length(d))
	for i in 1:length(d)
		result[i] = dmap(d[i])
	end
	return result
end
_rescale{T<:Number}(d::Vector{T}, scale::AxisScale{:lin}) = d #Optimization: Linear scale does not need re-scaling
_rescale{T<:IDataset}(input::T, xscale::AxisScale, yscale::AxisScale) =
	T(_rescale(input.x, xscale), _rescale(input.y, yscale))

function _rescale(input::IWaveform, xscale::AxisScale, yscale::AxisScale)
	ds = _rescale(input.ds, xscale, yscale)
	return IWaveform(input.id, ds, input.line, input.glyph, getextents(ds))
end

#Specialized on xscale/yscale for efficiency:
_rescale(inputlist::Vector{IWaveform}, xscale::AxisScale, yscale::AxisScale) =
	map((input)->_rescale(input, xscale, yscale), inputlist)

_rescale(inputlist::Vector{IWaveform}, axes::AxesRect) = _rescale(inputlist, axes.xscale, axes.yscale)

_rescale(pt::Point2D, xscale::AxisScale, yscale::AxisScale) =
	Point2D(datamap(DReal, xscale)(pt.x), datamap(DReal, yscale)(pt.y))
_rescale(pt::Point2D, axes::AxesRect) = _rescale(pt, axes.xscale, axes.yscale)

#Preprocess input dataset (rescale/reduce quantity of data/...):
#-------------------------------------------------------------------------------
#   (Updates display_data)
function preprocess_data(plot::Plot2D)
	#TODO: Find a way to preprocess x-vectors referencing same data only once?

	#Rescale data
	wfrmlist = _rescale(plot.data, plot.axes)

	#Update extents:
	plot.ext_data = rescale_rev(getextents(wfrmlist), plot.axes) #Update max extents
	setextents(plot, plot.ext) #Update extents, resolving any NaN fields.
	ext = getextents_xfrm(plot) #Read back extents, in transformed coordinates

	#Reduce data:
	plot.display_data = _reduce(wfrmlist, ext, plot.xres)
	plot.invalid_ddata = false

	#NOTE: Rescaling before data reduction is somewhat inefficient, but makes
	#      it easier to interpolate data in _reduce step.
	#TODO: Find a way to efficiently reduce before re-scaling???
end

#==High-level display functions
===============================================================================#
function update_ddata(plot::Plot2D)
	invalidate_extents(plot) #Always compute below:
	#TODO: Conditionnaly compute (only when data changed/added)?

	if plot.invalid_ddata
		preprocess_data(plot)
	end
end

#Last line


module Bioceananigans

export LightGrowth!
export MixedLayerDepth!

include("LightGrowth.jl")
include("MixedLayerDepth.jl")

end # end module


"""
Wrapper for minimization function
"""

## Minimization function
"""
Method to resize bounds if float is passed
"""
function _resize_bounds(lx, nx::Int)
    if length(lx) == 1 && nx > 1
        lx = lx*ones(nx)
    end
    return lx
end

"""
    minimize(func!::Function, x0::Vector, ng::Int; kwargs...)

Minimize function, common interface to IPOPT and SNOPT

Objective function structure:

```julia
function func!(g, x)
    <modify g>
    return <objective value to be minimized>
end
```

# Arguments
    - `func!::Function`: fitness function with signature `rosenbrock!(g, x)` and returning objective
    - `x0::Vector: initial guess
    - `ng::Int: number of constraints
    - `lx::Vector: lower bound on x
    - `ux::Vector: upper bound on x
    - `lg::Vector: lower bound on constraints
    - `ug::Vector: upper bound on constraints
    - `solver::String`: solver, "ipopt" or "snopt"
    - `options::Dict`: options used by solver
    - `sparsity::AbstractSparsityPattern`: sparsity pattern, defualt is DensePattern()
    - `derivatives::AbstractDiffMethod`: derivative method, ForwardFD() or CentralFD() or ForwardAD() or ReverseAD()
    - `outputfile::Boolean`: whether to create output file
    - `verbosity::Int`: verbosirty level, >= 0

# Returns
    - `Tuple`: xstar, fstar, info
"""
function minimize(func!::Function, x0::Vector, ng::Int; kwargs...)
    # unpack values
    lx = _assign_from_kwargs(Dict(kwargs), :lx, -Inf)
    ux = _assign_from_kwargs(Dict(kwargs), :ux, Inf)
    lg = _assign_from_kwargs(Dict(kwargs), :lg, -Inf)
    ug = _assign_from_kwargs(Dict(kwargs), :ug, 0.0)
    solver      = _assign_from_kwargs(Dict(kwargs), :solver, "ipopt")
    options     = _assign_from_kwargs(Dict(kwargs), :options,  Dict())
    sparsity    = _assign_from_kwargs(Dict(kwargs), :sparsity, DensePattern())
    derivatives = _assign_from_kwargs(Dict(kwargs), :derivatives, ForwardFD())
    outputfile  = _assign_from_kwargs(Dict(kwargs), :outputfile, false)
    verbosity   = _assign_from_kwargs(Dict(kwargs), :verbosity, 0)

    # snopt-specific settings
    lencw       = _assign_from_kwargs(Dict(kwargs), :lencw, 500)
    iSumm       = _assign_from_kwargs(Dict(kwargs), :iSumm, 6)   # print to screen if 6

    # initialize number of decision variables
    nx = length(x0)

    # resize bounds if input length is 1
    lx = _resize_bounds(lx, nx)
    ux = _resize_bounds(ux, nx)
    lg = _resize_bounds(lg, ng)
    ug = _resize_bounds(ug, ng)

    # create cache
    # if verbosity>0
    #     println("derivatives: $derivatives")
    # end
    cache = _create_cache(sparsity, derivatives, func!, nx, ng)

    # determine sparsity pattern
    rows, cols = _get_sparsity(sparsity, nx, ng)

    if cmp(solver, "ipopt") == 0
        xstar, fstar, info = minimize_ipopt(options, cache, x0, lx, ux, lg, ug, rows, cols, outputfile)
    elseif cmp(solver, "snopt") == 0
        xstar, fstar, info = minimize_snopt(options, cache, x0, lx, ux, lg, ug, rows, cols, outputfile, lencw, iSumm)
    end
    #print("typeof-cache: ")
    #println(typeof(cache))
    return xstar, fstar, info
end

"""
    _assign_from_kwargs(keyargs::Dict, keyname, default)

Utility to unpack kwargs
"""
function _assign_from_kwargs(keyargs::Dict, keyname, default)
    if haskey(keyargs, keyname)==true
        var = keyargs[keyname]
    else
        var = default
    end
    return var
end


# Conditionals

for op ∈ (:(==), :(!=), :(<), :(<=), :(>), :(>=) )
    rename = VECTOR_SYMBOLS[op]
    # rename_mask = Symbol(rename, :_mask)
    @eval begin
        # scalar versions handled in floating_point_arithmetic.jl
        # @inline $rename(s1::ScalarTypes, s2::ScalarTypes) = $op(s1,s2)
        @vectordef $rename function Base.$op(v1, v2) where {W,T}
            llvmwrap_bitmask(Val{$(QuoteNode(op))}(), extract_data(v1), extract_data(v2))
        end
        @vectordef $rename function Base.$op(v1, v2::T) where {W,T}
            $rename(extract_data(v1), v2)
        end
        @vectordef $rename function Base.$op(v1::T, v2) where {W,T}
            $rename(extract_data(v1), v2)
        end
    end
end
@inline visfinite(s::ScalarTypes) = isfinite(s)
@inline function visfinite(v1::Vec{W,T}) where {W,T<:FloatingTypes}
    U = uint_type(T)
    em = vbroadcast(Vec{W,U}, exponent_mask(T))
    iv = vreinterpret(Vec{W,U}, v1)
    vnot_equal(vand(iv, em), em)
end
@inline visfinite(v1::SVec) = SVec(visfinite(extract_data(v1)))
@inline Base.isfinite(v1::SVec) = SVec(visfinite(extract_data(v1)))

@inline visinf(s1::ScalarTypes) = isinf(s1)
@inline visinf(v1::Vec{W,T}) where {W,T<:FloatingTypes} = visequal(vabs(v1), vbroadcast(Vec{W,T},Inf))
@inline visinf(v1::SVec) = SVec(visinf(extract_data(v1)))
@inline Base.isinf(v1::SVec) = SVec(visinf(extract_data(v1)))

@inline visnan(s1::ScalarTypes) = isnan(s1)
@inline visnan(v1::Vec{W,T}) where {W,T<:FloatingTypes} = vnot_equal(v1, v1)
@inline visnan(v1::SVec) = SVec(visnan(extract_data(v1)))
@inline Base.isnan(v1::SVec) = SVec(visnan(extract_data(v1)))

@inline vissubnormal(s1::ScalarTypes) = issubnormal(s1)
@inline function vissubnormal(v1::Vec{W,T}) where {W,T<:FloatingTypes}
    U = uint_type(T)
    em = vbroadcast(Vec{W,U}, exponent_mask(T))
    sm = vbroadcast(Vec{W,U}, significand_mask(T))
    iv = vreinterpret(Vec{W,U}, v1)
    vand(visequal(vand(iv, em), vbroadcast(Vec{W,U}, 0)), vnot_equal(vand(iv, sm), vbroadcast(Vec{W,U}, 0)))
end
@inline vissubnormal(v1::SVec) = SVec(vissubnormal(extract_data(v1)))
@inline Base.issubnormal(v1::SVec) = SVec(vissubnormal(extract_data(v1)))

@inline signbit(s1::ScalarTypes) = signbit(s1)
@inline function vsignbit(v1::Vec{W,T}) where {W,T<:FloatingTypes}
    U = uint_type(T)
    sm = vbroadcast(Vec{W,U}, sign_mask(T))
    iv = vreinterpret(Vec{W,U}, v1)
    vnot_equal(vand(iv, sm), vbroadcast(Vec{W,U}, 0))
end
@inline function vsignbit(v1::SVec{W,T}) where {W,T<:FloatingTypes}
    SVec(vsignbit(extract_data(v1)))
end
@inline function Base.signbit(v1::SVec{W,T}) where {W,T<:FloatingTypes}
    SVec(vsignbit(extract_data(v1)))
end

@inline vifelse(c::Bool, x, y) = c ? x : y
# @inline vifelse(c::Bool, x, y) = ifelse(c, x, y)
@generated function vifelse(v1::Vec{W,Bool}, v2::Vec{W,T}, v3::Vec{W,T}) where {W,T}
    btyp = llvmtype(Bool)
    vbtyp = "<$W x $btyp>"
    abtyp = "[$W x $btyp]"
    typ = llvmtype(T)
    vtyp = "<$W x $typ>"
    atyp = "[$W x $typ]"
    decls = String[]
    instrs = String[]
    push!(instrs, "%cond = trunc $vbtyp %0 to <$W x i1>")
    push!(instrs, "%res = select <$W x i1> %cond, $vtyp %1, $vtyp %2")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Vec{W,T},
            Tuple{Vec{W,Bool}, Vec{W,T}, Vec{W,T}},
            v1, v2, v3)
    end
end
@inline function vifelse(v1::AbstractSIMDVector{W,Bool},
        v2::AbstractSIMDVector{W,T}, v3::AbstractSIMDVector{W,T}) where {W,T}
    SVec(vifelse(extract_data(v1), extract_data(v2), extract_data(v3)))
end
# @inline function vifelse(v1::VecOrProd{W,Bool},
#         v2::VecOrProd{W,T}, v3::VecOrProd{W,T}) where {W,T}
#     vifelse(extract_data(v1), extract_data(v2), extract_data(v3))
# end
# @inline function Base.ifelse(v1::SVec{W,Bool},
#         v2::SVec{W,T}, v3::SVec{W,T}) where {W,T}
#     SVec(vifelse(extract_data(v1), extract_data(v2), extract_data(v3)))
# end

@generated function vifelse(mask::U, v2::Vec{W,T}, v3::Vec{W,T}) where {W,T,U<:Unsigned}
    @assert 8sizeof(U) >= W
    btyp = llvmtype(Bool)
    vbtyp = "<$W x $btyp>"
    abtyp = "[$W x $btyp]"
    typ = llvmtype(T)
    vtyp = "<$W x $typ>"
    atyp = "[$W x $typ]"
    decls = String[]
    instrs = String[]

    mtyp_input = llvmtype(U)
    mtyp_trunc = "i$W"

    if mtyp_input == mtyp_trunc
        push!(instrs, "%cond = bitcast $mtyp_input %0 to <$W x i1>")
    else
        push!(instrs, "%condtrunc = trunc $mtyp_input %0 to $mtyp_trunc")
        push!(instrs, "%cond = bitcast $mtyp_trunc %condtrunc to <$W x i1>")
    end
    # push!(instrs, "%cond = trunc $vbtyp %0 to <$W x i1>")
    if T <: FloatingTypes && Base.libllvm_version >= v"9"
        push!(instrs, "%res = select fast <$W x i1> %cond, $vtyp %1, $vtyp %2")
    else
        push!(instrs, "%res = select <$W x i1> %cond, $vtyp %1, $vtyp %2")
    end
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Vec{$W,$T},
            Tuple{$U, Vec{$W,$T}, Vec{$W,$T}},
            mask, v2, v3)
    end
end

@vpromote vifelse 3

@inline vifelse(U::Unsigned, v2::AbstractSIMDVector, v3::AbstractSIMDVector) = SVec(vifelse(U, extract_data(v2), extract_data(v3)))
@inline vifelse(U::Unsigned, v2::Vec{W,T}, s::Union{T,Int}) where {W,T} = vifelse(U, v2, vbroadcast(Vec{W,T}, s))
@inline vifelse(U::Unsigned, v2::AbstractSIMDVector{W,T}, s::Union{T,Int}) where {W,T} = SVec(vifelse(U, extract_data(v2), vbroadcast(Vec{W,T}, s)))
@inline vifelse(U::Unsigned, s::Union{T,Int}, v2::Vec{W,T}) where {W,T} = vifelse(U, vbroadcast(Vec{W,T}, s), v2)
@inline vifelse(U::Unsigned, s::Union{T,Int}, v2::AbstractSIMDVector{W,T}) where {W,T} = SVec(vifelse(U, vbroadcast(Vec{W,T}, s), extract_data(v2)))

@inline function vifelse(f::F, m::Mask{W}, vargs::Vararg{<:Any,N}) where {F<:Function,W,N}
    vifelse(m, f(vargs...), @inbounds(vargs[N]))
end

@vectordef visodd function Base.isodd(v) where {W,T<:Integer}
    visequal(vand(v, one(T)), one(T))
end

# @inline vand(U::Unsigned, vb::Vec{W,Bool}) where {W} = U & vecbool_to_unsigned(vb)
# @inline vand(vb::Vec{W,Bool}, U::Unsigned) where {W} = U & vecbool_to_unsigned(vb)


#Defines the mission system that contains the reference trajectory or other mission related
#properties.

#Define Abstract Mission System for customisability
abstract type AbstractMissionSystem end

########### Template function ###########################
mutable struct MissionSystem{T <: Number} <: AbstractMissionSystem
    Active::Bool
    Ts::T
    MissionFunction::Function
    Tprev::T
end
MissionSystem() = MissionSystem(false, 0.0, x -> nothing, 0.0)
MissionSystem(Ts::Number) = MissionSystem(false, Float64(Ts), x->nothing, 0.0)
MissionSystem(Ts::Number, missfunc::Function) = MissionSystem(true, Float64(Ts), missfunc, 0.0)
MissionSystem(Active::Bool, Ts::Number, missfunc::Function) = MissionSystem(Active, Float64(Ts), missfunc, 0.0)

#Set the callback affect condition of the ControlSystem
#placeholder!
function affect!(int,dts::Type{<:AbstractMissionSystem})
    for cu in user_cache(int)
        cu.u = [dts.MissionFunction(int.u)...]
    end
  dts.Tprev = int.t
end


using Flux, Test
using Flux.CUDA
using Flux: cpu, gpu
using Statistics: mean
using LinearAlgebra: I, cholesky, Cholesky
using SparseArrays: sparse, SparseMatrixCSC, AbstractSparseArray

@testset "CUDA" begin
  x = randn(5, 5)
  cx = gpu(x)
  @test cx isa CuArray

  @test Flux.onecold(gpu([1.0, 2.0, 3.0])) == 3

  x = Flux.onehotbatch([1, 2, 3], 1:3)
  cx = gpu(x)
  @test cx isa Flux.OneHotMatrix && cx.indices isa CuArray
  @test (cx .+ 1) isa CuArray

  m = Chain(Dense(10, 5, tanh), Dense(5, 2), softmax)
  cm = gpu(m)

  @test all(p isa CuArray for p in params(cm))
  @test cm(gpu(rand(10, 10))) isa CuArray{Float32,2}

  xs = rand(5, 5)
  ys = Flux.onehotbatch(1:5,1:5)
  @test collect(cu(xs) .+ cu(ys)) ≈ collect(xs .+ ys)

  c = gpu(Conv((2,2),3=>4))
  x = gpu(rand(10, 10, 3, 2))
  l = c(gpu(rand(10,10,3,2)))
  @test gradient(x -> sum(c(x)), x)[1] isa CuArray

  c = gpu(CrossCor((2,2),3=>4))
  x = gpu(rand(10, 10, 3, 2))
  l = c(gpu(rand(10,10,3,2)))
  @test gradient(x -> sum(c(x)), x)[1] isa CuArray

end

@testset "onehot gpu" begin
  y = Flux.onehotbatch(ones(3), 1:2) |> gpu;
  @test (repr("text/plain", y); true)

  gA = rand(3, 2) |> gpu;
  @test gradient(A -> sum(A * y), gA)[1] isa CuArray

  # construct from CuArray
  x = [1, 3, 2]
  y = Flux.onehotbatch(x, 0:3)
  y2 = Flux.onehotbatch(x |> gpu, 0:3)
  @test y2.indices isa CuArray
  @test y2 |> cpu == y
end

@testset "onecold gpu" begin
  y = Flux.onehotbatch(ones(3), 1:10) |> gpu;
  l = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
  @test Flux.onecold(y) isa CuArray
  @test y[3,:] isa CuArray
  @test Flux.onecold(y, l) == ['a', 'a', 'a']
end

@testset "onehot forward map to broadcast" begin
  oa = OneHotArray(rand(1:10, 5, 5), 10) |> gpu
  @test all(map(identity, oa) .== oa)
  @test all(map(x -> 2 * x, oa) .== 2 .* oa)
end

@testset "restructure gpu" begin
  dudt = Dense(1,1) |> gpu
  p,re = Flux.destructure(dudt)
  foo(x) = sum(re(p)(x))
  @test gradient(foo, cu(rand(1)))[1] isa CuArray
end

@testset "GPU functors" begin
  @testset "Cholesky" begin
    M = 2.0*I(10) |> collect
    Q = cholesky(M)
    Q_gpu = Q |> gpu
    @test Q_gpu isa Cholesky{<:Any,<:CuArray}
    Q_cpu = Q_gpu |> cpu
    @test Q_cpu == cholesky(eltype(Q_gpu).(M))
  end

  @testset "isbits array types" begin
    struct SimpleBits
      field::Int32
    end

    @test gpu((;a=ones(1))).a isa CuVector{Float32}
    @test gpu((;a=['a', 'b', 'c'])).a isa CuVector{Char}
    @test gpu((;a=[SimpleBits(1)])).a isa CuVector{SimpleBits}
  end
end

@testset "gpu(cpu(x)) inside gradient" begin
  a = randn(Float32, 4, 4)
  ca = cu(a)

  # Trivial functions
  @test gradient(x -> sum(abs, gpu(x)), a)[1] isa Matrix
  @test gradient(x -> sum(gpu(x)), a)[1] isa Matrix
  @test_skip gradient(x -> sum(gpu(x)), a')[1] isa Matrix  # sum(::Adjoint{T,CuArray}) makes a Fill
  @test gradient(x -> sum(abs, cpu(x)), ca)[1] isa CuArray
  # This test should really not go through indirections and pull out Fills for efficiency
  # but we forcefully materialise. TODO: remove materialising CuArray here
  @test gradient(x -> sum(cpu(x)), ca)[1] isa CuArray # This involves FillArray, which should be GPU compatible
  @test gradient(x -> sum(cpu(x)), ca')[1] isa LinearAlgebra.Adjoint

  # Even more trivial: no movement
  @test gradient(x -> sum(abs, cpu(x)), a)[1] isa Matrix
  @test gradient(x -> sum(abs, cpu(x)), a')[1] isa LinearAlgebra.Adjoint
  @test gradient(x -> sum(cpu(x)), a)[1] isa typeof(gradient(sum, a)[1]) # FillArray
  @test gradient(x -> sum(abs, gpu(x)), ca)[1] isa CuArray
  @test_skip gradient(x -> sum(abs, gpu(x)), ca')[1] isa CuArray # KernelError: passing and using non-bitstype argument

  # More complicated, Array * CuArray is an error
  g0 = gradient(x -> sum(abs, (a * (a * x))), a)[1]
  @test g0 ≈ gradient(x -> sum(abs, cpu(ca * gpu(a * x))), a)[1]
  @test cu(g0) ≈ gradient(x -> sum(abs, gpu(a * cpu(ca * x))), ca)[1]
  @test gradient(x -> sum(gpu(cpu(x))), a)[1] isa Matrix
  @test gradient(x -> sum(gpu(cpu(x))), ca)[1] isa CuArray

  g4 = gradient(x -> sum(a * (a' * x)), a)[1]  # no abs, one adjoint
  @test g4 ≈ gradient(x -> sum(cpu(ca * gpu(a' * x))), a)[1]
  @test cu(g4) ≈ gradient(x -> sum(gpu(a * cpu(ca' * x))), ca)[1]

  # Scalar indexing of an array, needs OneElement to transfer to GPU
  # https://github.com/FluxML/Zygote.jl/issues/1005
  @test gradient(x -> cpu(2 .* gpu(x))[1], Float32[1,2,3]) == ([2,0,0],)
  @test gradient(x -> cpu(gpu(x) * gpu(x))[1,2], Float32[1 2 3; 4 5 6; 7 8 9]) == ([2 6 8; 0 2 0; 0 3 0],)
end

@testset "gpu(x) and cpu(x) on structured arrays" begin
  @test cpu(1:3) isa UnitRange
  @test cpu(range(1, 3, length = 4)) isa AbstractRange

  # OneElement isn't GPU compatible
  g1 = Zygote.OneElement(1, (2,3), axes(ones(4,5)))
  @test cpu(g1) isa Zygote.OneElement

  g2 = Zygote.Fill(1f0, 2)
  @test cpu(g2) isa Zygote.FillArrays.AbstractFill

  g3 = transpose(Float32[1 2; 3 4])
  @test parent(cpu(g3)) isa Matrix{Float32}

  @testset "Sparse Arrays" begin
    @test cpu(sparse(rand(3,3))) isa SparseMatrixCSC
    a = sparse(rand(3,3))
    @test cpu(a) === a
    @test gpu(sparse(rand(3,3))) isa CUDA.CUSPARSE.CuSparseMatrixCSC
  end

  # Check that gpu() converts these to CuArrays. This a side-effect of using the same functions
  # in gpu() as in the gradient of cpu(). A different design could avoid having gpu() used alone
  # move these, if that turns out to be desirable.
  @test gpu(g1) isa CuArray
  @test gpu(g1) ≈ cu(Matrix(g1))
  @test gpu(g2) isa CuArray
  @test gpu(g2) ≈ cu(Vector(g2))
  @test parent(gpu(g3)) isa CuArray
end


module CCDReduction

using Statistics
using LazyStack
using FITSIO
using DataFrames: DataFrame
using ResumableFunctions

export subtract_bias,
       subtract_bias!,
       subtract_overscan,
       subtract_overscan!,
       flat_correct,
       flat_correct!,
       trim,
       trimview,
       crop,
       cropview,
       combine,
       subtract_dark,
       subtract_dark!,
       fitscollection,
       arrays,
       filenames,
       ccds,
       CCDData,
       data,
       hdr

include("ccddata.jl")
include("methods.jl")
include("fits.jl")
include("collection.jl")

end


#####################################################
# Surya
# NREL
# November 2021
# NAERM Threats
# Testing of EGRET2SIIP module
#######################################################
# Loading the required packages
#######################################################
# Load EGRET2SIIP module here
#####################################################################################
# Read EGRET System JSON
#####################################################################################
import JSON
Location_1 = dirname(pwd())*"/test/DAY_AHEAD_Model_2020-07-01_2020-07-14.json";
Location_2 = dirname(pwd())*"/test/REAL_TIME_Model_2020-07-01_2020-07-14.json";
EGRET_json_DA = JSON.parsefile(Location_1);
EGRET_json_RT = JSON.parsefile(Location_2);
#####################################################################################
# Convert EGRET JSON to CSV in a format SIIP Tabular Data understands
#####################################################################################
location, base_MVA,rt_flag = EGRET2SIIP.parse_EGRET_JSON(EGRET_json_DA,EGRET_json_RT=EGRET_json_RT);

if (rt_flag)
    sys_DA,sys_RT = EGRET2SIIP.parse_tabular_data(location,base_MVA,rt_flag,ts_pointers_file ="CSV");
else
    sys_DA = EGRET2SIIP.parse_tabular_data(location,base_MVA,rt_flag,ts_pointers_file="CSV");
end

#           (or)
sys_DA, sys_RT = EGRET2SIIP.EGRET_TO_PSY(EGRET_json_DA,EGRET_json_RT=EGRET_json_RT,ts_pointers_file="CSV");


using Ramnet.Encoders

using ..Approximators: ContinousActionPolicy
using ..Buffers: MultiStepDynamicBuffer, add!

mutable struct ContinousFunctionalPG{OS,AS,T<:Real,O<:StaticArray{Tuple{OS},T,1},A<:StaticArray{Tuple{AS},T,1},E<:AbstractEncoder{T}} <: AbstractAgent{O,A}
    γ::Float64
    policy::ContinousActionPolicy{OS,AS,T,O,A,E}
    buffer::MultiStepDynamicBuffer{O,A}
    observation::Union{Nothing,O}
    done::Bool
    action::Union{Nothing,A}
    min_obs::Union{Nothing,O}
    max_obs::Union{Nothing,O}
    # rng::MersenneTwister
end

function ContinousFunctionalPG(::Type{O}, ::Type{A}, n, start_learning_rate, end_learning_rate, learning_rate_decay, epochs, discount, encoder::E, start_cov::Float64, end_cov::Float64, cov_decay::Int; seed::Union{Nothing,UInt} = nothing) where {OS,AS,T<:Real,O<:StaticArray{Tuple{OS},T,1},A<:StaticArray{Tuple{AS},T,1},E<:AbstractEncoder{T}}
    # !isnothing(seed) && seed < 0 && throw(DomainError(seed, "Seed must be non-negative"))
    # rng = isnothing(seed) ? MersenneTwister() : MersenneTwister(seed)

    n < 1 && throw(DomainError(n, "Tuple size must be at least 1"))

    policy = ContinousActionPolicy(O, A, n, encoder; start_cov, end_cov, cov_decay, start_learning_rate, end_learning_rate, learning_rate_decay, epochs, partitioner = :uniform_random, seed)
    buffer = MultiStepDynamicBuffer{O,A}()

    ContinousFunctionalPG{OS,AS,T,O,A,E}(discount, policy, buffer, nothing, false, nothing, nothing, nothing)
end

function ContinousFunctionalPG(env, encoder::E; tuple_size, start_learning_rate, end_learning_rate, learning_rate_decay, epochs, discount, start_cov, end_cov, cov_decay, seed = nothing) where {T<:Real,E<:AbstractEncoder{T},C<:AbstractMatrix{T}}
    # ContinousFunctionalPG{action_length(env),T,observation_type(env),action_type(env),E,C}(n, observation_length(env), η, γ, encoder, cov; seed)
    ContinousFunctionalPG(observation_type(env), action_type(env), tuple_size, start_learning_rate, end_learning_rate, learning_rate_decay, epochs, discount, encoder, start_cov, end_cov, cov_decay; seed)
end

function observe!(agent::ContinousFunctionalPG{OS,AS,T,O,A,E}, observation::O) where {OS,AS,T,O,A,E}
    agent.observation = observation
    agent.done = false

    # Finding the edges of the domain
    if isnothing(agent.min_obs)
        agent.min_obs = observation
    else
        agent.min_obs = min.(agent.min_obs, observation)
    end

    if isnothing(agent.max_obs)
        agent.max_obs = observation
    else
        agent.max_obs = max.(agent.max_obs, observation)
    end

    nothing
end

# TODO: All observe! methods are the same for all agents. Generalize.
# TODO: This method does not need to take in the action
function observe!(agent::ContinousFunctionalPG{OS,AS,T,O,A,E}, action::A, reward::Float64, observation::O, done::Bool) where {OS,AS,T,O,A,E}
    # REMEMBER: you are forcing reward clipping here!!
    add!(agent.buffer, agent.observation, agent.action, clamp(reward, -1.0, 1.0))

    agent.observation = observation
    agent.done = done

    # Finding the edges of the domain
    agent.min_obs = min.(agent.min_obs)
    agent.max_obs = max.(agent.max_obs)

    nothing
end

function select_action(agent::ContinousFunctionalPG{OS,AS,T,O,A,E}, observation::O; deterministic::Bool = false) where {OS,AS,T,O,A,E}
    return select_action(agent.policy, observation; deterministic)
end

function select_action!(agent::ContinousFunctionalPG{OS,AS,T,O,A,E}, observation::O) where {OS,AS,T,O,A,E}
    agent.action = select_action(agent, observation)

    return agent.action
end

function update!(agent::ContinousFunctionalPG)
    if agent.done
        G = 0.0
        for transition in Iterators.reverse(agent.buffer)
            G = transition.reward + agent.γ * G

            update!(agent.policy, transition.observation, transition.action, G)
        end

        reset!(agent.buffer)
    end

    nothing
end

function Agents.reset!(agent::ContinousFunctionalPG; seed::Union{Nothing,UInt} = nothing)
    reset!(agent.policy; seed)
    reset!(agent.buffer)

    agent.observation = nothing
    agent.action = nothing
    agent.done = false

    nothing
end


def foo# 3

deftype Foo (A -> B)

;; here is a comment
; here is another comment

deftype (Foo a b)
  (a -> b)

(qat :: A -> b)
def tdeftype (lambda a b c -> (a + b))

def (add a b)
  (a + b)

def wat [[1,2,3], [4,5,6]]

def (catstrs strs)
  (foldr f
         (head strs)
         (tail strs))

def strs ["aa", "bb"]

def (mymap f xs)
  if ((length xs) == 0)
      then
xs
  else
    ((f (head xs))
     : (mymap f (tail xs)))

def empty []

def getFile
  (readFile "./parse.js")

;;def fileLines
;;  (getFile >>=
;;   ((mapM_ putStrLn) . lines))

def (testUnary n)
  ((-n) + n)

def (foo bar)
  let {
    lol =  [1,
   (lambda qwerty blah ->
      [qerty, blah,
       (lambda whatever -> whatever)])]
  }
  if bar
    then [lol,(- 1.2),"lulz",lol]
  else if something
    then [,]
  else
    somethingelse

def (splitHelp acc xs ys)
  if (null xs)
    then ((reverse acc), ys)
  else if (null (tail xs))
    then ((reverse acc), ys)
  else
    (splitHelp ((head ys) : acc)
               (tail (tail xs))
               (tail ys))

def (splitxs xs)
  (splitHelp [] xs xs)

def r def
  {
    a = 4
  }
    a

def main
  let {
    (f a) -> a
    unary = (print (testUnary 6))
    splitted = def {
                xs = (fst (splitxs [12,3,4,56]))
                } (xs ++ [0,9])
      }
      if False
        then superduper
      else
        (unary +
         fileLines +
         (print splitted))

def blah (3 / 4)


function setTtlMode(rhd::FPGA,mode)

    value = 0;
    value += mode[1] ? 1 : 0;
    value += mode[2] ? 2 : 0;
    value += mode[3] ? 4 : 0;
    value += mode[4] ? 8 : 0;
    value += mode[5] ? 16 : 0;
    value += mode[6] ? 32 : 0;
    value += mode[7] ? 64 : 0;
    value += mode[8] ? 128 : 0;

    SetWireInValue(rhd,WireInTtlOut,value,0x000000ff)
    UpdateWireIns(rhd)
end

function manual_trigger(rhd,trigger,triggerOn)
    #Trigger 0-7
    SetWireInValue(rhd,WireInManualTriggers, (triggerOn ? 1 : 0) << trigger,1<<trigger);
    UpdateWireIns(rhd)
end

clearTtlOut(rhd::FPGA)=(SetWireInValue(rhd,WireInTtlOut, 0x0000);UpdateWireIns(rhd))

function getTtlIn(rhd::FPGA,ttlInArray)

    UpdateWireOuts(rhd)
    ttlIn=GetWireOutValue(rhd,WireOutTtlIn)
    for i=1:16
        ttlInArray[i] = 0
        if (ttlIn & (1 << (i-1))) > 0
            ttlInArray[i] = 1
        end
    end

    ttlInArray
end

function setTtlOut(rhd::FPGA,ttlOutArray)

    ttlOut=Int32(0)
    for i=1:16
        if ttlOutArray[i]>0
            ttlOut += (1<< (i-1))
        end
    end

    SetWireInValue(rhd,WireInTtlOut,ttlOut)
    UpdateWireIns(rhd)
    nothing
end

function sendTimePulse(fpga::FPGA,val::Bool)

    if val==true
        fpga.ttloutput += (1 << (16-1))
    else
        fpga.ttloutput -= (1 << (16-1))
    end

    SetWireInValue(fpga,WireInTtlOut,fpga.ttloutput)
    UpdateWireIns(fpga)
    nothing
end

#=
Fast settle function allows TTL inputs to black amplifier channels
Can be used to prevent amplifiers from saturation (such as during stimulation)
=#
function enableExternalFastSettle(rhd::FPGA,enable)

    SetWireInValue(rhd,WireInMultiUse, (enable ? 1 : 0))
    UpdateWireIns(rhd)
    if (OPEN_EPHYS)
        ActivateTriggerIn(rhd,TrigInExtFastSettle,0)
    else
        ActivateTriggerIn(rhd,TrigInConfig,6)
    end

    nothing
end

function setExternalFastSettleChannel(rhd::FPGA,channel)

    #error checking goes here

    SetWireInValue(rhd,WireInMultiUse,channel)
    UpdateWireIns(rhd)
    if (OPEN_EPHYS)
        ActivateTriggerIn(rhd,TrigInExtFastSettle,1)
    else
        ActivateTriggerIn(rhd,TrigInConfig,7)
    end
end


module FixedEffectModels

##############################################################################
##
## Dependencies
##
##############################################################################
import Base: size, copyto!, getindex, length, fill!, eltype, length, view, adjoint
import LinearAlgebra: mul!, rmul!, norm, Matrix, Diagonal, cholesky, cholesky!, Symmetric, Hermitian, rank, dot, eigen, axpy!, svd, I, Adjoint, diag, qr
import LinearAlgebra.BLAS: gemm!
import Statistics: mean, quantile
import Printf: @sprintf
import Distributions: ccdf, TDist, FDist, Chisq
import DataFrames: DataFrame, AbstractDataFrame, completecases, names!, ismissing, groupby, groupindices
import Combinatorics: combinations
using CategoricalArrays
using FillArrays
import StatsBase: coef, nobs, coeftable, vcov, predict, residuals, var, RegressionModel, model_response, stderror, confint, fit, CoefTable, dof_residual, r2, adjr2, deviance, mss, rss, islinear, response
import StatsModels: @formula,  Formula, ModelFrame, ModelMatrix, Terms, coefnames, evalcontrasts, check_non_redundancy!
using Reexport
@reexport using StatsBase
@reexport using StatsModels
@reexport using FixedEffects

##############################################################################
##
## Exported methods and types
##
##############################################################################

export reg,
partial_out,
allvars,
fes,
WeightFormula,

AbstractRegressionResult,
title,
top,
RegressionResult,
RegressionResultIV,
RegressionResultFE,
RegressionResultFEIV,

AbstractVcovFormula,
VcovSimpleFormula,
VcovRobustFormula,
VcovClusterFormula,
VcovFormula,

AbstractVcovMethod,
VcovMethod,
VcovSimpleMethod,
VcovWhiteMethod,
VcovClusterMethod,

vcov!,
shat!,
VcovData,

Model,
@model

##############################################################################
##
## Load files
##
##############################################################################
include("utils/model.jl")
include("utils/weights.jl")
include("utils/fixedeffects.jl")
include("utils/basecol.jl")
include("utils/tss.jl")

include("formula/formula_iv.jl")
include("formula/formula_fe.jl")

include("RegressionResult.jl")

include("vcov/types.jl")
include("vcov/vcovsimple.jl")
include("vcov/vcovrobust.jl")
include("vcov/vcovcluster.jl")
include("vcov/utils.jl")

include("reg.jl")
include("partial_out.jl")

end  # module FixedEffectModels


import Geometry

"""
    struct JacobianMatrix{T} <: AbstractMatrix{T}
Jacobian matrix. Block-diagonal local matrix represents intraelement operations,
off-diagonal represents connectivity between elements.

`diag[:,iS1,:,iS2,iK]` is a square matrix d(f(u[:,iS1,iK))/d(u[:,iS2,iK])
for states iS1, iS2, element iK. Designed so that `diag[:,:,:,:,iK]` can be
in-place reshaped for matrix multiplication.

`offdiag[iFN2,iS2,iN1,iS1,iF,iK1]` is a rectangular matrix
d(f(u[iN1,iS1,iK1])/d(u[iFN2,iS2,mesh.e2e[iF,iK1]]), or effect of neighbor iF on
element iK1. Again, can in-place reshape for matrix multiplication. Note that
for upper triangular portion of global matrix (`mesh.e2e[iF,iK1]<iK1`),
the matrix `offdiag[:,:,:,:,iF,iK1]` actually yields the transpose of
operator needed. That is, the lower triangular part represents the effect of
neighbor on face nodes of iK1, and the upper triangular part represents the
transpose of the effect of neighbor's face nodes on iK1. This format therefore
only needs one array to store all off-diagonals.
"""
struct JacobianMatrix{T} <: AbstractMatrix{T}
    diag::Array{T,5}
    offdiag::Array{T,6}
    mesh::Geometry.Mesh

    function JacobianMatrix(diag::Array{T,5}, offdiag::Array{T,6}, mesh::Geometry.Mesh) where {T}
        @assert size(diag,1) == size(diag,3) "Diagonal of Jacobian must be square!"
        @assert size(diag,2) == size(diag,4) "Number of Jacobian states must be consistent!"
        @assert size(offdiag,5) == Geometry.N_FACES "Number of Jacobian neighbors does not match geometry!"
        new{T}(diag,offdiag,mesh)
    end
end

"""
    size(A::JacobianMatrix)
Returns a tuple containing the dimensions of actual operator A
"""
function Base.size(A::JacobianMatrix)
    (m,ns,n,ns2,ne) = size(A.diag)
    return (m*ns*ne, n*ns2*ne)
end
"""
    getindex(A::JacobianMatrix, i::Integer)
Linear scalar indexing into Jacobian matrix
"""
@inline function Base.getindex(A::JacobianMatrix, i::Integer)
    (sz1,sz2) = size(A)
    j = rem1(i, sz1); i = div(i, sz2, RoundUp)
    return getindex(A,i,j)
end
"""
    getindex(A::JacobianMatrix{T}, i::Integer, j::Integer)
2-dimensional scalar indexing into JacobianMatrix. Not efficient.
"""
@inline function Base.getindex(A::JacobianMatrix{T}, i::Integer, j::Integer) where {T}
    @boundscheck checkbounds(Bool, A,i,j)
    (m,ns,n,ns2,ne) = size(A.diag)
    im = rem1(i ,m ); i2 = div(i ,m , RoundUp)
    is = rem1(i2,ns); i3 = div(i2,ns, RoundUp)
    ie = rem1(i3,ne)
    jn = rem1(j ,n ); j2 = div(j ,n , RoundUp)
    js = rem1(j2,ns); j3 = div(j2,ns, RoundUp)
    je = rem1(j3,ne)
    if ie==je
        # On diagonal
        if checkbounds(Bool,A.diag, im,is,jn,js,ie)
            return @inbounds A.diag[im,is,jn,js,ie]
        end
    elseif je ∈ A.mesh.e2e[:,ie]
        iF = findfirst(x->x==je,A.mesh.e2e[:,ie])
        if je < ie
            # Lower triangular part stored directly
            iFN2 = findfirst(x->x==jn,A.mesh.ef2n[:,iF])
            iFN = 0
            if !isnothing(iFN2)
                # We need the original face node
                iFN = mesh.nfn2fn[iFN2]
            else
                # Node is not on the face => no Jacobian dependence
                return zero(T)
            end
            return A.offdiag[iFN,is,jn,js,iF,ie]
        else
            # Upper triangular part stored transposed
            iFN = findfirst(x->x==im,A.mesh.ef2n[:,iF])
            iFN2 = 0
            if !isnothing(iFN)
                # We need the neighbor's face node
                iFN2 = mesh.nfn2fn[iFN]
            else
                # Node is not on the face => no Jacobian dependence
                return zero(T)
            end
            return A.offdiag[iFN2,js,im,is,iF,ie]
        end
    else
        # Non-neighboring elements => no Jacobian dependence
        return zero(T)
    end
end
"""
    setindex!(A::JacobianMatrix, v, i::Integer)
Linear scalar indexed assignment of `JacobianMatrix`
"""
@inline function Base.setindex!(A::JacobianMatrix, v, i::Integer)
    (sz1,sz2) = size(A)
    j = rem1(i, sz1); i = div(i, sz2, RoundUp)
    setindex!(A, v, i, j)
end
"""
    setindex!(A::JacobianMatrix{T}, v, i::Integer, j::Integer)
2-dimensional scalar assignment of `JacobianMatrix`. Not efficient.
"""
@inline function Base.setindex!(A::JacobianMatrix{T}, v, i::Integer, j::Integer) where {T}
    @boundscheck checkbounds(Bool, A,i,j)
    (m,ns,n,ns2,ne) = size(A.diag)
    im = rem1(i ,m ); i2 = div(i ,m , RoundUp)
    is = rem1(i2,ns); i3 = div(i2,ns, RoundUp)
    ie = rem1(i3,ne)
    jn = rem1(j ,n ); j2 = div(j ,n , RoundUp)
    js = rem1(j2,ns); j3 = div(j2,ns, RoundUp)
    je = rem1(j3,ne)
    if ie==je
        # On diagonal
        if checkbounds(Bool,A.diag, im,is,jn,js,ie)
            @inbounds setindex!(A.diag, v, im,is,jn,js,ie)
        end
    elseif je ∈ A.mesh.e2e[:,ie]
        iF = findfirst(x->x==je,A.mesh.e2e[:,ie])
        if je < ie
            # Lower triangular part stored directly
            iFN2 = findfirst(x->x==jn,A.mesh.ef2n[:,iF])
            iFN = 0
            if !isnothing(iFN2)
                # We need the original face node
                iFN = mesh.nfn2fn[iFN2]
            else
                if !iszero(v)
                    # Node is not on the face => no Jacobian dependence
                    throw(ArgumentError("Cannot set off-diagonal entry ($i,$j) to non-zero value $v"))
                else
                    return v
                end
            end
            return setindex!(A.offdiag, v, iFN,is,jn,js,iF,ie)
        else
            # Upper triangular part stored transposed
            iFN = findfirst(x->x==im,A.mesh.ef2n[:,iF])
            iFN2 = 0
            if !isnothing(iFN)
                # We need the neighbor's face node
                iFN2 = mesh.nfn2fn[iFN]
            else
                if !iszero(v)
                    # Node is not on the face => no Jacobian dependence
                    throw(ArgumentError("Cannot set off-diagonal entry ($i,$j) to non-zero value $v"))
                else
                    return v
                end
            end
            return setindex!(A.offdiag, v, iFN2,js,im,is,iF,ie)
        end
    else
        # Non-neighboring elements => no Jacobian dependence
        if !iszero(v)
            throw(ArgumentError("Cannot set off-diagonal entry ($i,$j) to non-zero value $v"))
        else
            return v
        end
    end
end

(==)(A::JacobianMatrix, B::JacobianMatrix) = (A.diag == B.diag && A.offdiag == B.offdiag)
(-)(A::JacobianMatrix) = JacobianMatrix(-A.diag, -A.offdiag, A.mesh)
function (+)(A::JacobianMatrix, B::JacobianMatrix)
    @assert A.mesh === B.mesh
    JacobianMatrix(A.diag+B.diag, A.offdiag+B.offdiag, A.mesh)
end
function (-)(A::JacobianMatrix, B::JacobianMatrix)
    @assert A.mesh === B.mesh
    JacobianMatrix(A.diag-B.diag,A.offdiag-B.offdiag,mesh)
end
# TODO: Matrix addition with LocalMatrix as well?
(*)(c::Number, A::JacobianMatrix) = JacobianMatrix(c*A.diag, c*A.offdiag, A.mesh)
(*)(A::JacobianMatrix, c::Number) = JacobianMatrix(c*A.diag, c*A.offdiag, A.mesh)
(/)(A::JacobianMatrix, c::Number) = JacobianMatrix(A.diag / c, A.offdiag / c, A.mesh)
(*)(A::JacobianMatrix, B::JacobianMatrix) = error("Too expensive an operation!")

" Store `faces` on `volume` array according to given mask: vol[mask] += face "
function distribute_face_to_nodes!(volume, faces, mask)
    volume_shaped = reshape(volume, size(volume,1), prod(size(volume)[2:end])))
    faces_shaped = reshape(faces, size(faces,1), prod(size(faces)[2:end]))
    for j = 1:size(volume_shaped,2)
        for i = 1:size(mask,1)
            volume_shaped[mask[i],j] += faces_shaped[i,j]
        end
    end
end
" Collect `volume` into `faces` array according to given mask: face = vol[mask] "
function collect_face_from_nodes!(faces, volume, mask)
    volume_shaped = reshape(volume, size(volume,1), prod(size(volume)[2:end])))
    faces_shaped = reshape(faces, size(faces,1), prod(size(faces)[2:end]))
    for j = 1:size(volume_shaped,2)
        for i = 1:size(mask,1)
            faces_shaped[i,j] = volume_shaped[mask[i],j]
        end
    end
end

function (*)(A::JacobianMatrix{T1}, x::AbstractVector{T2}) where {T1, T2}
    (m,ns,n,ns2,ne) = size(A.diag)
    b = zeros(typeof(zero(T1)*zero(T2)), m*ns*ne)
    b_shaped = reshape(b, m,ns,ne)
    b_mult   = reshape(b, m*ns,ne)
    x_shaped = reshape(x, n,ns2,ne)
    x_mult   = reshape(x, n*ns2,ne)
    faces_mult = zeros(typeof(zero(T1)*zero(T2)), A.mesh.n_face_nodes*ns)
    faces_shaped = reshape(faces, A.mesh.n_face_nodes, ns)
    for iK = 1:ne
        # Diagonal of A
        @views Adiag = reshape(A.diag[:,:,:,:,iK], m*ns, n*ns2)
        @views b_mult[:,iK] .+= Adiag*x_mult[:,iK]
        # Off-diagonal of A
        for iF = 1:Geometry.N_FACES
            nK = A.mesh.e2e[iF,iK]
            if nK > 0
                Aoff = reshape(@view A.offdiag[:,:,:,:,iF,iK], A.mesh.n_faces_nodes*ns,m*ns)
                nF = A.mesh.e2f[iF,iK]
                if nK < iK
                    # Lower triangular part stored directly
                    @views faces_mult .= Aoff*x_mult[:,iK]
                    @views distribute_face_to_nodes!(b_shaped[:,:,iK], faces_shaped, A.mesh.ef2n[:,iF])
                else
                    # Upper triangular part stored transposed
                    @views collect_face_from_nodes!(faces_shaped, x_shaped[:,:,nK], A.mesh.ef2n[:,nF])
                    @views b_mult[:,iK] .+= Aoff'*faces_mult
                end
            end
        end
    end
    return b
end
function (*)(A::JacobianMatrix{T1}, x::SolutionVector{T2}) where {T1, T2}
    (m,ns,n,ns2,ne) = size(A.diag)
    b_shaped = zeros(typeof(zero(T1)*zero(T2)), m,ns,ne)
    b_mult = reshape(b, m*ns,ne)
    x_mult = reshape(x.data, n*ns2,ne)
    faces_mult = zeros(typeof(zero(T1)*zero(T2)), A.mesh.n_face_nodes*ns)
    faces_shaped = reshape(faces, A.mesh.n_face_nodes, ns)
    @views for iK = 1:ne
        # Diagonal of A
        Adiag = reshape(A.diag[:,:,:,:,iK], m*ns, n*ns2)
        b_mult[:,iK] .+= Adiag*x_mult[:,iK]
        # Off-diagonal of A
        for iF = 1:Geometry.N_FACES
            nK = A.mesh.e2e[iF,iK]
            if nK > 0
                Aoff = reshape(A.offdiag[:,:,:,:,iF,iK], A.mesh.n_faces_nodes*ns,m*ns)
                nF = A.mesh.e2f[iF,iK]
                if nK < iK
                    # Lower triangular part stored directly
                    faces_mult .= Aoff*x_mult[:,iK]
                    distribute_face_to_nodes!(b_shaped[:,:,iK], faces_shaped, A.mesh.ef2n[:,iF])
                else
                    # Upper triangular part stored transposed
                    collect_face_from_nodes!(faces_shaped, x.data[:,:,nK], A.mesh.ef2n[:,nF])
                    b_mult[:,iK] .+= Aoff'*faces_mult
                end
            end
        end
    end
    return SolutionVector(b_shaped)
end
"""
    mul!(y::SolutionVector, A::JacobianMatrix, x::SolutionVector, α::Number, β::Number)
Return A x α + y β by overwriting y
"""
function LinearAlgebra.mul!(y::SolutionVector{T1}, A::JacobianMatrix, x::SolutionVector, α::Number, β::Number) where {T1}
    (m,ns,n,ns2,ne) = size(A.diag)
    (nx,ns2x,nex) = true_size(x)
    (my,nsy,ney) = true_size(y)
    @assert m == my && ns  == nsy  && ne == ney "JacobianMatrix $(size(A.diag)), SolutionVector $(true_size(y)) mismatched size"
    @assert n == nx && ns2 == ns2x && ne == nex "JacobianMatrix $(size(A.diag)), SolutionVector $(true_size(x)) mismatched size"
    y_shaped = reshape(y.data, m*ns*ne)
    y_mult   = reshape(y.data, m*ns,ne)
    x_shaped = reshape(x.data, n,ns2,ne)
    x_mult   = reshape(x.data, n*ns2,ne)
    faces_mult = zeros(T1, A.mesh.n_face_nodes*ns)
    faces_shaped = reshape(faces, A.mesh.n_face_nodes, ns)
    @views for iK = 1:ne
        # Diagonal of A
        Adiag = reshape(A.diag[:,:,:,:,iK], m*ns, n*ns2)
        # y[iK] = Adiag*x[iK]*α + y[iK]*β
        mul!(y_mult[:,iK], Adiag, x_mult[:,iK], α, β)
        # Off-diagonal of A
        for iF = 1:Geometry.N_FACES
            nK = A.mesh.e2e[iF,iK]
            if nK > 0
                Aoff = reshape(A.offdiag[:,:,:,:,iF,iK], A.mesh.n_faces_nodes*ns,m*ns)
                nF = A.mesh.e2f[iF,iK]
                if nK < iK
                    # Lower triangular part stored directly
                    faces_mult .= Aoff*x_mult[:,iK]
                    # y[iK] += Aoff*faces[iK]*α
                    distribute_face_to_nodes!(y_shaped[:,:,iK], faces_shaped, A.mesh.ef2n[:,iF])
                else
                    # Upper triangular part stored transposed
                    collect_face_from_nodes!(faces_shaped, x_shaped[:,:,nK], A.mesh.ef2n[:,nF])
                    # y[iK] += Aoff'*faces[nK]*α
                    mul!(y_mult[:,iK], Aoff', faces_mult, α, 1.0)
                end
            end
        end
    end
end

"""
    ldiv!(A::JacobianMatrix, x::SolutionVector)
In-place linear solve A\\x using numerical iterative solver based on
preconditioned GMRES.
TODO: Decide actual algorithm and create production-level version of code.
"""
function LinearAlgebra.ldiv!(A::JacobianMatrix, x::SolutionVector)
    println("my ldiv!")
    (m,ns,n,ns2,ne) = size(A.diag)
    (nx,ns2x,nex) = true_size(x)
    @assert nx == m == n && ns2x == ns2 == ns && ne == nex "JacobianMatrix A $(size(A.diag)) cannot \\ x $(true_size(x))"
    # TODO: program the solver here
    error("Not yet programmed!")
end
function \(A::LocalMatrix, x::SolutionVector)
    println("my \\")
    (m,ns,n,ns2,ne) = size(A.diag)
    (nx,ns2x,nex) = true_size(x)
    @assert nx == m == n && ns2x == ns2 == ns && ne == nex "JacobianMatrix A $(size(A.diag)) cannot \\ x $(true_size(x))"
    b = deepcopy(x)
    LinearAlgebra.ldiv!(A, b)
end

function L2_error!(u::SolutionVector, u_true::SolutionVector, M::JacobianMatrix)
    u_true -= u
    return L2_norm(u_true, M)
end
function L2_norm(u::SolutionVector, M::JacobianMatrix)
    Mu = M*u
    (n,ns,ne) = true_size(Mu)
    l2_norms = zeros(ns)
    for iK = 1:ne
        for iS = 1:ns
            @views l2_norms[iS] += u.data[:,iS,iK]'*Mu.data[:,iS,iK]
        end
    end
    return l2_norms
end


using CLIMA
using CLIMA.Mesh.Topologies: StackedCubedSphereTopology, cubedshellwarp, grid1d
using CLIMA.Mesh.Grids: DiscontinuousSpectralElementGrid
using CLIMA.Mesh.Filters
using CLIMA.DGmethods: DGModel, init_ode_state, VerticalDirection
using CLIMA.DGmethods.NumericalFluxes: Rusanov, CentralNumericalFluxGradient,
                                       CentralNumericalFluxDiffusive
using CLIMA.ODESolvers
using CLIMA.GeneralizedMinimalResidualSolver
using CLIMA.ColumnwiseLUSolver: ManyColumnLU
using CLIMA.VTK: writevtk, writepvtu
using CLIMA.GenericCallbacks: EveryXWallTimeSeconds, EveryXSimulationSteps
using CLIMA.PlanetParameters: planet_radius, day
using CLIMA.MoistThermodynamics: air_density, soundspeed_air, internal_energy
using CLIMA.Atmos: AtmosModel, SphericalOrientation,
                   DryModel, NoPrecipitation, NoRadiation, NoFluxBC,
                   ConstantViscosityWithDivergence,
                   vars_state, vars_aux,
                   Gravity, HydrostaticState, IsothermalProfile,
                   AtmosAcousticGravityLinearModel, AtmosLESConfiguration
using CLIMA.VariableTemplates: flattenednames

using MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test

const output_vtk = false

function main()
  CLIMA.init()
  ArrayType = CLIMA.array_type()

  mpicomm = MPI.COMM_WORLD
  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = Dict("DEBUG" => Logging.Debug,
                  "WARN"  => Logging.Warn,
                  "ERROR" => Logging.Error,
                  "INFO"  => Logging.Info)[ll]

  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))

  polynomialorder = 5
  numelem_horz = 10
  numelem_vert = 5

  timeend = 60 * 60
  # timeend = 33 * 60 * 60 # Full simulation
  outputtime = 60 * 60

  expected_result = Dict()
  expected_result[Float32] = 9.5064378310656000e+13
  expected_result[Float64] = 9.5073452847081828e+13

  for FT in (Float32, Float64)
    result = run(mpicomm, polynomialorder, numelem_horz, numelem_vert,
                 timeend, outputtime, ArrayType, FT)
    @test result ≈ expected_result[FT]
  end
end

function run(mpicomm, polynomialorder, numelem_horz, numelem_vert,
             timeend, outputtime, ArrayType, FT)

  setup = AcousticWaveSetup{FT}()

  vert_range = grid1d(FT(planet_radius), FT(planet_radius + setup.domain_height), nelem = numelem_vert)
  topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)

  grid = DiscontinuousSpectralElementGrid(topology,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = polynomialorder,
                                          meshwarp = cubedshellwarp)

  model = AtmosModel{FT}(AtmosLESConfiguration;
                         orientation=SphericalOrientation(),
                           ref_state=HydrostaticState(IsothermalProfile(setup.T_ref), FT(0)),
                          turbulence=ConstantViscosityWithDivergence(FT(0)),
                            moisture=DryModel(),
                              source=Gravity(),
                          init_state=setup)
  linearmodel = AtmosAcousticGravityLinearModel(model)

  dg = DGModel(model, grid, Rusanov(),
               CentralNumericalFluxDiffusive(), CentralNumericalFluxGradient())

  lineardg = DGModel(linearmodel, grid, Rusanov(),
                     CentralNumericalFluxDiffusive(),
                     CentralNumericalFluxGradient();
                     direction=VerticalDirection(),
                     auxstate=dg.auxstate)

  # determine the time step
  element_size = (setup.domain_height / numelem_vert)
  acoustic_speed = soundspeed_air(FT(setup.T_ref))
  dt_factor = 445
  dt = dt_factor * element_size / acoustic_speed / polynomialorder ^ 2
  # Adjust the time step so we exactly hit 1 hour for VTK output
  dt = 60 * 60 / ceil(60 * 60 / dt)
  nsteps = ceil(Int, timeend / dt)

  Q = init_ode_state(dg, FT(0))

  linearsolver = ManyColumnLU()

  odesolver = ARK2GiraldoKellyConstantinescu(dg, lineardg, linearsolver, Q;
                                             dt = dt, t0 = 0, split_nonlinear_linear=false)

  filterorder = 18
  filter = ExponentialFilter(grid, 0, filterorder)
  cbfilter = EveryXSimulationSteps(1) do
    Filters.apply!(Q, 1:size(Q, 2), grid, filter, VerticalDirection())
    nothing
  end

  eng0 = norm(Q)
  @info @sprintf """Starting
                    ArrayType       = %s
                    FT              = %s
                    polynomialorder = %d
                    numelem_horz    = %d
                    numelem_vert    = %d
                    dt              = %.16e
                    norm(Q₀)        = %.16e
                    """ "$ArrayType" "$FT" polynomialorder numelem_horz numelem_vert dt eng0

  # Set up the information callback
  starttime = Ref(now())
  cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s=false)
    if s
      starttime[] = now()
    else
      energy = norm(Q)
      runtime = Dates.format(convert(DateTime, now() - starttime[]), dateformat"HH:MM:SS")
      @info @sprintf """Update
                        simtime = %.16e
                        runtime = %s
                        norm(Q) = %.16e
                        """ gettime(odesolver) runtime energy
    end
  end
  callbacks = (cbinfo, cbfilter)

  if output_vtk
    # create vtk dir
    vtkdir = "vtk_acousticwave" *
      "_poly$(polynomialorder)_horz$(numelem_horz)_vert$(numelem_vert)" *
      "_dt$(dt_factor)x_$(ArrayType)_$(FT)"
    mkpath(vtkdir)

    vtkstep = 0
    # output initial step
    do_output(mpicomm, vtkdir, vtkstep, dg, Q, model)

    # setup the output callback
    cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do
      vtkstep += 1
      Qe = init_ode_state(dg, gettime(odesolver))
      do_output(mpicomm, vtkdir, vtkstep, dg, Q, model)
    end
    callbacks = (callbacks..., cbvtk)
  end

  solve!(Q, odesolver; numberofsteps=nsteps, adjustfinalstep=false, callbacks=callbacks)

  # final statistics
  engf = norm(Q)
  @info @sprintf """Finished
  norm(Q)                 = %.16e
  norm(Q) / norm(Q₀)      = %.16e
  norm(Q) - norm(Q₀)      = %.16e
  """ engf engf/eng0 engf-eng0
  engf
end

Base.@kwdef struct AcousticWaveSetup{FT}
  domain_height::FT = 10e3
  T_ref::FT = 300
  α::FT = 3
  γ::FT = 100
  nv::Int = 1
end

function (setup::AcousticWaveSetup)(bl, state, aux, coords, t)
  # callable to set initial conditions
  FT = eltype(state)

  r = norm(coords, 2)
  @inbounds λ = atan(coords[2], coords[1])
  @inbounds φ = asin(coords[3] / r)
  h = r - FT(planet_radius)

  β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))
  f = (1 + cos(FT(π) * β)) / 2
  g = sin(setup.nv * FT(π) * h / setup.domain_height)
  Δp = setup.γ * f * g
  p = aux.ref_state.p + Δp

  state.ρ = air_density(setup.T_ref, p)
  state.ρu = SVector{3, FT}(0, 0, 0)
  state.ρe = state.ρ * (internal_energy(setup.T_ref) + aux.orientation.Φ)
  nothing
end

function do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, testname = "acousticwave")
  ## name of the file that this MPI rank will write
  filename = @sprintf("%s/%s_mpirank%04d_step%04d",
                      vtkdir, testname, MPI.Comm_rank(mpicomm), vtkstep)

  statenames = flattenednames(vars_state(model, eltype(Q)))
  auxnames = flattenednames(vars_aux(model, eltype(Q)))
  writevtk(filename, Q, dg, statenames, dg.auxstate, auxnames)

  ## Generate the pvtu file for these vtk files
  if MPI.Comm_rank(mpicomm) == 0
    ## name of the pvtu file
    pvtuprefix = @sprintf("%s/%s_step%04d", vtkdir, testname, vtkstep)

    ## name of each of the ranks vtk files
    prefixes = ntuple(MPI.Comm_size(mpicomm)) do i
      @sprintf("%s_mpirank%04d_step%04d", testname, i - 1, vtkstep)
    end

    writepvtu(pvtuprefix, prefixes, (statenames..., auxnames...))

    @info "Done writing VTK: $pvtuprefix"
  end
end

main()


#__precompile__()

module EvoDuplexes

using BioSymbols
using BioSequences
using GenomicFeatures
using IntervalTrees
using SuffixArrays
using Gadfly
using Automa

using BufferedStreams
using Libz

using PyCall
using PyPlot

using ScikitLearn
using ScikitLearn.Utils: meshgrid

@sk_import ensemble: IsolationForest

@pyimport matplotlib.font_manager as fm
@pyimport scipy.stats as stats

import Automa
import Automa.RegExp: @re_str
import Compat: take!

importall BioSymbols
importall GenomicFeatures

include("../src/pairs.jl")
include("../src/energy.jl")
include("../src/rnaduplex.jl")
include("../src/intervals.jl")
include("../src/traverse.jl")
include("../src/mafreader.jl")
include("../src/gtrmodel.jl")
include("../src/newick.jl")
include("../src/suffix.jl")
include("../src/evoduplex.jl")
include("../src/io.jl")
include("../src/regions.jl")
include("../src/train.jl")

export RNADuplex,
       EvoDuplex,
       MAFSpecies,
       MAFRecord,
       MAFReader,
       reverse_complement!,
       readmaf!,
       read!,
       done,
       RNASuffixArray,
       RNADuplexArray,
       DuplexCollection,
       DuplexInterval,
       energy,
       npairs,
       traverse,
       collect,
       stitch,
       covariance,
       distance,
       score!,
       score,
       dinucleotide_entropy,
       extend_branches!,
       set_prob_mat!,
       parsenewick,
       GTR_SINGLE_Q,
       GTR_PAIRED_Q,
       isgzipped,
       randcode,
       brackets,
       loadbed,
       writebed,
       loadbedgraph,
       DistanceForest,
       train!,
       predict
end


using MLUtils
using MLUtils.Datasets
using MLUtils: RingBuffer, eachobsparallel
using SparseArrays
using Random, Statistics
using Test
using Transducers
using FoldsThreads: TaskPoolEx
using ChainRulesTestUtils: test_rrule
using Zygote: ZygoteRuleConfig
using ChainRulesCore: rrule_via_ad

showcompact(io, x) = show(IOContext(io, :compact => true), x)

# --------------------------------------------------------------------
# create some test data
Random.seed!(1335)
const X = rand(4, 15)
const y = repeat(["setosa","versicolor","virginica"], inner = 5)
const Y = permutedims(hcat(y,y), [2,1])
const Yt = hcat(y,y)
const yt = Y[1:1,:]
const Xv = view(X,:,:)
const yv = view(y,:)
const XX = rand(20,30,15)
const XXX = rand(3,20,30,15)
const vars = (X, Xv, yv, XX, XXX, y)
const tuples = ((X,y), (X,Y), (XX,X,y), (XXX,XX,X,y))
const Xs = sprand(10, 15, 0.5)
const ys = sprand(15, 0.5)
const X1 = hcat((1:15 for i = 1:10)...)'
const Y1 = collect(1:15)

struct EmptyType end

struct FallbackType end
Base.getindex(::FallbackType, i) = 1234
Base.length(::FallbackType) = 5678

struct CustomType end
MLUtils.numobs(::CustomType) = 15
MLUtils.getobs(::CustomType, i::Int) = i
MLUtils.getobs(::CustomType, i::AbstractVector) = collect(i)

struct CustomArray{T,N} <: AbstractArray{T,N}
    length::Int
    size::NTuple{N, Int}

    function CustomArray{T}(n...) where T
        N = length(n)
        new{T,N}(prod(n), Tuple(n))
    end
end
CustomArray(n...) = CustomArray{Float32}(n...)
Base.length(x::CustomArray) = x.length
Base.size(x::CustomArray) = x.size
Base.getindex(x::CustomArray{T}, i::Int) where T = T(1)
Base.getindex(x::CustomArray{T, 1}, i::AbstractVector{Int}) where T = CustomArray{T}(length(i))

struct CustomArrayNoView{T,N} <: AbstractArray{T,N}
    length::Int
    size::NTuple{N, Int}

    function CustomArrayNoView{T}(n...) where T
        N = length(n)
        new{T,N}(prod(n), Tuple(n))
    end
end
CustomArrayNoView(n...) = CustomArrayNoView{Float32}(n...)
Base.length(x::CustomArrayNoView) = x.length
Base.size(x::CustomArrayNoView) = x.size
Base.getindex(x::CustomArrayNoView{T}, i::Int) where T = T(1)
Base.getindex(x::CustomArrayNoView{T, 1}, i::AbstractVector{Int}) where T = CustomArrayNoView{T}(length(i))
Base.view(x::CustomArrayNoView, i...) = error("view not allowed")

struct CustomRangeIndex
    n::Int
end
Base.length(r::CustomRangeIndex) = r.n
Base.getindex(r::CustomRangeIndex, idx::Int) = idx
Base.getindex(r::CustomRangeIndex, idxs::UnitRange) = idxs

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

include("test_utils.jl")

# @testset "MLUtils.jl" begin

@testset "batchview" begin; include("batchview.jl"); end
@testset "eachobs" begin; include("eachobs.jl"); end
@testset "dataloader" begin; include("dataloader.jl"); end
@testset "folds" begin; include("folds.jl"); end
@testset "observation" begin; include("observation.jl"); end
@testset "obsview" begin; include("obsview.jl"); end
@testset "obstransform" begin; include("obstransform.jl"); end
@testset "randobs" begin; include("randobs.jl"); end
@testset "resample" begin; include("resample.jl"); end
@testset "splitobs" begin; include("splitobs.jl"); end
@testset "utils" begin; include("utils.jl"); end
@testset "eachobsparallel" begin; include("parallel.jl"); end

@testset "Datasets/datasets" begin; include("Datasets/datasets.jl"); end
@testset "Datasets/generators" begin; include("Datasets/generators.jl"); end

# end


@testset "Independent Bounds" begin
    m = BabyPOMDP()
    rs = MemorizingSource(1,1,MersenneTwister(37))
    sb = ScenarioBelief([1=>true], rs, 0, false)
    b = IndependentBounds(0.0, -1e-5)
    @test bounds(b, m, sb) == (0.0, -1e-5)
    b = IndependentBounds(0.0, -1e-5, consistency_fix_thresh=1e-5)
    @test bounds(b, m, sb) == (0.0, 0.0)
    b = IndependentBounds(0.0, -1e-4, consistency_fix_thresh=1e-5)
    @test bounds(b, m, sb) == (0.0, -1e-4)
end


"""
This function returns all the subtypes at the leaf nodes (i.e. which do not have any further subtypes) of a given type.
"""
function leaftypes(t::Any, level=1, leaves = Vector{Any}())
    for s in InteractiveUtils.subtypes(t)
        if length(InteractiveUtils.subtypes(s)) == 0
            push!(leaves, s)
        else
            leaftypes(s, level+1, leaves)
        end
    end
    return leaves
end

"""
This function gets the capacity market forward years for the simulation.
"""
function get_capacity_forward_years(sim::Union{AgentSimulation, AgentSimulationData})
    forward_years = 1
    if get_markets(sim)[:Capacity]
        capacity_mkt_params = DataFrames.DataFrame(CSV.File(joinpath(get_data_dir(get_case(sim)), "markets_data", "capacity_mkt_param.csv")))
        forward_years = capacity_mkt_params.forward_years[1]
    end
    return forward_years
end

"""
This function returns the size of the project (in MW)
based on the project type and size type.
"""
function size_in_MW(investor_dir::String,
                    projecttype::String,
                    sizetype::String)
    filename = joinpath(investor_dir, "sizedict.csv")
    sizedata = read_data(filename)

    size = sizedata[findfirst(x-> x == projecttype, sizedata.type),
                findfirst(x -> x == sizetype, names(sizedata))]
    return size
end

"""
This function returns the size of the project (in MW)
if already specified in MW.
"""
function size_in_MW(investor_dir::String,
                    projecttype::String,
                    size::Float64)
    return size
end

"""
This function returns the size of the PSY Device.
"""
get_device_size(device::PSY.ThermalStandard) = PSY.get_active_power_limits(device)[:max]
get_device_size(device::PSY.RenewableDispatch) = PSY.get_rating(device)
get_device_size(device::PSY.HydroEnergyReservoir) = PSY.get_active_power_limits(device)[:max]
get_device_size(device::PSY.GenericBattery) = PSY.get_rating(device)

"""
This function returns line rating for PSY Lines.
"""
get_line_rating(line::PSY.Line) = PSY.get_rate(line)
get_line_rating(line::PSY.HVDCLine) = line.active_power_limits_from[:max]

"""
This function removes leap days from the data if year is not leap year.
"""
function remove_leap_day!(df:: DataFrames.DataFrame, start_year::Int64)
    if start_year % 4 != 0
        filter!(row -> row["Month"] != 2 || row["Day"] != 29, df)
    end

    return
end

"""
This function finds the start and end years
for reading price data.
"""
function find_price_years(construction_year::Int64,
                          end_life_year::Int64,
                          iteration_year::Int64,
                          yearly_horizon::Int64)

    price_years = (start_year = max((construction_year -
                                    iteration_year + 1),
                                    1),
                   end_year = min((end_life_year -
                                  iteration_year + 1),
                                  (yearly_horizon)))

    return price_years
end

"""
This function finds the start and end years
for updating project profit and npv data.
"""
function find_update_years(construction_year::Int64,
                          end_life_year::Int64,
                          iteration_year::Int64,
                          yearly_horizon::Int64)

    update_years = (start_year = max(construction_year,
                                     iteration_year),
                    end_year = min(end_life_year,
                                   yearly_horizon + iteration_year - 1))

    return update_years
end

"""
This function returns the number of hours modeled in the market clearing problem.
"""
function get_num_hours(system::MarketClearingProblem{Z, T}) where {Z, T}
    num_hours = T

    return num_hours
end

"""
This function returns an O&M cost array filled with the fixed
annual O&M cost values from update year to the end life year.
"""
function create_OMcost_array(finance::Finance,
                            scenario_name::String,
                            project_life_end::Int64,
                            update_start_year::Int64,
                            iteration_year::Int64)

    OM_cost_array = zeros(size(get_scenario_profit(finance)[scenario_name][iteration_year], 2))

    for i in update_start_year:project_life_end
        OM_cost_array[i] = get_fixed_OM_cost(finance)
    end

    return OM_cost_array

end

"""
This function return a queue cost array filled with the queue cost values
corresponding to each year spent in the queue.
"""
function create_queue_cost_array(size_profit::Int64,
         decision_year::Int64,
         queue_time::Int64,
         queue_cost::Vector{Float64})

    queue_cost_array = zeros(size_profit)

    for i in decision_year:(decision_year + queue_time - 1)
        queue_cost_array[i] = queue_cost[i - decision_year + 1]
    end

    return queue_cost_array

end
"""
This function does nothing if the project is not an Option project.
"""
function update_lifecycle!(project::P,
                           iteration_year::Int64,
                           simulation_years::Int64) where P <: Project{<: BuildPhase}
    return
end

"""
This function updates the decision, construction and endlife years of Option projects.
"""
function update_lifecycle!(project::P,
                           iteration_year::Int64,
                           simulation_years::Int64) where P <: Project{Option}

        finance_data = get_finance_data(project)

        queue_time = length(get_queue_cost(finance_data))

        if queue_time + iteration_year <  simulation_years
            set_decision_year!(project, (iteration_year + 1))
            set_construction_year!(project, (get_decision_year(project) +
                                            queue_time +
                                            get_lag_time(finance_data)))

            end_life_year = get_construction_year(project) + get_life_time(finance_data)  - 1
            set_end_life_year!(project, end_life_year)

            set_effective_investment_cost!(project, get_investment_cost(finance_data)[queue_time + iteration_year + 1])

            products = get_products(project)
            finance_data.realized_profit = AxisArrays.AxisArray(zeros(length(products), end_life_year),
                                                    AxisArrays.Axis{:prod}(get_name.(products)),
                                                    AxisArrays.Axis{:year}(1:1:end_life_year))

            finance_data.annual_cashflow = zeros(end_life_year)

            scenario_profit = get_scenario_profit(finance_data)

            for scenario_name in keys(scenario_profit)
                rows = size(scenario_profit[scenario_name][iteration_year], 1)
                cols = get_end_life_year(project)

                set_expected_npv!(finance_data, zeros(cols))
                set_expected_utility!(finance_data, zeros(cols))

                finance_data.scenario_profit[scenario_name] = [AxisArrays.AxisArray(zeros(length(products), cols),
                                                     AxisArrays.Axis{:prod}(get_name.(products)),
                                                     AxisArrays.Axis{:year}(collect(1:1:cols)))
                                                     for y in 1:cols]
                set_scenario_npv!(finance_data, scenario_name, zeros(cols))
                set_scenario_utility!(finance_data, scenario_name, zeros(cols))
            end
        end

    return
end

"""
This function updates the total installed capacity each year.
"""
function update_installed_cap!(installed_capacity::Vector{Float64},
                               announced_projects::Vector{<: Project{<: BuildPhase}},
                               iteration_year::Int64,
                               simulation_years::Int64)

    for year in iteration_year:simulation_years
        installed_capacity[year] = 0.0
        for project in announced_projects
            if (get_construction_year(project) <= year) && (get_end_life_year(project) >= year)
                installed_capacity[year] += get_maxcap(project)
            end
        end
    end
    return installed_capacity
end

"""
This function does nothing if project is not in Existing phase.
"""
function extrapolate_profits!(project::Project{<:BuildPhase}, final_year::Int64)
    return
end

"""
This function extrapolates the profits of Existing projects beyond the simulation horizon.
"""
function extrapolate_profits!(project::Project{Existing}, final_year::Int64)
    end_life_year = get_end_life_year(project)
    finance_data = get_finance_data(project)
    if end_life_year > final_year
        for y in (final_year + 1):end_life_year
            set_annual_cashflow!(finance_data, y, get_annual_cashflow(finance_data)[final_year])
            for product in get_products(project)
                name = get_name(product)
                set_realized_profit!(finance_data, name, y, get_realized_profit(finance_data)[name, final_year])
            end
        end
    end
    return
end


using Random

abstract type Modifier end

struct Modifiers <: Modifier
    modifiers::Array{Modifier}
end

Modifiers(modifiers::Modifier ...) = Modifiers([modifiers...])

function apply!(modifiers::Modifiers, data::DataFrame, event::Event)
    for modifier in modifiers.modifiers
        apply!(modifier, data, event)
    end

    return nothing
end

struct NoiseModifier <: Modifier
    seed::Random.MersenneTwister
    attribute::Float64 # The fraction of attribute values to disturb
end

NoiseModifier(attribute::Float64, seed::Int) = NoiseModifier(Random.seed!(seed), attribute)
NoiseModifier(attribute::Float64) = NoiseModifier(Random.default_rng(), attribute)

function apply!(modifier::NoiseModifier, data::DataFrame, event::Event)
    return nothing
end

struct FilterModifier <: Modifier
    columns::Array{Symbol}

    function FilterModifier(columns)
        _columns = unique(columns)
        if length(_columns) != length(columns) @warn "There are duplicate columns." end
        return new(_columns)
    end
end

FilterModifier(columns::Symbol...) = FilterModifier([columns...])

function apply!(modifier::FilterModifier, data::DataFrame, event::Event)
    columns = Symbol[]
    for col in modifier.columns
        if !(col in propertynames(data))
            throw(ArgumentError("stream doesn't have the column $col"))
        else
            push!(columns, col)
        end
    end

    select!(data, columns)
    return nothing
end

struct AlterDataModifier <: Modifier
    alter!::Function
end

function apply!(modifier::AlterDataModifier, data::DataFrame, event::Event)
    modifier.alter!(data, event)

    return nothing
end


using Test, MPI

using ClimateMachine
using ClimateMachine.MPIStateArrays
using ClimateMachine.VariableTemplates: @vars, varsindex
using StaticArrays

ClimateMachine.init()
const ArrayType = ClimateMachine.array_type()
const mpicomm = MPI.COMM_WORLD

const V = @vars begin
    a::Float32
    b::SVector{3, Float32}
    c::SMatrix{3, 8, Float32}
    d::Float32
end

@testset "MPIStateArray varsindex" begin
    Q = MPIStateArray{Float32, V}(mpicomm, ArrayType, 4, 29, 8)
    @test Q.a === view(MPIStateArrays.realview(Q), :, 1:1, :)
    @test Q.b === view(MPIStateArrays.realview(Q), :, 2:4, :)
    @test Q.c === view(MPIStateArrays.realview(Q), :, 5:28, :)
    @test Q.d === view(MPIStateArrays.realview(Q), :, 29:29, :)

    P = similar(Q)
    @test P.a === view(MPIStateArrays.realview(P), :, 1:1, :)
    @test P.b === view(MPIStateArrays.realview(P), :, 2:4, :)
    @test P.c === view(MPIStateArrays.realview(P), :, 5:28, :)
    @test P.d === view(MPIStateArrays.realview(P), :, 29:29, :)

    A = MPIStateArray{Float32}(mpicomm, ArrayType, 4, 29, 8)
    @test_throws ErrorException A.a
end


## Functions to compute the reduced shape

# for reductions that expand 0 dims to 1
reduced_dims(a::AbstractArray, region) = reduced_dims(size(a), region)

# for reductions that keep 0 dims as 0
reduced_dims0(a::AbstractArray, region) = reduced_dims0(size(a), region)
reduced_dims{N}(siz::NTuple{N,Int}, d::Int, rd::Int) = (d == 1 ? tuple(rd, siz[d+1:N]...) :
                                                        d == N ? tuple(siz[1:N-1]..., rd) :
                                                        1 < d < N ? tuple(siz[1:d-1]..., rd, siz[d+1:N]...) :
                                                        siz)::typeof(siz)
reduced_dims{N}(siz::NTuple{N,Int}, d::Int) = reduced_dims(siz, d, 1)
reduced_dims0{N}(siz::NTuple{N,Int}, d::Int) = 1 <= d <= N ? reduced_dims(siz, d, (siz[d] == 0 ? 0 : 1)) : siz

function reduced_dims{N}(siz::NTuple{N,Int}, region)
    rsiz = [siz...]
    for i in region
        if 1 <= i <= N
            rsiz[i] = 1
        end
    end
    tuple(rsiz...)::typeof(siz)
end

function reduced_dims0{N}(siz::NTuple{N,Int}, region)
    rsiz = [siz...]
    for i in region
        if i <= i <= N
            rsiz[i] = (rsiz[i] == 0 ? 0 : 1)
        end
    end
    tuple(rsiz...)::typeof(siz)
end

function regionsize(a, region)
    s = 1
    for d in region
        s *= size(a,d)
    end
    s
end

###### Generic reduction functions #####

## initialization

for (Op, initfun) in ((:AddFun, :zero), (:MulFun, :one), (:MaxFun, :typemin), (:MinFun, :typemax))
    @eval initarray!{T}(a::AbstractArray{T}, ::$(Op), init::Bool) = (init && fill!(a, $(initfun)(T)); a)
end

for (Op, initval) in ((:AndFun, true), (:OrFun, false))
    @eval initarray!(a::AbstractArray, ::$(Op), init::Bool) = (init && fill!(a, $initval); a)
end

reducedim_initarray{R}(A::AbstractArray, region, v0, ::Type{R}) = fill!(similar(A,R,reduced_dims(A,region)), v0)
reducedim_initarray{T}(A::AbstractArray, region, v0::T) = reducedim_initarray(A, region, v0, T)

reducedim_initarray0{R}(A::AbstractArray, region, v0, ::Type{R}) = fill!(similar(A,R,reduced_dims0(A,region)), v0)
reducedim_initarray0{T}(A::AbstractArray, region, v0::T) = reducedim_initarray0(A, region, v0, T)

# TODO: better way to handle reducedim initialization
#
# The current scheme is basically following Steven G. Johnson's original implementation
#
function reducedim_init{T}(f, op::AddFun, A::AbstractArray{T}, region)
    if method_exists(zero, (Type{T},))
        x = evaluate(f, zero(T))
        z = zero(x) + zero(x)
        Tr = typeof(z) == typeof(x) && !isbits(T) ? T : typeof(z)
    else
        z = zero(sum(f, A))
        Tr = typeof(z)
    end
    return reducedim_initarray(A, region, z, Tr)
end

function reducedim_init{T}(f, op::MulFun, A::AbstractArray{T}, region)
    if method_exists(zero, (Type{T},))
        x = evaluate(f, zero(T))
        z = one(x) * one(x)
        Tr = typeof(z) == typeof(x) && !isbits(T) ? T : typeof(z)
    else
        z = one(prod(f, A))
        Tr = typeof(z)
    end
    return reducedim_initarray(A, region, z, Tr)
end

reducedim_init{T}(f, op::MaxFun, A::AbstractArray{T}, region) = reducedim_initarray0(A, region, typemin(evaluate(f, zero(T))))
reducedim_init{T}(f, op::MinFun, A::AbstractArray{T}, region) = reducedim_initarray0(A, region, typemax(evaluate(f, zero(T))))
reducedim_init{T}(f::Union(AbsFun,Abs2Fun), op::MaxFun, A::AbstractArray{T}, region) =
    reducedim_initarray(A, region, zero(evaluate(f, zero(T))))

reducedim_init(f, op::AndFun, A::AbstractArray, region) = reducedim_initarray(A, region, true)
reducedim_init(f, op::OrFun, A::AbstractArray, region) = reducedim_initarray(A, region, false)

# specialize to make initialization more efficient for common cases

typealias CommonReduceResult Union(Uint64,Uint128,Int64,Int128,Float32,Float64,Complex64,Complex128)

for (IT, RT) in ((:CommonReduceResult, :T), (:SmallSigned, :Int), (:SmallUnsigned, :Uint))
    @eval begin
        reducedim_init{T<:$IT}(f::Union(IdFun,AbsFun,Abs2Fun), op::AddFun, A::AbstractArray{T}, region) =
            reducedim_initarray(A, region, zero($RT))
        reducedim_init{T<:$IT}(f::Union(IdFun,AbsFun,Abs2Fun), op::MulFun, A::AbstractArray{T}, region) =
            reducedim_initarray(A, region, one($RT))
    end
end
reducedim_init(f::Union(IdFun,AbsFun,Abs2Fun), op::AddFun, A::AbstractArray{Bool}, region) =
    reducedim_initarray(A, region, 0)

## generic (map)reduction

has_fast_linear_indexing(a::AbstractArray) = false
has_fast_linear_indexing(a::Array) = true

function check_reducdims(R, A)
    # Check whether R has compatible dimensions w.r.t. A for reduction
    #
    # It returns an integer value value (useful for choosing implementation)
    # - If it reduces only along leading dimensions, e.g. sum(A, 1) or sum(A, (1, 2)),
    #   it returns the length of the leading slice. For the two examples above,
    #   it will be size(A, 1) or size(A, 1) * size(A, 2).
    # - Otherwise, e.g. sum(A, 2) or sum(A, (1, 3)), it returns 0.
    #
    lsiz = 1
    had_nonreduc = false
    for i = 1:ndims(A)
        sRi = size(R, i)
        sAi = size(A, i)
        if sRi == 1
            if sAi > 1
                if had_nonreduc
                    lsiz = 0  # to reduce along i, but some previous dimensions were non-reducing
                else
                    lsiz *= sAi  # if lsiz was set to zero, it will stay to be zero
                end
            end
        else
            sRi == sAi ||
                throw(DimensionMismatch("Reduction on array of size $(size(A)) with output of size $(size(R))"))
            had_nonreduc = true
        end
    end
    return lsiz
end

@ngenerate N typeof(R) function _mapreducedim!{T,N}(f, op, R::AbstractArray, A::AbstractArray{T,N})
    lsiz = check_reducdims(R, A)
    isempty(A) && return R
    @nextract N sizeR d->size(R,d)
    sizA1 = size(A, 1)

    if has_fast_linear_indexing(A) && lsiz > 16
        # use mapreduce_impl, which is probably better tuned to achieve higher performance
        nslices = div(length(A), lsiz)
        ibase = 0
        for i = 1:nslices
            @inbounds R[i] = mapreduce_impl(f, op, A, ibase+1, ibase+lsiz)
            ibase += lsiz
        end
    elseif size(R, 1) == 1 && sizA1 > 1
        # keep the accumulator as a local variable when reducing along the first dimension
        @nloops N i d->(d>1? (1:size(A,d)) : (1:1)) d->(j_d = sizeR_d==1 ? 1 : i_d) begin
            @inbounds r = (@nref N R j)
            for i_1 = 1:sizA1
                @inbounds v = evaluate(f, (@nref N A i))
                r = evaluate(op, r, v)
            end
            @inbounds (@nref N R j) = r
        end
    else
        # general implementation
        @nloops N i A d->(j_d = sizeR_d==1 ? 1 : i_d) begin
            @inbounds v = evaluate(f, (@nref N A i))
            @inbounds (@nref N R j) = evaluate(op, (@nref N R j), v)
        end
    end
    return R
end

mapreducedim!(f, op, R::AbstractArray, A::AbstractArray) = _mapreducedim!(f, op, R, A)

function mapreducedim!(f::Function, op, R::AbstractArray, A::AbstractArray)
    is(op, +) ? _mapreducedim!(f, AddFun(), R, A) :
    is(op, *) ? _mapreducedim!(f, MulFun(), R, A) :
    is(op, &) ? _mapreducedim!(f, AndFun(), R, A) :
    is(op, |) ? _mapreducedim!(f, OrFun(), R, A) :
    _mapreducedim!(f, op, R, A)
end

reducedim!{RT}(op, R::AbstractArray{RT}, A::AbstractArray) = mapreducedim!(IdFun(), op, R, A, zero(RT))

mapreducedim(f, op, A::AbstractArray, region, v0) = mapreducedim!(f, op, reducedim_initarray(A, region, v0), A)
mapreducedim{T}(f, op, A::AbstractArray{T}, region) = mapreducedim!(f, op, reducedim_init(f, op, A, region), A)

reducedim(op, A::AbstractArray, region, v0) = mapreducedim(IdFun(), op, A, region, v0)
reducedim(op, A::AbstractArray, region) = mapreducedim(IdFun(), op, A, region)

##### Specific reduction functions #####

for (fname, Op) in [(:sum, :AddFun), (:prod, :MulFun),
                    (:maximum, :MaxFun), (:minimum, :MinFun),
                    (:all, :AndFun), (:any, :OrFun)]

    fname! = symbol(string(fname, '!'))
    @eval begin
        $(fname!)(f::Union(Function,Func{1}), r::AbstractArray, A::AbstractArray; init::Bool=true) =
            mapreducedim!(f, $(Op)(), initarray!(r, $(Op)(), init), A)
        $(fname!)(r::AbstractArray, A::AbstractArray; init::Bool=true) = $(fname!)(IdFun(), r, A; init=init)

        $(fname)(f::Union(Function,Func{1}), A::AbstractArray, region) =
            mapreducedim(f, $(Op)(), A, region)
        $(fname)(A::AbstractArray, region) = $(fname)(IdFun(), A, region)
    end
end

for (fname, fbase, Fun) in [(:sumabs, :sum, :AbsFun),
                            (:sumabs2, :sum, :Abs2Fun),
                            (:maxabs, :maximum, :AbsFun),
                            (:minabs, :minimum, :AbsFun)]
    fname! = symbol(string(fname, '!'))
    fbase! = symbol(string(fbase, '!'))
    @eval begin
        $(fname!)(r::AbstractArray, A::AbstractArray; init::Bool=true) =
            $(fbase!)($(Fun)(), r, A; init=init)
        $(fname)(A::AbstractArray, region) = $(fbase)($(Fun)(), A, region)
    end
end

##### findmin & findmax #####

# Generate the body for a reduction function reduce!(f, Rval, Rind, A), using a comparison operator f
# Rind contains the index of A from which Rval was taken
function gen_findreduction_body(N, f::Function)
    F = Expr(:quote, f)
    quote
        (isempty(Rval) || isempty(A)) && return Rval, Rind
        for i = 1:$N
            (size(Rval, i) == size(A, i) || size(Rval, i) == 1) || throw(DimensionMismatch("Find-reduction on array of size $(size(A)) with output of size $(size(Rval))"))
            size(Rval, i) == size(Rind, i) || throw(DimensionMismatch("Find-reduction: outputs must be of the same size"))
        end
        @nexprs $N d->(sizeR_d = size(Rval,d))
        # If we're reducing along dimension 1, for efficiency we can make use of a temporary.
        # Otherwise, keep the result in Rval/Rind so that we traverse A in storage order.
        k = 0
        @inbounds if size(Rval, 1) < size(A, 1)
            @nloops $N i d->(d>1? (1:size(A,d)) : (1:1)) d->(j_d = sizeR_d==1 ? 1 : i_d) begin
                tmpRv = (@nref $N Rval j)
                tmpRi = (@nref $N Rind j)
                for i_1 = 1:size(A,1)
                    k += 1
                    tmpAv = (@nref $N A i)
                    if ($F)(tmpAv, tmpRv)
                        tmpRv = tmpAv
                        tmpRi = k
                    end
                end
                (@nref $N Rval j) = tmpRv
                (@nref $N Rind j) = tmpRi
            end
        else
            @nloops $N i A d->(j_d = sizeR_d==1 ? 1 : i_d) begin
                k += 1
                tmpAv = (@nref $N A i)
                if ($F)(tmpAv, (@nref $N Rval j))
                    (@nref $N Rval j) = tmpAv
                    (@nref $N Rind j) = k
                end
            end
        end
        Rval, Rind
    end
end

eval(ngenerate(:N, :(typeof((Rval,Rind))), :(_findmin!{T,N}(Rval::AbstractArray, Rind::AbstractArray, A::AbstractArray{T,N})), N->gen_findreduction_body(N, <)))
findmin!{R}(rval::AbstractArray{R}, rind::AbstractArray, A::AbstractArray; init::Bool=true) = _findmin!(initarray!(rval, typemax(R), init), rind, A)
findmin{T}(A::AbstractArray{T}, region) =
    isempty(A) ? (similar(A,reduced_dims0(A,region)), zeros(Int,reduced_dims0(A,region))) :
                  _findmin!(reducedim_initarray0(A, region, typemax(T)), zeros(Int,reduced_dims0(A,region)), A)

eval(ngenerate(:N, :(typeof((Rval,Rind))), :(_findmax!{T,N}(Rval::AbstractArray, Rind::AbstractArray, A::AbstractArray{T,N})), N->gen_findreduction_body(N, >)))
findmax!{R}(rval::AbstractArray{R}, rind::AbstractArray, A::AbstractArray; init::Bool=true) = _findmax!(initarray!(rval, typemin(R), init), rind, A)
findmax{T}(A::AbstractArray{T}, region) =
    isempty(A) ? (similar(A,reduced_dims0(A,region)), zeros(Int,reduced_dims0(A,region))) :
                  _findmax!(reducedim_initarray0(A, region, typemin(T)), zeros(Int,reduced_dims0(A,region)), A)