using ModelingToolkit, Unitful, OrdinaryDiffEq, DiffEqJump
using Test
MT = ModelingToolkit
@parameters τ [unit = u"ms"]
@variables t [unit = u"ms"] E(t) [unit = u"kJ"] P(t) [unit = u"MW"]
D = Differential(t)

@test MT.get_unit(t) == u"ms"
@test MT.get_unit(E) == u"kJ"
@test MT.get_unit(τ) == u"ms"

@parameters β [unit = u"°"] α [unit = u"°C"] γ [unit = 1u"s"]
@test_throws MT.ValidationError MT.get_unit(β)
@test_throws MT.ValidationError MT.get_unit(α)
@test_throws MT.ValidationError MT.get_unit(γ)

unitless = Unitful.unit(1)
@test MT.get_unit(0.5) == unitless
@test MT.get_unit(t) == u"ms"
@test MT.get_unit(P) == u"MW"
@test MT.get_unit(τ) == u"ms"

@test MT.get_unit(τ^-1) == u"ms^-1"
@test MT.equivalent(MT.get_unit(D(E)),u"MW")
@test MT.equivalent(MT.get_unit(E/τ), u"MW")
@test MT.get_unit(2*P) == u"MW"
@test MT.get_unit(t/τ) == unitless
@test MT.equivalent(MT.get_unit(P - E/τ),u"MW")
@test MT.equivalent(MT.get_unit(D(D(E))),u"MW/ms")

@test MT.get_unit(1.0^(t/τ)) == unitless
@test MT.get_unit(exp(t/τ)) == unitless
@test MT.get_unit(sin(t/τ)) == unitless
@test MT.get_unit(sin(1u"rad")) == unitless
@test MT.get_unit(t^2) == u"ms^2"

@test !MT.validate(E^1.5 ~ E^(t/τ))
@test MT.validate(E^(t/τ) ~ E^(t/τ))

eqs = [D(E) ~ P - E/τ
        0 ~ P]
@test MT.validate(eqs)
@named sys = ODESystem(eqs)
@named sys = ODESystem(eqs, t, [P, E], [τ])

@test !MT.validate(D(D(E)) ~ P)
@test !MT.validate(0 ~ P + E*τ)

#Unit-free
@variables x y z u
@parameters σ ρ β
eqs = [0 ~ σ*(y - x)]
@test MT.validate(eqs)

#Array variables
@variables t x[1:3,1:3](t)
D = Differential(t)
eqs = D.(x) .~ x
ODESystem(eqs,name=:sys)

# Array ops
using Symbolics: unwrap, wrap
using LinearAlgebra
@variables t
sts = @variables x[1:3](t) y(t)
ps = @parameters p[1:3] = [1, 2, 3]
D = Differential(t)
eqs = [
       collect(D.(x) ~ x)
       D(y) ~ norm(x)*y
      ]
ODESystem(eqs, t, [sts...;], [ps...;],name=:sys)

#= Not supported yet b/c iterate doesn't work on unitful array
# Array ops with units
@variables t [unit =u"s"]
sts = @variables x[1:3](t) [unit = u"kg"] y(t) [unit = u"kg"]
ps = @parameters b [unit = u"s"^-1]
D = Differential(t)
eqs = [
       collect(D.(x) ~ b*x)
       D(y) ~ b*norm(x)
      ]
ODESystem(eqs, t, [sts...;], [ps...;])

#Array variables with units
@variables t [unit = u"s"] x[1:3,1:3](t) [unit = u"kg"]
@parameters a [unit = u"s"^-1]
D = Differential(t)
eqs = D.(x) .~ a*x
ODESystem(eqs)
=#

#Difference equation with units
@parameters t [unit = u"s"] a [unit = u"s"^-1]
@variables x(t) [unit = u"kg"]
δ = Differential(t)
D = Difference(t; dt = 0.1u"s")
eqs = [
    δ(x) ~ a*x
]
de = ODESystem(eqs, t, [x, y], [a],name=:sys)

@parameters t
@variables y[1:3](t)
@parameters k[1:3]
D = Differential(t)

eqs = [D(y[1]) ~ -k[1]*y[1] + k[3]*y[2]*y[3],
       D(y[2]) ~  k[1]*y[1] - k[3]*y[2]*y[3] - k[2]*y[2]^2,
       0 ~  y[1] + y[2] + y[3] - 1]

@named sys = ODESystem(eqs,t,y,k)

# Nonlinear system
@parameters a [unit = u"kg"^-1]
@variables x [unit = u"kg"]
eqs = [
    0 ~ a*x
]
@named nls = NonlinearSystem(eqs, [x], [a])

# SDE test w/ noise vector
@parameters τ [unit = u"ms"] Q [unit = u"MW"]
@variables t [unit = u"ms"] E(t) [unit = u"kJ"] P(t) [unit = u"MW"]
D = Differential(t)
eqs = [D(E) ~ P - E/τ
      P ~ Q]

noiseeqs = [0.1u"MW",
            0.1u"MW"]
@named sys = SDESystem(eqs, noiseeqs, t, [P, E], [τ, Q])
# With noise matrix
noiseeqs = [0.1u"MW" 0.1u"MW"
            0.1u"MW" 0.1u"MW"]
@named sys = SDESystem(eqs,noiseeqs, t, [P, E], [τ, Q])

noiseeqs = [0.1u"MW" 0.1u"MW"
            0.1u"MW" 0.1u"s"]
@test !MT.validate(eqs,noiseeqs)

#Test non-trivial simplifications
@variables t [unit = u"s"] V(t) [unit = u"m"^3] L(t) [unit = u"m"]
@parameters v [unit = u"m/s"] r [unit =u"m"^3/u"s"]
D = Differential(t)
eqs = [D(L) ~ v,
       V ~ L^3]
@named sys = ODESystem(eqs)
sys_simple = structural_simplify(sys)

eqs = [D(V) ~ r,
       V ~ L^3]
@named sys = ODESystem(eqs)
sys_simple = structural_simplify(sys)

@variables V [unit = u"m"^3] L [unit = u"m"]
@parameters v [unit = u"m/s"] r [unit = u"m"^3/u"s"] t [unit = u"s"]
eqs = [V ~ r*t,
       V ~ L^3]
@named sys = NonlinearSystem(eqs, [V, L], [t, r])
sys_simple = structural_simplify(sys)

eqs = [L ~ v*t,
       V ~ L^3]
@named sys = NonlinearSystem(eqs, [V,L], [t,r])
sys_simple = structural_simplify(sys)

#Jump System
@parameters β [unit = u"(mol^2*s)^-1"] γ [unit = u"(mol*s)^-1"] t [unit = u"s"] jumpmol [unit = u"mol"]
@variables S(t) [unit = u"mol"] I(t) [unit = u"mol"] R(t) [unit = u"mol"]
rate₁   = β*S*I
affect₁ = [S ~ S - 1*jumpmol, I ~ I + 1*jumpmol]
rate₂   = γ*I
affect₂ = [I ~ I - 1*jumpmol, R ~ R + 1*jumpmol]
j₁      = ConstantRateJump(rate₁, affect₁)
j₂      = VariableRateJump(rate₂, affect₂)
js      = JumpSystem([j₁, j₂], t, [S, I, R], [β, γ],name=:sys)

affect_wrong = [S ~ S - jumpmol, I ~ I + 1]
j_wrong      = ConstantRateJump(rate₁, affect_wrong)
@test_throws MT.ValidationError JumpSystem([j_wrong, j₂], t, [S, I, R], [β, γ],name=:sys)

rate_wrong   = γ^2*I
j_wrong     = ConstantRateJump(rate_wrong, affect₂)
@test_throws MT.ValidationError JumpSystem([j₁, j_wrong], t, [S, I, R], [β, γ],name=:sys)

# mass action jump tests for SIR model
maj1 = MassActionJump(2*β/2, [S => 1, I => 1], [S => -1, I => 1])
maj2 = MassActionJump(γ, [I => 1], [I => -1, R => 1])
@named js3  = JumpSystem([maj1, maj2], t, [S,I,R], [β,γ])

#Test unusual jump system
@parameters β γ t
@variables S(t) I(t) R(t)

maj1 = MassActionJump(2.0, [0 => 1], [S => 1])
maj2 = MassActionJump(γ, [S => 1], [S => -1])
@named js4  = JumpSystem([maj1, maj2], t, [S], [β, γ])


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 1.0.5
#   kernelspec:
#     display_name: Julia 1.4.2
#     language: julia
#     name: julia-1.4
# ---

# ### MRI reconstruction demo
# Single-coil under-sampled data reconstructed two ways:
# * edge-preserving regularization (akin to corner-rounding anisotropic TV) using gradient-based methods
# * synthesis sparsity via 1-norm of wavelet detail coefficients using fast proximal methods
#
# Although this is a 2D demo, the k-space sampling pattern used here would be appropriate only for a 3D scan with phase encoding in the two dimensions shown here.  So this demo is akin to a single slice of a 3D scan with 2D undersampling.
#
# This demo also reproduces Figs. 1 & 3 in the survey paper
# "Optimization methods for MR image reconstruction"
# in Jan 2020 IEEE Signal Processing Magazine
# [http://doi.org/10.1109/MSP.2019.2943645]
#
# 2019-03-06 Jeff Fessler, University of Michigan

# load all packages needed for this demo
# https://github.com/JeffFessler/MIRT.jl:
using MIRT: ellipse_im, prompt, embed, Afft, diffl_map, Aodwt
using MIRTjim: jim
using MIRT: ncg, ogm_ls, pogm_restart
using LinearAlgebra
using Plots; default(markerstrokecolor=:auto)
#using LinearMapsAA
using FFTW
using Random: seed!

# ### Create (synthetic) data

# shepp-logan phantom
nx,ny = 192,256
Xtrue = ellipse_im(ny, oversample=2)[Int((ny-nx)/2+1):Int(ny-(ny-nx)/2),:]
#Xtrue = ellipse_im(192, ny=256, oversample=2) # too small
nx,ny = size(Xtrue)
jim(Xtrue, "true image")
#prompt()

# +
#savefig("xtrue.pdf")

# +
# generate noisy, under-sampled k-space data
M,N = nx,ny
seed!(0); sampfrac = 0.3; samp = rand(M,N) .< sampfrac; sig = 1
mod2 = (N) -> mod.((0:N-1) .+ Int(N/2), N) .- Int(N/2)
samp .|= (abs.(mod2(M)) .< Int(M/8)) * (abs.(mod2(N)) .< Int(N/8))' # fully sampled center
@show sum(samp) / (M*N) # under-sampling factor

ytrue = fft(Xtrue)[samp]
y = ytrue + sig * √(2) * randn(ComplexF32, size(ytrue)) # complex noise!
ysnr = 20 * log10(norm(ytrue) / norm(y-ytrue))
@show ysnr;
# -

# display k-space sampling and zero-filled data
logger = (x;min=-6) -> log10.(max.(abs.(x) / maximum(abs.(x)), (10.)^min))
p1 = jim(samp, fft0=true, title="k-space sampling")
#savefig(p1, "samp.pdf")
jim(:abswarn, false) # suppress warnings about showing magnitude
p2 = jim(logger(embed(ytrue,samp)), fft0=true, title="k-space |data|")
plot(p1,p2)

# ### Prepare to reconstruct
# Creating a system matrix (encoding matrix) and an initial image
# The system matrix is a `LinearMapAA` object, akin to a `fatrix` in Matlab MIRT.

# system model from MIRT
F = Afft(samp) # operator!

# initial image based on zero-filled reconstruction
nrmse = (x) -> norm(x - Xtrue) / norm(Xtrue)
X0 = 1/(M*N) * (F' * y)
@show nrmse(X0)
jim(X0, "|X0|: initial image")
#prompt()

# ### Case 1: Edge-preserving regularization
#
# The cost function here is
# $
# 1/2 \|A x - y\|^2 + \beta \; \psi(T x)
# $
# where $\psi$ is the Fair potential function
# and $T$ is first-order finite differences in 2D,
# again implemented as a `LinearMapAA` object.

# Cost function for edge-preserving regularization
A = F
delta = 0.1
reg = 0.01 * M * N
# finite differences sparsifying transform for anisotropic TV:
T = diffl_map((M,N), 1:2 ; T=ComplexF32) # LinearMapAO
pot = (z,del) -> del^2 * (abs(z)/del - log(1 + abs(z)/del)) # Fair potential function
cost = (x) -> 1/2 * norm(A * x - y)^2 + reg * sum(pot.(T*x, delta))
dpot = (z,del) -> z / (1 + abs(z)/del) # potential derivative

# Nonlinear CG recon using edge-preserving regularization
fun = (x,iter) -> (cost(x), nrmse(x), time())
niter = 50
if !@isdefined(Xcg)
    Xcg, out_cg = ncg([A, T],
        [u -> u - y, v -> reg * dpot.(v,delta)], [t -> 1, t -> reg],
        X0, niter=niter, fun=fun)
end
@show nrmse(Xcg)

jim(Xcg, "CG: edge-preserving regularization")
#prompt()

# +
#savefig("xcg.pdf")
# -

# Nesterov FGM recon using edge-preserving regularization
if !@isdefined(Xfgm)
    fun_fgm = (iter,xk,yk,rs) -> (cost(yk), nrmse(yk), rs)
    f_grad = (x) -> A'*(A*x - y) + reg * (T' * dpot.(T * x, delta))
    L_f = N*M + 8*reg
    Xfgm, out_fgm = pogm_restart(X0, cost, f_grad, L_f;
        mom=:fpgm, niter=niter, fun=fun_fgm)
    @show nrmse(Xfgm)
end
jim(Xfgm, "FGM recon with edge-preserving regularization")
#prompt()

# OGM line-search recon using edge-preserving regularization
if !@isdefined(Xogm)
    Xogm, out_ogm = ogm_ls([A, T],
        [u -> u - y, v -> reg * dpot.(v,delta)], [t -> 1, t -> reg],
        X0, niter=niter, ninner=20, fun=fun)
end
@show nrmse(Xogm)

jim(Xogm, "OGM recon with edge-preserving regularization")
#prompt()

p1 = jim(Xtrue, "true")
p2 = jim(X0, "X0: initial")
p3 = jim(Xcg, "CG recon")
p5 = jim(X0 - Xtrue, "X0 error")
p6 = jim(Xcg - Xtrue, "Xcg error")
plot(p1, p2, p3, p1, p5, p6)
#plot(p1, p2, p3, layout=(1,3))
#plot(p5, p6)
#prompt()

# POGM - suboptimal in smooth case
if false
    fun_pogm = (iter,xk,yk,rs) -> (cost(xk), nrmse(xk), rs)
    f_grad = (x) -> A'*(A*x - y) + reg * (T' * dpot.(T * x, delta))
    L_f = N*M + 8*reg
    Xpogm, out_pogm = pogm_restart(X0, cost, f_grad, L_f;
        mom=:pogm, niter=niter, fun=fun_pogm)
    @show nrmse(Xpogm)
end

# plot cost vs iteration
cost_cg = [out_cg[k][1] for k=1:niter+1]
cost_fgm = [out_fgm[k][1] for k=1:niter+1]
cost_ogm = [out_ogm[k][1] for k=1:niter+1]
cost_min = min(minimum(cost_cg), minimum(cost_ogm))
plot(xlabel="iteration k", ylabel="Relative Cost", ytick=[0,8e4])
scatter!(0:niter, cost_fgm .- cost_min, markershape=:square, label="Cost FGM")
scatter!(0:niter, cost_ogm .- cost_min, markershape=:utriangle, label="Cost OGM")
scatter!(0:niter, cost_cg  .- cost_min, label="Cost CG")
#prompt()

# +
#savefig("cost_ogm_cg.pdf")
#savefig("cost_fgm_ogm_cg.pdf")
# -

# plot nrmse vs iteration
nrmse_cg = [out_cg[k][2] for k=1:niter+1]
nrmse_fgm = [out_fgm[k][2] for k=1:niter+1]
nrmse_ogm = [out_ogm[k][2] for k=1:niter+1]
#nrmse_pogm = [out_pogm[k][2] for k=1:niter+1]
plot(xlabel="iteration k", ylabel="NRMSE")
scatter!(0:niter, nrmse_fgm, markershape=:square, label="NRMSE FGM")
scatter!(0:niter, nrmse_ogm, markershape=:utriangle, label="NRMSE OGM")
scatter!(0:niter, nrmse_cg, label="NRMSE CG")
#scatter!(0:niter, nrmse_pogm, label="NRMSE POGM")
#prompt()

#savefig("nrmse_ogm_cg.pdf")

#nrmse(x_pogm), nrmse(X0)

#?pogm_restart

# ### Case 2: Wavelet sparsity in synthesis form
#
# The cost function here is
# $
# 1/2 \|A x - y\|^2 + \beta \; \|W x\|_1
# $
# where
# and $W$ is an orthogonal discrete (Haar) wavelet transform,
# again implemented as a `LinearMapAA` object.
# Because $W$ is unitary,
# we make the change of variables
# $z = W x$
# and solve for $z$
# and then let $x = W' z$
# at the end.

# Orthogonal discrete wavelet transform (LinearMapAO)
W, scales, _ = Aodwt((M,N) ; T = eltype(T))
jim(real(W * Xtrue) .* (scales .> 0), "wavelet detail coefficients")
#prompt()

# Cost function for edge-preserving regularization
Az = F * W'
Fnullz = (z) -> 0 # not needed
f_gradz = (z) -> Az' * (Az * z - y)
f_Lz = M*N # single coil cartesian DFT
regz = 0.03 * M * N # oracle from Xtrue wavelet coefficients!
pot = (z,del) -> del^2 * (abs(z)/del - log(1 + abs(z)/del)) # Fair potential function
costz = (z) -> 1/2 * norm(Az * z - y)^2 + reg * norm(z,1) # 1-norm regularizer
soft = (z,c) -> sign(z) * max(abs(z) - c, 0) # soft thresholding
g_prox = (z,c) -> soft.(z, regz * c) # proximal operator
z0 = W * X0;

# #### Run ISTA=PGM and FISTA=FPGM and POGM, the latter two with adaptive restart
# See [http://doi.org/10.1007/s10957-018-1287-4]

# +
# Run proximal methods
function fun_ista(iter, xk_z, yk, is_restart)
    if mod(iter, 50) == 0
#        @show iter
    end
    xh = W' * xk_z
return (costz(xk_z), nrmse(xh), is_restart) # , psnr(xh)) # time()
end

function fun_fista(iter, xk, yk_z, is_restart)
    xh = W' * yk_z
    return (costz(yk_z), nrmse(xh), is_restart) # , psnr(xh)) # time()
end

z_ista, out_ista = pogm_restart(z0, Fnullz, f_gradz, f_Lz; mom=:pgm, niter=niter,
    restart=:none, restart_cutoff=0., g_prox=g_prox, fun=fun_ista)
Xista = W'*z_ista
@show nrmse(Xista)

z_fista, out_fista = pogm_restart(z0, Fnullz, f_gradz, f_Lz; mom=:fpgm, niter=niter,
    restart=:gr, restart_cutoff=0., g_prox=g_prox, fun=fun_fista)
Xfista = W'*z_fista
@show nrmse(Xfista)

z_pogm, out_pogm = pogm_restart(z0, Fnullz, f_gradz, f_Lz; mom=:pogm, niter=niter,
    restart=:gr, restart_cutoff=0., g_prox=g_prox, fun=fun_fista)
Xpogm = W'*z_pogm
@show nrmse(Xpogm)

jim(Xfista, "FISTA/FPGM")
jim(Xpogm, "POGM with ODWT")
#prompt()

# +
#savefig("xpogm_odwt.pdf")
# -

# plot cost vs iteration
cost_ista = [out_ista[k][1] for k=1:niter+1]
cost_fista = [out_fista[k][1] for k=1:niter+1]
cost_pogm = [out_pogm[k][1] for k=1:niter+1]
cost_min = min(minimum(cost_ista), minimum(cost_pogm))
plot(xlabel="iteration k", ylabel="Relative cost")
scatter!(0:niter, cost_ista  .- cost_min, label="Cost ISTA")
scatter!(0:niter, cost_fista .- cost_min, markershape=:square, label="Cost FISTA")
scatter!(0:niter, cost_pogm  .- cost_min, markershape=:utriangle, label="Cost POGM")
#prompt()

# +
#savefig("cost_pogm_odwt.pdf")
# -

# plot nrmse vs iteration
nrmse_ista = [out_ista[k][2] for k=1:niter+1]
nrmse_fista = [out_fista[k][2] for k=1:niter+1]
nrmse_pogm = [out_pogm[k][2] for k=1:niter+1]
plot(xlabel="iteration k", ylabel="NRMSE", ylim=[0.01,0.08])
scatter!(0:niter, nrmse_ista, label="NRMSE ISTA")
scatter!(0:niter, nrmse_fista, markershape=:square, label="NRMSE FISTA")
scatter!(0:niter, nrmse_pogm, markershape=:utriangle, label="NRMSE POGM")
#prompt()

# +
#savefig("nrmse_pogm_odwt.pdf")
# -

p1 = jim(Xtrue, "true")
p2 = jim(X0, "X0: initial")
p3 = jim(Xpogm, "POGM recon")
p5 = jim(X0 - Xtrue, "X0 error")
p6 = jim(Xpogm - Xtrue, "Xpogm error")
plot(p1, p2, p3, p1, p5, p6)
#plot(p1, p2, p3, layout=(1,3))
#plot(p5, p6)
#prompt()


function get_sizes(obj::Objective)
    n = size(obj.Q,1)
    m = size(obj.R,1)
    return n,m
end

function get_sizes(X::Matrix,U::Matrix)
    n,N = size(X)
    m = size(U,1)
    return n,m,N
end

function convertInf!(A::Matrix,infbnd=1.1e20)
    infs = isinf.(A)
    A[infs] = sign.(A[infs])*infbnd
    return nothing
end

function get_weights(method,N::Int)
    if method == :trapezoid
        weights = ones(N)
        weights[[1,end]] = 0.5
    elseif method == :hermite_simpson_separated
        weights = ones(N)*2/3
        weights[2:2:end] = 4/3
        weights[[1,end]] = 1/3
    end
    weights *= dt
    return weights
end

function get_bounds(obj::Objective,N::Int)
    n,m = get_sizes(obj)
    lb = zeros((n+m),N)
    ub = zeros((n+m),N)

    lb[1:n,:] .= obj.x_min
    ub[1:n,:] .= obj.x_max
    lb[n+(1:m),:] .= obj.u_min
    ub[n+(1:m),:] .= obj.u_max

    # Initial Condition
    lb[1:n,1] .= obj.x0
    ub[1:n,1] .= obj.x0

    # Terminal Constraint
    lb[1:n,N] .= obj.xf
    ub[1:n,N] .= obj.xf

    # Convert Infinite bounds
    convertInf!(lb)
    convertInf!(ub)
    return vec(lb), vec(ub)
end

"""
Stack state and controls for all time steps into a single vector of variables
Z = [X1,U1,X2,U2,..,] (all transposed)
"""
function packZ(X,U)
    n, N = size(X)
    m = size(U,1)
    Z = zeros(n+m,N)
    Z[1:n,:] .= X
    Z[n+1:end,1:end] .= U
    Z = vec(Z)
end

function unpackZ(Z, sze)
    n,m,N = sze
    Z = reshape(Z,n+m,N)
    X = Z[1:n,:]
    U = Z[n+1:end,:]
    return X,U
end

function get_pack(Z, obj)
    n,m = get_sizes(obj)
    return (n,m, length(Z)//(n*m))
end

function get_initial_state(obj::Objective, N::Int)
    n,m = get_sizes(obj)
    X0 = line_trajectory(obj.x0, obj.xf, N)
    U0 = zeros(m,N)
    Z0 = packZ(X0,U0)
end

"""
Evaluate Objective Value
"""
function cost(X,U,weights,obj::Objective)
    N = size(X,2)
    Qf = obj.Qf; xf = obj.xf;
    J = zeros(eltype(X),N)
    for k = 1:N
        J[k] = cost_k(X[:,k],U[:,k],obj)
    end
    J = weights'J
    J += 0.5*(X[:,N] - xf)'*Qf*(X[:,N] - xf)
end

function cost_k(x,u,obj::Objective)
    xf = obj.xf
    0.5*(x-xf)'obj.Q*(x-xf) + 0.5*u'obj.R*u
end

"""
Gradient of Objective
"""
function cost_gradient(X,U,weights,obj::Objective)
    Q = obj.Q; xf = obj.xf; R = obj.R; Qf = obj.Qf;
    n,m,N = get_sizes(X,U)
    grad_f = zeros(eltype(X),(n+m)*N)
    grad_f = reshape(grad_f, n+m, N)
    for k = 1:N
        grad_f[1:n,k] = weights[k]*Q*(X[:,k] - xf) # weight includes dt
        grad_f[n+1:end,k] = weights[k]*R*U[:,k]
    end
    grad_f[1:n,N] += Qf*(X[:,N] - xf)
    return vec(grad_f)
end

"""
Evaluate constraint values
[ dt*(f(x1)+f(x2))/2 - x2 + x1     ] Collocation Constraints
[ dt*(f(x2)+f(x3))/2 - x3 + x2     ]
[           ...                 ]
[ dt*(f(xN-1)+f(xN))/2 - xN + xN-1 ]
"""
function calc_constraints(X,U,method,dt,f::Function)
    n,m,N = get_sizes(X,U)
    g = zeros(eltype(X),(N-1)*n)
    g = reshape(g,n,N-1)
    if method == :trapezoid
        for k = 1:N-1
            # Collocation Constraints
            g[:,k] = dt*( f(X[:,k+1],U[:,k+1]) + f(X[:,k],U[:,k]) )/2 - X[:,k+1] + X[:,k]
        end
    elseif method == :hermite_simpson_separated
        fVal = zeros(eltype(X),n,N)
        for k = 1:N
            fVal[:,k] = f(X[:,k],U[:,k])
        end
        iLow = 1:2:N-1
        iMid = iLow + 1
        iUpp = iMid + 1

        midpoints =  - X[:,iMid] + (X[:,iLow]+X[:,iUpp])/2 + dt*(fVal[:,iLow] - fVal[:,iUpp])/8
        collocation = - X[:,iUpp] + X[:,iLow] + dt*(fVal[:,iLow] + 4*fVal[:,iMid] + fVal[:,iUpp])/6
        g[:,iLow] = collocation
        g[:,iMid] = midpoints
    end
    # reshape(g,n*(N-1),1)
    return vec(g)
end

"""
Constraint Jacobian
"""
function constraint_jacobian(X,U,dt,method,F::Function)
    n,m,N = get_sizes(X,U)
    Z = packZ(X,U)
    jacob_g = zeros((N-1)*n,N*(n+m))
    z = reshape(Z,n+m,N)

    if method == :trapezoid
        nm = n+m
        # First time step
        fz = F(z[:,1])
        jacob_g[1:n,1:n+m] .= dt*fz/2+eye(n,n+m)

        # Loop over time steps
        for k = 2:N-1
            off_1 = (k-1)*(n)
            off_2 = (k-1)*(n+m)
            # Calculate (n,n+m) Jacobian of both states and controls
            fz = F(z[:,k])
            jacob_g[off_1-n+(1:n),off_2+(1:n+m)] .= dt*fz/2 - eye(n,n+m)
            jacob_g[off_1 + (1:n),off_2+(1:n+m)] .= dt*fz/2 + eye(n,n+m)
        end

        # Last time step
        fz = F(z[:,N])
        jacob_g[end-n+1:end,end-n-m+1:end] = dt*fz/2-eye(n,n+m)

    elseif method == :hermite_simpson_separated
        nSeg = Int((N-1)/2)
        jacob_g = zeros(n*(N-1),(n+m)N)
        fz1 = F(z[:,1])
        Inm = Matrix(I,n,n+m)

        function calc_block(k)
            vals = zeros(2n,3(n+m))
            fz = F(z[:,k])
            vals[0+(1:n),(1:n+m)] .= dt*fz/6 + Inm
            vals[n+(1:n),(1:n+m)] .= dt*fz/8 + Inm/2
            fm = F(z[:,k+1])
            vals[0+(1:n),n+m+(1:n+m)] .= 2*dt*fm/3
            vals[n+(1:n),n+m+(1:n+m)] .= -Inm
            fz1 .= F(z[:,k+2])
            vals[0+(1:n),2(n+m)+(1:n+m)] .=  dt*fz1/6 - Inm
            vals[n+(1:n),2(n+m)+(1:n+m)] .= -dt*fz1/8 + Inm/2
            return vals
        end

        for i = 1:nSeg
            off_1 = 2(i-1)*(n)
            off_2 = 2(i-1)*(n+m)
            k = 2i-1

            jacob_g[off_1+(1:2n), off_2+(1:3(n+m))] = calc_block(k)
        end
    end

    return jacob_g
end


module Apriori

using DataFrames
using Combinatorics
using CSV

ATTRIBUTE_TYPE = Int32

function apriori(
    data::DataFrame,
    min_relative_support=0.2,
    min_confidence=0.3;
    X_in_antecedent=Vector{Symbol}()::Vector{Symbol},
    antecedent_in_X=Vector{Symbol}()::Vector{Symbol},
    X_in_consequent=Vector{Symbol}()::Vector{Symbol},
    consequent_in_X=Vector{Symbol}()::Vector{Symbol},
    )::Vector{Pair{Tuple{Vector{Symbol},Vector{Symbol}},Tuple{Float64,Float64}}}

    attributes = propertynames(data)::Array{Symbol,1}

    symbol2index = Dict([x => y for (x, y) in zip(attributes, one(ATTRIBUTE_TYPE):convert(ATTRIBUTE_TYPE, length(attributes)))])

    filters = __prepare_filters(symbol2index, X_in_antecedent, antecedent_in_X, X_in_consequent, consequent_in_X)

    frequent_itemsets = __apriori_frequent_itemsets(data, min_relative_support)

    rules = __apriori_rule_gen(frequent_itemsets, min_confidence, filters)
    map(x -> translate_rule(x, attributes), rules)
end

function apriori_frequent_itemsets(data::DataFrame, min_relative_support=0.2)
    attributes = propertynames(data)
    freq_itemsets = __apriori_frequent_itemsets(data, min_relative_support)
    symbol2index = Dict([x => y for (x, y) in zip(attributes, 1:ncol(data))])
    frequent_itemsets = map(x -> Set{Symbol}(getindex.(Ref(attributes), x[1])) => x[2], freq_itemsets)
end

function __apriori_frequent_itemsets(data::DataFrame, min_relative_support=0.2)::Vector{Pair{Vector{ATTRIBUTE_TYPE},Int64}}
    supp = floor(min_relative_support * nrow(data))

    frequent_itemsets = Vector{Pair{Vector{ATTRIBUTE_TYPE},Int64}}([[]=>nrow(data)])

    itemsets = [Vector{ATTRIBUTE_TYPE}([i]) for i in 1:ncol(data)]
    infrequents = Set{Vector{ATTRIBUTE_TYPE}}()

    while true
        itesmets_w_support = map(x -> x => support(data, x), itemsets)
        frequent_itesmets_w_support = filter(x -> x[2] > supp, itesmets_w_support)
        append!(frequent_itemsets, frequent_itesmets_w_support)

        if (length(frequent_itesmets_w_support) <= 1) break end

        itemsets = filter(x -> !any(subsets(x, length(x) - 1) .∈ (infrequents,)), merge_vectors(map(x -> x[1], frequent_itesmets_w_support)))
        infrequents = map(y -> y[1], filter(x -> x[2] <= supp, itesmets_w_support))
    end

    frequent_itemsets
end

function __prepare_filters(translate_dict,
    X_in_antecedent=Vector{Symbol}()::Vector{Symbol},
    antecedent_in_X=Vector{Symbol}()::Vector{Symbol},
    X_in_consequent=Vector{Symbol}()::Vector{Symbol},
    consequent_in_X=Vector{Symbol}()::Vector{Symbol},
    )

    t = [
        :x_in_ant => X_in_antecedent,
        :ant_in_x => antecedent_in_X,
        :x_in_con => X_in_consequent,
        :con_in_x => consequent_in_X
    ]

    @assert all(X_in_antecedent .∉ (X_in_consequent,))

    filters = Dict(filter(x -> length(x[2]) > 0, map(x -> x[1] => sort(translate_itemset(x[2], translate_dict)), t)))
end

function apriori_rule_gen(
    frequent_itemsets::Vector{Pair{Set{Symbol},supp_int}},
    attribute_names::Vector{Symbol},
    min_confidence=0.3;
    X_in_antecedent=Vector{Symbol}()::Vector{Symbol},
    antecedent_in_X=Vector{Symbol}()::Vector{Symbol},
    X_in_consequent=Vector{Symbol}()::Vector{Symbol},
    consequent_in_X=Vector{Symbol}()::Vector{Symbol},
    ) where {supp_int <: Signed}

    symbol2index = Dict([x => convert(ATTRIBUTE_TYPE, y) for (x, y) in zip(attribute_names, 1:length(attribute_names))])
    frequent_itemsets = map(x -> sort(Vector{ATTRIBUTE_TYPE}(getindex.(Ref(symbol2index), x[1]))) => x[2], frequent_itemsets)

    filters = __prepare_filters(symbol2index, X_in_antecedent, antecedent_in_X, X_in_consequent, consequent_in_X)

    rules = __apriori_rule_gen(frequent_itemsets, min_confidence, filters)
    result = map(x -> translate_rule(x, attribute_names), rules)
    result::Vector{Pair{Tuple{Vector{Symbol},Vector{Symbol}},Float64}}
end

function __apriori_rule_gen(
    frequent_itemsets::Vector{Pair{Vector{ATTRIBUTE_TYPE},supp_int}},
    min_confidence=0.3,
    filters=Dict()
    ) where {supp_int <: Signed}

    df = Dict(frequent_itemsets)

    dataset_size = df[[]]

    all_strong_rules = Vector{Pair{Tuple{Vector{ATTRIBUTE_TYPE},Vector{ATTRIBUTE_TYPE}},Tuple{Float64,Float64}}}()

    for (Z_set, Z_sup) in frequent_itemsets
        #if (length(Z_set) == 1) continue end

        Z = Z_set
        i = 1
        Y = subsets(Z, i)

        if (haskey(filters, :x_in_ant))
            X = filters[:x_in_ant]
            if (all(X .∈ (Z,)))
                Y = filter(x -> x[end] ∉ X, Y)
            else continue end
        end

        if (haskey(filters, :con_in_x))
            X = filters[:con_in_x]
            Y = filter(x -> x[end] ∈ X, Y)
        end

        if (haskey(filters, :ant_in_x))
            X = filters[:ant_in_x]
            X = filter(x -> x[1] ∈ Z, X)
            ZmX = filter(x -> x ∉ X, Z)
            if (length(ZmX) > 0 && length(X) > 0)
                sp = ((X, ZmX) => (Z_sup/dataset_size, Z_sup / df[X]))
                if (sp[2][2] >= min_confidence &&
                    (!haskey(filters, :x_in_ant) || all(filters[:x_in_ant] .∈ (sp[1][1],))) && # ghetto fix na inne filtry żeby obejmowały te wyjątki, ale nadal są powtórzenia wśród tych special cases
                    (!haskey(filters, :con_in_x) || all(sp[1][2] .∈ (filters[:con_in_x],))) &&
                    (!haskey(filters, :x_in_con) || all(filters[:x_in_con] .∈ (sp[1][2],))) &&
                    sp ∉ all_strong_rules)
                    append!(all_strong_rules,  [sp])
                end
            end
            Y = filter(x -> x[end] ∈ X, Y)
            Y = map(x -> vcat(ZmX, x), Y)
        end

        if (haskey(filters, :x_in_con))
            X = filters[:x_in_con]
            if (all(X .∈ (Z,))) # czy wszystkie elementy X są w Z
                ZmX = filter(x -> x ∉ X, Z)
                if (length(ZmX) > 0)
                    sp = (ZmX, X) => (Z_sup/dataset_size, Z_sup / df[ZmX] )
                    if (sp[2][2] >= min_confidence &&
                        (!haskey(filters, :x_in_ant) || all(filters[:x_in_ant] .∈ (sp[1][1],))) && # ghetto solution
                        (!haskey(filters, :con_in_x) || all(sp[1][2] .∈ (filters[:con_in_x],))) &&
                        (!haskey(filters, :ant_in_x) || all(sp[1][1] .∈ (filters[:ant_in_x],))) &&
                        sp ∉ all_strong_rules)
                        append!(all_strong_rules, [sp])
                    end
                end
                Y = filter(x -> x[end] ∉ X, Y)
                Y = map(x -> vcat(X, x), Y)
            else continue end
        end

        if (length(Y) > 0 &&  all(length.(Y) .< length(Z)))
            i = length(Y[1])
            while true
                strong_rules = filter(x -> x[2][2] >= min_confidence, map(x -> begin
                    ant = antecedent(Z, x)
                    (ant, x) => (Z_sup/dataset_size, Z_sup / df[ant] )
                end, Y))

                append!(all_strong_rules, strong_rules)
                i += 1
                if (i == length(Z_set) || length(strong_rules) == 0) break end

                Y = merge_vectors(map(x -> x[1][2], strong_rules))
            end
        end
    end
    all_strong_rules::Vector{Pair{Tuple{Vector{ATTRIBUTE_TYPE},Vector{ATTRIBUTE_TYPE}},Tuple{Float64,Float64}}}
end

function merge_vectors(itemsets::Vector{Vector{ATTRIBUTE_TYPE}})
    result = Vector{Vector{ATTRIBUTE_TYPE}}()
    itemsets_len = length(itemsets)::Int64
    for i in 1:itemsets_len
        itemset1 = itemsets[i]
        prefix = itemset1[1:end - 1]
        suffixes = Vector{Vector{ATTRIBUTE_TYPE}}()
        for j in i + 1:itemsets_len
            itemset2 = itemsets[j]
            if prefix != itemset2[1:end - 1]
                break
            end
            push!(suffixes, sort([itemset1[end], itemset2[end]]))
        end
        merged_itemsets = map(x -> vcat(prefix, x), sort(suffixes))
        append!(result, merged_itemsets)
    end
    result::Vector{Vector{ATTRIBUTE_TYPE}}
end

function translate_itemset(element::Vector{Symbol}, dict::Dict{Symbol,ATTRIBUTE_TYPE})
    Vector{ATTRIBUTE_TYPE}(getindex.(Ref(dict), element))
end

function translate_rule(rule::Pair{Tuple{Vector{ATTRIBUTE_TYPE},Vector{ATTRIBUTE_TYPE}},Tuple{Float64,Float64}}, dict::Vector{Symbol})
    (getindex.(Ref(dict), rule[1][1]), getindex.(Ref(dict), rule[1][2])) => rule[2]
end

function support(data::DataFrame, attrs::Vector{ATTRIBUTE_TYPE})
    nrow(data[reduce((x, y) -> x .& y, [data[!, attr] .== true for attr in attrs]),:])
end

function subsets(vec::Vector{ATTRIBUTE_TYPE}, len)
    Vector.(combinations([vec...], len))
end

function antecedent(fvec, vec)
    filter(x -> x ∉ vec, fvec)
end

function dummy_dataset(attrs, rows)
    @assert attrs <= 26 "TODO the attr label fun to generate more"
    attr_names = collect('a':'z')[1:attrs]
    DataFrame([Symbol(i) => rand(Bool, rows) for i in attr_names])
end

# Creates a semi-random dataset that exhibits a strong association rule
function dummy_dataset_biased(attrs, rows, relative_sup, antecedent, consequent, support, confidence=1)
    @assert attrs <= 26 "TODO the attr label fun to generate more"
    attr_names = collect('a':'z')[1:attrs]
    df = DataFrame([Symbol(i) => map(x->x<relative_sup, rand(rows)) for i in attr_names])
    # print(df)
    antecedent_indices = findall(col_name -> col_name in antecedent, attr_names)
    consequent_indices = findall(col_name -> col_name in consequent, attr_names)
    antecedent_count = nrow(filter(row -> all([row[index] for index in antecedent_indices]), df))
    rule_count = nrow(filter(row -> all([row[index] for index in vcat(antecedent_indices, consequent_indices)]), df))
    required_rule_count = rows * support

    remove_consequent_count = required_rule_count * (1-confidence)
    for row in eachrow(df)
        if !all([row[index] for index in vcat(antecedent_indices, consequent_indices)])
            for index in vcat(antecedent_indices, consequent_indices)
                row[index] = true
            end
            if remove_consequent_count>0
                row[consequent_indices[1]] = false
                remove_consequent_count -= 1
            end
            rule_count = rule_count + 1
        end
        if rule_count >= required_rule_count
            break
        end
    end
    return df
end

function mushroom_dataset()
    download("https://gist.githubusercontent.com/krynju/120e623b7b0277e45126baf8aa39edad/raw/f980a374b787cd7e539e821450e4de91f948b5fd/mushroom.csv", "mushroom.csv")
    d = CSV.read("mushroom.csv", DataFrame)
    select!(d, [(atr => (x -> (;[Symbol(atr * "=" * string(value)) => d[!,atr] .== value for value in unique(d[!,atr])]...)) => AsTable)  for atr in names(d)])
end

function pretty(rules::Vector{Pair{Tuple{Vector{Symbol},Vector{Symbol}},Tuple{Float64,Float64}}})
    DataFrame([(antecedent=join(r[1][1],", "), consequent=join(r[1][2],", "), support=r[2][1], confidence=r[2][2])  for r in rules])
end

export apriori, dummy_dataset, dummy_dataset_biased, apriori_rule_gen, apriori_frequent_itemsets, merge_vectors, mushroom_dataset, pretty

end


include("TestFunctions.jl")

# vim:filetype=julia:foldmethod=syntax


"""
Implementation of Koza-type (tree-based) Genetic Programming

The constructor takes following keyword arguments:

- `populationSize`: The size of the population
- `initialization`: A strategy for population initialization (default: `:half`)
  - Possible values: `:grow`, `:full`, `:half`
- `evaluate`: A collection of fitness cases required for evaluation of a fitness function

"""
@kwdef struct TreeGP{T} <: AbstractOptimizer
    populationSize::Integer = 50
    terminals::Vector{Symbol} = [:x, :rand]
    functions::Vector{Symbol} = [:+, :-, :*, :/, :^]
    mindepth::Int = 0
    maxdepth::Int = 10
    initialization::Symbol = :grow
    evaluate::AbstractVector{T} = zeros(T, 0)
end

"""
    initial_population(method::TreeGP, expr::{Expr,Nothing})

Initialize a random population of expressions derived from `expr`.
"""
function initial_population(method::GP, expr::{Expr,Nothing}))
    n = population_size(method)

    return [genrandexpr(method.terminals, method.functions,
                        mindepth=method.mindepth, maxdepth=method.maxdepth)
            for i in 1:n]
end

arity(sym) = 2
function genrandexpr(terms, funcs; maxdepth=2, mindepth=1, method=:grow)
    # println("depth: $maxdepth => $mindepth")
    tl = length(terms)
    fl = length(funcs)
    expr = if (maxdepth <= 0 || ( method == :grow && rand() < tl/(tl+fl) )) && mindepth <= 0
        termsym = rand(terms)
        term = termsym == :rand ? rand() : termsym
        # println("term: $term")
        term
    else
        func = rand(funcs)
        # println("func: $func, ")
        args = Any[]
        for i in 1:arity(func)
            arg = genrandexpr(terms, funcs,
                              maxdepth=maxdepth-1, mindepth=mindepth-1, method=method)
            # println("arg_$i: $arg ")
            push!(args, arg)
        end
        Expr(:call, func, args...)
    end
    return expr
end

terminals = [:x, :rand, 1, 2, 3, 4]
functions = [:+, :-, :*, :/, :^]

ex = genrandexpr(terminals, functions, maxdepth=3, mindepth=2)
exfun = eval(Expr(:->, :x, ex))
exfun(1)

depth(ex) = isa(ex, Expr) ? maximum( depth(e) for e in ex.args )+1 : 0
children(ex) = isa(ex, Expr) ? sum( children(e) for e in ex.args) : 1

children(ex)
depth(ex)
dump(ex)

function subtree(root, idx)
    m, left, right = if !isa(root, Expr) || length(root.args) < 2
        1, root, root
    else
        left, right = root.args[2:3] # left & right nodes
        children(left)+1, left, right
    end
    if m == idx
        return root
    elseif m > idx
        return subtree(left, idx)
    else
        return subtree(right, idx - m)
    end
end
randsubtree(ex) = subtree(ex, rand(1:children(ex)))

[subtree(ex, i) for i in 1:children(ex)]


function neural_ode(model,x,tspan, solver,
                    args...;kwargs...)
  p = destructure(model)
  dudt_(u::TrackedArray,p,t) = restructure(model,p)(u)
  dudt_(u::AbstractArray,p,t) = Flux.data(restructure(model,p)(u))
  prob = ODEProblem(dudt_,x,tspan,p)
  return diffeq_adjoint(p,prob,solver,args...;u0=x,kwargs...)
end

function neural_ode_rd(model,x,tspan,
                       args...;kwargs...)
  dudt_(u,p,t) = model(u)
  _x = x isa TrackedArray ? x : param(x)
  prob = ODEProblem(dudt_,_x,tspan)
  # TODO could probably use vcat rather than collect here
  solve(prob, args...; kwargs...) |> Tracker.collect
end

function neural_dmsde(model,x,mp,tspan,
                      args...;kwargs...)
  dudt_(u,p,t) = model(u)
  g(u,p,t) = mp.*u
  prob = SDEProblem(dudt_,g,param(x),tspan,nothing)
  # TODO could probably use vcat rather than collect here
  solve(prob, args...; kwargs...) |> Tracker.collect
end

# Flux Layer Interface
struct NeuralODE
    model
    tspan
    solver
    args
    kwargs
end

# Optional args and kwargs
NeuralODE(model,tspan,solver,kwargs)= NeuralODE(model,tspan,solver,(),kwargs)
NeuralODE(model,tspan,solver)= NeuralODE(model,tspan,solver,(),())

# Play nice with Flux
Flux.@treelike NeuralODE
Flux.params(n::NeuralODE) = params(n.model)

function (n::NeuralODE)(x)
    return neural_ode(n.model,x,n.tspan,n.solver,n.args...;n.kwargs...)
end


export Point, RouteSection

"""
    Radius Earth [m]

Source: en.wikipedia.org/wiki/Earth_radius
"""
const Rₑ_m = 6371008.7714

"Point type with latitude `ϕ` [deg] and longitude `λ` [deg]"
struct Point{T<:Float64}
    ϕ::T
    λ::T
end

(-)(x::Point, y::Point) = Point(x.ϕ - y.ϕ, x.λ - y.λ)
(rad2deg)(x::Point) = Point(rad2deg(x.ϕ), rad2deg(x.λ))
(deg2rad)(x::Point) = Point(deg2rad(x.ϕ), deg2rad(x.λ))

"RouteSection type with start position pos₁ [Point] and end position pos₂
[Point_deg]"
struct RouteSection{T}
    pos₁::T
    pos₂::T
end

"""
    normalize(value::Float64[, lower::Float64 = 0.0, upper::Float64 = 360.0])

Normalize a `value` to stay within the `lower` and `upper` limit.

Example:
normalize(370.0, lower=0.0, upper=360.0)
10.0
"""
function normalize(value::Float64, lower::Float64 = 0.0, upper::Float64 = 360.0)
    result = value
    if result > upper || value == lower
        value = lower + abs(value + upper) % (abs(lower) + abs(upper))
    end
    if result < lower || value == upper
        value = upper - abs(value - lower) % (abs(lower) + abs(upper))
    end
    value == upper ? result = lower : result = value
    return result
end

#TODO Get a more efficient way to normalize latitude
"""
    normalize(point::Point)

Normalize a `point`.
"""
function normalize(point::Point)
    return Point(asind(sind(point.ϕ)), normalize(point.λ, -180.0, 180.0))
end


#=
Popis
  Funkcia vypocita diameter grafu.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozena vlastnost vrchola :: String
  components : parameter pre prepinanie graf/max_komponent :: Bool

Navratova hodnota
  : diameter maximalneho komponentu :: Float64
  g_diameter : diameter grafu :: Float64

Poznamky
  Funkcia podla hodnoty parametra components vypocita hodnotu
  pre cely graf alebo pre najvacsi komponent.
=#

function graph_diameter(g::GenericGraph; dist_key::AbstractString="weight", components::Bool=false)

  if components
    comp_diameter = Dict{Array{ExVertex,1}, Float64}()
    comp_length = Int64[]

    for comp in connected_components(g)
      comp_diameter[comp] = maximum([vertex_eccentricity(v.index, g) for v in comp])
      push!(comp_length, length(comp))
    end

    max_comp = maximum(comp_length)

    return [d[2] for d=filter!(x -> length(x[1]) == max_comp, collect(zip(keys(comp_diameter), values(comp_diameter))))]
  else
    g_diameter = [vertex_eccentricity(v.index, g) for v in vertices(g)]

    return maximum(g_diameter)
  end

end


using Resample
using Test
using Dates
using DataFrames
using Statistics

times_long = DateTime(2019):Minute(5):DateTime(2019, 1, 31, 23, 55)
times_short = DateTime(2019):Hour(1):DateTime(2019, 1, 31, 23)

inds_long = eachindex(times_long)
inds_short = 1:12:length(times_long)

power_long = rand(length(times_long))
energy_long = power_long / 12

df = DataFrame(
    Index = inds_long,
    Time = times_long,
    Power = power_long,
    Energy = energy_long,
)

power_short_times = resample(power_long, times_long, times_short)
power_short_inds = resample(power_long, inds_long, inds_short)
power_short_step = resample(power_long, times_long, Hour(1))

energy_short_times = resample(energy_long, times_long, times_short, Sum())
energy_short_inds = resample(energy_long, inds_long, inds_short, Sum())
energy_short_step = resample(energy_long, times_long, Hour(1), Sum())

df_short_times = resample(df, :Time, times_short, [First(), Mean(), Sum()])
df_short_inds = resample(df, :Index, inds_short, [First(), Mean(), Sum()])
df_short_step = resample(df, :Time, Hour(1), [First(), Mean(), Sum()])

@testset "Resample" begin
    @test power_short_times == power_short_inds == power_short_step
    @test energy_short_times == energy_short_inds == energy_short_step
    @test df_short_times == df_short_inds == df_short_step

    @test mean(power_long) ≈ mean(power_short_times)
    @test sum(energy_long) ≈ sum(energy_short_times)
end


"""

# ancestor_graph

$(SIGNATURES)

Internal
"""
function ancester_graph(e::NamedArray)
  if sum(size(e)) == 0
    return(e)
  end
  indicator_matrix(Int.(round.(Int, inv(2 * I(size(e, 1)) - e))))
end

"""

## indicator_matrix

$(SIGNATURES)

Internal
"""
function indicator_matrix(e::NamedArray)
  for i in 1:size(e, 1)
    for j in 1:size(e, 2)
      e[i, j] = e[i, j] == 0 ? 0 : 1
    end
  end
  e
end

"""

# transitive_closure

$(SIGNATURES)

Internal
"""
function transitive_closure(a::NamedArray)
  size(a, 1) == 1 && return(a)
  r = copy(a)
  while true
    b = sign.( r * r )
    all( b == r ) && break
    r = b
  end
  r
end

"""

# `induced_covariance_graph`

$(SIGNATURES)

Internal
"""
function induced_covariance_graph(d::DAG, sel, c; debug=false)

  @assert all([c in d.vars for c in sel]) "Selection nodes are not among vertices."

  if isnothing(c)
    cs = []
  elseif typeof(c) == Symbol
    cs = [c]
  elseif length(c) > 0
    cs = c
  end

  @assert all([c in d.vars for c in cs]) "Conditioning nodes are not among vertices."
  #@assert !(length([cs in sel]) > 0) "Conditioning nodes in selected nodes."

  l = setdiff(d.vars, union(sel, cs))
  debug && println(l)
  l = union(l, sel)
  debug && println(l)
  r = union(sel, cs)
  debug && println(r)

  e = edge_matrix(d.a)                  # From adjacency matrix to edge matrix
  debug && println(e)
  al = transpose(StructuralCausalModels.ancester_graph(e[l, l]))
  debug && println(al)
  if length(cs) > 0
    trl = StructuralCausalModels.indicator_matrix( transpose(e[l, cs]) * al)
  else
    trl = al - al
  end
  debug && println(trl)
  dlr = StructuralCausalModels.indicator_matrix(I(length(l)) + transpose(trl) * trl)
  debug && println(dlr)
  cl = StructuralCausalModels.transitive_closure(dlr)
  debug && println(cl)
  out = StructuralCausalModels.indicator_matrix(( al * cl * transpose(al)))
  debug && println(out)
  out = out[sel, sel]
  debug && println(out)
  debug && println(adjacency_matrix(out))
  adjacency_matrix(out)

end

export
  induced_covariance_graph


include("meta.jl")

mutable struct Diffable #<: Real

  value::Float64
  precursors::Vector{Tuple{Float64,Diffable}} #gradient weight + the Diffable node that is dependent during a forward pass of the computation graph (antecedent from a backwards perspective, subsequent from a forward perspective)
  grad

  Diffable(a) = new(a, Vector{Tuple{Float64,Diffable}}(), 0.0)
end

function clear!(x::Diffable)

  x.precursors = similar(x.precursors, 0)
end

"""sadly 8x faster than a cooler one-liner"""
function backwards!(y::Diffable)

    if isempty(y.precursors) # terminal node
      y.grad = 1
    else
      y.grad = 0 # remove this to accumulate
      for (weight,precursor) in y.precursors
        y.grad += weight * backwards!(precursor)
      end
    end

  return y.grad
end

"""
generate conversion functions for a binary operator
e.g. `@create_binaryfns +`
yields
`+(a::Diffable, b) = +(a, Diffable(b))`
`+(a, b::Diffable) = +(Diffable(a), b)`
"""
macro create_conversions(operator)

  quote
    $(esc(operator))(a::Diffable, b::T) where T <: Real = $(esc(operator))(a, Diffable(b))
    $(esc(operator))(a::T, b::Diffable) where T <: Real = $(esc(operator))(Diffable(a), b)
  end
end

"""
generate base functions for a binary operator in terms of arguments `a` and `b`

calculate the result and push the result (weight) and partial derivative into the computation graph

example output:
```
function *(a::Diffable, b::Diffable)

  result = Diffable(a.value * b.value)
  push!(a.precursors, (b.value, result))
  push!(b.precursors, (a.value, result))

  return result
end
```
"""
macro create_binaryfn(operator, ∂a, ∂b, a = :a, b = :b)

  ∂a = replace_expr!(∂a, a => :($a.value), b => :($b.value))
  ∂b = replace_expr!(∂b, a => :($a.value), b => :($b.value))

  quote
    function $(esc(operator))(a::Diffable, b::Diffable)

      result = Diffable($(esc(operator))($a.value, $b.value))
      push!($(:(a.precursors)), ($∂a, result))
      push!($(:(b.precursors)), ($∂b, result))

      return result
    end

    @create_conversions($(esc(operator)))
  end
end

"""unary version"""
macro create_unaryfn(fn, ∂a, a = :a)

  ∂a = replace_expr!(∂a, a => :($a.value))

  quote
    function $(esc(fn))(a::Diffable)

      result = Diffable($(esc(fn))($a.value))
      push!($(:(a.precursors)), ($∂a, result))

      return result
    end
  end

end

# actually define the adjoints:

# binary functions
@create_binaryfn(+, 1.0, 1.0)
@create_binaryfn(-, 1.0, -1.0)
@create_binaryfn(*, b, a)
@create_binaryfn(^, b * a^(b-1), a^b * log(a)) #x^a, a^x

/(a::Diffable, b::Diffable) = a * (b^-1)
@create_conversions /

# unary functions
-(a::Diffable) = -1 * a

@create_unaryfn(inv, -a * a^-2) #used internally for calls to a^-1
@create_unaryfn(exp, result.value) #short-circuit recalculating exp(a)
@create_unaryfn(log, 1/a)
@create_unaryfn(sin, cos(a))
@create_unaryfn(cos, -sin(a))
@create_unaryfn(abs, sign(a)) #equivalent to |x|/x, but handles 0 to avoid random errors at the expense of correctness.

relu(x) = x > 0 ? x : 0
@create_unaryfn(relu, a > 0)