import Distributions: logpdf
import Random: rand
import Base: length
export Gaussian, normal_logpdf, StandardGaussian, GaussianPair, logpdf
export standard_to_natural, natural_to_standard
export standard_to_expectation, expectation_to_standard
export expectation_to_natural, natural_to_expectation
export standard_to_unconstrained, unconstrained_to_standard
export natural_to_unconstrained, unconstrained_to_natural
export expectation_to_unconstrained, unconstrained_to_expectation

"""
    Gaussian{Tm, Tchol_Σ<:Cholesky}

A (multivariate) Gaussian with mean vector `m` and cholesky factorisation `chol_Σ` of the
covariance matrix.
"""
struct Gaussian{Tm<:AbstractVector{<:Real}, Tchol_Σ<:Cholesky}
    m::Tm
    chol_Σ::Tchol_Σ
end

length(x::Gaussian) = length(x.m)

rand(rng::AbstractRNG, x::Gaussian) = x.m + x.chol_Σ.U' * randn(rng, length(x.m))
function rand(rng::AbstractRNG, x::Gaussian, S::Int)
    return x.m .+ x.chol_Σ.U' * randn(rng, length(x.m), S)
end

# Returns a `Vector` of length `size(Y, 2)` containing the logpdf of each column of `Y`
# under `x`.
function logpdf(x::Gaussian, Y::AbstractMatrix{<:Real})
    @assert length(x) == size(Y, 1)
    tmp = length(x) * log(2π) + logdet(x.chol_Σ)
    return .-(tmp .+ diag_Xt_invA_X(x.chol_Σ, Y .- x.m)) / 2
end
logpdf(x::Gaussian, y::AbstractVector{<:Real}) = logpdf(x, reshape(y, :, 1))[1]

# A vectorised logpdf for the normal distribution because Zygote's broadcasting sucks
# if you don't do things correctly :S Even differnetiating through this kind of sucks,
# which is a real shame.
function normal_logpdf(y, m, σ)
    return .-(log(2π) .+ 2 .* log.(σ) .+ ((y .- m) ./ σ).^2) ./ 2
end

"""
    StandardGaussian

`D`-dimensional multivariate Gaussian with zero mean and identity covariance.
"""
struct StandardGaussian
    D::Int
end

length(X::StandardGaussian) = X.D

rand(rng::AbstractRNG, X::StandardGaussian) = randn(rng, length(X))
rand(rng::AbstractRNG, X::StandardGaussian, S::Int) = randn(rng, length(X), S)

function logpdf(X::StandardGaussian, y::AbstractVector{<:Real})
    return -(X.D * log(2π) + sum(abs2, y)) / 2
end

function logpdf(X::StandardGaussian, Y::AbstractMatrix{<:Real})
    return .-(X.D * log(2π) .+ vec(sum(abs2, Y;  dims=1))) / 2
end

"""
    GaussianPair{Tx1<:Gaussian, Tx2<:Gaussian}

A pair of independent Gaussians `x1` and `x2`.
"""
struct GaussianPair{Tx1<:Gaussian, Tx2<:Gaussian}
    x1::Tx1
    x2::Tx2
end

length(x::GaussianPair) = length(x.x1) + length(x.x2)

function rand(rng::AbstractRNG, x::GaussianPair, S::Int)
    return (x1=rand(rng, x.x1, S), x2=rand(rng, x.x2, S))
end
rand(rng::AbstractRNG, x::GaussianPair) = (x1=rand(rng, x.x1), x2=rand(rng, x.x2))

function logpdf(x::GaussianPair, Y::NamedTuple{(:x1, :x2)})
    return logpdf(x.x1, Y.x1) + logpdf(x.x2, Y.x2)
end

#
# Constrained parametrisations.
#

"""
    natural_to_standard(θ₁::AbstractVector{<:Real}, θ₂::AbstractMatrix{<:Real})

Convert from natural to standard parametrisation.
"""
function natural_to_standard(θ₁::AbstractVector{<:Real}, θ₂::Symmetric{<:Real})
    S = Symmetric(inv(cholesky(-θ₂)) ./ 2)
    return S * θ₁, S
end

# forwards-mode rule for natural_to_standard. Avoids some repeated computation.
function frule(
    ::typeof(natural_to_standard),
    θ₁::AbstractVector{<:Real},
    θ₂::Symmetric{<:Real},
)
    C = cholesky(-θ₂)
    S = Symmetric(inv(C) ./ 2)
    m = S * θ₁
    function natural_to_standard_pushforward(Δself, Δθ₁, Δθ₂)
        ΔS = Symmetric((C \ Matrix((C \ Δθ₂')')) ./ 2)
        Δm = ΔS * θ₁ + S * Δθ₁
        return (Δm, ΔS)
    end
    return (m, S), natural_to_standard_pushforward
end

"""
    standard_to_natural(m::AbstractVector{<:Real}, S::Symmetric{<:Real})

Convert from standard to natural parametrisation.
"""
function standard_to_natural(m::AbstractVector{<:Real}, S::Symmetric{<:Real})
    C = inv(cholesky(S))
    return C * m, Symmetric(C ./ (-2))
end

"""
    expectation_to_standard(η₁::AbstractVector{<:Real}, η₂::Symmetric{<:Real})

Convert from expectation to standard parametrisation.
"""
function expectation_to_standard(η₁::AbstractVector{<:Real}, η₂::Symmetric{<:Real})
    return η₁, η₂ - Symmetric(η₁ * η₁')
end

"""
    standard_to_expectation(m::AbstractVector{<:Real}, S::Symmetric{<:Real})

Convert from standard to natural parametrisation.
"""
function standard_to_expectation(m::AbstractVector{<:Real}, S::Symmetric{<:Real})
    return m, Symmetric(S + m * m')
end

function frule(
    ::typeof(standard_to_expectation),
    m::AbstractVector{<:Real},
    S::Symmetric{<:Real},
)
    function standard_to_expectation_pushforward(Δself, Δm, ΔS::Symmetric{<:Real})
        Δη₁ = Δm
        Δη₂ = Symmetric(ΔS + m * Δm' + Δm * m')
        return Δη₁, Δη₂
    end
    return standard_to_expectation(m, S), standard_to_expectation_pushforward
end

"""
    expectation_to_natural(η₁::AbstractVector{<:Real}, η₂::Symmetric{<:Real})

Convert from expectation to natural parameters.
"""
function expectation_to_natural(η₁::AbstractVector{<:Real}, η₂::Symmetric{<:Real})
    return standard_to_natural(expectation_to_standard(η₁, η₂)...)
end

"""
    natural_to_expectation(θ₁::AbstractVector{<:Real}, θ₂::Symmetric{<:Real})

Convert from natural to expectation parametrisation.
"""
function natural_to_expectation(θ₁::AbstractVector{<:Real}, θ₂::Symmetric{<:Real})
    m, S = natural_to_standard(θ₁, θ₂)
    η₁, η₂ = standard_to_expectation(m, S)
    return η₁, η₂
end

function frule(
    ::typeof(natural_to_expectation),
    θ₁::AbstractVector{<:Real},
    θ₂::Symmetric{<:Real},
)
    (m, S), nat_to_std_pushforward = frule(natural_to_standard, θ₁, θ₂)
    (η₁, η₂), std_to_exp_pushforward = frule(standard_to_expectation, m, S)
    function natural_to_expectation_pushforward(Δself, Δθ₁, Δθ₂::Symmetric)
        Δm, ΔS = nat_to_std_pushforward(DoesNotExist(), Δθ₁, Δθ₂)
        return std_to_exp_pushforward(DoesNotExist(), Δm, ΔS)
    end
    return (η₁, η₂), natural_to_expectation_pushforward
end

#
# Unconstrained parametrisations.
#

"""
    standard_to_unconstrained(m::AbstractVector{<:Real}, S::Symmetric{<:Real})

Convert from standard to mean and unconstrained upper-triangular matrix.
"""
function standard_to_unconstrained(m::AbstractVector{<:Real}, S::Symmetric{<:Real})
    return m, to_positive_definite_softplus_inv(chol(S))
end

function frule(
    ::typeof(standard_to_unconstrained),
    m::AbstractVector{<:Real},
    S::Symmetric{<:Real},
)
    # Perform forwards-pass and retain information required to compute pushforward.
    U, pushforward_U = frule(chol, S)
    V, pushforward_V = frule(to_positive_definite_softplus_inv, U)

    function standard_to_unconstrained_pushforward(Δself, Δm, ΔS::Symmetric{<:Real})
        ΔU = pushforward_U(DoesNotExist(), ΔS)
        ΔV = pushforward_V(DoesNotExist(), ΔU)
        return (Δm, ΔV)
    end

    return (m, V), standard_to_unconstrained_pushforward
end

"""
    unconstrained_to_standard(m::AbstractVector{<:Real}, U::UpperTriangular{<:Real})

Convert from mean + unconstrained upper-triangular matrix to standard parametrisation.
"""
function unconstrained_to_standard(m::AbstractVector{<:Real}, U::UpperTriangular{<:Real})
    UA = to_positive_definite_softplus(U)
    return m, Symmetric(UA'UA)
end

"""
    natural_to_unconstrained(θ₁::AbstractVector{<:Real}, θ₂::AbstractMatrix{<:Real})

Convert from natural to mean and unconstrained upper-triangular matrix.
"""
function natural_to_unconstrained(θ₁::AbstractVector{<:Real}, θ₂::Symmetric{<:Real})
    return standard_to_unconstrained(natural_to_standard(θ₁, θ₂)...)
end

function frule(
    ::typeof(natural_to_unconstrained),
    θ₁::AbstractVector{<:Real},
    θ₂::Symmetric{<:Real},
)
    (m_, S), pushforward_standard = frule(natural_to_standard, θ₁, θ₂)
    (m, U), pushforward_unconstrained = frule(standard_to_unconstrained, m_, S)

    function pushforward_natural_to_unconstrained(Δself, Δθ₁, Δθ₂::Symmetric{<:Real})
        Δm_, ΔS = pushforward_standard(DoesNotExist(), Δθ₁, Δθ₂)
        return pushforward_unconstrained(DoesNotExist(), Δm_, ΔS)
    end
    return (m, U), pushforward_natural_to_unconstrained
end

"""
    unconstrained_to_natural(m::AbstractVector{<:Real}, U::AbstractMatrix{<:Real})

Convert from mean and unconstrained upper-triangular matrix to natural.
"""
function unconstrained_to_natural(m::AbstractVector{<:Real}, U::UpperTriangular{<:Real})
    return standard_to_natural(unconstrained_to_standard(m, U)...)
end

"""
    expectation_to_unconstrained(η₁::AbstractVector{<:Real}, η₂::AbstractMatrix{<:Real})

Convert from expecation parameters to mean and unconstrained upper-triangular matrix.
"""
function expectation_to_unconstrained(
    η₁::AbstractVector{<:Real},
    η₂::AbstractMatrix{<:Real},
)
    return standard_to_unconstrained(expectation_to_standard(η₁, η₂)...)
end

"""
    unconstrained_to_expectation(η₁::AbstractVector{<:Real}, η₂::AbstractMatrix{<:Real})

Convert from expecation parameters to mean and unconstrained upper-triangular matrix.
"""
function unconstrained_to_expectation(
    m::AbstractVector{<:Real},
    U::UpperTriangular{<:Real},
)
    return standard_to_expectation(unconstrained_to_standard(m, U)...)
end


# variational autoencoder with Gaussian observed and latent variables
mutable struct VAE{T}
    # encoder / recognizer
    We1::AbstractMatrix{T}  # encoder: layer 1 weight
    be1::AbstractVector{T}  # encoder: layer 1 bias
    We2::AbstractMatrix{T}  # encoder: layer 2 weight
    be2::AbstractVector{T}  # encoder: layer 2 bias
    We3::AbstractMatrix{T}  # encoder: layer 3 mu weight
    be3::AbstractVector{T}  # encoder: layer 3 mu bias
    We4::AbstractMatrix{T}  # encoder: layer 3 log(sigma^2) weight
    be4::AbstractVector{T}  # encoder: layer 3 log(sigma^2) bias
    # decoder / generator
    Wd1::AbstractMatrix{T}  # decoder: layer 1 weight
    bd1::AbstractVector{T}  # decoder: layer 1 bias
    Wd2::AbstractMatrix{T}  # decoder: layer 2 weight
    bd2::AbstractVector{T}  # decoder: layer 2 bias
    Wd3::AbstractMatrix{T}  # decoder: layer 3 weight
    bd3::AbstractVector{T}  # decoder: layer 3 bias
end

function Base.show(io::IO, m::VAE{T}) where T
    print(io, "VAE{$T}($(size(m.We1,2)), $(size(m.We1,1)), $(size(m.We2,1)), " *
          "$(size(m.We3,1)), $(size(m.Wd1,1)), $(size(m.Wd2,1)), $(size(m.Wd3,1)))")
end

VAE{T}(n_inp, n_he1, n_he2, n_z, n_hd1, n_hd2, n_out) where T =
    VAE{T}(
        # encoder
        xavier_init(n_he1, n_inp),
        zeros(n_he1),
        xavier_init(n_he2, n_he1),
        zeros(n_he2),
        xavier_init(n_z, n_he2),
        zeros(n_z),
        xavier_init(n_z, n_he2),
        zeros(n_z),
        # decoder
        xavier_init(n_hd1, n_z),
        zeros(n_hd1),
        xavier_init(n_hd2, n_hd1),
        zeros(n_hd2),
        xavier_init(n_out, n_hd2),
        zeros(n_out)
    )

# empty constructor is required by XGrad
VAE{T}() where T = VAE{T}(1, 1, 1, 1, 1, 1, 1)

function encode(m::VAE, x)
    he1 = softplus.(m.We1 * x .+ m.be1)
    he2 = softplus.(m.We2 * he1 .+ m.be2)
    mu = m.We3 * he2 .+ m.be3
    log_sigma2 = m.We4 * he2 .+ m.be4
    return mu, log_sigma2
end

function decode(m::VAE, z)
    hd1 = softplus.(m.Wd1 * z .+ m.bd1)
    hd2 = softplus.(m.Wd2 * hd1 .+ m.bd2)
    x_rec = logistic.(m.Wd3 * hd2 .+ m.bd3)
    return x_rec
end

# TODO: make cost with Gaussian decoder
# TODO: figure out why KLD is different from the one in papers
function vae_cost(m::VAE, eps, x)
    mu, log_sigma2 = encode(m, x)
    z = mu .+ sqrt.(exp.(log_sigma2)) .* eps
    x_rec = decode(m, z)
    # loss
    rec_loss = -sum(x .* log.(1e-10 + x_rec) + (1 - x) .* log.(1e-10 + 1 - x_rec), 1)
    KLD = -0.5 * sum(1 + log_sigma2 .- mu .^ 2 - exp.(log_sigma2), 1)
    cost = mean(rec_loss .+ KLD)
end


#This is a workspace for hyperspectral imaging methods

"""
    ACE(Background, X, Target)

Untested
"""
function ACE(Background, X, Target)
    @assert( length(size(Background)) < 4 )
    if length(size(Background)) == 3
        Background = reshape(Background, prod( size( Background )[1:2 ]), size( Background )[ 3 ]  )
    end
    mu = Statistics.mean(Background, dims = 1)
    mcent = Background .- mu
    covinv = Base.inv( ( 1.0 / size(Background)[1] ) .* (mcent' * mcent) )
    tmu = Target .- mu
    xmu = X .- mu
    numerator = tmu' * covinv * xmu
    denominator = ( tmu' * covinv * tmu ) * ( xmu' * covinv * xmu )
    return (numerator * numerator) / denominator
end

"""
    MatchedFilter(X, Target)

Untested

MatchedFilter is always superior to CEM. Xiurui Geng, Luyan Ji, Weitun Yang, Fuxiang Wang, Yongchao Zhao
https://arxiv.org/pdf/1612.00549.pdf
"""
function MatchedFilter(X, Target)
    @assert( length(size(X)) < 4 )
    if length(size(X)) == 3
        X = reshape(X, prod( size( X )[1:2 ]), size( X )[ 3 ]  )
    end
    mu = Statistics.mean(X, dims = 1)
    mcent = X .- mu
    covinv = Base.inv( ( 1.0 / size(X)[1] ) .* (mcent' * mcent) )
    tmu = Target .- mu
    xmu = X .- mu
    numerator = covinv * tmu
    denominator = ( tmu' * covinv * tmu )
    return numerator / denominator
end


# energy calculation using linked lists method
function utotal_parallel(box, DATA, first_atom, next_atom, nc)

  nthreads = Threads.nthreads()     # number of threads available
  Ut = zeros(nthreads)              # vector with nthreads positions

 Threads.@threads  for iat in 1:DATA.N                  # paralelization through the particles
    icell = trunc(Int64,box[iat][1]/DATA.cutoff) + 1
    jcell = trunc(Int64,box[iat][2]/DATA.cutoff) + 1
    for i in icell-1:icell+1
      id = Threads.threadid()
      for j in jcell-1:jcell+1
        iw, jw = wrapcell(nc, i, j)
        jat = first_atom[iw,jw]
        while jat > 0
          if jat > iat
            rij = pbcseparation(box[jat],box[iat],DATA.side)
            if rij <= 2.
              Ut[id] += upair(rij,DATA)
            end
          end
          jat = next_atom[jat]
        end
      end
    end
  end
  return sum(Ut)
end


include("../../ccode/mpiops.jl")
using ADOPT

mpi_init()

function kappa(x, y)
    return 1+x^2+y^2
end

function rhs(x, y)
    -2*x*(1 - x)*(x^2 + y^2 + 1) + 2*x*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) -
             2*y*(1 - y)*(x^2 + y^2 + 1) + 2*y*(-x*y*(1 - x) + x*(1 - x)*(1 - y))
end

function ufunc(x, y)
    x * (1-x) * y * (1-y)
end

mc = MPIConfig(300)
global_to_local, local_to_global = dofmap(mc)
X, Y = get_xy(mc)
f_local = rhs.(X, Y)
using Random; Random.seed!(233)
θ = Variable(fc_init([2,20,20,20,1]))
θ_shared = mpi_bcast(θ)
κ_local = abs(fc([X'[:] Y'[:]], [20,20,20,1], θ_shared) + 5.0)|>squeeze
κ_local = reshape(κ_local, (mc.n, mc.n))
u_local = poisson_solver(κ_local, f_local, mc)

@load "data/$(mpi_size())_$(mpi_rank()).jld2" U

loss = sum(mpi_sum((u_local - U)^2))
g = gradients(loss, θ)

sess = Session(); init(sess)

L = run(sess, loss)

if mpi_rank()==0
    @info "Initial loss = $L"
end

function calculate_loss(x)
    L = run(sess, loss, θ=>x)
    L
end

function calculate_gradients(G, x)
    G[:] = run(sess, g, θ=>x)
end

losses  = Float64[]
function step_callback(x)
    @info "Loss = $x"
    push!(losses, x)
end

initial_x = run(sess, θ)
options = Options()
result = ADOPT.mpi_optimize(calculate_loss, calculate_gradients, initial_x, LBFGS(), options; step_callback = step_callback)

if mpi_rank()==0
    minimizer = result.minimizer
    @save "result.jld2" result losses
end

if mpi_size()>1
    mpi_finalize()
end


workspace();
using BoxChopper;

box = [0 1; 0 1; 0 1]
r₀ = [0.5, 0.5, 0.5]
nout = [-2,-1,0]

nr₀ = nout⋅r₀
vcbits = BoxChopper.calc_vcbits(box, nout, nr₀)
nvc = count_ones(vcbits)  # number of vertices contained

assert(BoxChopper.isquadsect(vcbits))
@code_warntype BoxChopper.rvol_quadsect(box, nout, nr₀, vcbits)


# Simple example softmax function with associated utils
# Doesn't do a good job of making util function reusable
using cg
using MNIST

using ImageMagick
using Images

function softmax_regression(input_dim::Int, output_dim::Int, step_size::Float64)
    # Wow that's verbose - start overloading so constants become nodes automatically
    W1 = cg.variable(cg.zeros(cg.constant([output_dim, input_dim])), "W1")
    b1 = cg.variable(cg.zeros(cg.constant(output_dim)), "b1")

    input = cg.placeholder([input_dim], "input")
    label = cg.placeholder([output_dim], "label")
    unnormed = cg.broadcast("+", cg.matmul(W1, input), b1)
    prediction = cg.softmax_stable(unnormed)
    loss = cg.crossentropy(label, prediction)

    train = cg.sgd_optimizer(loss, [W1, b1], cg.constant(step_size))

    (input, label, prediction, loss, train)
end

function onehot(labels::AbstractArray; min::Int = 0, max::Int = 9)
    labels = round(Int, labels)
    result = zeros(Float64, max - min + 1, length(labels))
    for (i, label) in enumerate(labels)
        result[label - min + 1, i] = 1.0
    end
    result
end

function shuffle_data(data, labels)
    perm = shuffle(collect(1:size(labels, 2)))
    data[:, perm], labels[:, perm]
end

# labels = 1hot encoded
function train_steps(sess, batch_size, steps, data, labels)
    @assert size(data, 2) == size(labels, 2)
    @assert batch_size <= size(labels, 2)
    sdata, slabel = shuffle_data(data, labels)
    tsteps = 0
    bs = 1
    while tsteps < steps
        be = bs + batch_size - 1
        # Reshuffle and startover
        if be > size(labels, 2)
            sdata, slabel = shuffle(data, labels)
            bs = 1
            continue
        end
        tsteps += 1
        sess.values[input] = sdata[:, bs:be]
        sess.values[label] = slabel[:, bs:be]
        cg.interpret(sess, sgd)
    end
end

function compute_loss(session, loss, data, labels)
    session.values[input] = data
    session.values[label] = labels
    cg.interpret(session, loss)
end

function visualize_filters(filters, w, h)
    output = zeros(h, w * size(filters,1))
    for i in 1:size(filters,1)
        start = (i-1) * h
        filter = filters[i, :]
        high = maximum(filter)
        low = minimum(filter)
        scaled = (filter - low) / (high - low)
        output[:, start+1:start+w] = scaled
    end
    convert(Image, output)
end

function accuracy(data, labels)
    sess.values[input] = data
    result = cg.interpret(sess, output)
    correct = findmax(labels, 1)[2]
    prediction = findmax(result, 1)[2]
    sum(correct .== prediction) / length(correct)
end

# Scale the data to [0,1] to avoid overflow
test_raw = testdata()
train_raw = traindata()
test = (test_raw[1] / 255.0, test_raw[2])
train = (train_raw[1] / 255.0, train_raw[2])

input, label, output, loss, sgd = softmax_regression(784, 10, 0.01)
sess = cg.Session(sgd)

if !isinteractive()
    @show compute_loss(sess, loss, train[1], onehot(train[2]))
    @show compute_loss(sess, loss, test[1], onehot(test[2]))
    train_steps(100, 1000, train[1], onehot(train[2]))
    @show compute_loss(sess, loss, train[1], onehot(train[2]))
    @show compute_loss(sess, loss, test[1], onehot(test[2]))
end


#=

julia julia_callback.jl
  2.457 ms (100002 allocations: 10.68 MiB)

=#
using BenchmarkTools

function rober(u)
    k1 = 0.04
    k2 = 3e7
    k3 = 1e4
    y1, y2, y3 = u
    dy1 = -k1 * y1 + k3 * y2 * y3
    dy2 = k1 * y1 - k2 * y2 * y2 - k3 * y2 * y3
    dy3 = k2 * y2 * y2
    return [dy1, dy2, dy3]
end

rober([1.0, 0.0, 0.0])

function call_function(func)
    u = [1.0, 0.0, 0.0]
    for time in (0:1e5)
        u = func(u)
    end
    return u
end

@btime call_function(rober)


push!(LOAD_PATH, "./src")

# "Packages"
using Revise # reduce need for recompile
using Plots
using Documenter
using LinearAlgebra
using SparseArrays

# "User defined modules"
using CommonUtils
using Basis1D
using Basis2DQuad
using Basis3DHex
using UniformHexMesh

using SetupDG
using UnPack

N = 3
K1D = 8
CFL = .75;
T = .5 # endtime

VX,VY,VZ,EToV = uniform_hex_mesh(K1D,K1D,K1D)

rd = init_reference_hex(N,gauss_lobatto_quad(0,0,N))
md = init_mesh((VX,VY,VZ),EToV,rd)

# make domain periodic
@unpack Nfaces = rd
@unpack xf,yf,zf,K,mapM,mapP,mapB = md
LX = 2; LY = 2; LZ = 2
mapPB = build_periodic_boundary_maps(xf,yf,zf,LX,LY,LZ,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB
@pack! md = mapP

"Time integration"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N+1)*(N+2)*3/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

function rhs(u,rd::RefElemData,md::MeshData,params...)

    @unpack Dr,Ds,Dt,LIFT,Vf = rd
    @unpack rxJ,sxJ,txJ,ryJ,syJ,tyJ,rzJ,szJ,tzJ,J,nxJ,nyJ,nzJ,sJ = md
    @unpack mapP,mapB = md

    uM = Vf*u
    du = uM[mapP]-uM
    uflux = @. .5*(du*nxJ - abs(nxJ)*du)

    ur,us,ut = (A->A*u).((Dr,Ds,Dt))
    ux = @. rxJ*ur + sxJ*us + txJ*ut
    rhsu = ux + LIFT*uflux

    return @. -rhsu/J
end

# set initial conditions
@unpack x,y,z = md
u = @. sin(pi*x)
u = @. exp(-25*(x^2+y^2))
resQ = zeros(size(x))
for i = 1:Nsteps
    for INTRK = 1:5
        rhsQ = rhs(u,rd,md)
        @. resQ = rk4a[INTRK]*resQ + dt*rhsQ
        @. u    = u + rk4b[INTRK]*resQ
    end

    if i%10==0 || i==Nsteps
        println("Number of time steps: ", i, " out of ", Nsteps)
    end
end

"plotting nodes"
gr(aspect_ratio=1,legend=false,
   markerstrokewidth=0,markersize=2)

@unpack Vp = rd
xp,yp,zp = (x->Vp*x).((x,y,z))
vv = Vp*u
ids = map(x->x[1],findall(@. abs(zp[:])<1e-10))
(xp,yp,zp,vv) = (x->x[ids]).((xp,yp,zp,vv))
scatter(xp,yp,vv,zcolor=vv,camera=(0,90))


export JacobianFreeNewtonKrylovSolver, JacobianAction

"""
mutable struct JacobianAction{FT, AT}
    rhs!
    ϵ::FT
    Q::AT
    Qdq::AT
    Fq::AT
    Fqdq::AT
end

Solve for Frhs = F(q), the Jacobian action is computed

    ∂F(Q)      F(Q + eΔQ) - F(Q)
    ---- ΔQ ≈ -------------------
     ∂Q                e

...
# Arguments
- `rhs!`           : nonlinear operator F(Q)
- `ϵ::FT`          : ϵ used for finite difference, e = e(Q, ϵ)
- `Q::AT`          : cache for Q
- `Qdq::AT`        : container for Q + ϵΔQ
- `Fq::AT`         : cache for F(Q)
- `Fqdq::AT`       : container for F(Q + ϵΔQ)
...
"""
mutable struct JacobianAction{FT, AT}
    rhs!
    ϵ::FT
    Q::AT
    Qdq::AT
    Fq::AT
    Fqdq::AT
end

function JacobianAction(rhs!, Q, ϵ)
    return JacobianAction(
        rhs!,
        ϵ,
        similar(Q),
        similar(Q),
        similar(Q),
        similar(Q),
    )
end

"""
Approximates the action of the Jacobian of a nonlinear
form on a vector `ΔQ` using the difference quotient:

      ∂F(Q)      F(Q + e ΔQ) - F(Q)
JΔQ = ---- ΔQ ≈ -------------------
       ∂Q                e

Compute  JΔQ with cached Q and F(Q), and the direction  dQ
"""
function (op::JacobianAction)(JΔQ, dQ, args...)
    rhs! = op.rhs!
    Q = op.Q
    Qdq = op.Qdq
    ϵ = op.ϵ
    Fq = op.Fq
    Fqdq = op.Fqdq

    FT = eltype(dQ)
    n = length(dQ)
    normdQ = norm(dQ, weighted_norm)

    if normdQ > ϵ
        factor = FT(1 / (n * normdQ))
    else
        # initial newton step, ΔQ = 0
        factor = FT(1 / n)
    end

    β = √ϵ
    e = factor * β * norm(Q, 1, false) + β

    Qdq .= Q .+ e .* dQ

    rhs!(Fqdq, Qdq, args...)

    JΔQ .= (Fqdq .- Fq) ./ e

end

"""
update cached Q and F(Q) before each Newton iteration
"""
function update_Q!(op::JacobianAction, Q, args...)
    op.Q .= Q
    Fq = op.Fq

    op.rhs!(Fq, Q, args...)
end

"""
Solve for Frhs = F(Q), by finite difference

    ∂F(Q)      F(Q + eΔQ) - F(Q)
    ---- ΔQ ≈ -------------------
     ∂Q                e

     Q^n+1 = Q^n - dF/dQ(Q^{n})⁻¹ (F(Q^n) - Frhs)

     set ΔQ = F(Q^n) - Frhs
"""
mutable struct JacobianFreeNewtonKrylovSolver{FT, AT} <: AbstractNonlinearSolver
    # small number used for finite difference
    ϵ::FT
    # tolerances for convergence
    tol::FT
    # Max number of Newton iterations
    M::Int
    # Linear solver for the Jacobian system
    linearsolver
    # container for unknows ΔQ, which is updated for the linear solver
    ΔQ::AT
    # contrainer for F(Q)
    residual::AT
end

"""
JacobianFreeNewtonKrylovSolver constructor
"""
function JacobianFreeNewtonKrylovSolver(
    Q,
    linearsolver;
    ϵ = 1.e-8,
    tol = 1.e-6,
    M = 30,
)
    FT = eltype(Q)
    residual = similar(Q)
    ΔQ = similar(Q)
    return JacobianFreeNewtonKrylovSolver(
        FT(ϵ),
        FT(tol),
        M,
        linearsolver,
        ΔQ,
        residual,
    )
end

"""
JacobianFreeNewtonKrylovSolver initialize the residual
"""
function initialize!(
    rhs!,
    Q,
    Qrhs,
    solver::JacobianFreeNewtonKrylovSolver,
    args...,
)
    # where R = Qrhs - F(Q)
    R = solver.residual
    # Computes F(Q) and stores in R
    rhs!(R, Q, args...)
    # Computes R = R - Qrhs
    R .-= Qrhs
    return norm(R, weighted_norm)
end

"""
Solve for Frhs = F(Q), by finite difference

Q^n+1 = Q^n - dF/dQ(Q^{n})⁻¹ (F(Q^n) - Frhs)

set ΔQ = F(Q^n) - Frhs

...
# Arguments
- `rhs!`:  functor rhs!(Q) =  F(Q)
- `jvp!`:  Jacobian action jvp!(ΔQ)  = dF/dQ(Q) ⋅ ΔQ
- `preconditioner`: approximation of dF/dQ(Q)
- `Q` : Q^n
- `Qrhs` : Frhs
- `solver`: linear solver
...
"""
function donewtoniteration!(
    rhs!,
    jvp!,
    preconditioner::AbstractPreconditioner,
    Q,
    Qrhs,
    solver::JacobianFreeNewtonKrylovSolver,
    args...,
)

    FT = eltype(Q)
    ΔQ = solver.ΔQ
    ΔQ .= FT(0.0)

    # R(Q) == 0, R = F(Q) - Qrhs, where F = rhs!
    # Compute right-hand side for Jacobian system:
    # J(Q)ΔQ = -R
    # where R = Qrhs - F(Q), which is computed at the end of last step or in the initialize function
    R = solver.residual

    # R = F(Q^n) - Frhs
    # ΔQ = dF/dQ(Q^{n})⁻¹ (Frhs - F(Q^n)) = -dF/dQ(Q^{n})⁻¹ R
    iters =
        linearsolve!(jvp!, preconditioner, solver.linearsolver, ΔQ, -R, args...)

    # Newton correction Q^{n+1} = Q^n + dF/dQ(Q^{n})⁻¹ (Frhs - F(Q^n))
    Q .+= ΔQ

    # Compute residual norm and residual for next step
    rhs!(R, Q, args...)
    R .-= Qrhs
    resnorm = norm(R, weighted_norm)

    return resnorm, iters
end


module FiniteDifferencesQuasi

function __init__()
    @warn "The FiniteDifferencesQuasi.jl package has been deprecated in favour of JuliaApproximation/CompactBases.jl"
    nothing
end

import Base: eltype, axes, size, ==, getindex, checkbounds, copyto!, similar, show, step
import Base.Broadcast: materialize

using ContinuumArrays
import ContinuumArrays: Basis, ℵ₁, Derivative, Inclusion, @simplify

using QuasiArrays
import QuasiArrays: AbstractQuasiMatrix, QuasiAdjoint, MulQuasiArray,
    PInvQuasiMatrix, InvQuasiMatrix, QuasiDiagonal,
    BroadcastQuasiArray, SubQuasiArray

using IntervalSets

using LazyArrays

using LinearAlgebra
import LinearAlgebra: Matrix, dot
using SparseArrays
using BandedMatrices
using FillArrays

using Printf

include("restricted_bases.jl")
include("finite_differences.jl")
include("inner_products.jl")
include("operators.jl")
include("derivatives.jl")
include("densities.jl")

export AbstractFiniteDifferences, FiniteDifferences,
    RadialDifferences, StaggeredFiniteDifferences,
    NumerovFiniteDifferences,
    Derivative, dot, QuasiDiagonal

end # module


# Primal infeasible test problems
# Problems with m1 feasible equality constraints, one infeasible m2-sized second-order-cone constraint and one feasible psd constraint
# constrain s[m1] = t of [t;x] to be -1, which makes the second-order-cone constraint infeasible

nn = 1
rng = Random.MersenneTwister(1313)

@testset "Primal infeasible QP problems - Testset 3" begin
  for iii = 1:nn

    # choose size of problem
    n = rand(rng,10:50)
    m1 = rand(rng,2:10) #equality constraint
    m2 = rand(rng,3:10) # second order cone constraint
    r = rand(rng,4:10)
    m3 = r^2 # one sdp cone constraint of size r
    m = m1+m2+m3

    A = sprandn(rng,m,n,0.8)*50
    xtrue = rand(rng,n,1)*50
    s1 = zeros(m1,1)
    s2 = randn(rng,m2,1)
    s3 = vec(generate_pos_def_matrix(rng, r))
    s = [s1;s2;s3]
    b = A*xtrue+s

    # enforce the t in ||x|| < t of second order cone constraint to be t < 0 --> infeasible
    A[m1+1,:] = zeros(n)
    b[m1+1] = -1
    b = vec(b)

    # create dual feasible problem Px+q+A'y = 0, and y in K*
    P = generate_pos_def_matrix(rng, n)
    ytrue_1 = randn(rng,m1,1)*50
    ytrue_2 = randn(rng,m2-1,1)*50
    ytrue_2 = [norm(ytrue_2,2)+1;ytrue_2]
    ytrue_3 = vec(generate_pos_def_matrix(rng, r))
    ytrue = [ytrue_1;ytrue_2;ytrue_3]
    q = (-P*xtrue - A'*ytrue)[:]

    A1 = -A[1:m1,:]
    A2 = -A[m1+1:m1+m2,:]
    A3 = -A[m1+m2+1:end,:]
    b1 = b[1:m1,:]
    b2 = b[m1+1:m1+m2]
    b3 = b[m1+m2+1:end]

    cs1 = COSMO.Constraint(A1,b1,COSMO.ZeroSet)
    cs2 = COSMO.Constraint(A2,b2,COSMO.SecondOrderCone)
    cs3 = COSMO.Constraint(A3,b3,COSMO.PsdCone)

    settings = COSMO.Settings(max_iter=10000,eps_abs = 1e-5,eps_rel=1e-5)

    model = COSMO.Model()
    assemble!(model,P,q,[cs1;cs2;cs3], settings = settings)
    res = COSMO.optimize!(model);

    @test res.status == :Primal_infeasible
  end
end
nothing


using PPInterpolation, Test
using ForwardDiff

@testset "Akima" begin
    x = [3.0, 5.0, 6.0, 8.0, 9.0, 11.0, 12.0, 14.0, 15.0]
    f = [10.0, 10.0, 10.0, 10.0, 10.5, 15.0, 50.0, 60.0, 85.0]
    varieties = [C2Hyman89(), C2MP(), HuynRational(), VanLeer(), FritschButland(), Brodlie(), VanAlbada()]
    for v in varieties
        spline = makeCubicPP(x, f, PPInterpolation.SECOND_DERIVATIVE, 0.0, PPInterpolation.SECOND_DERIVATIVE, 0.0, v)
        nSample = 1000
        previousVal = spline(x[1])
        for i = 1:nSample
            ti = x[1] + (x[end] - x[1]) * (i) / (nSample)
            zval = spline(ti)
            if i > 2
                @test zval >= previousVal
            end
            previousVal = zval
        end
    end
    #     varieties = [ C2Hyman89(), C2MP(), HuynRational(), VanLeer(), FritschButland(), Brodlie(), VanAlbada() ]^C
    #  p = plot()
    #  for v in varieties
    #        spline = makeCubicPP(x,f,PPInterpolation.SECOND_DERIVATIVE,0.0,PPInterpolation.SECOND_DERIVATIVE,0.0,v)
    #        plot!(p, t, spline.(t), label=string(v))
    #        end
    #  plot!(legend=:top)

end

@testset "FrtschCarlson" begin
    x = [7.99, 8.09, 8.19, 8.7, 9.2, 10.0, 12.0, 15.0, 20.0]
    f = [0.0, 2.76429e-5, 4.37498e-2, 0.169183, 0.469428, 0.943740, 0.998636, 0.999919, 0.999994]

    varieties = [C2Hyman89(), C2MP(), HuynRational(), VanLeer(), FritschButland(), Brodlie(), VanAlbada()]
    for v in varieties
        spline = makeCubicPP(x, f, PPInterpolation.SECOND_DERIVATIVE, 0.0, PPInterpolation.SECOND_DERIVATIVE, 0.0, v)
        nSample = 1000
        previousVal = spline(x[1])
        for i = 1:nSample
            ti = x[1] + (x[end] - x[1]) * (i) / (nSample)
            zval = spline(ti)
            if i > 2
                @test zval >= previousVal
            end
            previousVal = zval
        end
    end
end
@testset "HaganCurveAuto" begin
    #for quintic: first pass compute cubic first derivative, second pass, use cubic b, and optimize second derivative.
    t = [0.0, 0.1, 1.0, 2.0, 3.0, 4.0, 9.0, 20.0, 30.0]
    y = [0.06, 0.06, 0.06, 0.05, 0.05, 0.05, 0.065, 0.06, 0.06]
    t = [0.0, 0.1, 1.0, 4.0, 9.0, 20.0, 30.0]
    y = [0.081, 0.081, 0.07, 0.044, 0.07, 0.04, 0.03]
    t = [0, 0.1, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0, 2.5, 3.0, 3.5]
    y = [0.012, 0.008, 0.0104, 0.00995, 0.010633333333333333, 0.010849999999999999, 0.010179999999999998, 0.009816666666666665, 0.008612499999999999, 0.007789999999999998, 0.007741666666666665, 0.007635714285714285]

    # t = [0.0, 0.1, 1, 2, 3, 4, 5, 6, 7, 30.0]
    # f = [0.01, 0.01, 0.01, 0.01, 0.01, 0.03, 0.03, 0.03, 0.03, 0.02]
    # y = zeros(length(t))
    # y[1] = f[1]
    # y[2:end] = (y[1:end-1] .* t[1:end-1] + (t[2:end]-t[1:end-1]) .* f[2:end]) ./ t[2:end]
    # y[3:end-2] = f[3:end-2] ./ t[3:end-2]
    # println(y)
    yt = zeros(length(t))
    yt[1:end] = y[1:end] .* t[1:end]
    spline = makeCubicPP(t, yt, PPInterpolation.SECOND_DERIVATIVE, 0.0, PPInterpolation.SECOND_DERIVATIVE, 0.0, C2())
    nSample = 1000
    for i = 1:nSample
        ti = (t[end]) * (i) / (nSample)
        val = evaluateDerivative(spline, ti)
        zval = spline(ti) / ti
        println(ti, " CubicPP ", val, " ", zval)
        autoval = ForwardDiff.derivative(spline, ti)
        @test isapprox(val, autoval, atol = 1e-15)
    end
    #AUTODIFF // where is autodiff useful (hagan wave read),  BigFloat
end


"""
		mlp_block(inplanes::Integer, hidden_planes::Integer, outplanes::Integer = inplanes;
							dropout = 0., activation = gelu)

Feedforward block used in many MLPMixer-like and vision-transformer models.

# Arguments
- `inplanes`: Number of dimensions in the input.
- `hidden_planes`: Number of dimensions in the intermediate layer.
- `outplanes`: Number of dimensions in the output - by default it is the same as `inplanes`.
- `dropout`: Dropout rate.
- `activation`: Activation function to use.
"""
function mlp_block(inplanes::Integer, hidden_planes::Integer, outplanes::Integer = inplanes;
									 dropout = 0., activation = gelu)
  Chain(Dense(inplanes, hidden_planes, activation), Dropout(dropout),
        Dense(hidden_planes, outplanes), Dropout(dropout))
end

"""
    gated_mlp(gate_layer, inplanes::Integer, hidden_planes::Integer,
							outplanes::Integer = inplanes; dropout = 0., activation = gelu)

Feedforward block based on the implementation in the paper "Pay Attention to MLPs".
([reference](https://arxiv.org/abs/2105.08050))

# Arguments
- `gate_layer`: Layer to use for the gating.
- `inplanes`: Number of dimensions in the input.
- `hidden_planes`: Number of dimensions in the intermediate layer.
- `outplanes`: Number of dimensions in the output - by default it is the same as `inplanes`.
- `dropout`: Dropout rate.
- `activation`: Activation function to use.
"""
function gated_mlp_block(gate_layer, inplanes::Integer, hidden_planes::Integer,
												 outplanes::Integer = inplanes; dropout = 0., activation = gelu)
	@assert hidden_planes % 2 == 0 "`hidden_planes` must be even for gated MLP"
	return Chain(Dense(inplanes, hidden_planes, activation),
							 Dropout(dropout),
							 gate_layer(hidden_planes),
							 Dense(hidden_planes ÷ 2, outplanes),
							 Dropout(dropout))
end
gated_mlp_block(::typeof(identity), args...; kwargs...) = mlp_block(args...; kwargs...)


module Follow
using Plots
using BoundaryIntegralMethod
using BoundaryIntegralMethod.DirectSweep
using JLD

function followResonance(k0, b0, changeFn, steps; convergence = 1e-6, n=200, outputDir=nothing, fileBase="",
    maxDist=nothing, plotFn = nothing, startStep=1, namingFn = nothing, callbackFn = nothing)
  plots = []
  kfs = []
  bs = []
  b = b0
  k = k0
  if outputDir == nothing
    outputDir = "output/follow-$k0"
  end
  if fileBase != ""
    fileBase = fileBase+"-"
  end
  mkpath(outputDir)
  if plotFn == nothing
    plotFn = (b, xf, kf, baseName) -> exportTotalField(b, xf, kf; outFile="$baseName.h5", xRange=(-2,2), yRange=(-2,2), dx=.01, dy=.01);
  end

  for i in startStep:steps
    try
      (ks, xs) = directSweep(b, n, k)
      ks = ks[1:6]
      # Try closest resonance
      @debug ks
      p = sortperm(abs.(imag.(ks).-imag(k)))
      @info "Choosing from ks:"
      [@info "$i: $(ks[i])" for i in 1:length(ks)]
      # candidateK = ks[p[1]]
      # candidateX = xs[p[1],:]
      candidateK = ks[1]
      candidateX = xs[1,:]
      @info "Sortperm would be $(p[1]): $(ks[p[1]]), abs would be 1: $(ks[1])"
      @info "Choosing 1: $(ks[1])"

      (kf, xf) = residualInverseIteration(25, convergence, b, n, candidateK, candidateX)
      @info "At step $i, converged to $kf"
      paramName = if (namingFn != nothing)
        "-$(namingFn(b))"
      else
        ""
      end

      baseName= "$outputDir/$fileBase$i$paramName-$kf"
      field = plotFn(b, xf, kf, baseName)
      p = heatmap(clamp.(abs.(field), 0, .6), cbar=false, aspect_ratio=1, ticks=nothing, title="$i$paramName-k=$kf")
      savefig("$baseName.png")
      if maxDist != nothing && maxDist < abs(k-kf)
        @error "Stopping after $i steps, max distance between wavenumber exceeded in a step ($k -> $kf)"
        break
      end
      push!(kfs, kf)
      push!(plots, p)
      push!(bs, b)
      writeDatafile("$outputDir/$(fileBase)data.jld", kfs, bs)
      if callbackFn != nothing
        callbackFn(b, kf, xf, "$fileBase$i$paramName", baseName)
      end
      k = kf
      b = changeFn(b, i)
    catch e
      @error "Stopping after $i steps, caught error" exception=e
      break
    end
  end
  (kfs, plots)
end

function writeDatafile(name, kfs, bs)
  jldopen(name, "w") do file
    write(file, "ks/full", kfs)
    write(file, "ks/Re", real(kfs))
    write(file, "ks/Im", imag(kfs))
    write(file, "boundaries", bs)
  end
end

end


using Test
using SDPT3

@testset "Semidefinite Programming example" begin
    @testset "Min offdiag" begin
        blk = ["s" 2.0]
        c = [0.0 1.0
             1.0 0.0]
        A = [1.0 0.0 0.0
             0.0 0.0 1.0]
        b = [1.0
             1.0]
        obj, X, y, Z, info, runhist = sdpt3(blk, [Matrix(A')], [c], b)
        tol = 1e-5
        @test obj[1] ≈ -2 atol=tol rtol=tol
        @test obj[2] ≈ -2 atol=tol rtol=tol
        @test X[1] ≈ [ 1.0 -1.0
                      -1.0  1.0] atol=tol rtol=tol
        @test y ≈ -ones(2) atol=tol rtol=tol
        @test Z[1] ≈ ones(2, 2) atol=tol rtol=tol
        @test info["termcode"] == 0.0
    end
    @testset "Min trace" begin
        blk = ["s" 2.0]
        c = [1.0 0.0
             0.0 1.0]
        A = [0.0 1/√2 0.0]
        b = [1.0]
        obj, X, y, Z, info, runhist = sdpt3(blk, [Matrix(A')], [c], b)
        tol = 1e-5
        @test obj[1] ≈ 2 atol=tol rtol=tol
        @test obj[2] ≈ 2 atol=tol rtol=tol
        @test X[1] ≈ ones(2, 2) atol=tol rtol=tol
        @test y ≈ [2.0] atol=tol rtol=tol
        @test Z[1] ≈ [ 1.0 -1.0
                      -1.0  1.0] atol=tol rtol=tol
        @test info["termcode"] == 0.0
    end
end


#! /usr/bin/julia

# Rosetta Code, Detect division by zero

function isdefinite{T<:Number}(n::T)
    !isequal(n, NaN) && abs(n) != Inf
end

for n in {1, 1//1, 1.0, 1im, 0}
    d = n/0
    print("Divding ", n, " by 0 ")
    if isdefinite(d)
        println("results in ", d, ".")
    else
        println("yields an indefinite value (", d, ").")
    end
end


function check_model_form(model)

  if typeof(model) <:Blanchard_Kahn_Form
    assessment = check_blanchard_kahn_form(model)
  elseif typeof(model) <:Klein_Form
    assessment = check_klein_form(model)
  elseif typeof(model) <:Binder_Pesaran_Form
    assessment = check_binder_pesaran_form(model)
  elseif typeof(model) <:Sims_Form
    assessment = check_sims_form(model)
  elseif typeof(model) <:Gomme_Klein_Form
    assessment = check_gomme_klein_form(model)
  elseif typeof(model) <:Lombardo_Sutherland_Form
    assessment = check_lombardo_sutherland_form(model)
  end

  return assessment

end

function check_blanchard_kahn_form(model::Blanchard_Kahn_Form)

  assessment = "Pass"

  if model.nx < 1
    assessment = "Fail: Too few predetermined variables"
  elseif model.ny < 1
    assessment = "Fail: Too few non-predetermined variables"
  elseif model.nx+model.ny != size(model.a,1)
    assessment = "Fail: A is not conformable with nx and ny"
  elseif size(model.a,1) != size(model.a,2)
    assessment = "Fail: A is not square"
  elseif size(model.a,1) != size(model.c,1)
    assessment = "Fail: A and C must have the same number of rows"
  elseif size(model.c,2) > model.nx
    assessment = "Fail: There are more innovations than predetermined variables"
  elseif size(model.sigma,1) != size(model.sigma,2)
    assessment = "Fail: SIGMA must be square"
  elseif size(model.c,2) != size(model.sigma,1)
    assessment = "Fail: C and SIGMA are not conformable"
  end

  return assessment

end

function check_klein_form(model::Klein_Form)

  assessment = "Pass"

  if model.nx < 1
    assessment = "Fail: Too few predetermined variables"
  elseif model.ny < 1
    assessment = "Fail: Too few non-predetermined variables"
  elseif model.nx+model.ny != size(model.a,1)
    assessment = "Fail: A is not conformable with nx and ny"
  elseif size(model.a,1) != size(model.a,2)
    assessment = "Fail: A is not square"
  elseif size(model.b,1) != size(model.b,2)
    assessment = "Fail: B is not square"
  elseif size(model.a) != size(model.b)
    assessment = "Fail: A and B are not conformable"
  elseif size(model.a,1) != size(model.c,1)
    assessment = "Fail: A and C must have the same number of rows"
  elseif size(model.c,2) > model.nx
    assessment = "Fail: There are more innovations than predetermined variables"
    elseif size(model.sigma,1) != size(model.sigma,2)
    assessment = "Fail: SIGMA must be square"
  elseif size(model.c,2) != size(model.sigma,1)
    assessment = "Fail: C and SIGMA are not conformable"
  end

  return assessment

end

function check_binder_pesaran_form(model::Binder_Pesaran_Form)

  assessment = "Pass"

  if size(model.a,1) != size(model.a,2)
    assessment = "Fail: A is not square"
  elseif size(model.a) != size(model.a1)
    assessment = "Fail: A and A1 are not conformable"
  elseif size(model.a) != size(model.b)
    assessment = "Fail: A and B are not conformable"
  elseif size(model.a,1) != size(model.c,1)
    assessment = "Fail: A and C must have the same number of rows"
  elseif size(model.sigma,1) != size(model.sigma,2)
    assessment = "Fail: SIGMA must be square"
  elseif size(model.c,2) != size(model.sigma,1)
    assessment = "Fail: C and SIGMA are not conformable"
  end

  return assessment

end

function check_sims_form(model::Sims_Form)

  assessment = "Pass"

  if size(model.gamma0,1) != size(model.gamma0,2)
    assessment = "Fail: GAMMA0 is not square"
  elseif size(model.gamma1,1) != size(model.gamma1,2)
    assessment = "Fail: GAMMA1 is not square"
  elseif size(model.gamma0) != size(model.gamma1)
    assessment = "Fail: GAMMA0 and GAMMA1 are not conformable"
  elseif size(model.gamma0,1) != size(model.c,1)
    assessment = "Fail: GAMMA0 and C must have the same number of rows"
  elseif size(model.c,2) != 1
    assessment = "Fail: C must be a vector"
  elseif size(model.gamma0,1) != size(model.psi,1)
    assessment = "Fail: GAMMA0 and PSI must have the same number of rows"
  elseif size(model.gamma0,1) > size(model.pi,1)
    assessment = "Fail: GAMMA0 and PI must have the same number of rows"
  elseif size(model.sigma,1) != size(model.sigma,2)
    assessment = "Fail: SIGMA must be square"
  elseif size(model.psi,2) != size(model.sigma,1)
    assessment = "Fail: PSI and SIGMA are not conformable"
  end

  return assessment

end

function check_gomme_klein_form(model::Gomme_Klein_Form)

  assessment = "Pass"

  if model.nx < 1
    assessment = "Fail: Too few predetermined variables"
  elseif model.ny < 1
    assessment = "Fail: Too few non-predetermined variables"
  elseif model.nx+model.ny != size(model.derivs,1)
    assessment = "Fail: Need first derivative for all model equations"
  elseif 2*size(model.derivs,1) != size(model.derivs,2)
    assessment = "Fail: The dimensions of the first-derivative matrix are incorrect"
  elseif size(model.derivs,2) != size(model.hessians,2)
    assessment = "Fail: The derivative matrix and the hessian matrix are not conformable"
  elseif size(model.hessians,1) != (model.nx+model.ny)*size(model.hessians,2)
    assessment = "Fail: The dimensions of the hessian matrix are incorrect"
  elseif size(model.eta,1) != model.nx
    assessment = "Fail: Eta has too few rows"
  elseif size(model.eta,2) != size(model.sigma,1)
    assessment = "Fail: Eta has too few columns"
  elseif size(model.sigma,1) != size(model.sigma,2)
    assessment = "Fail: SIGMA must be square"
  end

  return assessment

end

function check_lombardo_sutherland_form(model::Lombardo_Sutherland_Form)

  assessment = "Pass"

  if model.nx < 1
    assessment = "Fail: Too few predetermined variables"
  elseif model.ny < 1
    assessment = "Fail: Too few non-predetermined variables"
  elseif model.nx+model.ny != size(model.derivs,1)
    assessment = "Fail: Need first derivative for all model equations"
  elseif 2*size(model.derivs,1) != size(model.derivs,2)
    assessment = "Fail: The dimensions of the first-derivative matrix are incorrect"
  elseif size(model.derivs,2) != size(model.hessians,2)
    assessment = "Fail: The derivative matrix and the hessian matrix are not conformable"
  elseif size(model.hessians,1) != (model.nx+model.ny)*size(model.hessians,2)
    assessment = "Fail: The dimensions of the hessian matrix are incorrect"
  elseif size(model.eta,1) != model.nx
    assessment = "Fail: Eta has too few rows"
  elseif size(model.eta,2) != size(model.sigma,1)
    assessment = "Fail: Eta has too few columns"
  elseif size(model.sigma,1) != size(model.sigma,2)
    assessment = "Fail: SIGMA must be square"
  end

  return assessment

end


using LittleScienceTools
using LittleScienceTools.Measuring
# using LittleScienceTools.Random
# using LittleScienceTools.Roots
# using LittleScienceTools.Ising
# using Compat
# using Erdos
using Test
using Statistics
using LinearAlgebra
using Random
using KahanSummation

@testset "Measuring" begin
    include("Measuring/observable.jl")
    include("Measuring/obstable.jl")
end

# @testset "Ising" begin
#     include("Ising/runtests.jl")
# end

# @testset "Roots" begin
#     include("Roots/Roots.jl")
# end

# @testset "Random" begin
# #     include("Random/parisi_rapuano.jl")
#     include("Random/discrete_distribution.jl")
#     include("Random/Random.jl")
# end

# @testset "utils" begin
#     include("utils.jl")
# end


# Advent of Code 2021, Day 5: Hydrothermal Venture
# (c) 2022 TillW

struct Point
    x::Int; y::Int
end

struct Line
    first::Point; last::Point
end

function getCoords(txt::AbstractString)::Point
    parts = split(txt, ",")
    Point(parse(Int,parts[1]), parse(Int,parts[2]))
end

function readInput(filename::String)
    lines = Line[]
    input = readlines(filename)
    for txtLine in input
        txtPoints = split(txtLine, " -> ");
        line = Line(getCoords(txtPoints[1]), getCoords(txtPoints[2]))
        lines = [lines;line]
    end
    lines
end

function addLine(area::Array{Int}, line::Line, ignoreDiagonal::Bool)
    if (line.first.x == line.last.x || line.first.y == line.last.y)
        for x = min(line.first.x, line.last.x):max(line.first.x, line.last.x)
            for y = min(line.first.y, line.last.y):max(line.first.y, line.last.y)
                area[x+1,y+1] += 1
            end
        end
    elseif !ignoreDiagonal
        left = line.first.x < line.last.x ? line.first : line.last
        right = left == line.last ? line.first : line.last
        slope = left.y < right.y ? 1 : -1
        for i = 0:(right.x - left.x)
            area[left.x+1+i, left.y+1+(i*slope)] += 1
        end
    end
end

function printArea(area::Array{Int})
    for i = 1:length(area)
        print(area[i] == 0 ? "." : area[i])
        if (i % size(area, 1) == 0) println() end
    end
end

function solvePuzzle(inputFilename::String, ignoreDiagonalLines::Bool)
    lines = readInput(inputFilename)
    maxx = maximum([map((e) -> e.first.x, lines); map((e) -> e.last.x, lines)])
    maxy = maximum([map((e) -> e.first.y, lines); map((e) -> e.last.y, lines)])
    area = zeros(Int, maxx+1, maxy+1)
    for line in lines
        addLine(area, line, ignoreDiagonalLines)
    end
    count((x) -> x > 1, area)
end

function solvePuzzle1(inputFileName::String) solvePuzzle(inputFileName, true) end
function solvePuzzle2(inputFileName::String) solvePuzzle(inputFileName, false) end

print("Number of overlaps: ")
print(solvePuzzle1("input_05.txt"))
print("/")
println(solvePuzzle2("input_05.txt"))


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# https://discourse.julialang.org/t/solve-ode-with-many-different-initial-conditions/66384

# %%
using DifferentialEquations
using StaticArrays
using Base.Threads
using PyPlot
using Parameters

# %%
# define system parameters
Ip = 0.5
N  = 2
w  = 0.057
T  = 2π/w
F0 = 0.075
a = 1.0

# combine the system parameters into a single variable
p = (; Ip, N, w, T, F0, a,)

# define laser pulse
function Fx(p, t)
    @unpack F0, N, w, T = p
    F0 * cos(w*t/(2N))^2 * cos(w*t) * (abs(t) < N*T/2)
end
function Fy(p, t)
    @unpack F0, N, w, T = p
    F0 * cos(w*t/(2N))^2 * sin(w*t) * (abs(t) < N*T/2)
end

# define electron trajectory
function traj(u, p, t)
    @unpack a = p
    r3i = (u[1]^2 + u[2]^2 + a)^(-1.5)
    du1 = u[3]
    du2 = u[4]
    du3 = -u[1]*r3i - Fx(p, t)
    du4 = -u[2]*r3i - Fy(p, t)
    SVector(du1, du2, du3, du4)
end

# define rate
function W_adk(p, F, kd)
    @unpack Ip = p
    exp(-(2.0*(kd^2+2Ip)^1.5)/3F)
end

# define result box
Px  = range(-2, 2, length = 200)
Py  = range(-2, 2, length = 200)

# define simulation parameters
simparam = (
    Nt = 500,
    Nkd = 100,
    kd_max = 2.0,
)

# %%
function calcprob(p, Px, Py, simparam)
    @unpack Ip, N, w, T, F0, a = p
    Px_num, dPx = length(Px), step(Px)
    Py_num, dPy = length(Py), step(Py)
    @unpack Nt, Nkd, kd_max = simparam

    Prob_nthreads = zeros(Px_num, Py_num, nthreads())
    Trajs = [ODEProblem(traj, @SVector(zeros(4)), (-N*T/2, N*T/2), p) for _ in 1:nthreads()]

    @threads for tr in range(-0.9N*T/2, 0.9N*T/2, length = Nt)
        tid = threadid()
        Fxtr = Fx(p, tr)
        Fytr = Fy(p, tr)
        Ftr = hypot(Fxtr, Fytr)
        phi = atan(-Fytr, -Fxtr)
        r0 = Ip/Ftr
        x0 = r0*cos(phi)
        y0 = r0*sin(phi)
        tspan = (tr, N*T/2)

        for kd in range(-kd_max, kd_max, length = Nkd)
            kx0 = kd*cos(phi + 0.5π)
            ky0 = kd*sin(phi + 0.5π)
            rate = W_adk(p, Ftr, kd)

            u0 = SVector(x0, y0, kx0, ky0)
            Traj = remake(Trajs[tid]; u0=u0, tspan=tspan)
            sol = solve(Traj, Tsit5(), reltol=1e-6, save_everystep=false, save_start=false)
            (x, y, px, py) = sol.u[end]

            E_inf = 0.5*(px^2 + py^2) - 1/(x^2 + y^2 + a)
            E_inf >= zero(E_inf) || continue
            pxidx = round(Int, (px - Px[1])/dPx) + 1
            pyidx = round(Int, (py - Py[1])/dPy) + 1
            checkbounds(Bool, Prob_nthreads, pxidx, pyidx, tid) || continue
            @inbounds Prob_nthreads[pxidx, pyidx, tid] += rate
        end
    end
    Prob = reshape(sum(Prob_nthreads, dims=3), size(Prob_nthreads)[1:2])
end

# %%
# plotting
function plotprob(Px, Py, Prob)
    fig, ax = plt.subplots()
    ax.imshow(Prob', extent=(Px[1], Px[end], Py[1], Py[end]), origin="lower")
    ax.set_xlim(Px[1], Px[end])
    ax.set_ylim(Py[1], Py[end])
    ax.set_xlabel("\$p_x\$ (a.u.)")
    ax.set_ylabel("\$p_y\$ (a.u.)")
    #plt.show()
end

# %%
@time Prob = calcprob(p, Px, Py, simparam)
@time Prob = calcprob(p, Px, Py, simparam)
@time Prob = calcprob(p, Px, Py, simparam)
flush(stdout)
plotprob(Px, Py, Prob)

# %%

0018/Solve ODE with many different initial conditions Part 2.jl


export julia_expression

# to make a variable name from operators
varname(A::BaseOperator) = string(A.name,BaseOpType_expr[Int(A.t)])
varname(A::BaseOpProduct) = Symbol(varname.(A.v)...)
varname(A::Union{ExpVal,Corr}) = varname(A.ops)

# for sum indices use a valid and short symbol that cannot be expressed as a single char, i.e., no risk of collisions with "real" index names
# (the "standard" symbol '#' is not a variable name, so it becomes var"#")
indexpr(ind::QuIndex) = issumindex(ind) ? Symbol(:s̄,subscript(ind.num)) : (isintindex(ind) ? ind.num : Symbol(ind))
indexpr(A) = indexpr.(indices(A))

# the varnames optional (and ignored) argument allows to define a closure getexpr(A) = julia_expression(A,my_names)
function julia_expression(A::Param,varnames=nothing)
    # the parsing here allows "parameters" that are functions
    # e.g., with names such as "f(t)"
    x = Meta.parse(string(A.name))
    # is it a function call?
    if x isa Expr
        @assert isempty(indices(A))
    else
        x = :( $x[$(indexpr(A)...)] )
    end
    A.state == 'c' ? :( conj($x) ) : x
end
julia_expression(A::δ,varnames=nothing) = :( I[$(indexpr(A)...)] )
function julia_expression(A::Union{ExpVal,Corr},varnames=nothing)
    A = unalias(A)
    x = varname(A)
    if varnames === nothing || x ∈ varnames
        :( $x[$(indexpr(A)...)] )
    else
        # <A> = <A^†>^*
        Ap = A'
        x = varname(Ap)
        x ∈ varnames || throw(ValueError("passed allowed varnames = $(varnames), but neither $(varname(A)) nor $x are valid."))
        # since conjugation changes the order, we have to use the indices of A after conjugation
        :( conj($x[$(indexpr(Ap)...)]) )
    end
end
num_expr(s::Number) = :( $(isreal(s) ? Float64(s) : ComplexF64(s)) )
function julia_expression(A::QuTerm,varnames=nothing,s=1)
    isempty(A) && return num_expr(s)
    isempty(A.bares) || throw(ArgumentError("Cannot convert term $A with bare operators to a julia expression."))
    sexpr = isone(s) ? [] : num_expr(s)
    exprs = [sexpr; julia_expression.(A.δs); julia_expression.(A.params);
             julia_expression.(A.expvals,(varnames,)); julia_expression.(A.corrs,(varnames,))]
    length(exprs)==1 ? exprs[1] : :( *($(exprs...)) )
end
function julia_expression(A::QuExpr,varnames=nothing)
    isempty(A) && return 0
    exprs = [julia_expression(term,varnames,s) for (term,s) in sort!(Tuple.(collect(A.terms)))]
    length(exprs)==1 ? exprs[1] : :( +($(exprs...)) )
end


## First pass at math functions. There are more domain restrictions to be implemented.

### SetPrecision

@sjseealso_group(N, SetPrecision, GetPrecison, BigFloatInput, BigIntInput, BI, BF)

@mkapprule SetPrecision :nargs => 1
@sjdoc SetPrecision """
    SetPrecison(n)

Set the precsion of BigFloat numbers to `n` decimal digits. If `N` does not give the result you
want, you can use `SetPrecision`.
"""
@doap function SetPrecision(n::Real)
    setprecision(round(Int,3.322*n))
    return n
end

### GetPrecision

@mkapprule GetPrecision :nargs => 0
@sjdoc GetPrecision """
    GetPrecision(n)

Get the precsion of BigFloat numbers.
"""
@doap function GetPrecision()
    return precision(BigFloat)
end

"""
    evalmath(expr::Expr)

Evaluate `expr`. This is used for writing Symata interfaces
to math functions. Currently, it is equivalent to `eval`.
"""
evalmath(expr::Expr) = Symata.eval(expr::Expr)

## We also need to use these to convert Symata expressions to Julia
## This function writes apprules for math functions. Usally dispatches floating point
## args to Julia functions. Some have a fall through to SymPy
## These tuples have 1,2, or 3 members. Symbols are for Julia, Symata, and SymPy
## If only one member is present, the second is constructed by putting an inital capital on the first.
## If only one or two members are present, we do not fall back to SymPy.
## The last list has tuples with 2 args for which we do not use any Julia function.

# Typical symbols for Julia, Symata, SymPy
function mtr(sym::Symbol)
    s = string(sym)
    sjf = uppercasefirst(s)
    (sym, Symbol(sjf), sym)
end

# We choose not to implement :cis. Julia has it. sympy and mma, not

# List of not yet implemented sympy functions from keys(sympy.functions)
# :E1,:Ei,:FallingFactorial,:Heaviside, :Id,
# :Max,:Min,:Piecewise,
#  :Ynm_c,:Znm,
# :__builtins__,:__doc__,:__file__,:__name__,:__package__,:__path__,
# :acos,:acosh,:acot,:acoth,:acsc,:adjoint,
#  :assoc_laguerre,:assoc_legendre,
#  :bspline_basis, :bspline_basis_set, :ceiling,:chebyshevt_root,
# :chebyshevu_root,:combinatorial
#  :elementary,
#  :exp_polar, :factorial2,:ff
# :floor, :hyper,:im, :jacobi_normalized, :jn, :jn_zeros,
# :periodic_argument,:piecewise_fold,:polar_lift,
# :principal_branch,:re,:real_root,:rf,:root, :special,
# :transpose,:unbranched_argument,:yn

# TODO: FIX: sympy zeta may take two args
# two arg both float or complex : zeta(s,z)  (with domain restrictions)
# BellB should be split into BellB and BellY

# TODO: Implement rewrite
# sympy.airyai(x)[:rewrite](sympy.hyper)
# PyObject -3**(2/3)*x*hyper((), (4/3,), x**3/9)/(3*gamma(1/3)) + 3**(1/3)*hyper((), (2/3,), x**3/9)/(3*gamma(2/3))

# TODO: rename these arrays. Name by number of args plus identifying number.
# ie. two_args1, two_args2. Then explain in comment the particulars. A naming
# scheme to capture all possibilities is to cumbersome.

#  mtr(:exp)  <-- trying with and without this
   const single_arg_float_complex =   # check, some of these can't take complex args
    [ mtr(:sin), mtr(:cos), mtr(:tan), (:sind,:SinD), (:cosd,:CosD),(:tand,:TanD),
         (:sinpi,:SinPi), (:cospi,:CosPi), mtr(:sinh), mtr(:cosh),
         mtr(:tanh), (:acos,:ArcCos,:acos), (:asin,:ArcSin,:asin),
         (:atan,:ArcTan,:atan),(:acosd,:ArcCosD), (:asind,:ArcSinD),
         (:atand,:ArcTanD), mtr(:sec), mtr(:csc), mtr(:cot),(:secd,:SecD),(:csc,:CscD),(:cot,:CotD),
         (:asec,:ArcSec),(:acsc,:ArcCsc),(:acot,:ArcCot),  # (:acotd,:ArcCotD)
         (:csch,), mtr(:coth,),(:asinh,:ASinh,:asinh), (:acosh,:ACosh, :acosh),(:atanh,:ATanh, :atanh),
         (:asech,:ArcSech),(:acsch,:ArcCsch),(:acoth,:ArcCoth, :acoth),
         (:cosc,),
         (:log1p,),(:exp2,),(:exp10,),(:expm1,),(:abs2,),
         mtr(:erfc), mtr(:erfi),(:erfcx,),(:dawson,),(:real,:Re, :re),(:imag,:Im, :im),
         (:angle,:Arg,:arg),  (:lgamma, :LogGamma, :loggamma),
         (:lfact,:LogFactorial), mtr(:digamma), mtr(:trigamma),
         (:airyai,:AiryAi,:airyai),
         (:airybi,:AiryBi,:airybi),(:airyaiprime,:AiryAiPrime,:airyaiprime),(:airybiprime,:AiryBiPrime,:airybiprime),
         (:besselj0,:BesselJ0),(:besselj1,:BesselJ1),(:bessely0,:BesselY0),(:bessely1,:BesselY1),
         (:zeta,:Zeta,:zeta)
#         (:eta,:DirichletEta,:dirichlet_eta)
         ]

# (:log,),   removed from list above, because it must be treated specially (and others probably too!)

const single_arg_float_int_complex = [(:sinc,)]

const single_arg_float = [(:erfcinv,:InverseErfc,:erfcinv),(:invdigamma,:InverseDigamma)]

# FIXME  Ceiling should probably return integer type. So it can't be included in generic code here (or punt to sympy for everything)
    const single_arg_float_int = [(:signbit,:SignBit), (:ceil, :Ceiling, :ceiling), (:floor, :Floor, :floor)]

#  (:factorial,:Factorial, :factorial), <--- we do this by hand

    const single_arg_int = [(:isqrt,:ISqrt),(:ispow2,:IsPow2),(:nextpow2,:NextPow2),(:prevpow2,:PrevPow2),
                      (:isprime,:PrimeQ), (:(~), :BitNot)
                        ]

    const two_arg_int = [(:binomial,:Binomial,:binomial)
                         ]

# Do NDigits by hand for now!
@mkapprule NDigits :nargs => 1:2
@doap NDigits(n::T,b::V) where {T<:Integer,V<:Integer} = ndigits(n, base=b)
@doap NDigits(n::T) where {T<:Integer} = ndigits(n)
@doap NDigits(n) = mx
@doap NDigits(n,b) = mx

# TODO
#  Two or more args
#   (:gcd,:GCD), (:lcm, :LCM), Also sympy does these for polynomials

# Complicated!
    const one_or_two_args1 = [(:polygamma, :PolyGamma, :polygamma)]

    const two_arg_float_and_float_or_complex =
     [
      (:besselj,:BesselJ, :besselj), (:besseljx,:BesselJx), (:bessely,:BesselY,:bessely),
      (:besselyx,:BesselYx), (:hankelh1,:HankelH1,:hankel1), (:hankelh1x,:HankelH1x),
      (:hankelh2,:HankelH2,:hankel2), (:hankelh2x,:HankelH2x), (:besseli,:BesselI, :besseli),
      (:besselix,:BesselIx), (:besselk,:BesselK,:besselk), (:besselkx,:BesselKx),
       (:atan2,:ArcTan2,:atan2)
      ]

    const  two_arg_float = [ (:beta,),(:lbeta,:LogBeta),(:hypot,)]

    ## There are no Julia functions for these (or at least we are not using them).
    ## First symbol is for  Symata, Second is SymPy

# TODO: Migrate functions out this non-specific list
    const no_julia_function = [  (:ExpIntegralE, :expint),
                           (:GegenbauerC, :gegenbauer), (:HermiteH, :hermite),
                           (:BellB, :bell), (:BernoulliB, :bernoulli), (:CatalanNumber, :catalan),
                           (:EulerE, :euler), (:Subfactorial, :subfactorial), (:Factorial2, :factorial2),
                           (:FactorialPower, :FallingFactorial), (:Pochhammer, :RisingFactorial),
 (:LeviCivita, :LeviCivita),
(:KroneckerDelta, :KroneckerDelta),  (:HypergeometricPFQ, :hyper), (:FactorSquareFree, :sqf), (:MellinTransform, :mellin_transform),
(:InverseMellinTransform, :inverse_mellin_transform), (:SineTransform, :sine_transform), (:InverseSineTransform, :inverse_sine_transform),
(:CosineTransform, :cosine_transform), (:InverseCosineTransform, :inverse_cosine_transform),
(:HankelTransform, :hankel_transform), (:InverseHankelTransform, :inverse_hankel_transform)
]
# LeviCivita is different than in mathematica

# Mobius needs filtering
const no_julia_function_one_arg = [ (:EllipticK, :elliptic_k), (:HeavisideTheta,:Heaviside),
                                    (:CosIntegral, :Ci), (:SinIntegral, :Si),
                                    (:FresnelC, :fresnelc), (:FresnelS, :fresnels), (:DiracDelta, :DiracDelta),
                                    (:LogIntegral, :Li), (:Ei, :Ei), (:ExpandFunc, :expand_func), (:Denominator, :denom),
                                    (:MoebiusMu, :mobius), (:EulerPhi, :totient), (:Divisors, :divisors), (:DivisorCount, :divisor_count),
                                     (:SinhIntegral, :Shi),(:CoshIntegral, :Chi) ]

# elliptic_e take more than ints!
const no_julia_function_one_or_two_int = [ (:HarmonicNumber, :harmonic) , (:EllipticE, :elliptic_e) ]

# FIXME: DivisorSigma has arg order reversed in Mma
const no_julia_function_two_args = [(:LegendreP, :legendre_poly), (:EllipticF, :elliptic_f),
                                    (:ChebyshevT, :chebyshevt), (:ChebyshevU, :chebyshevu),
                                    (:Cyclotomic, :cyclotomic_poly), (:SwinnertonDyer, :swinnerton_dyer_poly),
                                    (:SphericalBesselJ, :jn), (:SphericalBesselY, :yn), (:DivisorSigma, :divisor_sigma)
#                                    (:PolyLog, :polylog), (:SphericalBesselJ, :jn), (:SphericalBesselY, :yn), (:DivisorSigma, :divisor_sigma)
                                    ]

const no_julia_function_two_or_three_args = [ (:EllipticPi, :elliptic_pi), (:LaguerreL, :laguerre_poly)]

const no_julia_function_three_args = [ (:LerchPhi, :lerchphi) ]

const no_julia_function_four_args = [ (:JacobiP, :jacobi), (:SphericalHarmonicY, :Ynm) ]

function make_math()

# Note: in Mma and Julia, catalan and Catalan are Catalan's constant. In sympy catalan is the catalan number
# can't find :stirling in sympy

    for x in no_julia_function_one_arg
        sjf = x[1]
        eval(macroexpand(Symata, :( @mkapprule $sjf)))
        write_sympy_apprule(x[1],x[2],1)
    end

    for x in no_julia_function_two_args
        sjf = x[1]
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 2)))
        write_sympy_apprule(x[1],x[2],2)
    end

    for x in no_julia_function_four_args
        nargs = 4
        sjf = x[1]
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => $nargs)))
        write_sympy_apprule(x[1],x[2],nargs)
    end

    for x in no_julia_function_three_args
        nargs = 3
        sjf = x[1]
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => $nargs)))
        write_sympy_apprule(x[1],x[2],nargs)
    end

    for x in no_julia_function_two_or_three_args
        sjf = x[1]
        eval(macroexpand(Symata, :( @mkapprule $sjf)))
        write_sympy_apprule(x[1],x[2],2)
        write_sympy_apprule(x[1],x[2],3)
    end

    for x in no_julia_function_one_or_two_int
        sjf = x[1]
        eval(macroexpand(Symata, :( @mkapprule $sjf)))
        write_sympy_apprule(x[1],x[2],1)
        write_sympy_apprule(x[1],x[2],2)
    end

    # TODO: update this code
    for x in no_julia_function
        set_up_sympy_default(x...)
        clear_attributes(x[1]) ## FIXME: Listable was set in previous line. Now we unset it. We need a better system for this.
        set_sysattributes(x[1])
    end

    # Ok, this works. We need to clean it up
    for x in single_arg_float_complex
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 1 )))
        write_julia_numeric_rule(jf,sjf,"AbstractFloat")
        write_julia_numeric_rule(jf,sjf,"CAbstractFloat")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],1)
        end
        set_attribute(Symbol(sjf),:Listable)
    end

    for x in single_arg_float
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 1 )))
        write_julia_numeric_rule(jf,sjf,"AbstractFloat")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],1)
        end
        set_attribute(Symbol(sjf),:Listable)
    end

    for x in single_arg_float_int
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 1 )))
        write_julia_numeric_rule(jf,sjf,"Real")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],1)
        end
    end

    # This is all numbers, I suppose
    for x in single_arg_float_int_complex
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 1 )))
        write_julia_numeric_rule(jf,sjf,"Real")
        write_julia_numeric_rule(jf,sjf,"CReal")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],1)
        end
    end

    for x in single_arg_int
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 1 :argtypes => [Integer] )))
        write_julia_numeric_rule(jf,sjf,"Integer")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],1)
        end
    end

    for x in two_arg_int
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 2 )))
        write_julia_numeric_rule(jf,sjf,"Integer", "Integer")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],2)
        end
    end

    # Mma allows one arg, as well
    for x in one_or_two_args1
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf )))
        write_julia_numeric_rule(jf,sjf,"Integer", "AbstractFloat")
        write_julia_numeric_rule(jf,sjf,"Integer", "CAbstractFloat")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],2)
        end
    end

    for x in two_arg_float
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 2 )))
        write_julia_numeric_rule(jf,sjf,"AbstractFloat", "AbstractFloat")
    end

    for x in two_arg_float_and_float_or_complex
        jf,sjf = only_get_sjstr(x...)
        eval(macroexpand(Symata, :( @mkapprule $sjf :nargs => 2 )))
        write_julia_numeric_rule(jf,sjf,"AbstractFloat", "AbstractFloat")
        write_julia_numeric_rule(jf,sjf,"AbstractFloat", "CAbstractFloat")
        write_julia_numeric_rule(jf,sjf,"Integer", "AbstractFloat")
        write_julia_numeric_rule(jf,sjf,"Integer", "CAbstractFloat")
        if length(x) == 3
            write_sympy_apprule(sjf,x[3],2)
        end
    end
end

Typei(i::Int)  =  "T" *string(i)
complex_type(t::AbstractString) = "Complex{" * t * "}"
function write_julia_numeric_rule(jf, sjf, types...)
    annot = join(AbstractString[ "T" * string(i) * "<:" *
                                 (types[i][1] == 'C' ? types[i][2:end] : types[i])  for i in 1:length(types)], ", ")
    protargs = join(AbstractString[ "x" * string(i) * "::" *
                                (types[i][1] == 'C' ? complex_type(Typei(i)) : Typei(i))
                                    for i in 1:length(types)], ", ")
    callargs = join(AbstractString[ "x" * string(i) for i in 1:length(types)], ", ")
#    appstr = "do_$sjf{$annot}(mx::Mxpr{:$sjf},$protargs) = $jf($callargs)"
    appstr = "do_$sjf(mx::Mxpr{:$sjf},$protargs) where {$annot} = $jf($callargs)"
    eval(Meta.parse(appstr))
end

function only_get_sjstr(jf,sjf,args...)
    return jf, sjf
end

function only_get_sjstr(jf)
    st = string(jf)
    sjf = Symbol(uppercasefirst(st))
    return jf, sjf
end

function get_sjstr(jf,sjf)
    do_common(sjf)
    return jf, sjf
end

function get_sjstr(jf)
    st = string(jf)
    sjf = String(uppercasefirst(st))
    do_common(sjf)
    return jf, sjf
end

## FIXME: This is outdated. Some of this is handled in @mkapprule
# Handle functions that do *not* fall back on SymPy
function do_common(sjf)
    aprs = "Symata.apprules(mx::Mxpr{:$sjf}) = do_$sjf(mx,margs(mx)...)"
    aprs1 = "do_$sjf(mx::Mxpr{:$sjf},x...) = mx"
    evalmath(Meta.parse(aprs))
    evalmath(Meta.parse(aprs1))
    set_attribute(Symbol(sjf),:Protected)
    set_attribute(Symbol(sjf),:Listable)
end

# Handle functions that fall back on SymPy
function get_sjstr(jf, sjf, sympyf)
    set_up_sympy_default(sjf, sympyf)
    return jf, sjf
end

# Faster if we don't do interpolation
function write_sympy_apprule(sjf, sympyf, nargs::Int)
    callargs = Array{AbstractString}(undef, 0)
    sympyargs = Array{AbstractString}(undef, 0)
    for i in 1:nargs
        xi = "x" * string(i)
        push!(callargs, xi)
        push!(sympyargs, "sjtopy(" * xi * ")")
        end
    cstr = join(callargs, ", ")
    sstr = join(sympyargs, ", ")
    aprpy = "function do_$sjf(mx::Mxpr{:$sjf},$cstr)
               try
                 (sympy[:$sympyf]($sstr) |> pytosj)
               catch e
                 showerror(stdout, e)
                 mx
               end
            end"
    evalmath(Meta.parse(aprpy))
end

function set_up_sympy_default(sjf, sympyf)
    aprs = "Symata.apprules(mx::Mxpr{:$sjf}) = do_$sjf(mx,margs(mx)...)"
    aprs1 = "function do_$sjf(mx::Mxpr{:$sjf},x...)
               try
                 (sympy[:$sympyf](map(sjtopy,x)...) |> pytosj)
               catch
                   mx
               end
           end"
    evalmath(Meta.parse(aprs))
    evalmath(Meta.parse(aprs1))
    set_attribute(Symbol(sjf),:Protected)
    set_attribute(Symbol(sjf),:Listable)
end

make_math()

#### Gamma function

@mkapprule Gamma :nargs => 1:2

@doap Gamma(a) = sjgamma(a)
@doap Gamma(a,z) = sjgamma(a,z)

set_pytosj(:gamma, :Gamma)
set_pytosj(:uppergamma, :Gamma)
set_sjtopy(:Gamma, :sympy_gamma)

@sjdoc Gamma """
    Gamma(a)

the gamma function.

    Gamma(a,z)

the upper incomplete Gamma function.
"""

### GammaRegularized

## Implementation is not complete
@mkapprule GammaRegularized :nargs => 2:3

@doap function GammaRegularized(a,z::Number)
    if z == 0
       return  _gammaregularizedzero(mx,a)
    end
   return  _gammaregularizednozero(mx,a,z)
end

function _gammaregularizedzero(mx,a::Real)
    a > 0 ? 1 : a == 0 ? 0 : ComplexInfinity
end

function _gammaregularizedzero(mx,a::Complex)
    ra = real(a)
    ra > 0 ? 1 : ra == 0 ? Indeterminate : ComplexInfinity
end

function _gammaregularizedzero(mx,a)
    mx
end

function _gammaregularizednozero(mx,a::Real,z::Number)
    a < 0 && return zero(a) ## need to consider type of z as well
    mxpr(:Gamma,a,z)/mxpr(:Gamma,a)
end

function _gammaregularizednozero(mx,a::Number,z::Number)
    mxpr(:Gamma,a,z)/mxpr(:Gamma,a)
end

function _gammaregularizednozero(mx,a,z)
    mx
end

### Erf

@mkapprule Erf :nargs => 1:2

@doap Erf(x) = sjerf(x)
@doap Erf(x,y) = sjerf(x,y)

set_pytosj(:erf, :Erf)
set_pytosj(:erf2, :Erf)
set_sjtopy(:Erf, :sympy_erf)

### InverseErf

@mkapprule InverseErf :nargs => 1:2

@doap InverseErf(x::T) where {T<:AbstractFloat} = x < 1 && x > -1 ? erfinv(x) : mx
@doap InverseErf(x) = pytosj(sympy[:erfinv](sjtopy(x)))
@doap InverseErf(x,y) = pytosj(sympy[:erf2inv](sjtopy(x),sjtopy(y)))

register_sjfunc_pyfunc(:InverseErf,:erfinv)
register_only_pyfunc_to_sjfunc(:InverseErf,:erf2inv)

### IntegerDigits

@sjdoc IntegerDigits """
    IntegerDigits(n,[, base][, pad])

return an array of the digits of `n` in the given `base`,
optionally padded with zeros to `pad` characters.

In contrast to Julia, more significant digits are at lower indexes.
"""
@mkapprule IntegerDigits  :nargs => 1:3
set_sysattributes(:IntegerDigits, :Listable)

_intdig(args) = MList(reverse!(args)...) |> setfixed

@doap IntegerDigits(n::Integer) = digits(n) |> _intdig
@doap IntegerDigits(n::Integer,b::Integer) = digits(n, base=convert(Int, b)) |> _intdig
@doap IntegerDigits(n::Integer,b::Integer,p::Integer) = digits(n, base=convert(Int, b), convert(Int, p)) |> _intdig

@mkapprule Log :nargs => 1:2

# sjlog defined in wrappers.jl
@doap Log(x,y) = sjlog(x,y)
@doap Log(x) = sjlog(x)

### N

@sjdoc N """
    N(expr)

try to give a the numerical value of `expr`.

    N(expr,p)

try to give `p` decimal digits of precision.

Sometimes `N` does not give the number of digits requested. In this case, you can use `SetPrecision`.
"""

## TODO: call evalf(expr,n) on sympy functions to get arbitrary precision numbers.
## Eg. N(LogIntegral(4),30) does not give correct result.
## ... a list of heads all of which should be converted and then passed to evalf.
## N needs to be rewritten
function apprules(mx::Mxpr{:N})
    outer_N(margs(mx)...)
end

function outer_N(expr)
    do_N(expr)
end

function outer_N(expr, p)
    if p > 16
        pr = precision(BigFloat)
        dig = round(Int,p*3.322)
        set_mpmath_dps(dig)     # This is for some SymPy math. But, it is not clear when this will help us
        setprecision( BigFloat, dig) # new form
        res = meval(do_N(expr,p))
        setprecision(BigFloat, pr)
        restore_mpmath_dps()
        return res
    else
        do_N(expr)
    end
end

do_N(x::Float64,dig) = x  # Mma does this
do_N(x,dig) = x
do_N(x) = x
do_N(n::T) where {T<:AbstractFloat} = n
do_N(n::T) where {T<:Real} = float(n)
do_N(n::Complex{T}) where {T<:Real} = complex(float(real(n)), float(imag(n)))
do_N(n::T,p::V) where {T<:Real,V<:Integer}  = float_with_precision(n,p)

# p is number of decimal digits of precision
# Julia doc says set_bigfloat_precision uses
# binary digits, but it looks more like decimal.

function float_with_precision(x,p)
        res = BigFloat(x)
end

# These rely on fixed-point evaluation to continue with N, this is not efficient
# We need to do it all here.

function _make_N_body(d1,d2)
    quote begin
        get_attribute(mx,:NHoldAll) && return mx;
        len = length(mx)
        args = margs(mx)
        nargs = newargs(len)
        start::Int = 1
        if get_attribute(mx,:NHoldFirst)
            nargs[1] = args[1]
            for i in 2:len
                nargs[i] = $d2
            end
        elseif get_attribute(mx,:NHoldRest)
            nargs[1] = $d1
            for i in 2:len
                nargs[i] = args[i]
            end
        else
            @inbounds for i in 1:len
                nargs[i] = $d2
            end
        end
        res = mxpr(mhead(mx),nargs)
    end
  end
end

function make_Mxpr_N()
    _Nbody = _make_N_body(:(do_N(args[1])), :(do_N(args[i])) )
    @eval begin
        function do_N(mx::Mxpr)
            $_Nbody
            res
        end
    end

    _Nbody = _make_N_body(:(do_N(args[1],p)), :(do_N(args[i],p)) )
    @eval begin
        function do_N(mx::Mxpr,p::T) where T<:Integer
            $_Nbody
            res
        end
    end
end

make_Mxpr_N()

@sjdoc GoldenRatio """
    GoldenRatio

equal to `(1+Sqrt(5))/2`.
"""

# We need to use dispatch as well, not conditionals
function do_N(s::SJSym)
    if s == :Pi || s == :π
        return float(MathConstants.pi)
    elseif s == :E
        return float(MathConstants.e)
    elseif s == :EulerGamma
        return float(MathConstants.eulergamma)
    elseif s == :GoldenRatio
        return float(MathConstants.golden)
    elseif s == :Catalan
        return float(MathConstants.catalan)
    end
    return s
end

function do_N(s::SJSym,pr::T) where T<:Integer
    if s == :Pi || s == :π
        return float_with_precision(MathConstants.pi, pr)
    elseif s == :E
        return float_with_precision(MathConstants.e, pr)
    elseif s == :EulerGamma
        return float_with_precision(MathConstants.eulergamma, pr)
    elseif s == :GoldenRatio
        return float_with_precision(MathConstants.golden, pr)
    elseif s == :Catalan
        return float_with_precision(MathConstants.catalan, pr)
    end
    return s
end

@sjdoc Precision """
    Precision(x)

get the precision of a floating point number `x` in decimal digits

    Precision()

get the current precision of `BigFloat` arithmetic

    Precision(Binary => True)
    Precision(x, Binary => True)

give precision as defined by the effective number of bits in the mantissa.
"""

@mkapprule Precision :nodefault => true, :options => Dict( :Binary => false )

#apprules(mx::Mxpr{:Precision}) = do_Precision(mx,margs(mx)...)

# TODO: mkapprule should write the correct default rule if keywords are present
do_Precision(mx::Mxpr{:Precision},args...; kws...) = mx

function do_Precision(mx::Mxpr{:Precision},x::AbstractFloat; Binary=false)
    Binary ? precision(x) : round(Int,precision(x) * log10(2))
end

function do_Precision(mx::Mxpr{:Precision}; Binary=false)
    pr = precision(BigFloat)
    Binary ? pr : round(Int,pr * log10(2))
end

@sjdoc Re """
    Re(z)

return the real part of `z`.
"""

@sjdoc Im """
    Im(z)

return the imaginary part of `z`.
"""

# Using Julia would be faster in some cases for ReIm and AbsArg
@sjdoc ReIm """
    ReIm(z)

returns `[Re(z),Im(z)]`.
"""

@mkapprule ReIm
@doap ReIm(z) = mxpr(:List,mxpr(:Re,z), mxpr(:Im,z))

# FIXME Abs(I-3) returns 10^(1/2) instead of (2^(1/2))*(5^(1/2)). i.e. does not evaluate to a fixed point
@sjdoc AbsArg """
    AbsArg(z)

return `[Abs(z),Arg(z)]`.
"""

@mkapprule AbsArg
@doap AbsArg(z) = mxpr(:List,mxpr(:Abs,z), mxpr(:Arg,z))

# Re and Im code below is disabled so it does not interfere with SymPy.
# Mma allows complex numbers of mixed Real type. So does Julia: Complex{Real}(1//3, 3.0)
# Implementation not complete. eg  Im(a + I *b) --> Im(a) + Re(b)

# do_Re{T<:Real}(mx::Mxpr{:Re}, x::Complex{T}) = real(x)
# do_Re(mx::Mxpr{:Re}, x::Real) = x

# function do_Re(mx::Mxpr{:Re}, m::Mxpr{:Times})
#     f = m[1]
#     return is_imaginary_integer(f) ? do_Re_imag_int(m,f) : mx
# end

# # dispatch on type of f. Maybe this is worth something.
# function do_Re_imag_int(m,f)
#     nargs = copy(margs(m))
#     popfirst!(nargs)
#     if length(nargs) == 1
#         return mxpr(:Times,-imag(f),mxpr(:Im,nargs))
#     else
#         return mxpr(:Times,-imag(f),mxpr(:Im,mxpr(:Times,nargs)))
#     end
# end

# do_Im{T<:Real}(mx::Mxpr{:Im}, x::Complex{T}) = imag(x)
# do_Im(mx::Mxpr{:Im}, x::Real) = zero(x)

# function do_Im(mx::Mxpr{:Im}, m::Mxpr{:Times})
#     f = m[1]
#     return is_imaginary_integer(f) ? do_Im_imag_int(m,f) : mx
# end

# function do_Im_imag_int(m,f)
#     nargs = copy(margs(m))
#     popfirst!(nargs)
#     if length(nargs) == 1
#         return mxpr(:Times,imag(f),mxpr(:Re,nargs))
#     else
#         return mxpr(:Times,imag(f),mxpr(:Re,mxpr(:Times,nargs)))
#     end
# end

#### Complex

@sjdoc Complex """
    Complex(a,b)

return a complex number when `a` and `b` are Reals

The conversion is done when the expression is parsed, so it is much faster than `a + I*b`.
"""

# Complex with two numerical arguments is converted at parse time. But, the
# arguments may evaluate to numbers only at run time, so this is needed.
# mkapprule requires that the first parameter do_Complex be annotated with the Mxpr type.

@mkapprule Complex
do_Complex(mx::Mxpr{:Complex},a::T,b::T) where {T<:Number} = complex(a,b)
do_Complex(mx::Mxpr{:Complex},a::T) where {T<:Number} = complex(a)

@sjdoc Rational """
    Rational(a,b), or a/b

return a `Rational` for `Integer`s `a` and `b`.  This is done when the
expression is parsed, so it is much faster than `a/b`.
"""

# Same here. But we need to use mdiv to reduce rationals to ints if possible.
@mkapprule Rational
do_Rational(mx::Mxpr{:Rational},a::Number,b::Number) = mdiv(a,b)

#### Rationalize

@sjdoc Rationalize """
    Rationalize(x)

return a `Rational` approximation of `x`.

    Rationalize(x,tol)

return an approximation differing from `x` by no more than `tol`.
"""

@mkapprule Rationalize
do_Rationalize(mx::Mxpr{:Rationalize},x::AbstractFloat) = rationalize(x)
do_Rationalize(mx::Mxpr{:Rationalize},x::AbstractFloat,tol::Number) = rationalize(x,tol=float(tol))
function do_Rationalize(mx::Mxpr{:Rationalize},x::Symbolic)
    r = doeval(mxpr(:N,x))  # we need to redesign do_N so that we can call it directly. See above
    return isa(r,AbstractFloat) ? do_Rationalize(mx,r) : x
end
function do_Rationalize(mx::Mxpr{:Rationalize},x::Symbolic,tol::Number)
    ndig = round(Int,-log10(tol))      # This is not quite correct.
    r = doeval(mxpr(:N,x,ndig))  # we need to redesign do_N so that we can call it directly. See above.
    return isa(r,AbstractFloat) ? do_Rationalize(mx,r,tol) : x
end
do_Rationalize(mx::Mxpr{:Rationalize},x) = x

#### Numerator

@sjdoc Numerator """
    Numerator(expr)

return the numerator of `expr`.
"""

apprules(mx::Mxpr{:Numerator}) = do_Numerator(mx::Mxpr{:Numerator},margs(mx)...)
do_Numerator(mx::Mxpr{:Numerator},args...) = mx

do_Numerator(mx::Mxpr{:Numerator},x::Rational) = numerator(x)
do_Numerator(mx::Mxpr{:Numerator},x) = x
function do_Numerator(mx::Mxpr{:Numerator},x::Mxpr{:Power})
    find_numerator(x)
end

function do_Numerator(mx::Mxpr{:Numerator},m::Mxpr{:Times})
    nargsn = newargs()
    args = margs(m)
    for i in 1:length(args)
        arg = args[i]
        res = find_numerator(arg)
        if res != 1
            push!(nargsn,res)
        end
    end
    if isempty(nargsn)
        return 1  # which one ?
    end
    if length(nargsn) == 1
        return nargsn[1]
    end
    return mxpr(mhead(m),nargsn)
end

find_numerator(x::Rational) = numerator(x)
find_numerator(x::Mxpr{:Power}) = pow_sign(x,expt(x)) > 0 ? x : 1

function pow_sign(x, texpt::Mxpr{:Times})
    fac = texpt[1]
    return is_Number(fac) && fac < 0 ? -1 : 1
end

pow_sign(x, texpt::Number) = texpt < 0 ? -1 : 1
pow_sign(x, texpt) = 1
find_numerator(x) = x

### Chop

@mkapprule Chop  :nargs => 1:2

@sjdoc Chop """
    Chop(expr)

set small floating point numbers in `expr` to zero.

    Chop(expr,eps)

set floating point numbers smaller in magnitude than `eps` in `expr` to zero.
"""
do_Chop(mx::Mxpr{:Chop}, x) = zchop(x)
do_Chop(mx::Mxpr{:Chop}, x, zeps) = zchop(x,zeps)
const chop_eps = 1e-14
zchop(x::T, eps=chop_eps) where {T<:AbstractFloat} = abs(x) > eps ? x : 0  # don't follow abstract type
zchop(x::T, eps=chop_eps) where {T<:Number} = x
zchop(x::Complex{T}, eps=chop_eps) where {T<:AbstractFloat} = complex(zchop(real(x),eps),zchop(imag(x),eps))
zchop(a::T, eps=chop_eps) where {T<:AbstractArray} = (b = similar(a); for i in 1:length(a) b[i] = zchop(a[i],eps) end ; b)
zchop!(a::T, eps=chop_eps) where {T<:AbstractArray} = (for i in 1:length(a) a[i] = zchop(a[i],eps) end ; a)
zchop(x::Expr,eps=chop_eps) = Expr(x.head,zchop(x.args)...)
zchop(x) = x
zchop(x,eps) = x

function zchop(mx::T) where T<:Mxpr
    nargs = similar(margs(mx))
    for i in 1:length(nargs) nargs[i] = zchop(mx[i]) end
    mxpr(mhead(mx), nargs)
end

function zchop(mx::T,zeps) where T<:Mxpr
    nargs = similar(margs(mx))
    for i in 1:length(nargs) nargs[i] = zchop(mx[i],zeps) end
    mxpr(mhead(mx), nargs)
end

### Conjugate

@mkapprule Conjugate :nargs => 1

@doap Conjugate(z::Number) = conj(z)
@doap Conjugate(z::Mxpr{:Plus}) = mxpr(:Plus, mapmargs( x -> mxpr(:Conjugate, x), margs(z))...)
@doap Conjugate(z::Mxpr) = z |> sjtopy |> sympy[:conjugate] |> pytosj

### Exp

# The parser normally takes care of this,
# But, when converting expressions from Sympy, we get Exp, so we handle it here.
# TODO: don't throw away other args. (but, we hope this is only called with returns from Sympy)
apprules(mx::Mxpr{:Exp}) = mxpr(:Power,:E,mx[1])

# TODO: handle 3rd argument
#### Mod
@mkapprule  Mod  :nargs => 2
do_Mod(mx::Mxpr{:Mod}, x::T, y::V) where {T<:Integer, V<:Integer} = mod1(x,y)

#### DivRem
@mkapprule  DivRem  :nargs => 2
do_DivRem(mx::Mxpr{:DivRem}, x::T, y::V) where {T<:Integer, V<:Integer} = mxprcf(:List,divrem(x,y)...)

#### Abs

@sjdoc Abs """
    Abs(z)

the absolute value of `z`.
"""

@mkapprule Abs :nargs => 1

@doap Abs(n::T) where {T<:Number} = mabs(n)

@doap function Abs(x)
    x |> sjtopy |> sympy[:Abs] |> pytosj
end

# The following code may be useful, but it is not complete
# # Abs(x^n) --> Abs(x)^n  for Real n
function do_Abs(mx::Mxpr{:Abs}, pow::Mxpr{:Power})
     doabs_pow(mx,base(pow),expt(pow))
end

# SymPy does not do this one
doabs_pow(mx,b,e::T) where {T<:Real} = mxpr(:Power,mxpr(:Abs,b),e)

function doabs_pow(mx,b,e)
    mx[1] |> sjtopy |> sympy[:Abs] |> pytosj
end

# do_Abs(mx::Mxpr{:Abs},prod::Mxpr{:Times}) = do_Abs(mx,prod,prod[1])

# #doabs(mx,prod,s::Symbol)

# function do_Abs{T<:Number}(mx::Mxpr{:Abs},prod,f::T)
#     f >=0 && return mx
#     if f == -1
#         return doabsmone(mx,prod,f)
#     end
#     args = copy(margs(prod))
#     args[1] = -args[1]
#     return mxpr(:Abs,mxpr(:Times,args))
# end

# function doabsmone{T<:Integer}(mx,prod,f::T)
#     args = copy(margs(prod))
#     popfirst!(args)
#     if length(args) == 1
#         return mxpr(:Abs,args)
#     else
#         return mxpr(:Abs,mxpr(:Times,args))
#     end
# end

# TODO Fix canonical routines so that 1.0 * a is not simplifed to a
function doabsmone(mx,prod,f::T) where T<:Real
    args = copy(margs(prod))
    popfirst!(args)
    if length(args) == 1
        res = mxpr(:Times,one(f),mxpr(:Abs,args))
    else
        res = mxpr(:Times,one(f),mxpr(:Abs,mxpr(:Times,args)))
    end
    return res
end

#### Sign

@mkapprule Sign :nargs => 1

do_Sign(mx::Mxpr{:Sign}, x::Number) = sign_number(x)

function sign_number(z::Complex{T}) where T<:Real
    av = mpow(real(z*conj(z)),-1//2)
    av == 1 ? z : mxprcf(:Times, z,  av)
end

sign_number(x::Real) = convert(Int,sign(x))

function do_Sign(mx::Mxpr{:Sign}, x)
    x |> sjtopy |> sympy[:sign] |> pytosj
end

# SymPy does not do Sign((I+1)*a), so we catch this case here.
function do_Sign(mx::Mxpr{:Sign}, x::Mxpr{:Times})
    length(x) == 0 && return 1
    x1 = x[1]
    if isa(x1,Number)
        mxpr(:Times, sign_number(x1), mxpr(:Sign , mxpr(:Times,x[2:end]...)))
    else
        x |> sjtopy |> sympy[:sign] |> pytosj
    end
end

# do_Sign(mx::Mxpr{:Sign}, x) = mx

# # Maybe each symbol should be a type. Like irrational or Mxpr
# function do_Sign(mx::Mxpr{:Sign}, x::SJSym)
#     x == :Pi && return 1
#     x == :E && return 1
#     x == :EulerGamma && return 1
#     return mx
# end

#### LowerGamma

@mkapprule LowerGamma

do_LowerGamma(mx::Mxpr{:LowerGamma}, a, z) =  mxpr(:Gamma,a) - mxpr(:Gamma,a,z)

@sjdoc LowerGamma """
    LowerGamma(a,z)

the lower incomplete gamma function.
"""

#### Norm

# TODO: Implement p norm
@mkapprule Norm :nargs => 1

@doap Norm(x::Number) = Abs(x)

@doap function Norm(x::Mxpr{:List})
    vectorq(x) || return mx # TODO error message
    local sum0 = 0
    for i in 1:length(x)
        sum0 += mpow(Abs(x[i]),2)
    end
    mpow(sum0,1//2)
end

@sjdoc ExpandFunc """
    ExpandFunc(expr)

rewrite some multi-parameter special functions using simpler functions.
"""

@sjdoc I """
    I

the imaginary unit.
"""

@sjdoc E """
    E

the base of the natural logarithm.
"""

@sjdoc Pi """
    Pi

the trigonometric constant `π`.
"""

@sjdoc BellB """
    BellB(n)

gives the `n^{th}` Bell number, `B_n`.

    BellB(n, x)

gives the `n^{th}` Bell polynomial, `B_n(x)`.

    BellB(n, k, [x1, x2, ... x_{n-k+1}])
"""

### Identity

@sjdoc Identity """
    Identity(expr)

return `expr`.
"""

@mkapprule Identity

@doap Identity(x) = x

### Total

@mkapprule Total

@sjdoc Total """
    Total(list)

give the sum of elements in `list`
"""

@doap Total(x::AbstractArray) = sum(x)

## We could try some things to speed this up... check that all elements are numbers ?
## with all() it takes about as much time to check as it does to sum. Maybe use try/catch
#@doap Total(x::Mxpr{:List}) = mxpr(:Plus,margs(x))
@doap Total(x::Mxpr{:List}) = _apply_plus(x)

### CosPi, SinPi

## for non-float arguments
## Not sure it is good to do this: N(CosPi(100)) will be less accurate than CosPi(N(100))
@doap CosPi(x) = mxpr(:Cos, mmul(x,Pi))
@doap SinPi(x) = mxpr(:Sin, mmul(x,Pi))

@sjdoc CosPi """
    CosPi(x)

returns `cos(π x)`. This is more accurate than `Cos(Pi*x)` for
floating point `x`.
"""

@sjdoc SinPi """
    SinPi(x)

returns `sin(π x)`. This is more accurate than `Sin(Pi*x)` for
floating point `x`.
"""

### ProductLog

@mkapprule ProductLog :nargs => 1:2

@doap function ProductLog(z::FloatRC)
    LambertW.lambertw(z)
end

@doap function ProductLog(z)
    pytosj(sympy[:LambertW](sjtopy(z)))
end

@doap function ProductLog(z::FloatRC, k::Number)
    LambertW.lambertw(z,k)
end

@doap function ProductLog(z,k)
    pytosj(sympy[:LambertW](sjtopy(z), sjtopy(k)))
end

### DirichletEta

## Note: sympy does not automatically transform eta to a function of zeta.
## However, Mma transforms it immediately
## However, if we do ex = DirichletEta(s) and then substitute 1 for s, we get indeterminate.
## So, one could argue for only converting DirichletEta(s) explicitly

@sjdoc DirichletEta """
    DirichletEta(s)

is the Dirichlet eta function.
"""

@mkapprule DirichletEta  :nargs => 1
@doap DirichletEta(x::FloatRC) = SpecialFunctions.eta(x)
@doap DirichletEta(x::Number) = _dirichlet_eta(x)
@doap DirichletEta(x::Integer) = x == 1 ? Log2 : _dirichlet_eta(x)
@doap DirichletEta(x) = _dirichlet_eta(x)
_dirichlet_eta(x::Number) = (1 - mpow(2, (1 - x))) * Zeta(x)
_dirichlet_eta(x) = mminus(1, mpow(2, mminus(1, x))) * Zeta(x)

### PolyLog

@sjdoc PolyLog """
    PolyLog(s,z)

is the polylogarithm function.
"""

@mkapprule PolyLog :nargs => 2

@doap PolyLog(s, z) = _polylog(mx,s,z)

_polyzero(z) = mmul(z , mpow(mminus(1,z),-1))

function _polylog(mx,s::Integer,z)
    s == 1 && return -Log(mminus(1,z))
    s == 0 && return _polyzero(z)
    if s < 0
        u = gensym()
        r = _polyzero(u)
        for i in 1:(-s)
            r = doeval(mxpr(:Factor, u*mxpr(:D,r,u)))
        end
        return replaceall(r, mxpr(:Rule, u, z))
    end
    if z == 1//2
        if s == 2
            return (Pisq - 6*(mpow(Log2,2)))/12
        elseif s == 3
            return (4*Log2^3  - 2*Pisq*Log2 + 21*Zeta(3))/24
        end
    end
    if s == 2
        z == I && return mmul(I,:Catalan) - Pisq/48
        z == MinusI && return -mmul(I,:Catalan) - Pisq/48
    end
    _polylog1(mx,s,z)
end

_polylog(mx,s,z) = _polylog1(mx,s,z)

function _polylog1(mx,s,z::Integer)
    z == 1 && return Zeta(s)
    z == -1 && return -DirichletEta(s)
    _polylog2(mx,s,z)
end

_polylog1(mx,s,z) = _polylog2(mx,s,z)

function _polylog2(mx,s1::Number,z1::Number)
    (s,z) = promote(s1,z1)
    polylognum(mx,s,z)
end

_polylog2(mx,s,z) = mx

polylognum(mx,s::FloatRC,z::FloatRC) = sympy[:polylog](s,z) |> pytosj
polylognum(mx,s,z) = mx

### ExpPolar

## sympy returns this sometimes. I am not sure what it is, but it usually just hinders
## desired evaluation. We could also handle this as a direct translation in sympy.jl.
## There may be some reason to preserve this.

@mkapprule ExpPolar
@doap ExpPolar(x) = Exp(x)

### Factorial

@sjdoc Factorial """
    Factorial(n)

is the factorial function.
"""
@mkapprule Factorial :nargs => 1
@doap Factorial(n::SJSym) = mx
@doap Factorial(n) = sjfactorial(n)

# function sjfactorial(mx::Mxpr{:DirectedInfinity})
#     _sjfactorial_inf(mx,mx[1])
# end

_sjfactorial_inf(mx, x, v) = mx
_sjfactorial_inf(mx,x,v::Real) = v > 0 ? Infinity : Indeterminate
_sjfactorial_inf(mx,x,v::Complex) = real(v) == 0 ? 0 : mx

@doap function Factorial(x::Mxpr{:DirectedInfinity})
    length(x) == 0 && return mx
    _sjfactorial_inf(mx,x,x[1])
end

### Sqrt

## we must have square root for when this is constructed programatically. eg. Map(Sqrt,[4])
## Sqrt is not always read and translated!
@mkapprule Sqrt :nargs => 1
@doap Sqrt(x) = mpow(x,1//2)

### CubeRoot

@mkapprule CubeRoot :nargs => 1

@doap CubeRoot(x::AbstractFloat) = cbrt(x)
@doap function CubeRoot(x::T) where T<:Union{Integer,Rational}
    x == 0 && return 0
    x > 0 && return mpow(x,1//3)
    return -mpow(-x,1//3)
end

@doap function CubeRoot(x::Complex)
    @symwarn("CubeRoot::preal: The parameter $x should be real valued.")
    return mx
end

@doap function CubeRoot(x::Mxpr)
    return mxpr(:Surd,x,3)
end

@doap function CubeRoot(x::Symbol)
    return mxpr(:Surd,x,3)
end

### Surd

@mkapprule Surd :nargs => 2

@doap function Surd(x,n::Integer)
    if iseven(n)
        @symwarn("Surd::noneg: Surd is not defined for even roots of negative values.")
        return mx
    end
    return _surd(mx,x,n)
end

function _surd(mx,x::Real,n)
    x >= 0 && return mpow(x,1//n)
    mminus(mpow(-x,1//n))
end

function _surd(mx,x::Complex,n)
    @symwarn("Surd::preal: The parameter I should be real valued")
    mx
end

_surd(mx,x,n) = mx

function _surd(mx,x::TimesT,n)
    length(x) < 2 && return mx
    isa(x[1],Real) && return mxpr(:Surd,x[1],n) * mxpr(:Surd, mxpr(:Times, x[2:end]...),n)
    mx
end

### Power

@mkapprule Power :nargs => 2

do_Power(mx::Mxpr{:Power}, b::Symbolic, n::Integer) = n == 1 ? b : n == 0 ? one(n) : mx

@doap Power(b::SJSym, expt) = b == :E ? dopowerE(mx, expt) : mx

@doap Power(b::SJSym, expt::Integer) = expt == 1 ? b : expt == 0 ? one(expt) :
    b == :E ? dopowerE(mx, expt) : mx

dopowerE(mx, expt::T) where {T<:AbstractFloat} = exp(expt)
dopowerE(mx, expt::Complex{T}) where {T<:AbstractFloat} = exp(expt)

function dopowerE(mx, expt)
    syexpt = sjtopy(expt)
    syres = sympy[:exp](syexpt)
    res = pytosj(syres)
    if is_Mxpr(res, :Exp)
        res[1] == expt && return mx
        return mxpr(:Power, :E, margs(res)...)
    end
    res
end

# Don't handle this yet.
do_Power(mx::Mxpr{:Power},   b::Complex{T},expt::Rational{V}) where {T<:Integer,V<:Integer} = mx
do_Power(mx::Mxpr{:Power},   b::Complex{T},expt::Complex{Rational{V}}) where {T<:Integer,V<:Integer} = mx

do_Power(mx::Mxpr{:Power},   b::T,expt::V) where {T<:Number,V<:Number} = mpow(b,expt)

# For some reason, we need this integer rule. For instance for (a^2)^2 --> a^4
do_Power(mx::Mxpr{:Power}, b::Mxpr{:Power}, exp::T) where {T<:Integer} = mpow(base(b), mmul(exp,expt(b)))
do_Power(   mx::Mxpr{:Power}, b::Mxpr{:Power}, exp::T) where {T<:Real} = mpow(base(b), mmul(exp,expt(b)))

do_Power(mx::Mxpr{:Power},   b::Mxpr{:Power}, exp) = is_Number(expt(b)) ? mpow(base(b), mmul(expt(b),exp)) : mx

do_Power(mx::Mxpr{:Power},   b::SJSym,expt::Complex{T}) where {T<:AbstractFloat} = b == :E ? exp(expt) : mx

do_Power(mx::Mxpr{:Power},   b::Mxpr{:DirectedInfinity},expt::T) where {T<:Integer} = mpow(b,expt)
do_Power(mx::Mxpr{:Power},   b::Mxpr{:DirectedInfinity},expt::T) where {T<:Number} = mpow(b,expt)

function do_Power(mx::Mxpr{:Power},b::T,expt::V) where {T<:Integer, V<:Rational}
    mpow(b,expt)
end

do_Power(mx,b,e) = mx

### // (Rational)

apprules(mx::Mxpr{://}) = makerat(mx,mx[1],mx[2])
makerat(mx::Mxpr{://},n::T,d::T) where {T<:Number} = n//d
makerat(mx,n,d) = mx

### BI

@sjdoc BI """
    BI(n)

convert the number `n` to type `BigInt`. Symata neither
detects integer overflow, nor automatically promote integers to `BigInt`.
But, a literal integer will automatically be given a large enough storage type without using
`BI`.
"""

@sjseealso_group(BI,BF,Big,BigIntInput)

### BF

@sjdoc BF """
    BF(n), or BF"n"

convert the number, or string `n` to a BigFloat. Symata neither
detects overflow, nor automatically promotes types from fixed to arbitrary precision. The form
`BF"n"` is more efficient, being a julia macro that converts the string "n" upon parsing.
"""

@sjdoc Big """
    Big(n)

convert `n` a maximum precision representation, typically
`BigInt` or `BigFloat`.
"""

apprules(mx::Mxpr{:BI}) = dobigint(mx,mx[1])
dobigint(mx,x) = mx
dobigint(mx,x::T) where {T<:Number} = BigInt(x)
dobigint(mx,x::T) where {T<:AbstractString} = Meta.parse(BigInt,x)

apprules(mx::Mxpr{:BF}) = dobigfloat(mx,mx[1])
dobigfloat(mx,x) = mx
dobigfloat(mx,x::T) where {T<:Number} = BigFloat(x)
dobigfloat(mx,x::T) where {T<:AbstractString} = Meta.parse(BigFloat,x)

### Big

apprules(mx::Mxpr{:Big}) = do_Big(mx,mx[1])
do_Big(mx,x) = mx
do_Big(mx,x::T) where {T<:Number} = big(x)

### Minus

apprules(mx::Mxpr{:Minus}) = is_Number(mx[1]) ? -mx[1] : mmul(-1, mx[1])