text:
test_gage_min_distance(
    [1., 0., 0.],
    [0., 1., 0.],
    sqrt(2.),
    "simple 3D"
)


text:
using JuMP, Gurobi, Ipopt
using DataFrames, JLD, CSVFiles
using MatpowerCases

### Include the file with input data functions
include("NetworkDataType.jl")
include("NetworkLoad.jl")
include("DLMP.jl")

testsystem = "AP15busDN"
#testsystem = "case141_OPF"

buses, lines, generators, datamat = NetworkLoad(testsystem)
lineset = 1:length(lines)
busset = 1:length(buses)
genset = 1:length(generators)
buses[1].Vmax = 1
buses[1].Vmin = 1

windset = Int[]
for g in genset
   if datamat["gen"][g,2] == "Wind"
      push!(windset, generators[g].gindex)
   end
end

root = -1
for g in genset
   if generators[g].gtype=="Root"
      root = g
   end
end
if root == -1
   display("No generator at root node")
   generators[1].gtype = "Root"
   root = 1
end

temp = Generator
for g in genset
   if generators[g].gtype == "Root"
      temp = Generator(length(genset)+1, "DG", generators[g].location, generators[g].Pg, generators[g].Qg, generators[g].Qmax,
         generators[g].Qmin, generators[g].Vg, generators[g].mBase, generators[g].status, generators[g].Pmax, generators[g].Pmin,
         generators[g].cost, generators[g].SUcost, generators[g].SDcost, generators[g].RU, generators[g].RD, generators[g].UPtime, generators[g].DNtime)
      #push!(generators, temp)
   end
end
genset = 1:length(generators)

########### define Utility generator set & Peer generator set #########
gensetU = root
gensetP = setdiff(genset, gensetU)
#######################################################################

B_g = []
for g in genset
   if ~any(x->x==generators[g].location, B_g)
      push!(B_g, generators[g].location)
      buses[generators[g].location].Pd = 0
   end
end
B_d = setdiff(busset,B_g)

B_gn = Array{Array{Int64}}(length(buses),1)
for ii in 1:length(buses)
   B_gn[ii] = Int64[]
end
for g in genset
   push!(B_gn[generators[g].location], g)
end
############
partlevel = 0.0
gam = partlevel*ones(busset)
############
m = Model(solver = GurobiSolver())

@variable(m, v[b in busset] >= 0) # variable for voltage square, voltage^2
@variable(m, a[l in lineset] >= 0) # variable for current square, current^2
@variable(m, fp[l in lineset])
@variable(m, fq[l in lineset])
@variable(m, pg[g in genset])
@variable(m, qg[g in genset])
@variable(m, OC >= 0)
@variable(m, OCgen[g in genset] >= 0)
# P2P variable
@variable(m, p[n in busset])
@variable(m, pnm[n in busset,k in busset])

@objective(m, Max, -OC)

# P2P constraints
@constraint(m, BiDirection[n in busset, k in busset], pnm[n,k] + pnm[k,n] == 0)
@constraint(m, SignG[n in B_g, k in busset], pnm[n,k] >= 0)
@constraint(m, SignD[n in B_d, k in busset], pnm[n,k] <= 0)
@constraint(m, NodalBalance[n in busset], p[n] == sum(pnm[n,k] for k in busset) )

@constraint(m, genPower[n in B_g], sum(pg[g] for g in setdiff(B_gn[n], gensetU)) == p[n])
@constraint(m, demand[n in B_d], - gam[n] * buses[n].Pd == p[n])

@constraint(m, OperatingCost, OC == sum(OCgen[g] for g in genset))

#### Polynomial Generation Cost
#@constraint(m, GenCost[g in genset], OCgen[g] == sum(generators[g].cost[k]*pg[g]^(length(generators[g].cost)-k) for k in 1:length(generators[g].cost)))
#### Linear Generation Cost
@constraint(m, GenCost[g in genset], OCgen[g] == generators[g].cost[2]*pg[g] + generators[g].cost[3])

#@constraint(m, LineCapFW[l = lineset], (fp[l])^2 + (fq[l])^2 <= ( lines[l].u)^2 )
@constraint(m, LineCapFW[l = 1:length(lineset)], norm([fp[l], fq[l]]) <= lines[l].u )

#@constraint(m, LineCapBW[l = lineset], (fp[l] - a[l]*lines[l].r)^2 + (fq[l] - a[l]*lines[l].x)^2 <= ( lines[l].u)^2 )
@constraint(m, LineCapBW[l = 1:length(lineset)], norm([fp[l] - a[l]*lines[l].r, (fq[l] - a[l]*lines[l].x)]) <= lines[l].u )

@constraint(m, BetweenNodes[l = lineset], v[lines[l].fbus] - 2*(lines[l].r*fp[l] + lines[l].x*fq[l]) + a[l]*(lines[l].r^2 + lines[l].x^2) == v[lines[l].tbus])

#@constraint(m, SOCP[l = lineset], ((fp[l])^2 + (fq[l])^2) <= v[lines[l].fbus] * a[l] )
@constraint(m, SOCP[l = lineset], norm([2*fp[l], 2*fq[l], v[lines[l].fbus]-a[l]]) <= v[lines[l].fbus] + a[l])

if testsystem == "AP15busDN" || testsystem == "IEEE33busDN" || testsystem == "ISONE8busTN" || testsystem == "IEEE33busDN2"
   @constraint(m, v[1] == 1) #Voltage constraint for root node
end

@constraint(m, PBalance[b = busset], sum(fp[l] for l in buses[b].outline)  - sum(fp[l] - lines[l].r*a[l] for l in buses[b].inline) - sum(pg[g] for g in B_gn[b]) + buses[b].Pd + v[b]*buses[b].Gs == 0)
@constraint(m, QBalance[b = busset], sum(fq[l] for l in buses[b].outline)  - sum(fq[l] - lines[l].x*a[l] for l in buses[b].inline) - sum(qg[g] for g in B_gn[b])+ buses[b].Qd - v[b]*buses[b].Bs == 0)

@constraint(m, Vmax[b = busset], v[b] <= buses[b].Vmax^2) #upper limit constraint for voltage square
@constraint(m, Vmin[b = busset], v[b] >= buses[b].Vmin^2) #lower limit constraint for voltage square

@constraint(m, PGmin[g = genset], pg[g] >= generators[g].Pmin)
@constraint(m, PGmax[g = genset], pg[g] <= generators[g].Pmax)
@constraint(m, QGmin[g = genset], qg[g] >= generators[g].Qmin)
@constraint(m, QGmax[g = genset], qg[g] <= generators[g].Qmax)

#print(m)
show(m)
println()
tic()
status = solve(m)
show(m)
println()
simulation_time = toc()
println()
println("Julia v",VERSION)
println("JuMP v",Pkg.installed("JuMP"))
println("Gurobi v",Gurobi.getlibversion())
#print("CPLEX v",CPLEX.version())

obj_value = getobjectivevalue(m)
v = getvalue(v)[:]
a = getvalue(a)[:]
fp = getvalue(fp)[:]
fq = getvalue(fq)[:]
pg = getvalue(pg)[:]
qg = getvalue(qg)[:]

OC = getvalue(OC)
OCgen = getvalue(OCgen)[:]

pgb = zeros(length(buses))
qgb = zeros(length(buses))
for b in 1:length(buses)
   if ~isempty(B_gn[b])
      pgb[b] = sum(pg[g] for g in B_gn[b])
      qgb[b] = sum(qg[g] for g in B_gn[b])
   end
end
pnm = getvalue(pnm)[:,:]
pnm[(abs.(pnm).<1e-3)]=0
result_pnm = pnm

DLMPInfo = DataFrame()
NodeInfo = DataFrame()
LineInfo = DataFrame()

dispatch = pg
SMP = generators[root].cost[end-1]
status1, DLMP, pg, NodeInfo, LineInfo, DLMPInfo = calculateDLMP(dispatch, buses, generators, lines, SMP, gensetP, gensetU);

cn = zeros(length(buses), length(buses))
for n in 1:length(buses), k in 1:length(buses)
    cn[n,k] = (DLMP[k]-DLMP[n])/2
end

line_marginrate = zeros(length(lines),1)
for ll in 1:length(lines)
    line_marginrate[ll] = round.((lines[ll].u  - sqrt((fp[ll])^2 + (fq[ll])^2))/ lines[ll].u,3)
end
line_loading = 1-line_marginrate
bus_voltage = sqrt.(v)

v_marginrate = round.(0.1-abs.(1-sqrt.(v)),3)/0.1

Loading = DataFrame(line_loading=vec(line_loading))
Vmargin = DataFrame(v_marginrate=vec(v_marginrate))

loading = zeros(length(lines))
for l in 1:length(lines)
   loading[l] = LineInfo[:flow][l] / lines[l].u
end

flow = LineInfo[:flow]
voltage = sqrt.(NodeInfo[:v])
##########
EP_P = zeros(length(busset))
for m in B_d
   EP_P[m] = sum(sum(generators[g].cost[end-1]*pnm[n,m] for g in setdiff(B_gn[n],gensetU)) for n in B_g if ~isempty(setdiff(B_gn[n],gensetU)))
end
EP_U = zeros(length(busset))
for m in B_d
   EP_U[m] = (1-gam[m])*buses[m].Pd*DLMP[m]
end
NUC = zeros(length(busset))
for m in B_d
   NUC[m] = sum(pnm[n,m]*(DLMP[m]-DLMP[n]) for n in B_g)
end
UR = EP_U + NUC
CP = EP_U + EP_P + NUC

GC = OCgen
sum_GC_P = sum(GC[g] for g in gensetP)
sum_GC_U = sum(GC[g] for g in gensetU)
sum_DGR = sum(GC[g] for g in gensetP)
sum_DGP = sum_DGR - sum_GC_P
UP = sum(UR) - sum_GC_U

UP = sum(UR) - sum_GC_U
CashFlow_mat = DataFrame(EP_P=EP_P,EP_U=EP_U,NUC=NUC,UR=UR,CP=CP)
CashFlow_sum = DataFrame(sum_CP=sum(CP), sum_EP_P=sum(EP_P), sum_EP_U=sum(EP_U), sum_NUC=sum(NUC), sum_DGR=sum_DGR, sum_DGP=sum_DGP, UP=UP)

println("partlevel = ""$partlevel")
println("status = ""$status")
println("max loading = ""$(maximum(loading)*100)""%")
println("max voltage = ""$(round(maximum(sqrt.(v)),3))")
println("min voltage = ""$(round(minimum(sqrt.(v)),3))")
println("max DLMP = ""$(round(maximum(DLMP),3))")
println("min DLMP = ""$(round(minimum(DLMP),3))")

method = "SC"
voltagefile = "$method$(length(buses))""_""$(Int(100*partlevel))""_voltage"".csv"
loadingfile = "$method$(length(buses))""_""$(Int(100*partlevel))""_loading"".csv"
DLMPfile = "$method$(length(buses))""_""$(Int(100*partlevel))""_DLMP"".csv"
cashflowmatfile = "$method$(length(buses))""_""$(Int(100*partlevel))""_cfmat"".csv"
cashflowsumfile = "$method$(length(buses))""_""$(Int(100*partlevel))""_cfsum"".csv"
# CSVFiles.save(voltagefile, DataFrame(voltage=voltage))
# CSVFiles.save(loadingfile, DataFrame(loading=100*loading))
# CSVFiles.save(DLMPfile, DataFrame(DLMP=DLMP))
# CSVFiles.save(cashflowmatfile, CashFlow_mat)
# CSVFiles.save(cashflowsumfile, CashFlow_sum)

# CSVFiles.save("LineInfo.csv",LineInfo)
##########
#resultfile = "output/$method$(length(buses))""_""$(Int(100*partlevel))"".jld"
resultfile = "$method$(length(buses))""_""$(Int(100*partlevel))"".jld"
# @save "$resultfile"

#######################################
Erho = zeros(length(buses),1)
Ecn = zeros(length(buses),1)

Tvol = sqrt.(sum(pnm,2).^2)
for n in 1:length(busset)
   if Tvol[n] != 0
      Ecn[n] = sum((DLMP[m]-DLMP[n])*pnm[n,m]/2 for m in 1:length(busset)) / Tvol[n]
   end
end
for n in 1:length(busset)
   if Tvol[n] != 0
      Erho[n] = sum((DLMP[m]+DLMP[n])*pnm[n,m]/2 for m in 1:length(busset)) / sum(pnm[n,m] for m in 1:length(busset))
   end
end

NUC_total = zeros(length(buses),1)
for n in 1:length(busset)
   if Tvol[n] != 0
      NUC_total[n] = sum((DLMP[m]-DLMP[n])*pnm[n,m]/2 for m in 1:length(busset))
   end
end
println("NUC total = ",sum(NUC_total))
println("Tvol = ", sum(Tvol))

ext: jl


text:
module Curves

export Curve, interpolate, apply, concat, drop_duplicates, firstpoint, lastpoint
export ItpLinear, ItpConstant, EtpFlat, EtpLine # type constants referring to Interpolations.jl

export Tenor, get_days, get_tenor, @t_str

using Interpolations
using RecipesBase

#=
The Curve object is intended to be immutable, i.e. each operation on it creates a new Curve instance.
In principle it would be possible to change points of the arrays curve.x or curve.y without creating a new Curve instance,
but this should be avoided because it may introduce inconsistencies (especially when using log-interpolation).
In order to have both extrapolation and interpolation, a nested extrapolation(interpolation) object (using Interpolations.jl)
is created.
The x and y arrays are passed to the Interpolations.jl interpolator directly if log-interpolation is switched off. In this
case, the interpolation object contains a pointer to the same array, thus there are no duplicates of the x and y arrays
in memory. The arrays x and y are mainly kept directly in the Curve struct (and not only inside the interpolation object)
to avoid problems with type inference.
However, if log interpolation is activated for one axis, a second array is created for the axis containing
the log values of the original array and stored in the interpolation object. This results in double memory usage, but
allows quick access to both the axis values (via curve.x and curve.y) and interpolated values.
Note that Interpolation.jl does not support log-interpolation (yet?), therefore it needs to be implemented here explicitly.
=#

include("tenors.jl")

# Type constants referring to Interpolations.jl

const ItpLinear = Gridded ∘ Linear
const ItpConstant = Gridded ∘ Constant

const EtpFlat = Flat
const EtpLine = Line

# Basic definition
abstract type AbstractCurve end

struct Curve{Tx <: AbstractVector, Ty <: AbstractVector,
        Titp <: Union{Interpolations.AbstractInterpolation, Nothing}} <: AbstractCurve
    "data points x-axis"
    x:: Tx
    "data points y-axis"
    y:: Ty
    "Interpolations.jl extrapolation(interpolation) object"
    etp:: Titp
    "is x-axis logarithmic?"
    logx:: Bool
    "is y-axis logarithmic?"
    logy:: Bool
end

"""
    Curve(x:: AbstractVector, y:: AbstractVector; method=ItpLinear(), extrapolation=EtpFlat(), logx=false, logy=false, sort=true)

Standard curve constructor.
Creates the interpolation/extrapolation object of the curve instance.
The points `x` and `y` do not need to be sorted, this is done in the Curve constructor.

Interpolation/extrapolation details can be changed using the keyword arguments, defaults are:

* linear interpolation
* constant extrapolation
* no logarithmic axes

Note that for method ˋGridded(x)ˋ must be used so that the x-grid can be non-uniform.

Valid choices for ˋmethodˋ are:

* ˋItpLinear()ˋ, corresponds to Interpolation.jl ˋGridded(Linear())ˋ - default
* ˋItpConstant()ˋ, corresponds to Interpolation.jl ˋGridded(Constant())ˋ

Valid choices for ˋextrapolationˋ are:

* ˋEtpFlat()ˋ, corresponds to Interpolation.jl ˋFlat()ˋ - default
* ˋEtpLine()ˋ, corresponds to Interpolation.jl ˋLine()ˋ

If the curve consists only of a single point, always constant extrapolation is used.

* `sort=true`: per default, the input points are sorted and duplicate x-values are removed. `sort=true` is unsafe and intended to be used for Curves.jl internal operations only, where it can be guaranteed that the points are sorted and not duplicate.

"""
function Curve(x:: AbstractVector, y:: AbstractVector;
        method=ItpLinear(), extrapolation=EtpFlat(), logx=false, logy=false, sort=true)
    length(x) == length(y) || error("length of x and y arrays must match")
    if length(x) == 1
        return Curve(x, y, nothing, logx, logy)
    else
        if sort
            x, y = uniquexy(x, y)
            perm = sortperm(x)
            x = x[perm]
            y = y[perm]
        end
        return Curve(x, y, extrapolate(interpolate(logx ? (log.(x),) : (x,), logy ?  log.(y) : y, method), extrapolation), logx, logy)
    end
end

Base.show(io:: IO, c:: Curve) = print(io, "x = $(c.x), y = $(c.y), logx = $(c.logx), logy = $(c.logy)")

"helper function to get Interpolations.jl interpolation method"
getitpm(c1:: Curve) = isnothing(c1.etp) ? nothing : c1.etp.itp.it

"helper function to get Interpolations.jl extrapolation method"
getetpm(c1:: Curve) = isnothing(c1.etp) ? nothing : c1.etp.et

"""
    Curve(c1:: Curve; method=getitpm(c1), extrapolation=getetpm(c1), logx=c1.logx, logy=c1.logy, sort=true)

Copy constructor to generate a new curve from an existing one.
The interpolation / extrapolation parameters can be changed.
"""
Curve(c1:: Curve; method=getitpm(c1), extrapolation=getetpm(c1), logx=c1.logx, logy=c1.logy, sort=true) =
    Curve(c1.x, c1.y, method=method, extrapolation=extrapolation, logx=logx, logy=logy, sort=sort)

"""
    Curve(x:: AbstractVector{<: AbstractString}, y; kwargs...)

Construct Curve objects from an array of tenor strings as x-axis.
"""
Curve(x:: AbstractVector{<: AbstractString}, y:: AbstractVector; kwargs...) = Curve(Tenor.(x), y; kwargs...)

"""
    Curve(x:: AbstractVector{Tenor}, y:: AbstractVector; offset:: Real = 0, kwargs...)

Construct Curve objects from an array of Tenor objects strings as x-axis.

With the `offset` keyword argument the points on the x-axis, defined by the given tenors, can be shifted.
This could be e.g. used to take a spot lag of financial instruments into account.
Note that the shift is always in calendar days, a spot lag given in business days must be converted to calendar days beforehand.
"""
Curve(x:: AbstractVector{Tenor}, y:: AbstractVector; offset:: Real = 0, kwargs...) = Curve(get_days.(x) .+ offset, y; kwargs...)

Curve(x:: Real, y:: Real; kwargs...) = Curve([x], [y]; kwargs...)
Curve(x:: Tenor, y:: Real; kwargs...) = Curve([get_days(x)], [y]; kwargs...)

Base.Broadcast.broadcastable(q:: Curve) = Ref(q) # treat it as a scalar in broadcasting

@recipe plot(c::Curve) = c.x, c.y

include("curve_functions.jl")

end # module


text:
using Pio3d

function prepare_ircad(id=1)
    # fn = io3d.datasets.join_path("medical", "orig", "3Dircadb1.1", "MASKS_DICOM", "liver", get_root=True)
    pth = Pio3d.datasets_join_path("medical/orig/3Dircadb1.$id/MASKS_DICOM/liver")

    datap = Pio3d.read3d(pth)
    data3d_full = datap["data3d"]
    # round(size(data3d_full, 1) / target_size1)
    # return data3d_full
    data3d_int8 = convert(Array{Int8, 3}, data3d_full .> 0)
    @info "reading Ircad $id done"

    return data3d_int8
end

# im = prepare_ircad(2)


text:
using LowRankModels
using Distributions

p = Params(0.00001,min_stepsize=0.00000000001,max_iter=5000)

# This is just to create a low-rank representation of some count data, I don't care how accurate it is.
m,n = 100, 50;
k = 2;
A = rand(Poisson(2), m, n); # Data is Poisson with mean 2
losses = convert(Array{Loss,1}, fill(PoissonLoss(10),n));
rx, ry = QuadReg(), QuadReg();
g_pre = GLRM(A, losses, rx, ry, k, scale=false, offset=false);

# let's check a different syntax works, too
g_pre = GLRM(A, PoissonLoss(), rx, ry, k, scale=false, offset=false);
X_real, Y_real, ch = fit!(g_pre, params=p);

# Now we do the actual model using the perfect data and try to recapture it.
A_real = impute(losses, X_real'*Y_real);

g = GLRM(A_real, losses, rx, ry, k, scale=true, offset=true);
X, Y, ch = fit!(g, params=p);
U = X'*Y;
A_imputed = impute(losses, X'*Y);

@show error_metric(g)
errors(Domain[l.domain for l in losses], losses, U, A_real);


text:
module Elsa

using Requires

# function __init__()
#     # @require Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" include("diagonalizers/pardiso.jl")
#     # @require KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" include("diagonalizers/krylovkit.jl")
#     @require Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97" include("diagonalizers/arpack.jl")
#     # @require ArnoldiMethod = "ec485272-7323-5ecc-a04f-4719b315124d" include("diagonalizers/arnoldimethod.jl")
#     # @require IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" include("diagonalizers/iterativesolvers.jl")
# end

using StaticArrays, NearestNeighbors, SparseArrays, LinearAlgebra, OffsetArrays,
      ProgressMeter, LinearMaps, Random

using SparseArrays: getcolptr

export sublat, bravais, lattice, dims, hopping, onsite, hamiltonian, randomstate,
       mul!, supercell, unitcell, semibounded, bloch, bloch!, optimize!, similarmatrix,
       sites, bandstructure, marchingmesh, defaultmethod, bands, vertices, states,
       flatten, wrap, transform!, combine,
       momentaKPM, dosKPM, averageKPM, densityKPM

export LatticePresets, RegionPresets, HamiltonianPresets

export LinearAlgebraPackage, ArpackPackage, KrylovKitPackage

export @SMatrix, @SVector, SMatrix, SVector

export ishermitian, I

const NameType = Symbol
const nametype = Symbol

const TOOMANYITERS = 10^8

include("iterators.jl")
include("presets.jl")
include("lattice.jl")
include("model.jl")
include("field.jl")
include("hamiltonian.jl")
include("state.jl")
include("mesh.jl")
include("diagonalizer.jl")
include("bandstructure.jl")
include("KPM.jl")
include("convert.jl")
include("tools.jl")

end


text:
"""
    $SIGNATURES

Transfer contents of `src::AbstractDFG` variables `syms::Vector{Symbol}` to `dest::AbstractDFG`.

Notes
- Reads, `dest` := `src`, for all `syms`
"""
function transferUpdateSubGraph!(dest::AbstractDFG,
                                 src::AbstractDFG,
                                 syms::Vector{Symbol}=union(ls(src)...),
                                 logger=ConsoleLogger()  )
  #
  with_logger(logger) do
    @info "transferUpdateSubGraph! -- syms=$syms"

    # TODO add with DFG v0.4
    # DFG.updateGraphSolverData!(src, dest, syms)
    for sym in syms
      vari = DFG.getVariable(src, sym)
      rc = size(solverData(vari).val)
      # TODO -- reduce to DFG functions only
      pp = getKDE(vari)
      rc2 = size(getPoints(pp))
      @info "sym=$sym, mem size of val=$rc and $(rc2)"
      updateFullVertData!(dest, vari, updateMAPest=true)
    end
  end
  nothing
end

"""
    $SIGNATURES

Build a new subgraph from `fgl<:AbstractDFG` containing all variables and factors
associated with `cliq`.  Additionally add the upward message prior factors as
needed for belief propagation (inference).

Notes
- `cliqsym::Symbol` defines the cliq where variable appears as a frontal variable.
- `varsym::Symbol` defaults to the cliq frontal variable definition but can in case a
  separator variable is required instead.
"""
function buildCliqSubgraphDown(fgl::AbstractDFG, treel::BayesTree, cliqsym::Symbol, varsym::Symbol=cliqsym)
  @warn "Obsolete, buildCliqSubGraph*() is no longer in use"
  # build a subgraph copy of clique
  cliq = whichCliq(treel, cliqsym)
  syms = getCliqAllVarIds(cliq)
  subfg = buildSubgraphFromLabels(fgl,syms)

  # add upward messages to subgraph
  msgs = getCliqParentMsgDown(treel, cliq)
  addMsgFactors!(subfg, msgs)
  return subfg
end


text:
module test_suite_applications

using DomainSets, BasisFunctions, FrameFun
using Test

FE = FrameFun
BA = BasisFunctions

## Settings

# Test fourier extensions for all parameters
Extensive = false

# Show matrix vector product timings
const show_mv_times = false
total_mv_allocs = 0
total_mv_time = 0.0

########
# Auxiliary functions
########

function delimit(s::AbstractString)
    println()
    println("############")
    println("# ",s)
    println("############")
end

#######
# Testing -- Accuracy
#######

function test_differential_equations_1d()
    @testset "diff 1D dirichlet" for solverstyle in (AZStyle(), DirectStyle())
        B = Fourier(101) → -1..1
        Dom = Interval(-0.5,0.5)
        # Set up Boundary conditions
        BC = DirichletBC(x->0)
        # Set up Differential equation
        f(x) = x
        Diff = differentiation(B)*differentiation(B)
        DE = DiffEquation(B, Dom, Diff, f, (BC,))
        # Actually solve the differential equation
        F = solve(DE, solverstyle=solverstyle, directsolver=:qr)
        sol(x) = x^3/6 - x/24
        error = abserror(sol,F)
        @test (error < sqrt(eps(real(codomaintype(B))))*10)
        error = abserror(f,F'')
        @test (error < sqrt(eps(real(codomaintype(B))))*100)
    end
    @testset "diff 1D mixed" for solverstyle in (AZStyle(), DirectStyle())
        B = Fourier(101) → -2..2
        Dom = Interval(-1.0,1.0)
        # Set up Boundary conditions
        BC1 = DirichletBC(x->0, Interval(-2.0,0.0))
        BC2 = NeumannBC(x->0, Interval(-2.0,0.0))
        # Set up Differential equation
        f(x) = x
        Diff = differentiation(B)*differentiation(B)
        DE = DiffEquation(B, Dom, Diff, f, (BC1,BC2))
        # Actually solve the differential equation
        F = solve(DE, solverstyle=solverstyle, directsolver=:qr)
        sol(x) = x^3/6 - x/2 -1/3
        error = abserror(sol,F)
        @test (error < sqrt(eps(real(codomaintype(B))))*10)
        error = abserror(f,F'')
        @test (error < sqrt(eps(real(codomaintype(B))))*100)
    end
end

function test_differential_equations_2d()
    # @testset "diff 2D dirichlet" for solverstyle in (AZStyle(), DirectStyle())
    @testset "diff 2D dirichlet" for solverstyle in (AZStyle(),)
        B = (Fourier(11) → -1..1)^2
        Dom = disk(0.8)
        # Set up Boundary conditions
        df = (x,y)->x-y
        BC = DirichletBC(df,DomainSets.euclideanspace(Val{2}()))
        # Set up Differential equation
        f(x,y) = 0
        Diff = differentiation(B,(2,0))+differentiation(B,(0,2))
        DE = DiffEquation(B, Dom, Diff, f, (BC,))
        # Actually solve the differential equation
        F = solve(DE, solverstyle=solverstyle, directsolver=:qr)
        error = abserror(df,F)
        @test (error < 0.3)
        error = abserror(f,∂x(∂x(F))+∂y(∂y(F)))
        @test (error < 0.3)
    end
    # @testset "diff 2D Neumann" for solverstyle in (AZStyle(), DirectStyle())
    # @testset "diff 2D Neumann" for solverstyle in (AZStyle(),)
    #     B = (Fourier(11) → -1..1)^2
    #     Dom = disk(0.8)
    #     # Set up Boundary conditions
    #     df = (x,y)->x-y
    #     BC = NeumannBC(df,DomainSets.euclideanspace(Val{2}()))
    #     # Set up Differential equation
    #     f(x,y) = 0
    #     Diff = differentiation(B,(2,0))+differentiation(B,(0,2))
    #     DE = DiffEquation(B, Dom, Diff, f, (BC,))
    #     # Actually solve the differential equation
    #     F = solve(DE, solverstyle=solverstyle, directsolver=:qr)
    #     # We should find a way to check the boundary condition more easily
    #     berror = sum(abs.(element(operator(DE),2,1)*coefficients(Diff*F)-element(FrameFun.rhs(DE),2)))/length(element(FrameFun.rhs(DE),2))
    #     @test berror < 0.6
    #     error = abserror(f,∂x(∂x(F))+∂y(∂y(F)))
    #     @test (error < 0.3)
    # end
end

# We just test the accuracy, not the smoothing properties
function test_smoothing_1d()
    @testset "Smoothing $(name(Basis(10)))" for Basis in (ChebyshevT,)
        basis = Basis(101)→ -1.0..1.0
        dom = -0.5..0.5
        f = exp
        fscale(dict, i) = 10.0^-4+i+i^2+i^3
        F = Fun(f, basis, dom; solverstyle = AZSmoothStyle(), scaling=fscale)
        @test (abserror(f,F) < 100*sqrt(eps(Float64)))
    end
end

function test_smoothing_2d()
    @testset "Smoothing $(name(Basis(10)))" for Basis in (ChebyshevT,)
        basis = (Basis(20)→(-1.0..1.0))⊗(Basis(20)→(-1.0..1.0))
        dom = disk(0.5)
        f = (x,y) -> exp(x*y)
        fscale(dict, I) = 10.0^-4+100*(I[1]^2+I[2]^2)
        F = Fun(f, basis, dom; solverstyle = AZSmoothStyle(), scaling=fscale)
        @test (abserror(f,F) < 100*sqrt(sqrt(eps(Float64))))
    end
end

test_smoothing_1d()

test_smoothing_2d()

test_differential_equations_1d()

test_differential_equations_2d()

if show_mv_times
    println("Total bytes in MV products:\t$total_mv_allocs")
    println("Total time in MV products:\t$total_mv_time")
end

end


text:
#==============================================================================
    Code for solving the Hamiltonian Jacboi Bellman for
	   an basic model: a neoclassical growth model: ρV(k) = max_{c} U(c) + V'(k)[F(k)-δk - c]
	   Where s(k) = F(k) - δk - c(k) and c(k) = U'^{-1}(V'(k))

	Translated Julia code from Matlab code by Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm

		Updated to julia 1.0.0
==============================================================================#

using Distributions, Plots

σ= 2.0 #
ρ = 0.05 #the discount rate
δ = 0.05 # the depreciation rate
A = 1.0
α= 0.3

k_ss = (α*A/(ρ+δ))^(1/(1-α))

H= 150
k_min = 0.001*k_ss
k_max = 2.0*k_ss

k = LinRange(k_min, k_max, H)
k = convert(Array, k) # create grid for a values
dk = (k_max-k_min)/(H-1)

maxit = 1000
ε = 10e-6

dVf, dVb = [zeros(H,1) for i=1:2]

#initial guess for V
v0 = (A.*k.^α).^(1-σ)/(1-σ)/ρ
v= v0

dist=[]

for n=1:maxit
	V=v

    # forward difference
	dVf[1:H-1] = (V[2:H]-V[1:H-1])/dk
	dVf[H] = 0

	# backward difference
	dVb[2:H] = (V[2:H]-V[1:H-1])/dk
	dVb[1] = 0 # the boundary condition

	I_concave = dVb .> dVf

    # consumption and savings with forward difference
    cf = dVf.^(-1/σ)
    μ_f = A.*k.^α - δ.*k -cf

    # consumption and savings with backward difference
    cb = dVb.^(-1/σ)
    μ_b = A.*k.^α - δ.*k -cb

	c0 = A.*k.^α - δ.*k
    dV0 = c0.^(-σ)

    # Now to make a choice between forward and backward difference
    If = μ_f .> 0
    Ib = μ_b .< 0
    I0 = 1.0.-If-Ib
    Ib[1] = false
    If[1] = true
    Ib[H] = true
    If[H] = false

    global dV_Upwind= dVf.*If + dVb.*Ib + dV0.*I0

    global c = dV_Upwind.^(-1/σ)
    V_change = c.^(1-σ)/(1-σ) + dV_Upwind.*(A.*k.^α - δ.*k-c) -ρ.*V

	# update
	Δ = .9*dk/(findmax(A.*k.^α- δ.*k-c)[1])
	global v = v + Δ*V_change
	push!(dist,findmax(abs.(V_change))[1])
	if dist[n] .< ε
		println("Value Function Converged Iteration=")
		println(n)
		break
	end

end

plot(dist, grid=false,
		xlabel="Iteration", ylabel="||V^{n+1} - V^n||",
		ylims=(-0.001,0.030),
		legend=false, title="")
png("Convergence")

v_err = c.^(1-σ)/(1-σ) + dV_Upwind.*(A.*k.^α - δ.*k -c) - ρ.*v

plot(k, v_err, grid=false,
		xlabel="k", ylabel="Error in the HJB equation",
		xlims=(k_min,k_max),
		legend=false, title="")
png("HJB_error")

plot(k, v, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k_min,k_max),
		legend=false, title="")
png("Value_function_vs_k")

plot(k, c, grid=false,
		xlabel="k", ylabel="c(k)",
		xlims=(k_min,k_max),
		legend=false, title="")
png("c(k)_vs_k")

# approximation at the borrowing constraint
k_dot = (A.*k.^α - δ.*k -c)

plot(k, k_dot, grid=false,
		xlabel="k", ylabel="s(k)",
		xlims=(k_min,k_max), title="", label="s(k)", legend=:topright)
plot!(k, zeros(H,1), label="", line=:dash)
png("stateconstraint")


text:
"""
    struct ModifyConstraintNotAllowed{F<:AbstractFunction, S<:AbstractSet,
                                      C<:AbstractFunctionModification} <: NotAllowedError
        constraint_index::ConstraintIndex{F, S}
        change::C
        message::String
    end

An error indicating that the constraint modification `change` cannot be applied
to the constraint of index `ci`.
"""
struct ModifyConstraintNotAllowed{
    F<:AbstractFunction,
    S<:AbstractSet,
    C<:AbstractFunctionModification,
} <: NotAllowedError
    constraint_index::ConstraintIndex{F,S}
    change::C
    message::String
end
function ModifyConstraintNotAllowed(
    ci::ConstraintIndex{F,S},
    change::AbstractFunctionModification,
    message = "",
) where {F<:AbstractFunction,S<:AbstractSet}
    return ModifyConstraintNotAllowed{F,S,typeof(change)}(ci, change, message)
end
function throw_modify_not_allowed(ci::ConstraintIndex, args...)
    return throw(ModifyConstraintNotAllowed(ci, args...))
end

function operation_name(err::ModifyConstraintNotAllowed{F,S}) where {F,S}
    return "Modifying the constraints $(err.constraint_index) with $(err.change)"
end

"""
    struct ModifyObjectiveNotAllowed{C<:AbstractFunctionModification} <: NotAllowedError
        change::C
        message::String
    end

An error indicating that the objective modification `change` cannot be applied
to the objective.
"""
struct ModifyObjectiveNotAllowed{C<:AbstractFunctionModification} <:
       NotAllowedError
    change::C
    message::String
end
function ModifyObjectiveNotAllowed(change::AbstractFunctionModification)
    return ModifyObjectiveNotAllowed(change, "")
end
function throw_modify_not_allowed(::ObjectiveFunction, args...)
    return throw(ModifyObjectiveNotAllowed(args...))
end

function operation_name(err::ModifyObjectiveNotAllowed)
    return "Modifying the objective function with $(err.change)"
end

"""
## Constraint Function

    modify(model::ModelLike, ci::ConstraintIndex, change::AbstractFunctionModification)

Apply the modification specified by `change` to the function of constraint `ci`.

An [`ModifyConstraintNotAllowed`](@ref) error is thrown if modifying
constraints is not supported by the model `model`.

### Examples

```julia
modify(model, ci, ScalarConstantChange(10.0))
```

## Objective Function

    modify(model::ModelLike, ::ObjectiveFunction, change::AbstractFunctionModification)

Apply the modification specified by `change` to the objective function of
`model`. To change the function completely, call `set` instead.

An [`ModifyObjectiveNotAllowed`](@ref) error is thrown if modifying
objectives is not supported by the model `model`.

### Examples

```julia
modify(model, ObjectiveFunction{ScalarAffineFunction{Float64}}(), ScalarConstantChange(10.0))
```
"""
function modify end

function modify(
    model::ModelLike,
    ci::ConstraintIndex,
    change::AbstractFunctionModification,
)
    return throw_modify_not_allowed(ci, change)
end

function modify(
    model::ModelLike,
    attr::ObjectiveFunction,
    change::AbstractFunctionModification,
)
    return throw_modify_not_allowed(attr, change)
end