using PyPlot
include("../src/kepler_solver.jl")

# Define a constant of 1/3:
const third = 1.0/3.0

function test_elliptic()
# This routine runs a test of the kep_elliptic function in kepler_solver_elliptic.jl
# Define the central force constant in terms of AU and days:
k = (2.0*pi/365.25)^2
# Initial position at 1 AU:
x0 = [0.0,1.0,0.0]
#r0 = norm(x0)
r0 = sqrt(x0[1]*x0[1]+x0[2]*x0[2]+x0[3]*x0[3])
# Circular velocity:
vcirc = sqrt(k/r0)
# Define initial velocity at apastron:
v0 = [0.98*vcirc,0.0,0.0]  # The eccentricity is about ~2(1-v0/vcirc).
h = 18.0 # 18-day timesteps

s0::Float64 = 0.0
x = zeros(Float64,3)
v = zeros(Float64,3)
s::Float64 = 0.0
t = 0.0
ssave = zeros(2)
#nsteps = 10000000
nsteps = 1000000
#nsteps = 1000
xsave=zeros(Float64,12,nsteps)  # First is time (1,:); next three are position (2-4,:); next three velocity (5-7,:); then r (8,:), drdt (9,:); then beta (10,:); finally s & ds (11:12,:)
# Create a variable to store the state at the end of a step:
state=zeros(Float64,12)

xsave[1,1]=0.0
for j=1:3
  xsave[j+1,1]=x0[j]
  xsave[j+4,1]=v0[j]
end
# Save beta:
xsave[8,1]=r0
beta0 = 2.0*k/r0-dot(v0,v0)
xsave[10,1]=beta0
#@inbounds for i=2:nsteps
for i=2:nsteps
#  x,v,r,drdt,s,beta,iter=kep_elliptic(x0,v0,r0,k,h,beta0,s0)
#  iter = kep_elliptic!(x0,v0,r0,k,h,beta0,s0,state)
  iter = kep_ell_hyp!(x0,v0,r0,k,h,beta0,s0,state)
  s = state[11]
  ds = state[12]
  if iter > 2
   println("iter: ",iter," ds: ",ds)
  end
# Increment time by h:
  state[1] += h
  for j=1:12
    xsave[j,i]=state[j]
  end
  if i >= 4
# Quadratic prediction for s of next step:
    s0 = 3.0*s - 3.0*ssave[1] + ssave[2]
    ssave[2] = ssave[1]
    ssave[1] = s
  else
    s0 = s
    ssave[2]=ssave[1]
    ssave[1]=s
  end
# Now, proceed to next step:
  for j=1:3
    x0[j]=state[1+j]
    v0[j]=state[4+j]
  end
  r0=state[8]
  beta0=state[10]
#  println(i,x,v)
end
#plot(xsave[1],xsave[2])
return xsave
end


# Simulation
# Assuming all bodies connected in a tree by joints

# include("RigidBody.jl")
# include("OrientationConversion.jl")
# include("Joint.jl")
# include("Force.jl")
# include("extF.jl")

using LinearAlgebra
using Revise
using DifferentialEquations
using StaticArrays
using BenchmarkTools

global GravityInInertial = [0.0,0.0,-9.806];

struct solSim
    # Structure to store results of Simulation for each rigid body
    r::Matrix{Float64}# where T<:Real
    v::Matrix{Float64}# where T<:Real
    β::Matrix{Float64}# where T<:Real
    βdot::Matrix{Float64}# where T<:Real
end

"""
simulate(tEnd,tSpan,j;g,extFVec)

tEnd: Simulation End time. (Float64) \\
tSpan: Result saved at every specified tSpan. (Float64) \\
j: Tuple of Joints. Length need not be specified. \\
g: Acceleration due to gravity. Default value: [0.0;0.0;-9.806] (Vector{Float64}) \\
extFVec: Vector of extForces acting on each rigid body. (Vector{extForces}) \\

Tolerance of ode solver set. RelTol = 1e-10, AbsTol = 1e-10. \\
Tsit5() solver used. \\
Returns a tuple containing tSim and vector of solutions in solSim form.
"""
function simulate(tEnd::Float64,tInt::Float64,j::Vector{Joint};
                  g::Vector{Float64}=[0.0;0.0;-9.806]);# where T<:Real#,extFVec::Vector{extForces}=Vector{extForces}(undef,1))
    # Ellipsis (...) to facilitate supplying variable number of arguments to the function
    # Initial condition (all bodies connected)
    global GravityInInertial = g
    # global extFList = extFVec
    X0 = Vector{Float64}(j[1].RB1.x)
    for k = 1:length(j)
        append!(X0,j[k].RB2.x)
    end
    # println("X0 type = ", typeof(X0))
    # println("tEnd type = ", typeof(tEnd))
    # Declaring the ODE Problem as per DifferentialEquations convention
    tSpanODE = (0.0, tEnd);

    prob = ODEProblem(mainDynODEDiff!,X0,tSpanODE,j)
    # sol = solve(prob, saveat = tSpan, RK4(), reltol=1e-10, abstol=1e-10)
    # sol = solve(prob, Tsit5(), reltol=1e-10, abstol=1e-10)#, saveat = 0:tSpan:tEnd)
    sol = solve(prob, Tsit5(), saveat=tInt, reltol=1e-10, abstol=1e-10)
    # sol = solve(prob, DP5(), reltol=1e-10, abstol=1e-10)

    solFinal = Vector{solSim}(undef,length(j))
    for i=1:length(j)
        rSol = Matrix{Float64}(undef,(length(sol.t),3))
        vSol = Matrix{Float64}(undef,(length(sol.t),3))
        βSol = Matrix{Float64}(undef,(length(sol.t),4))
        βdotSol = Matrix{Float64}(undef,(length(sol.t),4))
        for t=1:length(sol.t)
            rSol[t,:] = (sol.u[t][14*i+1:14*i+3])
            vSol[t,:] = (sol.u[t][14*i+8:14*i+10])
            βSol[t,:] = (sol.u[t][14*i+4:14*i+7])
            βdotSol[t,:] = (sol.u[t][14*i+11:14*i+14])
        end
    solFinal[i] = solSim(rSol,vSol,βSol,βdotSol)
    end
    return (sol.t, solFinal)
end

function mainDynODEDiff!(dX,X, j::Vector{Joint}, t::Float64)# where T<:Real
    dX[:] = mainDynODEDiff(X,j,t);
end

function mainDynODEDiff(X::Vector{T}, j::Vector{Joint}, t::Float64)::Vector{T} where T<:Real
    # ODE function to be used as per DifferentialEquations convention
    # Create extForcesList storing extForces for each rigid body
    # Create ForceConstraints Array storing constraint forces acting on each rigid body
    global GravityInInertial
    # Update RigidBodies
    updateRigidBody!(j[1].RB1,X[1:14]) # update inertial frame (not moving)
    for k=1:length(j)
        updateRigidBody!(j[k].RB2,X[14*k+1:14*(k+1)])
    end

    extFListCopy = externalFTotal(t,j); # Generates all the external forces explicitly specified, through joint actions or directly.

    U = genU(extFListCopy) # Construct an input matrix from external forces and torques supplied.

    ## Inplace
    dX = zeros(length(X));
    fxdot!(dX, X, U,j,GravityInInertial);

    ## Not inplace
    # dX = fxdot(X,U,j,GravityInInertial);

    # println("t = $t")
    return dX
end

function fxdot!(dX, X::Vector{T}, U::Matrix{S}, j::Vector{Joint}, GravityInInertial::Vector{Float64}) where {T<:Real,S<:Real}
    dX[:] = fxdot(X,U,j,GravityInInertial);
end

function fxdot(X::Vector{T},U::Matrix{S},j::Vector{Joint},GravityInInertial::Vector{Float64}) where {T<:Real,S<:Real}
    nJ = length(j); # Number of joints

    unconstrF, constrF = Constraint(X, U, j, GravityInInertial)

    xdot = Vector{Union{T,S}}(undef,14*(nJ+1));
    xdot[1:14] = zeros(14);

    for k=1:nJ
        xb = X[14*k+1:14*(k+1)]

        unconstrainedF_rb = unconstrF[:,j[k].RB2.bodyID]
        Fc_rb = constrF[:,j[k].RB2.bodyID]
        TotalForce = unconstrainedF_rb[1:3] + Fc_rb[1:3]
        TotalMoment = unconstrainedF_rb[4:7] + Fc_rb[4:7]
        genF = [TotalForce; TotalMoment]; # generalized force acting on RB
        genM = Array(genMatM(X, j[k].RB2)) # generalized mass matrix of RB
        xdot[14*k+1:14*k+7] = X[14*k+8:14*k+14]#xb[8:10]
        xdot[14*k+8:14*k+14] = inv(genM)*genF;
    end

    return xdot
end

function externalFTotal(t::T, j::Vector{Joint})::Vector{extForces} where T<:Real
    # Returns externally applied forces in total, not including gravity. Includes Forces from joints and other explicitly specified external forces.
    extFList = extF(t,j)

    # Generate forces from actuated joints on each body
    # ForceJoints = Matrix{Float64}(undef,6,2*length(j))
    # for k=1:length(j)
    #     ForceJoints[:,2*k-1], ForceJoints[:,2*k] = genJointF(t,j[k])
    # end

    # Add ForceJoints to extFList
    extFListCopy = deepcopy(extFList)
    # for k=1:length(j)
    #     extFListCopy[j[k].RB1.bodyID].Forces =
    #     vcat(extFListCopy[j[k].RB1.bodyID].Forces, reshape(ForceJoints[1:3,2*k-1],(1,3)))
    #     extFListCopy[j[k].RB1.bodyID].Positions =
    #     vcat(extFListCopy[j[k].RB1.bodyID].Positions, reshape(j[k].pos1,(1,3)))
    #     extFListCopy[j[k].RB1.bodyID].Torques =
    #     vcat(extFListCopy[j[k].RB1.bodyID].Torques, reshape(ForceJoints[4:6,2*k-1],(1,3)))
    #
    #     extFListCopy[j[k].RB2.bodyID].Forces =
    #     vcat(extFListCopy[j[k].RB2.bodyID].Forces, reshape(ForceJoints[1:3,2*k],(1,3)))
    #     extFListCopy[j[k].RB2.bodyID].Positions =
    #     vcat(extFListCopy[j[k].RB2.bodyID].Positions, reshape(j[k].pos2,(1,3)))
    #     extFListCopy[j[k].RB2.bodyID].Torques =
    #     vcat(extFListCopy[j[k].RB2.bodyID].Torques, reshape(ForceJoints[4:6,2*k],(1,3)))
    # end
    return extFListCopy
end

function genU(extFList::Vector{extForces})::Matrix{Float64}
    nB = length(extFList); # number of bodies
    u = zeros(6,nB); # Each column in u contains a vector of [F;τ] acting on the corresponding rigid body
    for b=1:nB
        extF_b = extFList[b]
        u[1:3,b] = sum(extF_b.Forces,dims = 1) # Sum up forces acting on body
        for n in 1:size(extF_b.Forces,1) # Sum up moments acting on body
            u[4:6,b] += cross(extF_b.Positions[n,:], extF_b.Forces[n,:])
        end
        u[4:6,b] += sum(extF_b.Torques, dims = 1)[:] # Sum up torques acting on body
    end
    return u
end

function getXU_0(j::Vector{Joint})
    X0 = Vector{Float64}(j[1].RB1.x)
    for k = 1:length(j)
        append!(X0,j[k].RB2.x)
    end

    extFListCopy = externalFTotal(0.0,j); # Generates all the external forces explicitly specified, through joint actions or directly.

    U0 = genU(extFListCopy) # Construct an input matrix from external forces and torques supplied.

    return X0, U0
end

function checkRevJoint(solQuad::solSim, solProp1::solSim, rjCube1::Array{Float64}, rjProp1::Array{Float64})
    ## Function that outputs relevant quantities constrained under a rotational joint
    tLen = size(solQuad.r,1)
    ωCube = Matrix{Float64}(undef,tLen,3);
    ωProp1 = Matrix{Float64}(undef,tLen,3);
    ωProp1InCube = Matrix{Float64}(undef,tLen,3);
    ωRel = Matrix{Float64}(undef,tLen,3);
    jointLoc1 = Matrix{Float64}(undef,length(tSim),3);
    for i=1:tLen
        dcm = quat2dcm(solQuad.β[i,:])
        dcm2 = quat2dcm(solProp1.β[i,:])
        ωCube[i,:] = angVel(solQuad.β[i,:],solQuad.βdot[i,:])
        ωProp1[i,:] = angVel(solProp1.β[i,:],solProp1.βdot[i,:])
        ωProp1InCube[i,:] = quat2dcm(solQuad.β[i,:])*transpose(quat2dcm(solProp1.β[i,:]))*ωProp1[i,:]
        ωRel[i,:] = ωProp1InCube[i,:] - ωCube[i,:]
        jointLoc1[i,:] = solQuad.r[i,:] + transpose(dcm)*rjCube1 - solProp1.r[i,:] - transpose(dcm2)*rjProp1;
        # jointLoc1[i,:] = solProp1.r[i,:] + transpose(dcm2)*rjProp1;
    end
    return ωCube, ωProp1, ωRel, jointLoc1
end

function checkRevJointIn(solPend1::solSim, rj2::Array{Float64})
    ## Function that outputs relevant quantities constrained under a rotational joint when first body is inertial frame
    tLen = size(solPend1.r,1)
    jointLoc1 = Matrix{Float64}(undef,tLen,3)
    ωSol = Matrix{Float64}(undef,tLen,3)
    for i=1:tLen
        tmp_dcm = quat2dcm(solPend1.β[i,:])
        ωSol[i,:] = angVel(solPend1.β[i,:],solPend1.βdot[i,:])
        jointLoc1[i,:] = solPend1.r[i,:] + transpose(tmp_dcm)*rj2
    end
    return ωSol, jointLoc1
end

function joinBodies!(j1::Vector{Joint}, j2::Vector{Joint}, b1_id::Integer, b2_id::Integer, rj1::Vector{T}, rj2::Vector{T}; type::String="Weld", axis::Vector{T}=[0.0;0.0;1.0], k::T=0.0, rL::T=0.0, jointForce::Vector{T} = zeros(T,3), jointTorque::Vector{T} = zeros(T,3))  where T<:Real
    # join the bodies b1 given by b1_id in j1 and b2_id in j2 at rj1 and rj2 respectively
    # Create new joint as a weld joint, unless specified
    len_j1 = length(j1); # old length of j1
    RB1 = body_in_jointVec(j1,b1_id);
    RB2 = body_in_jointVec(j2,b2_id);

    newJ = Joint(RB1, RB2, rj1, rj2, type = type, axis = axis, k = k, rL = rL, jointForce = jointForce, jointTorque = jointTorque);
    push!(j1,newJ)

     # Update bodyIDs of all bodies present in joint tree j2
    for i=1:length(j2)
        j2[i].RB2.bodyID += len_j1;
        push!(j1,j2[i])
    end
    filterInertialJoint!(j1); # This removes the connection of the body in joint tree j2 that was connected to the inertial frame (removes loop inside joint trees)
end


# algorithm type
struct ButuceaComte end

"""
    ButuceaComteOptions(; bandwidth = :auto)

The Butucea Comte estimator of linear functional in the deconvolution model. Given a `LinearEBayesTarget`
with characteristic function ``\\psi^*`` and samples convolved with Gaussian noise (with ``\\phi`` Standard Gaussian pdf),
it is defined as follows

```math
\\hat{L}_{\\text{BC},h_n} = \\frac{1}{2 \\pi n}\\sum_{i=1}^n \\int_{-1/h_n}^{1/h_n} \\exp(it Z_k) \\frac{\\psi^*(-t)}{\\varphi^*(t)}dt
```
``h`` is the bandwidth, the option `:auto` will pick it automatically.

## Reference:
  >Butucea, C. and Comte, F., 2009.
  >Adaptive estimation of linear functionals in the convolution model and applications.
  >Bernoulli, 15(1), pp.69-98.
"""
Base.@kwdef struct ButuceaComteOptions{S}
    bandwidth::S = :auto
end

fit(bcopt::ButuceaComteOptions, Zs) = bcopt #do nothing

struct ButuceaComteEstimator{EBT<:LinearEBayesTarget}
    target::EBT
    h::Float64
    f_real #ApproxFun based
    f_imag #ApproxFun based
end

# Make theoretical results of BC practical.
Base.@kwdef struct ButuceaComteNoiseDensityParams{T <: Real}
    γ::T
    α::T
    ρ::T
    κ0upper::T
    κ0lower::T
end

function ButuceaComteNoiseDensityParams(Z::StandardNormalSample)
    ButuceaComteNoiseDensityParams(γ = 0.0, ρ = 2.0, α = 0.5, κ0upper=1.0, κ0lower=1.0)
end

Base.@kwdef struct ButuceaComteSobolevParams{T<:Real}
   b::T
   L::T        # constant
   a::T = 0.0
   r::T = 0.0
end

Base.@kwdef struct ButuceaComteTargetParams{T<:Real}
   B::T        #polynomial decay (1+x^2)^B
   A::T
   R::T
   Cψ::T       #constant
end

function ButuceaComteTargetParams(target::PriorDensityTarget)
    ButuceaComteTargetParams(A=0.0, B=0.0, R=0.0, Cψ=1.0)
end

# only accounting for b,B for now
function butucea_comte_squared_bias_bound(noise::ButuceaComteNoiseDensityParams,
                                  sobolev::ButuceaComteSobolevParams,
                                  target::ButuceaComteTargetParams,
                                  h::Real)
    R = target.R
    r = sobolev.r

    a = target
    B = target.B
    b = sobolev.b

    L = sobolev.L
    Cψ = target.Cψ

    L*Cψ/π/(2b + 2B -1)*h^(1-2b-2B)
end

# only do case PriorDensityTarget for now
function butucea_comte_unit_var_proxy(h::Real)
    1/π^2*exp(h^2/2)/h^2
end

function default_bandwidth(::Type{ButuceaComteEstimator},
                           target::MarginalDensityTarget{<:StandardNormalSample},
                           n)
    sqrt(log(n))
end

default_bandwidth(::Type{ButuceaComteEstimator}, target::LFSRNumerator, n) = sqrt(log(n)*2/3)

default_bandwidth(::Type{ButuceaComteEstimator}, target::PosteriorMeanNumerator, n) = sqrt(log(n))

function default_bandwidth(::Type{ButuceaComteEstimator},
                           target::PriorDensityTarget,
                           n)
    # should make this also depend on noise distribution etc ...
    # π*m_n  = (log(n)/(2α + 1))^{1/ρ}
    #gaussian case α=1//2, ρ=2
    sqrt(log(n)/2)
end

function ButuceaComteEstimator(target::EBayesTarget; n::Integer)
    h = default_bandwidth(ButuceaComteEstimator, target, n)
    ButuceaComteEstimator(target, h)
end

function ButuceaComteEstimator(target::EBayesTarget, h::Float64)
    #todo: remove hardcoding of Normal
    _f(t) = cf(target,-t)/cf(Normal(),t)
    _set = Interval(-h,h)
    _f_real = Fun(real ∘ _f, _set)
    _f_imag = Fun(imag ∘ _f, _set)
    ButuceaComteEstimator(target, h, _f_real, _f_imag)
end

function (cb::ButuceaComteEstimator)(z)
    real_part = fourier(cb.f_real, z)
    imag_part = fourier(cb.f_imag, z)

    real(real_part + im*imag_part)/2/π
    #h = cb.h
    ##target = cb.target
    #f(t) = real(exp(im*t*z)*cf(target,-t)/cf(Normal(),t))
    #res, _ = quadgk(f, -h, +h) res/2π
end

function DiscretizedAffineEstimator(mhist, cb::ButuceaComteEstimator)
    DiscretizedAffineEstimator(mhist, z -> cb(z))
end

"""
    estimate(target::LinearEBayesTarget, bcopt::ButuceaComteOptions, Zs)

Estimate `target` based on samples `Zs` with the Butucea-Comte estimator with
[`ButuceaComteOptions`](@ref) bcopt.
"""
function estimate(target::LinearEBayesTarget, bcopt::ButuceaComteOptions,
                  Zs::AbstractVector{<:StandardNormalSample})
    (bcopt.bandwidth == :auto) || error("only :auto bandwidth for now")
    bcest = ButuceaComteEstimator(target;  n=nobs(Zs))
    mean(bcest.(response.(Zs)))
end

#function estimate(Xs,cb::ComteButucea)
#    mean(cb.Q.(Xs))
#end

#struct DiscretizedAffineEstimator{MH<:MCEBHistogram}
#    mhist::MH
#    Q::Vector{Float64}
#    Qo::Float64 #offset
#end

#default_bandwidth(::Type{ComteButucea}, n) = sqrt(log(n))
#default_bandwidth(::Type{ComteButucea}, target, n) = default_bandwidth(ComteButucea, n)


# Reversible jump MCMC for axis-aligned multivariate normal mixtures, using a conjugate prior.
module MVNaaRJ

include("RandomNumbers.jl")
using .RandomNumbers

using Statistics
using SpecialFunctions
lgamma_(x) = logabsgamma(x)[1]
lbeta_(a,b) = (logabsbeta(a,b))[1]

const Data = Array{Float64,1}
mutable struct Theta
    mu::Array{Float64,1}     # means
    lambda::Array{Float64,1} # precisions
end

# Include generic code
include("generic.jl")

logsumexp(a,b) = (m = max(a,b); m == -Inf ? -Inf : log(exp(a-m) + exp(b-m)) + m)

lognormpdf(x,m,v) = -0.5*(x-m)*(x-m)/v - 0.5*log(2*pi*v)
logmvnpdf(x,m,v,c) = (l=0.0; for j=1:length(x); l += lognormpdf(x[j],m[c,j],v[c,j]); end; l)

mutable struct Hyperparameters
    log_pk_vec::Array{Float64,1} # log of prior on the number of components
    gamma::Float64 # Dirichlet parameter
    d::Int64    # dimension
    # m[c,i]|v[c,i] ~ N(m,v[c,i]/r)
    m::Float64  # prior mean of m's
    r::Float64  # factor for prior precision of m's
    # v[c,i] ~ InvGamma(a,b)
    a::Float64  # prior shape of v's
    b::Float64  # prior scale of v's
end

function construct_hyperparameters(options)
    x = options.x
    n = length(x)
    d = length(x[1])
    mu = mean(x)
    v = mean([xi.*xi for xi in x]) - mu.*mu  # sample variance
    @assert(all(abs.(mu) .< 1e-10) && all(abs.(v .- 1.0) .< 1e-10), "Data must be normalized to zero mean, unit variance.")
    m = 0.0
    r = 1.0
    a = 1.0
    b = 1.0
    gamma = options.gamma
    lpk = eval(Meta.parse(options.log_pk))
    log_pk_fn(k) = Base.invokelatest(lpk,k)
    log_pk_vec = Float64[log_pk_fn(k) for k = 1:options.k_max]
    return Hyperparameters(log_pk_vec,gamma,d,m,r,a,b)
end

mutable struct List
    first::Int64
    next::Array{Int64,1}
    List(maxlen) = (l=new(); l.first=0; l.next=zeros(Int64,maxlen); return l)
end
insert!(l::List,i) = (l.next[i] = l.first; l.first = i) # Note: i must not already be in the list.
pop!(l::List) = (i = l.first; @assert(i!=0,"list capacity exceeded"); l.first = l.next[i]; i)
function remove!(l::List,i) # Note: i must already be in the list.
    if i==l.first
        l.first = l.next[i]
    else
        j = l.first
        while j != 0
            if i==l.next[j]; l.next[j] = l.next[i]; break; end
            j = l.next[j]
        end
    end
end
#pr(l::List) = (i=l.first; while i!=0; print(i," "); i=l.next[i]; end; println())

mutable struct IndexList
    first::Array{Int64,1}
    next::Array{Int64,1}
    IndexList(kmax,n) = (l=new(); l.first=zeros(Int64,kmax); l.next=zeros(Int64,n); return l)
end
insert!(l::IndexList,c,i) = (l.next[i] = l.first[c]; l.first[c] = i)
clear!(l::IndexList,c) = (l.first[c] = 0)

# sample from categorical distn p on list
function randp(p,list::List)
    u = rand()
    i = list.first
    s = p[i]
    while u > s
        i = list.next[i]
        s += p[i]
    end
    return i
end

# sample from uniform distn on list of k items
function randunif(k,list::List)
    u = rand()*k
    i = list.first
    s = 1.0
    while u > s
        i = list.next[i]
        s += 1.0
    end
    return i
end

# update assignments z
function update_assignments!(x,z,w,m,v,clist,log_p,p)
    for i = 1:length(x)
        c,log_s = clist.first,-Inf
        while c != 0
            log_p[c] = log(w[c]) + logmvnpdf(x[i],m,v,c)
            log_s = logsumexp(log_s,log_p[c])
            c = clist.next[c]
        end
        c = clist.first
        while c != 0
            p[c] = exp(log_p[c] - log_s)
            c = clist.next[c]
        end
        z[i] = randp(p,clist)
    end
end

# recompute t,s,sx,ssx,ilist
function recompute_statistics!(x,z,s,sx,ssx,ilist,clist,H)
    c = clist.first
    while c != 0
        s[c] = 0
        clear!(ilist,c)
        for j = 1:H.d; sx[c,j] = 0.0; ssx[c,j] = 0.0; end
        c = clist.next[c]
    end
    t = 0
    for i = 1:length(x)
        c = z[i]
        if s[c]==0; t+=1; end
        insert!(ilist,c,i)
        s[c] += 1
        for j = 1:H.d
            sx[c,j] += x[i][j]
            ssx[c,j] += x[i][j]*x[i][j]
        end
    end
    return t
end

# update weights, means, and variances
function update_parameters!(w,m,v,s,sx,ssx,clist,H)
    # update weights (unnormalized for now)
    c = clist.first
    w_sum = 0.0
    while c != 0
        w[c] = RandomNumbers.gamma(H.gamma+s[c],1)
        w_sum += w[c]
        c = clist.next[c]
    end

    c = clist.first
    while c != 0
        A = H.a + 0.5*(s[c]+1)
        for j = 1:H.d
            # update mean
            L = (H.r + s[c])/v[c,j]
            M = (H.r*H.m + sx[c,j]) / (H.r + s[c])
            m[c,j] = randn()/sqrt(L) + M

            # update variance
            B = H.b + 0.5*(ssx[c,j] - 2*m[c,j]*sx[c,j] + s[c]*m[c,j]*m[c,j]) + 0.5*H.r*(m[c,j]-H.m)*(m[c,j]-H.m)
            # if (A<=0 || B<=0); println(A," ",B)
                # println(H.b)
                # println(0.5*(ssx[c,j] - 2*m[c,j]*sx[c,j] + s[c]*m[c,j]*m[c,j]))
                # println(0.5*H.r*(m[c,j]-H.m)*(m[c,j]-H.m))
            # end
            v[c,j] = RandomNumbers.inverse_gamma(A,B)
        end

        # normalize weights
        w[c] = w[c]/w_sum
        c = clist.next[c]
    end
end

function log_prior_ratio(k,w,m,v,cm,c1,c2,s1,s2,H)
    ssm,slv,siv = 0.0,0.0,0.0
    for j = 1:H.d
        ssm += (m[c1,j]-H.m)^2/v[c1,j] + (m[c2,j]-H.m)^2/v[c2,j] - (m[cm,j]-H.m)^2/v[cm,j]
        slv += log(v[c1,j]*v[c2,j]/v[cm,j])
        siv += 1/v[c1,j] + 1/v[c2,j] - 1/v[cm,j]
    end
    log_rprior = (H.log_pk_vec[k+1] - H.log_pk_vec[k] # prior on k
               + (H.gamma-1)*log(w[c1]*w[c2]/w[cm]) - lbeta_(H.gamma,k*H.gamma) # prior on w
               + s1*log(w[c1]) + s2*log(w[c2]) - (s1+s2)*log(w[cm]) # prior on z
               + 0.5*H.d*log(H.r/(2*pi)) - 0.5*H.r*ssm - 0.5*slv # prior on m
               + H.d*H.a*log(H.b) - H.d*lgamma_(H.a) + (-H.a-1)*slv - H.b*siv) # prior on v
    return log_rprior
end

function log_Jacobian(w,v,u1,u2,cm,H)
    log_J = log(w[cm]) - 1.5*H.d*log(u1*(1-u1))
    for j = 1:H.d; log_J += log(1-u2[j]*u2[j]) + 1.5*log(v[cm,j]); end
    return log_J
end

# propose split-merge move
function split_merge!(k,t,x,z,w,m,v,s,sx,ssx,clist,cfree,ilist, zs,u2,u3,H)
    if ((k==1) || (rand()<0.5)) # propose split
        # choose component indices
        cm = randunif(k,clist)
        c1 = pop!(cfree)
        c2 = pop!(cfree)
        # sample u's
        log_pu = 0.0
        u1 = RandomNumbers.beta(2,2); log_pu += log(6*u1*(1-u1))
        for j = 1:H.d
            u2[j] = RandomNumbers.beta(2,2); log_pu += log(6*u2[j]*(1-u2[j]))
            u3[j] = rand(); # log_pu += 0.0
        end
        # split w, m, and v
        w[c1],w[c2] = w[cm]*u1, w[cm]*(1-u1)
        for j = 1:H.d
            m[c1,j] = m[cm,j] - u2[j]*sqrt(v[cm,j]*w[c2]/w[c1])
            m[c2,j] = m[cm,j] + u2[j]*sqrt(v[cm,j]*w[c1]/w[c2])
            v[c1,j] = u3[j]*(1-u2[j]*u2[j])*v[cm,j]*w[cm]/w[c1]
            v[c2,j] = (1-u3[j])*(1-u2[j]*u2[j])*v[cm,j]*w[cm]/w[c2]
        end
        # split z and simultaneously compute the log likelihood ratio
        log_pzs = 0.0
        log_rlik = 0.0
        s1,s2 = 0,0
        i = ilist.first[cm]
        while i != 0
            l1,l2,lm = logmvnpdf(x[i],m,v,c1),logmvnpdf(x[i],m,v,c2),logmvnpdf(x[i],m,v,cm)
            p1 = 1/(1 + exp(l2-l1)*w[c2]/w[c1])
            if rand() < p1
                zs[i] = c1
                s1 += 1
                log_pzs += log(p1)
                log_rlik += l1 - lm
            else
                zs[i] = c2
                s2 += 1
                log_pzs += log(1-p1)
                log_rlik += l2 - lm
            end
            i = ilist.next[i]
        end
        # log of the ratio of move probabilities (choice of split or merge)
        log_rmove = (k==1 ? log(0.5) : 0.0)

        # log prior ratio
        log_rprior = log_prior_ratio(k,w,m,v,cm,c1,c2,s1,s2,H)

        # log of the Jacobian
        log_J = log_Jacobian(w,v,u1,u2,cm,H)

        # acceptance probability
        p_accept = min(1,exp(log_rlik + log_rprior + log_rmove - log_pu - log_pzs + log_J))

        # accept or reject
        if rand() < p_accept
            # k,t,z,s,w,m,v,sx,ssx,clist,cfree,ilist
            k += 1
            t += Int((s1>0) && (s2>0))
            s[c1],s[c2] = s1,s2
            for j=1:H.d; sx[c1,j]=0.0; sx[c2,j]=0.0; ssx[c1,j]=0.0; ssx[c2,j]=0.0; end
            clear!(ilist,c1); clear!(ilist,c2)
            i = ilist.first[cm]
            while i != 0
                c = zs[i]
                z[i] = c
                for j = 1:H.d
                    sx[c,j] += x[i][j]
                    ssx[c,j] += x[i][j]*x[i][j]
                end
                next = ilist.next[i]
                insert!(ilist,c,i)
                i = next
            end
            insert!(clist,c1); insert!(clist,c2)
            remove!(clist,cm); insert!(cfree,cm)
        else
            insert!(cfree,c1); insert!(cfree,c2)
        end
    else # propose merge
        # choose component indices
        c1 = randunif(k,clist); remove!(clist,c1)
        c2 = randunif(k-1,clist); insert!(clist,c1)
        cm = pop!(cfree)
        # merge w, m, and v
        w[cm] = w[c1] + w[c2]
        for j = 1:H.d
            m[cm,j] = (w[c1]*m[c1,j] + w[c2]*m[c2,j])/w[cm]
            v[cm,j] = (w[c1]*(m[c1,j]^2 + v[c1,j]) + w[c2]*(m[c2,j]^2 + v[c2,j]))/w[cm] - m[cm,j]^2
        end
        # compute u's
        log_pu = 0.0
        u1 = w[c1]/w[cm]; log_pu += log(6*u1*(1-u1))
        for j = 1:H.d
            u2[j] = (m[cm,j] - m[c1,j])/sqrt(v[cm,j]*w[c2]/w[c1]); log_pu += log(6*u2[j]*(1-u2[j])*(0<=u2[j]<=1))
            u3[j] = (v[c1,j]*w[c1])/(v[cm,j]*w[cm]*(1-u2[j]*u2[j])); log_pu += (0<=u3[j]<=1 ? 0.0 : -Inf)
        end
        # merge z and compute the log likelihood ratio
        log_pzs = 0.0
        log_rlik = 0.0
        i = ilist.first[c1]
        while i != 0
            l1,l2,lm = logmvnpdf(x[i],m,v,c1),logmvnpdf(x[i],m,v,c2),logmvnpdf(x[i],m,v,cm)
            p1 = 1/(1 + exp(l2-l1)*w[c2]/w[c1])
            log_pzs += log(p1)
            log_rlik += lm - l1
            i = ilist.next[i]
        end
        i = ilist.first[c2]
        while i != 0
            l1,l2,lm = logmvnpdf(x[i],m,v,c1),logmvnpdf(x[i],m,v,c2),logmvnpdf(x[i],m,v,cm)
            p1 = 1/(1 + exp(l2-l1)*w[c2]/w[c1])
            log_pzs += log(1-p1)
            log_rlik += lm - l2
            i = ilist.next[i]
        end
        # log of the ratio of move probabilities (choice of split or merge)
        log_rmove = 0.0

        # log prior ratio
        log_rprior = -log_prior_ratio(k-1,w,m,v,cm,c1,c2,s[c1],s[c2],H)

        # log of the Jacobian
        log_J = -log_Jacobian(w,v,u1,u2,cm,H)

        # acceptance probability
        p_accept = min(1,exp(log_rlik + log_rprior + log_rmove + log_pu + log_pzs + log_J))

        # accept or reject
        if rand() < p_accept
            # k,t,z,w,m,v,s,sx,ssx,clist,cfree,ilist
            k -= 1
            t -= Int((s[c1]>0) && (s[c2]>0))
            s[cm] = s[c1] + s[c2]
            for j = 1:H.d
                sx[cm,j] = sx[c1,j] + sx[c2,j]
                ssx[cm,j] = ssx[c1,j] + ssx[c2,j]
            end
            clear!(ilist,cm)
            for i in [ilist.first[c1],ilist.first[c2]]
                while i != 0
                    z[i] = cm
                    next = ilist.next[i]
                    insert!(ilist,cm,i)
                    i = next
                end
            end
            insert!(clist,cm)
            remove!(clist,c1); remove!(clist,c2)
            insert!(cfree,c1); insert!(cfree,c2)
        else
            insert!(cfree,cm)
        end
    end
    return k,t
end

function birth_death!(k,t,n,w,m,v,s,sx,ssx,clist,cfree,ilist,H)
    k0 = k-t # number of empty components
    if rand() < 0.5 # propose birth of an empty component
        c = pop!(cfree)
        w[c] = RandomNumbers.beta(1,k)
        for j = 1:H.d
            v[c,j] = RandomNumbers.inverse_gamma(H.a,H.b)
            m[c,j] = randn()*sqrt(v[c,j]/H.r) + H.m
        end
        log_A = (H.log_pk_vec[k+1] - H.log_pk_vec[k]
              + (n+k*(H.gamma-1))*log(1-w[c]) + (H.gamma-1)*log(w[c])
              + lbeta_(1,k) - lbeta_(H.gamma,k*H.gamma) + log(k+1) - log(k0+1))
        p_accept = min(1,exp(log_A))

        # accept or reject
        if rand() < p_accept
            # k,t,z,w,m,v,s,sx,ssx,clist,cfree,ilist
            k += 1
            s[c] = 0
            for j = 1:H.d; sx[c,j] = 0.0; ssx[c,j] = 0.0; end
            insert!(clist,c)
            clear!(ilist,c)
        else
            insert!(cfree,c)
        end
    elseif (k0 > 0) # propose death of an empty component
        # randomly choose an empty component
        u = rand()*k0
        c = clist.first
        S = Int(s[c]==0)
        while u > S
            c = clist.next[c]
            S += Int(s[c]==0)
        end
        # compute acceptance probability
        log_A = (H.log_pk_vec[k] - H.log_pk_vec[k-1]
              + (n+(k-1)*(H.gamma-1))*log(1-w[c]) + (H.gamma-1)*log(w[c])
              + lbeta_(1,k-1) - lbeta_(H.gamma,(k-1)*H.gamma) + log(k) - log(k0))
        p_accept = min(1,exp(-log_A))

        # accept or reject
        if rand() < p_accept
            # k,t,z,w,m,v,s,sx,ssx,clist,cfree,ilist
            k -= 1
            remove!(clist,c)
            insert!(cfree,c)
        end
    end
    return k,t
end

function sampler(options,n_total,n_keep)
    H = construct_hyperparameters(options)
    x,n = options.x,options.n
    kmax = min(n,options.k_max)
    @assert(n==length(x))
    d = H.d # dimension

    # State
    # k,t,z,s,w,m,v,sx,ssx,cfirst,cnext,cfree,ilist
    k = 1 # number of components
    t = 1 # number of clusters (non-empty components)
    z = zeros(Int64,n) # z[i] = component assignment of datapoint i
    w = zeros(kmax) # w[c] = mixture weight of component c
    m = zeros(kmax,d) # m[c,:] = means of component c
    v = zeros(kmax,d) # v[c,:] = variances of component c
    s = zeros(Int64,kmax) # s[c] = # of points assigned to component c
    sx = zeros(kmax,d) # sx[c,:] = sum of the x's in component c
    ssx = zeros(kmax,d) # ssx[c,:] = sum of the squares of the x's in component c
    clist = List(kmax) # component list
    cfree = List(kmax) # list of unused components
    ilist = IndexList(kmax,n) # lists of indices assigned to each component

    # initialize with a single component
    z[:] .= 1
    s[1] = n
    m[1,:] = mean(x) # initialize to the sample mean
    v[1,:] = mean([(xi-vec(m[1,:])).^2 for xi in x]) # initialize to the sample variance
    sx[1,:] = sum(x)
    ssx[1,:] = sum([xi.^2 for xi in x])
    insert!(clist,1)
    for c=kmax:-1:2; insert!(cfree,c); end
    for i = 1:n; insert!(ilist,1,i); end

    # temporary variables
    p = zeros(kmax) # temporary variable for randp
    log_p = zeros(kmax) # temporary variable
    zs = zeros(Int64,n) # temporary assignments for split-merge
    u2,u3 = zeros(d),zeros(d) # temporary variables for split-merge

    # iterations to record
    keepers = round.(Int,range(round(Int,n_total/n_keep),stop=n_total,length=n_keep))
    keep_index = 0

    # records
    #k_r = zeros(Int64,n_total)
    t_r = zeros(Int8,n_total); @assert(kmax < 2^7)
    N_r = zeros(Int16,kmax+3,n_total); @assert(n < 2^15)
    z_r = zeros(Int8,n,n_keep); @assert(kmax < 2^7)
    theta_r = Array{Theta}(undef,kmax+3,n_keep)

    for iteration = 1:n_total
        # update weights, means, and variances
        update_parameters!(w,m,v,s,sx,ssx,clist,H)

        # update assignments z
        update_assignments!(x,z,w,m,v,clist,log_p,p)
        # recompute t,s,sx,ssx,ilist based on the new z
        t = recompute_statistics!(x,z,s,sx,ssx,ilist,clist,H)

        # split/merge
        k,t = split_merge!(k,t,x,z,w,m,v,s,sx,ssx,clist,cfree,ilist, zs,u2,u3,H)

        # birth/death
        k,t = birth_death!(k,t,n,w,m,v,s,sx,ssx,clist,cfree,ilist,H)

        # record-keeping
        #k_r[iteration] = k
        t_r[iteration] = t
        for i = 1:n
            N_r[z[i],iteration] += 1
        end
        if iteration==keepers[keep_index+1]
            keep_index += 1
            for i = 1:n; z_r[i,keep_index] = z[i]; end
            c = clist.first
            while c != 0
                theta_r[c,keep_index] = Theta(m[c,:][:], 1.0./v[c,:][:])
                c = clist.next[c]
            end
        end
    end

    return t_r,N_r,z_r,theta_r,keepers
end

end # module


@testset "PrimeIdealsSet" begin
  Qx, x = FlintQQ["x"]
  K, a = NumberField(x - 1, "a")
  O = maximal_order(K)

  S = @inferred PrimeIdealsSet(O, 2, 100)
  @test @inferred length(collect(S)) == 25

  S = @inferred PrimeIdealsSet(O, 2, 100, indexdivisors = false, ramified = false, degreebound = 100, coprimeto = 3*5*7)
  @test @inferred length(collect(S)) == 22

  K, a = NumberField(x^2 - 5, "a")
  O = maximal_order(K)

  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100))
  @test @inferred length(collect(S)) == 35
  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), indexdivisors = false)
  @test @inferred length(collect(S)) == 34
  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), indexdivisors = false, ramified = false)
  @test @inferred length(collect(S)) == 33
  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), indexdivisors = false, ramified = false, degreebound = 1)
  @test @inferred length(collect(S)) == 20

  K, a = NumberField(x^5 - x + 1, "a")
  O = maximal_order(K)

  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100))
  @test @inferred length(collect(S)) == 51

  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), degreebound = 1)
  @test @inferred length(collect(S)) == 18
  SS = collect(S)

  S = @inferred PrimeIdealsSet(O, fmpz(2), -1, degreebound = 1)
  z = 1
  T = NfOrdIdl[]
  for P in S
    push!(T, P)
    if z == 18
      break
    end
    z += 1
  end
  @test T == SS

  P = prime_decomposition(O, 2)[1][1]
  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), coprimeto = P)
  @test @inferred length(collect(S)) == 50

  el = Hecke.find_elem_of_valuation_1(P, P^2)
  @test valuation(el, P) == 1

  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), coprimeto = 2)
  @test @inferred length(collect(S)) == 49

  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), coprimeto = fmpz(6))
  @test @inferred length(collect(S)) == 48

  S = @inferred PrimeIdealsSet(O, fmpz(2), fmpz(100), coprimeto = O(30))
  @test @inferred length(collect(S)) == 47

  @test_throws ErrorException PrimeIdealsSet(O, fmpz(-1), fmpz(1))
  @test_throws ErrorException PrimeIdealsSet(O, fmpz(1), -2)
  @test_throws ErrorException PrimeIdealsSet(O, fmpz(1), 2, coprimeto = "bla")
end

@testset "PrimeIdealsSet" begin
  l = Hecke.primes_up_to(20)
  for i = 2:length(l)
    d = l[i]
    K = quadratic_field(d)[1]
    OK = maximal_order(K)
    p = next_prime(fmpz(10)^70)
    while length(prime_decomposition_type(OK, p)) != 2
      p = next_prime(p)
    end
    P = prime_decomposition(OK, p)[1][1]
    @test valuation(P.gen_two, P) == 1
  end

  _, x = PolynomialRing(QQ, cached = false)
  K = number_field(x^4 + 2*x^3 - 35*x^2 - 36*x + 5, "a", cached = false)[1]
  OK = maximal_order(K)
  @assert length(prime_decomposition_type(OK, 3)) == 2
  @assert length(prime_decomposition_type(OK, 5)) == 4
end


struct Params_SS{ComplexArray,Total,KF}
    # boundary conditions
    periodic::Bool
    # normalization
    total::Total

    # algorithm-specific
    # gradient kernel
    kernelfactor::KF
    # fft buffers
    complex_surface::ComplexArray
end

function Params_SS(img; periodic::Bool=true, kernelfactor=KernelFactors.sobel)
    box = size(img)
    complex_box = periodic ? box : box .* 2

    # scale factor
    total = cnt_total(img, periodic)

    p = Params_SS(
        periodic,
        total,
        kernelfactor,
        similar(img, ComplexF32, complex_box)
    )
    cf_type = periodic ? :periodic_point_point : :central_symmetry
    p, cf_type
end

"""
    correllation_function!(res, img, params::Params_SS)

Compute surface-surface correlation function in positive directions
"""
function correllation_function!(res, img, params::Params_SS)
    ix = CartesianIndices(img)

    f = params.complex_surface .= 0

    v_f = view(f, ix)

    v_f .= gradient_norm(img, params.kernelfactor)

    self_correlation!(f)

    res .= real.(v_f) ./ params.total
end

"""
    surfsurf(image; periodic = false)

Calculate `Fss(x)` (surface-surface) correlation function map
for the N-dimensional image and return a `CFMap` object.

The `image` contains the probability of the voxel being in the correct phase.

# Examples
```jldoctest
julia> surfsurf([1 0; 0 1]; periodic=true).result
2×2 Matrix{Float32}:
 0.125  0.125
 0.125  0.125
```
"""
function surfsurf(image; periodic::Bool=false, kernelfactor=KernelFactors.sobel)
    corr_function_map(image, Params_SS; periodic, kernelfactor)
end

function surfsurf(image, phase::Int; periodic::Bool=false, kernelfactor=KernelFactors.sobel)
    one_phase_img = image .== phase
    corr_function_map(one_phase_img, Params_SS; periodic, kernelfactor)
end


function diag_full(A, X0; kwargs...)
    Neig = size(X0, 2)
    Afull = Hermitian(Array(A))
    E = eigen(Afull)
    X = E.vectors[:, 1:Neig]
    λ = E.values[1:Neig]
    (λ=λ, X=X,
     residual_norms=zeros(Neig),
     iterations=0,
     converged=true,
     n_matvec=0)
end


# tensor.jl
#
# Tensor provides a dense implementation of an AbstractTensor type without any
# symmetry assumptions, i.e. it describes tensors living in the full tensor
# product space of its index spaces.

#++++++++++++++
# Tensor type:
#++++++++++++++
# Type definition and constructors:
#-----------------------------------
immutable Tensor{S,T,N} <: AbstractTensor{S,ProductSpace,T,N}
    data::Array{T,N}
    space::ProductSpace{S,N}
    function Tensor(data::Array{T},space::ProductSpace{S,N})
        if length(data)!=dim(space)
            throw(DimensionMismatch("data not of right size"))
        end
        if promote_type(T,eltype(S))!=eltype(S)
            warn("For a tensor in $(space), the entries should not be of type $(T)")
        end
        return new(reshape(data,map(dim,space)),space)
    end
end

# Show method:
#-------------
function Base.show{S,T,N}(io::IO,t::Tensor{S,T,N})
    print(io," Tensor ∈ $T")
    print(io,"[")
    for n=1:N
        n==1 || print(io, " ⊗ ")
        show(io,space(t,n))
    end
    println(io,"]:")
    Base.showarray(io,t.data;header=false)
end

# Basic methods for characterising a tensor:
#--------------------------------------------
space(t::Tensor,ind::Integer)=t.space[ind]
space(t::Tensor)=t.space

# General constructors
#---------------------
# with data
tensor{T<:Real,N}(data::Array{T,N})=Tensor{CartesianSpace,T,N}(data,mapreduce(CartesianSpace,⊗,size(data)))
function tensor{T<:Complex}(data::Array{T,1})
    warning("for complex array, consider specifying Euclidean index spaces")
    Tensor{ComplexEuclideanSpace,T,1}(data,⊗(ComplexSpace(size(data,1))))
end
function tensor{T<:Complex}(data::Array{T,2})
    warning("for complex array, consider specifying Euclidean index spaces")
    Tensor{ComplexEuclideanSpace,T,2}(data,ComplexSpace(size(data,1))⊗ComplexSpace(size(data,2))')
end

tensor{S,T,N}(data::Array{T},P::ProductSpace{S,N})=Tensor{S,T,N}(data,P)
tensor(data::Array,V::ElementarySpace)=tensor(data,⊗(V))

# without data
tensor{T}(::Type{T},P::ProductSpace)=tensor(Array(T,dim(P)),P)
tensor{T}(::Type{T},V::IndexSpace)=tensor(T,⊗(V))
tensor(V::Union{ProductSpace,IndexSpace})=tensor(Float64,V)

Base.similar{S,T,N}(t::Tensor{S},::Type{T},P::ProductSpace{S,N}=space(t))=tensor(similar(t.data,T,dim(P)),P)
Base.similar{S,T}(t::Tensor{S},::Type{T},V::S)=similar(t,T,⊗(V))
Base.similar{S,N}(t::Tensor{S},P::ProductSpace{S,N}=space(t))=similar(t,eltype(t),P)
Base.similar{S}(t::Tensor{S},V::S)=similar(t,eltype(t),V)

Base.zero(t::Tensor)=tensor(zero(t.data),space(t))

Base.zeros{T}(::Type{T},P::ProductSpace)=tensor(zeros(T,dim(P)),P)
Base.zeros{T}(::Type{T},V::IndexSpace)=zeros(T,⊗(V))
Base.zeros(V::Union{ProductSpace,IndexSpace})=zeros(Float64,V)

Base.rand{T}(::Type{T},P::ProductSpace)=tensor(rand(T,dim(P)),P)
Base.rand{T}(::Type{T},V::IndexSpace)=rand(T,⊗(V))
Base.rand(V::Union{ProductSpace,IndexSpace})=rand(Float64,V)

Base.eye{S<:ElementarySpace,T}(::Type{T},::Type{ProductSpace},V::S)=tensor(eye(T,dim(V)),V⊗dual(V))
Base.eye{S<:ElementarySpace}(::Type{ProductSpace},V::S)=eye(Float64,ProductSpace,V)

Base.eye{S<:ElementarySpace,T}(::Type{T},P::ProductSpace{S,2})=(P[1]==dual(P[2]) ? eye(T,ProductSpace,P[1]) : throw(SpaceError("Cannot construct eye-tensor when second space is not the dual of the first space")))
Base.eye{S<:ElementarySpace}(P::ProductSpace{S,2})=eye(Float64,P)

# tensors from concatenation
function tensorcat{S}(catind, X::Tensor{S}...)
    catind = collect(catind)
    isempty(catind) && error("catind should not be empty")
    # length(unique(catdims)) != length(catdims) && error("every dimension should appear only once")

    nargs = length(X)
    numindX = map(numind, X)

    all(n->(n == numindX[1]), numindX) || throw(SpaceError("all tensors should have the same number of indices for concatenation"))

    numindC = numindX[1]
    ncatind = setdiff(1:numindC,catind)
    spaceCvec = Array(S, numindC)
    for n = 1:numindC
        spaceCvec[n] = space(X[1],n)
    end
    for i = 2:nargs
        for n in catind
            spaceCvec[n] = directsum(spaceCvec[n], space(X[i],n))
        end
        for n in ncatind
            spaceCvec[n] == space(X[i],n) || throw(SpaceError("space mismatch for index $n"))
        end
    end
    spaceC = ⊗(spaceCvec...)
    typeC = mapreduce(eltype, promote_type, X)
    dataC = zeros(typeC, map(dim,spaceC))

    offset = zeros(Int,numindC)
    for i=1:nargs
        currentdims=ntuple(n->dim(space(X[i],n)),numindC)
        currentrange=[offset[n]+(1:currentdims[n]) for n=1:numindC]
        dataC[currentrange...] = X[i].data
        for n in catind
            offset[n]+=currentdims[n]
        end
    end
    return tensor(dataC,spaceC)
end

# Copy and fill tensors:
#------------------------
function Base.copy!(tdest::Tensor,tsource::Tensor)
    # Copies data of tensor tsource to tensor tdest if compatible
    space(tdest)==space(tsource) || throw(SpaceError("tensor spaces don't match"))
    copy!(tdest.data,tsource.data)
    return tdest
end
Base.fill!{S,T}(tdest::Tensor{S,T},value::Number)=fill!(tdest.data,convert(T,value))

# Vectorization:
#----------------
Base.vec(t::Tensor)=vec(t.data)
# Convert the non-trivial degrees of freedom in a tensor to a vector to be passed to eigensolvers etc.

# Conversion and promotion:
#---------------------------
Base.full(t::Tensor)=t.data

Base.promote_rule{S,T1,T2,N}(::Type{Tensor{S,T1,N}},::Type{Tensor{S,T2,N}})=Tensor{S,promote_type(T1,T2),N}
Base.promote_rule{S,T1,T2,N1,N2}(::Type{Tensor{S,T1,N1}},::Type{Tensor{S,T2,N2}})=Tensor{S,promote_type(T1,T2)}
Base.promote_rule{S,T1,T2}(::Type{Tensor{S,T1}},::Type{Tensor{S,T2}})=Tensor{S,promote_type(T1,T2)}

Base.promote_rule{S,T1,T2,N}(::Type{AbstractTensor{S,ProductSpace,T1,N}},::Type{Tensor{S,T2,N}})=AbstractTensor{S,ProductSpace,promote_type(T1,T2),N}
Base.promote_rule{S,T1,T2,N1,N2}(::Type{AbstractTensor{S,ProductSpace,T1,N1}},::Type{Tensor{S,T2,N2}})=AbstractTensor{S,ProductSpace,promote_type(T1,T2)}
Base.promote_rule{S,T1,T2}(::Type{AbstractTensor{S,ProductSpace,T1}},::Type{Tensor{S,T2}})=AbstractTensor{S,ProductSpace,promote_type(T1,T2)}

Base.convert{S,T,N}(::Type{Tensor{S,T,N}},t::Tensor{S,T,N})=t
Base.convert{S,T1,T2,N}(::Type{Tensor{S,T1,N}},t::Tensor{S,T2,N})=copy!(similar(t,T1),t)
Base.convert{S,T}(::Type{Tensor{S,T}},t::Tensor{S,T})=t
Base.convert{S,T1,T2}(::Type{Tensor{S,T1}},t::Tensor{S,T2})=copy!(similar(t,T1),t)

Base.float{S,T<:AbstractFloat}(t::Tensor{S,T})=t
Base.float(t::Tensor)=tensor(float(t.data),space(t))

Base.real{S,T<:Real}(t::Tensor{S,T})=t
Base.real(t::Tensor)=tensor(real(t.data),space(t))
Base.complex{S,T<:Complex}(t::Tensor{S,T})=t
Base.complex(t::Tensor)=tensor(complex(t.data),space(t))

for (f,T) in ((:float32,    Float32),
              (:float64,    Float64),
              (:complex64,  Complex64),
              (:complex128, Complex128))
    @eval (Base.$f){S}(t::Tensor{S,$T}) = t
    @eval (Base.$f)(t::Tensor) = tensor(($f)(t.data),space(t))
end

# Basic algebra:
#----------------
function Base.conj!(t1::Tensor,t2::Tensor)
    space(t1)==conj(space(t2)) || throw(SpaceError())
    copy!(t1.data,t2.data)
    conj!(t1.data)
    return t1
end

# transpose inverts order of indices, compatible with graphical notation
function Base.transpose!(tdest::Tensor,tsource::Tensor)
    space(tdest)==space(tsource).' || throw(SpaceError())
    N=numind(tsource)
    TensorOperations.tensorcopy!(tsource.data,1:N,tdest.data,reverse(1:N))
    return tdest
end
function Base.ctranspose!(tdest::Tensor,tsource::Tensor)
    space(tdest)==space(tsource)' || throw(SpaceError())
    N=numind(tsource)
    TensorOperations.tensorcopy!(tsource.data,1:N,tdest.data,reverse(1:N))
    conj!(tdest.data)
    return tdest
end

Base.scale!{S,T,N}(t1::Tensor{S,T,N},t2::Tensor{S,T,N},a::Number)=(space(t1)==space(t2) ? scale!(t1.data,t2.data,a) : throw(SpaceError());return t1)
Base.scale!{S,T,N}(t1::Tensor{S,T,N},a::Number,t2::Tensor{S,T,N})=(space(t1)==space(t2) ? scale!(t1.data,a,t2.data) : throw(SpaceError());return t1)

Base.LinAlg.axpy!(a::Number,x::Tensor,y::Tensor)=(space(x)==space(y) ? Base.LinAlg.axpy!(a,x.data,y.data) : throw(SpaceError()); return y)

-(t::Tensor)=tensor(-t.data,space(t))
+(t1::Tensor,t2::Tensor)= space(t1)==space(t2) ? tensor(t1.data+t2.data,space(t1)) : throw(SpaceError())
-(t1::Tensor,t2::Tensor)= space(t1)==space(t2) ? tensor(t1.data-t2.data,space(t1)) : throw(SpaceError())

# Scalar product and norm: only valid for EuclideanSpace
Base.dot{S<:EuclideanSpace}(t1::Tensor{S},t2::Tensor{S})= (space(t1)==space(t2) ? dot(vec(t1),vec(t2)) : throw(SpaceError()))
Base.vecnorm{S<:EuclideanSpace}(t::Tensor{S})=vecnorm(t.data) # frobenius norm

# Indexing
#----------
# # linear indexing using ProductBasisVector
# Base.getindex{S,T,N}(t::Tensor{S,T,N},b::ProductBasisVector{N,S})=getindex(t.data,Base.to_index(b))
# Base.setindex!{S,T,N}(t::Tensor{S,T,N},value,b::ProductBasisVector{N,S})=setindex!(t.data,value,Base.to_index(b))

@generated function Base.getindex{G,T,N}(t::Tensor{AbelianSpace{G},T,N},s::NTuple{N,G})
    quote
        @nexprs $N n->(r_n=to_range(s[n],space(t,n)))
        return @ncall $N slice t.data r
    end
end

Base.setindex!{G,T,N}(t::Tensor{AbelianSpace{G},T,N},v::Array,s::NTuple{N,G})=(length(v)==length(t[s]) ? copy!(t[s],v) : throw(DimensionMismatch()))
Base.setindex!{G,T,N}(t::Tensor{AbelianSpace{G},T,N},v::Number,s::NTuple{N,G})=fill!(t[s],v)

# Tensor Operations
#-------------------
TensorOperations.scalar(t::Tensor)=iscnumber(space(t)) ? t.data[1] : throw(SpaceError("Not a scalar"))

function TensorOperations.add!{CA}(α, A::Tensor, ::Type{Val{CA}}, β, C::Tensor, indCinA)
    # Implements C = β*C + α*permute(op(A))
    NA = numind(A)

    spaceA = CA == :C ? conj(space(A)) : space(A)

    for i = 1:NA
        spaceA[indCinA[i]] == space(C,i) || throw(SpaceError("incompatible index spaces of tensors"))
    end

    if NA == 0 #scalars
        C.data[1] = β*C.data[1] + α*A.data[1]
    elseif indCinA == collect(1:NA) #trivial permutation
        if β == 0
            scale!(copy!(C,A), α)
        else
            Base.LinAlg.axpy!(α, A, scale!(C, β))
        end
    else
        TensorOperations.add!(α, A.data, Val{CA}, β, C.data, indCinA)
    end

    return C
end

function TensorOperations.trace!{CA}(α, A::Tensor, ::Type{Val{CA}}, β, C::Tensor, indCinA, cindA1, cindA2)
    NA = numind(A)
    NC = numind(C)

    spaceA = CA == :C ? conj(space(A)) : space(A)

    for i = 1:NC
        spaceA[indCinA[i]] == space(C,i) || throw(SpaceError("space mismatch"))
    end

    for i = 1:div(NA-NC, 2)
        spaceA[cindA1[i]] == dual(spaceA[cindA2[i]]) || throw(SpaceError("space mismatch"))
    end

    TensorOperations.trace!(α, A.data, Val{CA}, β, C.data, indCinA, cindA1, cindA2)

    return C
end

function TensorOperations.contract!{CA,CB,ME}(α, A::Tensor, ::Type{Val{CA}}, B::Tensor, ::Type{Val{CB}}, β, C::Tensor, oindA, cindA, oindB, cindB, indCinoAB, ::Type{Val{ME}}=Val{:BLAS})
    # check size compatibility
    spaceA = CA == :C ? conj(space(A)) : space(A)
    spaceB = CB == :C ? conj(space(B)) : space(B)
    spaceC = space(C)

    cspaceA = spaceA[cindA]
    cspaceB = spaceB[cindB]

    ospaceA = spaceA[oindA]
    ospaceB = spaceB[oindB]

    ospaceAB = ospaceA ⊗ ospaceB

    for i = 1:length(cspaceA)
        cspaceA[i] == dual(cspaceB[i]) || throw(SpaceError("A and B have incompatible index spaces"))
    end

    for i in 1:length(indCinoAB)
        spaceC[i] == ospaceAB[indCinoAB[i]] || throw(SpaceError("C has incompatible index space"))
    end

    if length(indCinoAB) == 0
        Cscal = pointer_to_array(pointer(C.data, 1), ())
        TensorOperations.contract!(α, A.data, Val{CA}, B.data, Val{CB}, β, Cscal, oindA, cindA, oindB, cindB, indCinoAB, Val{ME})
    else
        TensorOperations.contract!(α, A.data, Val{CA}, B.data, Val{CB}, β, C.data, oindA, cindA, oindB, cindB, indCinoAB, Val{ME})
    end

    return C
end

function TensorOperations.similar_from_indices{T,CA}(::Type{T}, indices, A::Tensor, ::Type{Val{CA}}=Val{:N})
    spaceA = CA == :C ? conj(space(A)) : space(A)
    return similar(A, T, spaceA[indices])
end

function TensorOperations.similar_from_indices{T,CA,CB}(::Type{T}, indices, A::Tensor, B::Tensor,
                                                        ::Type{Val{CA}}=Val{:N}, ::Type{Val{CB}}=Val{:N})
    spaceA = CA == :C ? conj(space(A)) : space(A)
    spaceB = CB == :C ? conj(space(B)) : space(B)
    spaceAB = spaceA ⊗ spaceB
    return similar(A, T, spaceAB[indices])
end

# Methods below are only implemented for CartesianTensor or ComplexTensor:
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
typealias ComplexTensor{T,N} Tensor{ComplexSpace,T,N}
typealias CartesianTensor{T,N} Tensor{CartesianSpace,T,N}

# Index methods
#---------------
@eval function insertind{S}(t::Tensor{S},ind::Int,V::S)
    N=numind(t)
    0<=ind<=N || throw(IndexError("index out of range"))
    iscnumber(V) || throw(SpaceError("can only insert index with c-number index space"))
    spacet=space(t)
    newspace=spacet[1:ind] ⊗ V ⊗ spacet[ind+1:N]
    return tensor(t.data,newspace)
end
@eval function deleteind(t::Tensor,ind::Int)
    N=numind(t)
    1<=ind<=N || throw(IndexError("index out of range"))
    iscnumber(space(t,ind)) || throw(SpaceError("can only delete index with c-number index space"))
    spacet=space(t)
    newspace=spacet[1:ind-1] ⊗ spacet[ind+1:N]
    return tensor(t.data,newspace)
end

for (S,TT) in ((CartesianSpace,CartesianTensor),(ComplexSpace,ComplexTensor))
    @eval function fuseind(t::$TT,ind1::Int,ind2::Int,V::$S)
        N=numind(t)
        ind2==ind1+1 || throw(IndexError("only neighbouring indices can be fused"))
        1<=ind1<=N-1 || throw(IndexError("index out of range"))
        fuse(space(t,ind1),space(t,ind2),V) || throw(SpaceError("index spaces $(space(t,ind1)) and $(space(t,ind2)) cannot be fused to $V"))
        spacet=space(t)
        newspace=spacet[1:ind1-1]*V*spacet[ind2+1:N]
        return tensor(t.data,newspace)
    end
    @eval function splitind(t::$TT,ind::Int,V1::$S,V2::$S)
        1<=ind<=numind(t) || throw(IndexError("index out of range"))
        fuse(V1,V2,space(t,ind)) || throw(SpaceError("index space $(space(t,ind)) cannot be split into $V1 and $V2"))
        spacet=space(t)
        newspace=spacet[1:ind-1]*V1*V2*spacet[ind+1:N]
        return tensor(t.data,newspace)
    end
end

# Factorizations:
#-----------------
for (S,TT) in ((CartesianSpace,CartesianTensor),(ComplexSpace,ComplexTensor))
    @eval function svd!(t::$TT,n::Int,::NoTruncation)
        # Perform singular value decomposition corresponding to bipartion of the
        # tensor indices into the left indices 1:n and remaining right indices,
        # thereby destroying the original tensor.
        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)
        newdim=length(F[:S])
        newspace=$S(newdim)
        U=tensor(F[:U],leftspace*newspace')
        S=tensor(diagm(F[:S]),newspace*newspace')
        V=tensor(F[:Vt],newspace*rightspace)
        return U,S,V,abs(zero(eltype(t)))
    end

    @eval function svd!(t::$TT,n::Int,trunc::Union{TruncationDimension,TruncationSpace})
        # Truncate rank corresponding to bipartition into left indices 1:n
        # and remain right indices, based on singular value decomposition,
        # thereby destroying the original tensor.
        # Truncation parameters are given as keyword arguments: trunctol should
        # always be one of the possible arguments for specifying truncation, but
        # truncdim can be replaced with different parameters for other types of tensors.

        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        dim(trunc) >= min(leftdim,rightdim) && return svd!(t,n)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)
        sing=F[:S]

        # compute truncation dimension
        if eps(trunc)!=0
            sing=F[:S]
            normsing=vecnorm(sing)
            truncdim=0
            while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
                truncdim+=1
            end
            truncdim=min(truncdim,dim(trunc))
        else
            truncdim=dim(trunc)
        end
        newspace=$S(truncdim)

        # truncate
        truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
        U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
        S=tensor(diagm(sing[1:truncdim]),newspace*newspace')
        V=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
        return U,S,V,truncerr
    end

    @eval function svd!(t::$TT,n::Int,trunc::TruncationError)
        # Truncate rank corresponding to bipartition into left indices 1:n
        # and remain right indices, based on singular value decomposition,
        # thereby destroying the original tensor.
        # Truncation parameters are given as keyword arguments: trunctol should
        # always be one of the possible arguments for specifying truncation, but
        # truncdim can be replaced with different parameters for other types of tensors.

        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)

        # find truncdim based on trunctolinfo
        sing=F[:S]
        normsing=vecnorm(sing)
        truncdim=0
        while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
            truncdim+=1
        end
        newspace=$S(truncdim)

        # truncate
        truncerr=vecnorm(sing[(truncdim+1):end])/normsing
        U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
        S=tensor(diagm(sing[1:truncdim]),newspace*newspace')
        V=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
        return U,S,V,truncerr
    end

    @eval function leftorth!(t::$TT,n::Int,::NoTruncation)
        # Create orthogonal basis U for indices 1:n, and remainder C for right
        # indices, thereby destroying the original tensor.
        # Decomposition should be unique, such that it always returns the same
        # result for the same input tensor t. UC = QR is fastest but only unique
        # after correcting for phases.
        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        data=reshape(t.data,(leftdim,rightdim))
        if leftdim>rightdim
            newdim=rightdim
            tau=zeros(eltype(data),(newdim,))
            Base.LinAlg.LAPACK.geqrf!(data,tau)

            C=zeros(eltype(t),(newdim,newdim))
            for j in 1:newdim
                for i in 1:j
                    @inbounds C[i,j]=data[i,j]
                end
            end
            Base.LinAlg.LAPACK.orgqr!(data,tau)
            U=data
            for i=1:newdim
                tau[i]=sign(C[i,i])
            end
            scale!(U,tau)
            scale!(conj!(tau),C)
        else
            newdim=leftdim
            C=data
            U=eye(eltype(data),newdim)
        end

        newspace=$S(newdim)
        return tensor(U,leftspace*newspace'), tensor(C,newspace*rightspace), abs(zero(eltype(t)))
    end

    @eval function leftorth!(t::$TT,n::Int,trunc::Union{TruncationDimension,TruncationSpace})
        # Truncate rank corresponding to bipartition into left indices 1:n
        # and remain right indices, based on singular value decomposition,
        # thereby destroying the original tensor.
        # Truncation parameters are given as keyword arguments: trunctol should
        # always be one of the possible arguments for specifying truncation, but
        # truncdim can be replaced with different parameters for other types of tensors.

        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        dim(trunc) >= min(leftdim,rightdim) && return leftorth!(t,n)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)
        sing=F[:S]

        # compute truncation dimension
        if eps(trunc)!=0
            sing=F[:S]
            normsing=vecnorm(sing)
            truncdim=0
            while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
                truncdim+=1
            end
            truncdim=min(truncdim,dim(trunc))
        else
            truncdim=dim(trunc)
        end
        newspace=$S(truncdim)

        # truncate
        truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
        U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
        C=tensor(scale!(sing[1:truncdim],F[:Vt][1:truncdim,:]),newspace*rightspace)
        return U,C,truncerr
    end

    @eval function leftorth!(t::$TT,n::Int,trunc::TruncationError)
        # Truncate rank corresponding to bipartition into left indices 1:n
        # and remain right indices, based on singular value decomposition,
        # thereby destroying the original tensor.
        # Truncation parameters are given as keyword arguments: trunctol should
        # always be one of the possible arguments for specifying truncation, but
        # truncdim can be replaced with different parameters for other types of tensors.

        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)

        # compute truncation dimension
        sing=F[:S]
        normsing=vecnorm(sing)
        truncdim=0
        while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
            truncdim+=1
        end
        newspace=$S(truncdim)

        # truncate
        truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
        U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
        C=tensor(scale!(sing[1:truncdim],F[:Vt][1:truncdim,:]),newspace*rightspace)

        return U,C,truncerr
    end

    @eval function rightorth!(t::$TT,n::Int,::NoTruncation)
        # Create orthogonal basis U for right indices, and remainder C for left
        # indices. Decomposition should be unique, such that it always returns the
        # same result for the same input tensor t. CU = LQ is typically fastest but only
        # unique after correcting for phases.
        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        data=reshape(t.data,(leftdim,rightdim))
        if leftdim<rightdim
            newdim=leftdim
            tau=zeros(eltype(data),(newdim,))
            datat=transpose(data)
            Base.LinAlg.LAPACK.geqrf!(datat,tau)

            C=zeros(eltype(t),(newdim,newdim))
            for j in 1:newdim
                for i in 1:j
                    @inbounds C[j,i]=datat[i,j]
                end
            end
            Base.LinAlg.LAPACK.orgqr!(datat,tau)
            Base.transpose!(data,datat)
            U=data

            for i=1:newdim
                tau[i]=sign(C[i,i])
            end
            scale!(tau,U)
            scale!(C,conj!(tau))
        else
            newdim=rightdim
            C=data
            U=eye(eltype(data),newdim)
        end

        newspace=$S(newdim)
        return tensor(C,leftspace ⊗ dual(newspace)), tensor(U,newspace ⊗ rightspace), abs(zero(eltype(t)))
    end

    @eval function rightorth!(t::$TT,n::Int,trunc::Union{TruncationDimension,TruncationSpace})
        # Truncate rank corresponding to bipartition into left indices 1:n
        # and remain right indices, based on singular value decomposition,
        # thereby destroying the original tensor.
        # Truncation parameters are given as keyword arguments: trunctol should
        # always be one of the possible arguments for specifying truncation, but
        # truncdim can be replaced with different parameters for other types of tensors.

        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        dim(trunc) >= min(leftdim,rightdim) && return rightorth!(t,n)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)
        sing=F[:S]

        # compute truncation dimension
        if eps(trunc)!=0
            sing=F[:S]
            normsing=vecnorm(sing)
            truncdim=0
            while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
                truncdim+=1
            end
            truncdim=min(truncdim,dim(trunc))
        else
            truncdim=dim(trunc)
        end
        newspace=$S(truncdim)

        # truncate
        truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
        C=tensor(scale!(F[:U][:,1:truncdim],sing[1:truncdim]),leftspace*newspace')
        U=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
        return C,U,truncerr
    end

    @eval function rightorth!(t::$TT,n::Int,trunc::TruncationError)
        # Truncate rank corresponding to bipartition into left indices 1:n
        # and remain right indices, based on singular value decomposition,
        # thereby destroying the original tensor.
        # Truncation parameters are given as keyword arguments: trunctol should
        # always be one of the possible arguments for specifying truncation, but
        # truncdim can be replaced with different parameters for other types of tensors.

        N=numind(t)
        spacet=space(t)
        leftspace=spacet[1:n]
        rightspace=spacet[n+1:N]
        leftdim=dim(leftspace)
        rightdim=dim(rightspace)
        data=reshape(t.data,(leftdim,rightdim))
        F=svdfact!(data)

        # compute truncation dimension
        sing=F[:S]
        normsing=vecnorm(sing)
        truncdim=0
        while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
            truncdim+=1
        end
        newspace=$S(truncdim)

        # truncate
        truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
        C=tensor(scale!(F[:U][:,1:truncdim],sing[1:truncdim]),leftspace*newspace')
        U=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
        return C,U,truncerr
    end
end

# Methods below are only implemented for CartesianMatrix or ComplexMatrix:
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
typealias ComplexMatrix{T} ComplexTensor{T,2}
typealias CartesianMatrix{T} CartesianTensor{T,2}

function Base.pinv(t::Union{ComplexMatrix,CartesianMatrix})
    # Compute pseudo-inverse
    spacet=space(t)
    data=copy(t.data)
    leftdim=dim(spacet[1])
    rightdim=dim(spacet[2])

    F=svdfact!(data)
    Sinv=F[:S]
    for k=1:length(Sinv)
        if Sinv[k]>eps(Sinv[1])*max(leftdim,rightdim)
            Sinv[k]=one(Sinv[k])/Sinv[k]
        end
    end
    data=F[:V]*scale(F[:S],F[:U]')
    return tensor(data,spacet')
end

function Base.eig(t::Union{ComplexMatrix,CartesianMatrix})
    # Compute eigenvalue decomposition.
    spacet=space(t)
    spacet[1] == spacet[2]' || throw(SpaceError("eigenvalue factorization only exists if left and right index space are dual"))
    data=copy(t.data)

    F=eigfact!(data)

    Lambda=tensor(diagm(F[:values]),spacet)
    V=tensor(F[:vectors],spacet)
    return Lambda, V
end

function Base.inv(t::Union{ComplexMatrix,CartesianMatrix})
    # Compute inverse.
    spacet=space(t)
    spacet[1] == spacet[2]' || throw(SpaceError("inverse only exists if left and right index space are dual"))

    return tensor(inv(t.data),spacet)
end


export fista

mutable struct FISTA{matT,vecT} <: AbstractLinearSolver
  A::matT
  reg::Regularization
  x::vecT
  xᵒˡᵈ::vecT
  res::vecT
  res_norm::Float64
  res_norm_old::Float64
  ρ::Float64
  t::Float64
  tᵒˡᵈ::Float64
  iterations::Int64
  relTol::Float64
  normalizeReg::Bool
  regFac::Float64
end

"""
    FISTA(A::matT, x::vecT=zeros(eltype(A),size(A,2))
          ; reg=nothing, regName=["L1"], λ=[0.0], kargs...) where {matT,vecT}

creates a `FISTA` object for the system matrix `A`.

# Arguments
* `A`                       - system matrix
* `x::vecT`                 - Array with the same type and size as the solution
* (`reg=nothing`)           - Regularization object
* (`regName=["L1"]`)        - name of the Regularization to use (if reg==nothing)
* (`λ=[0.0]`)               - Regularization paramter
* (`ρ::Float64=1`)          - step size for gradient step
* (`t::Float64=1.0`)        - parameter for predictor-corrector step
* (`relTol::Float64=1.e-5`) - tolerance for stopping criterion
* (`iterations::Int64=50`)  - maximum number of iterations
"""
function FISTA(A::matT, x::vecT=zeros(eltype(A),size(A,2)); reg=nothing, regName=["L1"]
              , λ=[0.0]
              , ρ::Float64=1.0
              , t::Float64=1.0
              , relTol::Float64=eps()
              , iterations::Int64=50
              , normalizeReg::Bool=false
              , kargs...) where {matT,vecT}

  if reg == nothing
    reg = Regularization(regName, λ, kargs...)
  end

  xᵒˡᵈ = similar(x)
  res = similar(x,size(A,1))

  return FISTA(A,vec(reg)[1],x,xᵒˡᵈ,res,0.0,0.0,ρ,t,t,iterations,relTol,normalizeReg,1.0)
end

"""
    init!(it::FISTA{matT,T}, b::vecT
              ; A::matT=solver.A
              , x::vecT=similar(b,0)
              , t::Float64=1.0) where T

(re-) initializes the FISTA iterator
"""
function init!(solver::FISTA{matT,vecT}, b::vecT
              ; A::matT=solver.A
              , x::vecT=similar(b,0)
              , t::Float64=1.0
              ) where {matT,vecT}

  solver.A = A
  if isempty(x)
    if !isempty(b) #!iszero(b)
      solver.x[:] .= adjoint(A) * b
    else
      solver.x[:] .= zeros(T,size(A,2))
    end
  else
    solver.x[:] .= x
  end
  solver.xᵒˡᵈ[:] .= solver.x  # this could also be zero
  solver.res[:] .= A*solver.x-b
  solver.res_norm_old = 0.0
  solver.res_norm = norm(solver.res)
  solver.t = t
  solver.tᵒˡᵈ = t
  # normalization of regularization parameters
  if solver.normalizeReg
    solver.regFac = norm(b,1)/length(b)
  else
    solver.regFac = 1.0
  end
end

"""
    solve(solver::FISTA, b::Vector)

solves an inverse problem using FISTA.

# Arguments
* `solver::FISTA`                 - the solver containing both system matrix and regularizer
* `b::vecT`                     - data vector
* `A::matT=solver.A`            - operator for the data-term of the problem
* (`startVector::vecT=similar(b,0)`)  - initial guess for the solution
* (`solverInfo=nothing`)          - solverInfo object

when a `SolverInfo` objects is passed, the residuals are stored in `solverInfo.convMeas`.
"""
function solve(solver::FISTA, b::vecT; A::matT=solver.A, startVector::vecT=similar(b,0), solverInfo=nothing, kargs...) where {matT,vecT}
  # initialize solver parameters
  init!(solver, b; A=A, x=startVector)

  # log solver information
  solverInfo != nothing && storeInfo(solverInfo,solver.x,solver.res_norm)

  # perform FISTA iterations
  for (iteration, item) = enumerate(solver)
    solverInfo != nothing && storeInfo(solverInfo,solver.x,solver.res_norm)
  end

  return solver.x
end

"""
  iterate(it::FISTA{matT,vecT}, iteration::Int=0) where {matT,vecT}

performs one fista iteration.
"""
function iterate(solver::FISTA{matT,vecT}, iteration::Int=0) where {matT,vecT}
  if done(solver, iteration) return nothing end

  # gradient step
  solver.xᵒˡᵈ[:] .= solver.x
  solver.x[:] .= solver.x - solver.ρ* (solver.A' * solver.res)

  # proximal map
  solver.reg.prox!(solver.x, solver.regFac*solver.ρ*solver.reg.λ; solver.reg.params...)

  # predictor-corrector update
  solver.tᵒˡᵈ = solver.t
  solver.t = (1. + sqrt(1. + 4. * solver.tᵒˡᵈ^2)) / 2.
  solver.x[:] .= solver.x + (solver.tᵒˡᵈ-1)/solver.t*(solver.x-solver.xᵒˡᵈ)

  # update residual
  # solver.res .= solver.A*solver.x-solver.b
  solver.res .+= solver.A*(solver.x-solver.xᵒˡᵈ)  # this relies on a proper initialization in init()
  solver.res_norm_old = solver.res_norm
  solver.res_norm = norm(solver.res)

  # return the residual-norm as item and iteration number as state
  return solver.res_norm, iteration+1
end

@inline converged(solver::FISTA) = ( abs(solver.res_norm-solver.res_norm_old)/solver.res_norm_old < solver.relTol )

@inline done(solver::FISTA,iteration::Int) = converged(solver) || iteration>=solver.iterations


################################################################################
#
#          EllCrv/EllCrvFin.jl : Elliptic curves over finite fields
#
################################################################################

################################################################################
#
#  Submodules
#
################################################################################

include("EllCrv.jl")
include("../src/Misc/BigIntInterface.jl")
include("../src/Misc/PolyFac.jl")

################################################################################
#
#  Imports
#
################################################################################

using Random

################################################################################
#
#  Exports
#
################################################################################

export order, order_via_schoof, rand_point, crypto_curve

################################################################################
#
#  Random Point Generator
#
################################################################################

function issquare(x::Union{AbstractAlgebra.gfelem{Int32}, AbstractAlgebra.FinFieldElem})
    k = parent(x)
    S, t = AbstractAlgebra.PolynomialRing(k, "t", cached = false)

    f = t^2 - x

    # Distinct Degree Factorization for now.
    fac = factor_ddf(f)

    p = first(keys(fac.fac))

    if fac[p] == 2
        root = -AbstractAlgebra.coeff(p, 0)
        return true, k(root)
    elseif length(fac) == 2
        root = -coeff(p, 0)
        return true, k(root)
    else
        return false, zero(k)
    end
end

function rand_point(E::EllCrv)
    k = E.base_field
    (a4, a6) = a_invars(E)

    if E.short == false
        error("Not implemented for long form.")
    end

    while true

        x = rand(k)
        square = x^3 + a4*x + a6

        a = AbstractAlgebra.issquare(square)
        if a[1] == true
            y = a[2]
            P = E((x, y))
            return P
        end
    end
end

################################################################################
#
#  Elliptic Curve Generator over Finite Fields (Naive implementation at first.)
#  https://pdfs.semanticscholar.org/c60c/46a9ee6f90958f609f464723f6bf92835feb.pdf
#
################################################################################

function crypto_curve(r0::BigInt, k0::BigInt, h0::BigInt)
    if prod(r0, k0, h0) <= 0
        error("r0, k0, h0 must be positive integers.")
    end

    if r0 < 2^BigInt(159) || k0 < 4 || h0 < 200
        error("Invalid input.")
    end

    (d, p, r, k) = find_prime(r0, k0, h0)

    (E, G) = find_curve(d, p, r, k)

    return (d, p, r, k, E, G)
end

function find_prime(r0::BigInt, k0::BigInt, h0::BigInt)
    print("Enter a (rational) prime p: ")
    input = readline()
    p = parse(BigInt, chomp(input))

    m = 50

    if p == 0
        find_prime_delta_fixed(r0, k0, h0)
    elseif is_prime(p, m) == true && floor(log2(p)) == floor(log2(r0*k0))
        find_discriminant(p, r0, k0, h0)
    else
        error("Invalid prime.")
    end
end

function find_prime_delta_fixed(r0::BigInt, k0::BigInt, h0::BigInt)
    print("Enter a discriminant: ")
    input = readline()
    d = parse(BigInt, chomp(input))

    abs_d = abs(d)

    if k0 >= 4
        if mod(d, 8) == 1 || mod(abs_d, 8) == 7
            b = BigInt(floor(log2(4*r0)))
            i0 = 2^b
            i1 = 2^(b + 1) - 1
            interval = collect(BigInt, i0:i1)

            if find_prime_1_mod_8(r0, 4, d, interval)
                return true
            else
                error("Unable to find prime.")
            end
        else
            error("Get discriminant congruent to 1 mod 8,
            or its norm congruent to 7 mod 8.")
        end
    elseif k0 >= 2
        if mod(d, 16) == 8 || 12
            b = BigInt(floor(log2(r0*k0)))
            i0 = 2^b
            i1 = 2^(b + 1) - 1
            interval = collect(BigInt, i0:i1)

            if find_prime_0_mod_4(r0, k0, d, interval)
                return true
            else
                error("Unable to find prime.")
            end
        else
            error("Get discriminant congruent to 8 or 12 modulo 16.")
        end
    else
        if mod(d, 8) == 5
            b = BigInt(floor(log2(r0)))
            i0 = 2^b
            i1 = 2^(b + 1) - 1
            interval = collect(BigInt, i0:i1)

            if find_prime_5_mod_8(r0, 1, d, interval)
                return true
            else
                error("Unable to find prime.")
            end
        else
            error("Get discrimint congruent to 5 mod 8.")
        end
    end
end

function find_prime_0_mod_4(r0::BigInt, k0::BigInt, d::BigInt,
    interval::Array{BigInt, 1})
    # Number of successive tested pairs (t, y).
    # If not successful after T tries, init a new pair.
    const T = 2000

    # Bit of 2 in binary representation of t.
    t_bit_1 = Int

    abs_delta = assign(d)

    ###
end


using Base.Test
using QuantumOptics

mutable struct SpectralanalysisTestOperator <: Operator
end

@testset "spectralanalysis" begin

srand(0)

sprandop(b) = sparse(DenseOperator(b, rand(Complex128, length(b), length(b))))

# Test diagonalization
@test_throws ArgumentError eigenstates(SpectralanalysisTestOperator())
@test_throws bases.IncompatibleBases eigenstates(DenseOperator(GenericBasis(3), GenericBasis(4)))
@test_throws ArgumentError eigenenergies(SpectralanalysisTestOperator())
@test_throws bases.IncompatibleBases eigenenergies(DenseOperator(GenericBasis(3), GenericBasis(4)))

# Test hermitian diagonalization
b = GenericBasis(5)
a = randoperator(b)
H = (a+dagger(a))/2
U = expm(1im*H)
d = [-3, -2.6, -0.1, 0.0, 0.6]
D = DenseOperator(b, diagm(d))
Dsp = sparse(D)
@test eigenenergies(D) ≈ d
@test eigenenergies(Dsp, 3) ≈ d[1:3]

R = U*D*dagger(U)
Rsp = sparse(R)
@test eigenenergies((R+dagger(R))/2) ≈ d
@test eigenenergies((Rsp+dagger(Rsp))/2, 3) ≈ d[1:3]

E, states = eigenstates((R+dagger(R))/2, 3)
Esp, states_sp = eigenstates((Rsp+dagger(Rsp))/2, 3)
for i=1:3
    @test E[i] ≈ d[i]
    @test Esp[i] ≈ d[i]
    v = U.data[1,i]/states[i].data[1]
    @test states[i].data*v ≈ U.data[:,i]
    v = U.data[1,i]/states_sp[i].data[1]
    @test states_sp[i].data*v ≈ U.data[:,i]
end

# Test nonhermitian diagonalization
b = GenericBasis(5)
a = randoperator(b)
H = (a+dagger(a))/2
U = expm(1im*H)
d = [-3+0.2im, -2.6-0.1im, -0.1+0.5im, 0.0, 0.6+0.3im]
D = DenseOperator(b, diagm(d))
Dsp = sparse(D)
@test eigenenergies(D; warning=false) ≈ d
@test eigenenergies(Dsp, 3; warning=false) ≈ d[1:3]

R = U*D*dagger(U)
Rsp = sparse(R)
@test eigenenergies(R; warning=false) ≈ d
@test eigenenergies(Rsp, 3; warning=false) ≈ d[1:3]

E, states = eigenstates(R, 3; warning=false)
Esp, states_sp = eigenstates(Rsp, 3; warning=false)
for i=1:3
    @test E[i] ≈ d[i]
    @test Esp[i] ≈ d[i]
    v = U.data[1,i]/states[i].data[1]
    @test states[i].data*v ≈ U.data[:,i]
    v = U.data[1,i]/states_sp[i].data[1]
    @test states_sp[i].data*v ≈ U.data[:,i]
end

# Test simdiag
spinbasis = SpinBasis(1//2)
sx = sigmax(spinbasis)
sy = sigmay(spinbasis)
sz = sigmaz(spinbasis)
twospinbasis = spinbasis ⊗ spinbasis
Sx = full(sum([embed(twospinbasis, i, sx) for i=1:2]))/2.
Sy = full(sum([embed(twospinbasis, i, sy) for i=1:2]))/2.
Sz = full(sum([embed(twospinbasis, i, sz) for i=1:2]))/2.
Ssq = Sx^2 + Sy^2 + Sz^2
d, v = simdiag([Sz, Ssq])
@test d[1] == [-1.0, 0, 0, 1.0]
@test d[2] ≈ [2, 0.0, 2, 2]
@test_throws ErrorException simdiag([Sx, Sy])

threespinbasis = spinbasis ⊗ spinbasis ⊗ spinbasis
Sx3 = full(sum([embed(threespinbasis, i, sx) for i=1:3])/2.)
Sy3 = full(sum([embed(threespinbasis, i, sy) for i=1:3])/2.)
Sz3 = full(sum([embed(threespinbasis, i, sz) for i=1:3])/2.)
Ssq3 = Sx3^2 + Sy3^2 + Sz3^2
d3, v3 = simdiag([Ssq3, Sz3])
dsq3_std = eigenenergies(full(Ssq3))
@test diagm(dsq3_std) ≈ v3'*Ssq3.data*v3

fockbasis = FockBasis(4)
@test_throws ErrorException simdiag([Sy3, Sz3])
@test_throws ErrorException simdiag([full(destroy(fockbasis)), full(create(fockbasis))])

end # testset


function Convmtx(w,nr)
	w  = vec(w)
	nw = length(w)
	nd = nw + nr - 1
	cm = zeros(nd, nr)
	for i=1:nr
		cm[i:i+nw-1, i] = w
	end
	cm = sparse(cm)
	return cm
end


#
###
# AbstractDomain interface #TODO
###

"""
args:
    AbstractDomain{T,D}
    direction < D
ret:
    Bool
"""
function isperiodic end

"""
args:
    AbstractDomain
    direction
ret:
    Tuple of end points
"""
function end_points end

"""
check if domain end points match

args:
    AbstractDomain
    direction
ret:
    Bool
"""
function domains_match end

###
# IntervalDomain
###

""" 1D interval """
struct IntervalDomain{T<:Number,Ttag} <: AbstractDomain{T,1}
    x0::T
    x1::T
    isperiodic::Bool
    tags::Ttag

    function IntervalDomain(x0::Number, x1::Number, isperiodic, tags)
        T = promote_type(eltype.((x0, x1))...)
        new{T,typeof(tags)}(T(x0), T(x1), isperiodic, tags)
    end
end

function IntervalDomain(x0 = -1e0, x1 = 1e0,
                        isperiodic=false, tags = (nothing, nothing)
                       )
    IntervalDomain(x0, x1, isperiodic, tags)
end

isperiodic(dom::IntervalDomain) = dom.isperiodic
end_points(dom::IntervalDomain) = (dom.x0, dom.x1)

function domains_match(int1::IntervalDomain, int2::IntervalDomain)
    bools = isapprox.(end_points.((int1, int2))...)
    *(bools...)
end

###
# Box Domain
###

""" D-dimensional logically reectangular domain """
struct BoxDomain{T,D,Ti} <: AbstractDomain{T,D}
    intervals::Ti

    function BoxDomain(intervals::IntervalDomain...)
        T = promote_type(eltype.(intervals)...)
        D = length(intervals)
        new{T,D,typeof(intervals)}(intervals)
    end
end

isperiodic(box::BoxDomain, dir::Integer) = box.intervals[dir] |> isperiodic
end_points(box::BoxDomain, dir::Integer) = box.intervals[dir] |> end_points
end_points(box::BoxDomain) = end_points.(box.intervals)

Base.:*(int1::IntervalDomain, int2::IntervalDomain) = BoxDomain(int1, int2)
Base.:*(box1::BoxDomain, int2::IntervalDomain) = BoxDomain(box1.intervals..., int2)
Base.:*(int1::IntervalDomain, box2::BoxDomain) = BoxDomain(int1, box2.intervals...)
Base.:*(box1::BoxDomain, box2::BoxDomain) = BoxDomain(box1.intervals..., box2.intervals...)

###
# DeformedDomain
###

"""
Deform D-dimensional domain via mapping

args:
   -reference domain
   -mapping function
        (x1,...,xD) = map(r1, ..., rD)
   -isrescaling - do optimizations if
    mapping is a simple rescaling
        x1 = a1 + λ1 * x1(r1), ...,
        xD = aD + λD * xD(rD)
   -isseparable - do optimizations if
    mapping is separable
        x1 = x1(r1), ...,
        xD = xD(rD)
"""
struct DeformedDomain{T,D,Tdom<:AbstractDomain{T,D}, Tm} <: AbstractDomain{T,D}
    domain::Tdom
    mapping::Tm
#   isrescaling::Bool
    isseparable::Bool
end

function deform(domain, mapping = nothing, isseparable = false)
    DeformedDomain(domain, mapping, isseparable)
end

function domains_match(box1::BoxDomain{<:Number, D1}, box2::BoxDomain{<:Number, D2}) where{D1,D2}
    if D1 != D2
        @warn "dimension mismatch between $box1, $box2"
        return false
    end

    ret = true
    for i=1:D1
        ret * domains_match(box1.intervals[i], box2.intervals[i])
    end

    ret
end

# TODO - make a struct for mappings
# struct Deformations{D,Tmap} <: AbstractDeformation{D}
#   mapping::Tmap
#   isrescaling::Bool
#   isseparable::Bool
# end

function map_from_ref(domain, ref_domain;D=D) # TODO

    if domains_match(domain, ref_domain)
        return domain
    end

    xe = domain.end_points

    mapping = domain.mapping ==! nothing ? domain.mapping : (r -> r)
    mapping = mapping ∘ map
end

###
# Conveniences
###

function unit_box(D, args...) # TODO not exactly unit haha
    interval = IntervalDomain(-true, true, args...)
    BoxDomain((interval for i=1:D)...)
end

function polar2D(r, θ)
    x = @. r * cos(θ)
    y = @. r * sin(θ)
    x, y
end

function annulus_2D(r0, r1)
    intR = IntervalDomain(r0, r1, false, (:Inner, :Outer))
    intθ = IntervalDomain(-π,  π, true , (:Periodic, :Periodic))

    dom = intR * intθ

    deform(dom, polar2D)
end

###
# Transfinite Interpolation
###

function gordon_hall end # TODO

#


using Test
using SortingLab

x = [1,2,missing, 3]

@test isequal(sort(x), fsort(x))


@testset "Testing arma.jl" begin

    # set up
    phi = [.95, -.4, -.4]
    theta = zeros(3)
    sigma = .15
    lp = ARMA(phi, theta, sigma)

    # test simulate
    sim = simulation(lp, ts_length=250)
    @test length(sim) == 250

    # test impulse response
    imp_resp = impulse_response(lp, impulse_length=75)
    @test length(imp_resp) == 75

    @testset "test constructors" begin
        phi = 0.5
        theta = 0.4
        sigma = 0.15

        a1 = ARMA(phi, theta, sigma)
        a2 = ARMA([phi;], theta, sigma)
        a3 = ARMA(phi, [theta;], sigma)

        for nm in fieldnames(typeof(a1))
            @test getfield(a1, nm) == getfield(a2, nm)
            @test getfield(a1, nm) == getfield(a3, nm)
        end
    end

    @testset "test autocovariance" begin
        θ = 0.5
        σ = 0.15
        ma1 = ARMA(Float64[], [θ], σ)
        ac = autocovariance(ma1; num_autocov=5)

        # first is the variance. equal to (1 + θ^2) sigma^2
        @test isapprox(ac[1], (1+θ^2)*σ^2; atol=1e-3)

        # second should be θ σ^2
        @test isapprox(ac[2], θ*σ^2; atol=1e-3)

        # all others should be 0
        @test isapprox(ac[3:end], fill!(similar(ac[3:end]), zero(eltype(ac[3:end]))); atol=1e-3)
    end

end  # testset


using MultipleScattering
using Plots; pyplot()

function hankel_order_convergence(m=[0,1,2,3,4,5,6,7,8,9,10], volfrac = 0.1,
    radius = 1.0, maxtime = 40.0, k_arr=collect(LinRange(0.01,1.0,100)) )

    listener_position = [-10.0,0.0]
    shape = TimeOfFlight(listener_position,maxtime)

    seed = MersenneTwister(1).seed
    particles = random_particles(volfrac, radius, shape; seed = seed)

    simulations = Vector{FrequencySimulation{Float64}}(undef,length(m))

    for i = eachindex(m)
        simulations[i] = FrequencySimulation(particles, k_arr; seed=seed, hankel_order=m[i])
    end

    return simulations
end

function plot_hankel_order_convergence(simulations)
    responses = Vector{Vector{Complex{Float64}}}(undef,length(simulations))
    m = Vector{Int64}(undef,length(simulations))

    # labels = Matrix{String}(undef,1,0)
    for i = eachindex(simulations)
        responses[i] = reshape(simulations[i].response, size(simulations[i].response,1))
        m[i] = simulations[i].hankel_order
        labels = [labels "m = $(m[i])"]
    end

    error = [r .- responses[end] for r in responses[1:end-1]]
    integrated_error = norm.(error).*map(m->((m.k_arr[end]-m.k_arr[1])/length(m.k_arr)),simulations[1:end-1])

    colors = reshape(LinRange(RGB(0.6,1,0.6),RGB(0,0.4,0),length(m)),1,length(m))
    realcolors = reshape(LinRange(RGB(0.6,0.6,1),RGB(0,0,0.4),length(m)),1,length(m))
    imagcolors = reshape(LinRange(RGB(1,0.6,0.6),RGB(0.4,0,0),length(m)),1,length(m))

    absvec(v) = abs.(v)
    plot(
        plot(simulations[end],0.5),
        plot(simulations[1].k_arr, [real(responses),imag(responses)],
             labels=[ map(c -> "real "*c,labels)  map(c -> "imag "*c,labels)],
             xguide ="Wavenumber (k)", yguide ="Response", linecolor=[realcolors imagcolors]
        ),
        plot(simulations[1].k_arr, absvec.(error),
             yscale=:log10, labels=labels, linecolor=colors,
             xguide ="Wavenumber (k)", yguide ="Absolute error",
        ),
        plot(m[1:end-1], integrated_error,
             yscale=:log10, legend=false,
             xguide ="Hankel order", yguide ="\$L^2\$ integrated error",
        )
    )

end


#################################################################
# Randomized estimators of elementary linear algebraic quantities
#################################################################

export rcond, reigmax, reigmin, rnorm, rnorms

@doc doc"""
Randomized Gaussian matrices normalized by column

Input:
    `el`: element type
    `m`: number of rows
    `n`: number of columns or nothing
    `normalize`: whether or not to normalize (default: `true`)

Output:
    `Ω`: a matrix of dimensions `m` x `n` containing Gaussian random numbers of
    type `el`.
""" ->
function randnn(el, m::Int, normalize::Bool=true)
    if el <: Real
        Ω = randn(m)
    elseif el <: Complex
        Ω = randn(m) + im*randn(m)
    else
        throw(ValueError("Unsupported element type: $el"))
    end
    normalize ? Ω/norm(Ω) : Ω
end
function randnn(el, m::Int, n::Int, normalize::Bool=true)
    if el <: Real
        Ω = randn(m, n)
    elseif el <: Complex
        Ω = randn(m, n) + im*randn(m, n)
    else
        throw(ValueError("Unsupported element type: $el"))
    end
    normalize || return Ω
    for i=1:n
        Ω[:, i] /= norm(sub(Ω, :, i))
    end
    Ω
end

@doc doc"""
Randomized matrix norm estimator

Computes a probabilistic upper bound on the norm of a matrix `A`.

\cite[Lemma 4.1]{Halko2011} states (with slight notational change) that

‖A‖ ≤ α √(2/π) maxᵢ ‖Aωᵢ‖

with probability $$p=α^{-r}$$.

Inputs:
    `A`: Matrix whose norm to estimate
    `r`: Number of matrix-vector products to compute
    `p`: Probability of upper bound failing (default: 0.05)

Output:
    Estimate of ‖A‖.

See also:
    `rnorms()` for a different estimator that uses
    premultiplying by both `A` and `A'`
""" ->
function rnorm(A, r::Int, p::Real=0.05)
    @assert 0<p≤1
    α = p^(-1.0/r)
    m, n = size(A)
    Ω = randnn(eltype(A), n, r, false)
    AΩ = A*Ω
    mx = maximum([norm(sub(AΩ, :, j)) for j=1:r])
    α * √(2/π) * mx
end

@doc doc"""
Randomized matrix norm estimator using `A'A`

Computes a probabilistic upper bound on the norm of a matrix `A`.

\cite[Appendix]{Liberty2007} states (with minor change in notation) that

ρ = √(‖(A'A)ʲω‖/‖(A'A)ʲ⁻¹ω‖)

which is an estimate of the spectral norm of A produced by j
steps of the power method starting with normalized ω, is a lower
bound on the true norm by a factor

ρ ≤ α ‖A‖

with probability greater than $$1 - p$$, where
$$p = 4\sqrt{n/(j-1)} α^{-2j}$$.

Inputs:
    `A`: Matrix whose norm to estimate
    `j`: Number of power iterations to perform. (Default: 1)
    `p`: Probability of upper bound failing. (Default: 0.05)
    `At` (optional keyword): Transpose of `A`. (Default: `A'`)

Output:
    Estimate of ‖A‖.

Reference:
    Appendix of \cite{Liberty2007}.

    @article{Liberty2007,
        authors = {Edo Liberty and Franco Woolfe and Per-Gunnar Martinsson
	    and Vladimir Rokhlin and Mark Tygert},
        title = {Randomized algorithms for the low-rank approximation of matrices},
        journal = {Proceedings of the National Academy of Sciences},
	volume = {104},
	issue = {51},
	year = {2007},
	pages = {20167--20172},
        doi  = {10.1073/pnas.0709640104}
    }

Comment:
    see `rnorm()` for a different estimator that does not require
    premultiplying by `A'`
""" ->
function rnorms(A, j::Int=1, p::Real=0.05; At = A')
    @assert 0<p≤1
     m, n = size(A)
    α = ((j-1)/n*(p/4)^2)^(-1/(4j))

    Ωold = Ω = randnn(eltype(A), n)
    for i=1:j #Power iterations
        Ω, Ωold = At*(A*Ω), Ω
    end
    ρ = √(norm(Ω)/norm(Ωold))
    α*ρ
end

@doc doc"""
Randomized condition number estimator

Inputs:
    `A`: Matrix whose condition number to estimate.
         Must be square and support premultiply (`A*⋅`) and solve (`A\⋅`)
    `k`: Number of power iterations to run. (Default: 1, recommended: `k ≤ 3`)
    `p`: Probability that estimate fails to hold as an upper bound
       (Default: 0.05)

Output:

    The interval `(x, y)` which contains `κ(A)` with probability $$1 - p$$.

Implementation note:

    \cite{Dixon1983} originally describes this as a computation that
    can be done by computing the necessary number of power iterations given p
    and the desired accuracy parameter θ=y/x. However, these bounds were only
    derived under the assumptions of exact arithmetic. Empirically, k≥4 has
    been seen to result in incorrect results in that the computed interval does
    not contain the true condition number. This implemention therefore makes `k`
    an explicitly user-controllable parameter from which to infer the accuracy
    parameter and hence the interval containing κ(A).

Reference:
    \cite[Theorem 2]{Dixon1983}

    @article{Dixon1983,
        author = {Dixon, John D},
        doi = {10.1137/0720053},
        journal = {SIAM Journal on Numerical Analysis},
        number = {4},
        pages = {812--814},
	title = {Estimating Extremal Eigenvalues and Condition Numbers of
		Matrices},
	volume = {20},
        year = {1983}
    }
""" ->
function rcond(A, k::Int=1, p::Real=0.05)
    @assert 0<p≤1
    m, n = size(A)
    @assert m==n
    θ = (8n/(π*p^2))^(1/k)
    x = randnn(eltype(A), n)
    for i=1:k
        x = A*x
    end
    y = randnn(eltype(A), n)
    for i=1:k
        y = A\y
    end
    φ = ((x⋅x)*(y⋅y))^(1/(2k))
    (φ, θ*φ)
end

@doc doc"""
Randomized maximal eigenvalue estimator

Inputs:

    `A`: Matrix whose maximal eigenvalue to estimate.
         Must be square and support premultiply (`A*⋅`)
    `k`: Number of power iterations to run. (Default: 1, recommended: k ≤ 3)
    `p`: Probability that estimate fails to hold as an upper bound
       (Default: 0.05)

Output:

    The interval `(x, y)` which contains the maximal eigenvalue of `A` with
    probability $$1 - p$$.

Reference:

    \cite[Corollary of Theorem 1]{Dixon1983}.
""" ->
function reigmax(A, k::Int=1, p::Real=0.05)
    @assert 0<p≤1
    m, n = size(A)
    @assert m==n
    θ = (2n/(π*p^2))^(1/k)
    y = x = randnn(eltype(A), n)
    for i=1:k
        x = A*x
    end
    φ = y⋅x
    (φ, θ*φ)
end

@doc doc"""
Randomized minimal eigenvalue estimator

Inputs:

    `A`: Matrix whose minimal eigenvalue to estimate.
         Must be square and support backslash (`A\⋅`)
    `k`: Number of power iterations to run. (Default: 1, recommended: k ≤ 3)
    `p`: Probability that estimate fails to hold as an upper bound
         (Default: 0.05)

Output:

    The interval `(x, y)` which contains the minimal eigenvalue of `A` with
    probability $$1 - p$$.

Reference:

    \cite[Corollary of Theorem 1]{Dixon1983}.
""" ->
function reigmin(A, k::Int=1, p::Real=0.05)
    @assert 0<p≤1
    m, n = size(A)
    @assert m==n
    θ = (2n/(π*p^2))^(1/k)
    y = x = randnn(eltype(A), n)
    for i=1:k
        x = A\x
    end
    φ = y⋅x
    (inv(θ*φ), inv(φ))
end

@doc doc"""
A subsampled random Fourier transform

Parameter:

    l :: Number of vectors to return
""" ->
immutable srft{T<:Integer}
    l :: T
end

(*)(A::AbstractMatrixFcn, B::srft) = error("method only defined to avoid ambiguity. If you need this method please open a pull request")

@doc doc"""
Applies a subsampled random Fourier transform to the columns of `A`

Inputs:

    `A`: A matrix to transform
    `Ω`: A `srft` type

Output:
    `B`: A matrix of dimensions size(A,1) x Ω.l

Reference:

    \[Equation 4.6]{Halko2011}
""" ->
function *(A, Ω::srft)
    m, n = size(A)
    B = A*Diagonal(exp(2π*im*rand(n))/√Ω.l)
    B = vcat([fft(A[i,:]) for i=1:m]...) #Factor of √n cancels out
    B[:, randperm(n)[1:Ω.l]]
end


using Printf
using Logging
using Oceananigans

Logging.global_logger(OceananigansLogger())

function simulate_lid_driven_cavity(; Re, N, end_time)
    topology = (Flat, Bounded, Bounded)
    domain = (y=(0, 1), z=(0, 1))
    grid = RectilinearGrid(topology=topology, size=(N, N); domain...)

    v_bcs = FieldBoundaryConditions(top = ValueBoundaryCondition(1),
                                    bottom = ValueBoundaryCondition(0))

    w_bcs = FieldBoundaryConditions(north = ValueBoundaryCondition(0),
                                    south = ValueBoundaryCondition(0))

    model = NonhydrostaticModel(
                       grid = grid,
                   buoyancy = nothing,
                    tracers = nothing,
                   coriolis = nothing,
        boundary_conditions = (v=v_bcs, w=w_bcs),
                    closure = ScalarDiffusivity(ν=1/Re)
    )

    u, v, w = model.velocities
    ζ = Field(∂y(w) - ∂z(v))

    fields = (; v, w, ζ)
    global_attributes = Dict("Re" => Re)
    output_attributes = Dict("ζ" => Dict("longname" => "vorticity", "units" => "1/s"))
    field_output_writer =
        NetCDFOutputWriter(model, fields, filepath="lid_driven_cavity_Re$Re.nc", schedule=TimeInterval(0.1),
                           global_attributes=global_attributes, output_attributes=output_attributes)

    max_Δt = 0.25 * model.grid.Δyᵃᶜᵃ^2 * Re / 2  # Make sure not to violate diffusive CFL.
    wizard = TimeStepWizard(cfl=0.1, Δt=1e-6, max_change=1.1, max_Δt=max_Δt)

    cfl = AdvectiveCFL(wizard)
    dcfl = DiffusiveCFL(wizard)

    simulation = Simulation(model, Δt=wizard, stop_time=end_time, progress=print_progress,
                            iteration_interval=20, parameters=(cfl=cfl, dcfl=dcfl))

    simulation.output_writers[:fields] = field_output_writer

    run!(simulation)

    return simulation
end

function print_progress(simulation)
    model = simulation.model
    cfl, dcfl = simulation.parameters

    # Calculate simulation progress in %.
    progress = 100 * (model.clock.time / simulation.stop_time)

    # Find maximum velocities.
    vmax = maximum(abs, model.velocities.v)
    wmax = maximum(abs, model.velocities.w)

    i, t = model.clock.iteration, model.clock.time
    @info @sprintf("[%06.2f%%] i: %d, t: %.3f, U_max: (%.2e, %.2e), CFL: %.2e, dCFL: %.2e, next Δt: %.2e",
                   progress, i, t, vmax, wmax, cfl(model), dcfl(model), simulation.Δt.Δt)

    return nothing
end

simulate_lid_driven_cavity(Re=100,   N=128, end_time=15)
simulate_lid_driven_cavity(Re=400,   N=128, end_time=20)
simulate_lid_driven_cavity(Re=1000,  N=128, end_time=25)
simulate_lid_driven_cavity(Re=3200,  N=128, end_time=50)
simulate_lid_driven_cavity(Re=5000,  N=256, end_time=50)
simulate_lid_driven_cavity(Re=7500,  N=256, end_time=75)
simulate_lid_driven_cavity(Re=10000, N=256, end_time=100)


using Base
using Test
using LinearAlgebra
using Surrogates

@testset "Radial" begin

    @testset "1D" begin

        #WITHOUT ADD_POINT
        x = [1.0,2.0,3.0]
        y = [4.0,5.0,6.0]
        a = 0
        b = 4
        my_rad = RadialBasis(x,y,a,b,z->norm(z),1)
        est = my_rad(3.0)
        @test est ≈ 7.875

        #WITH ADD_POINT, adding singleton
        x = [1.0,2.0,3.0]
        y = [4.0,5.0,6.0]
        a = 0
        b = 4
        my_rad = RadialBasis(x,y,a,b,z->norm(z),1)
        add_point!(my_rad,4.0,10.0)
        est = my_rad(3.0)
        @test est ≈ 6.499999999999991

        #WITH ADD_POINT, adding more

        x = [1.0,2.0,3.0]
        y = [4.0,5.0,6.0]
        a = 0
        b = 4
        my_rad = RadialBasis(x,y,a,b,z->norm(z),1)
        est_rad = my_rad(3.0)
        add_point!(my_rad,[3.2,3.3,3.4],[8.0,9.0,10.0])
        est_rad2 = my_rad(3.0)
        @test est_rad2 ≈ 6.49416593818781

    end
    @testset "ND" begin
        #WITHOUT ADD_POINT

        x = [(1.0,2.0,3.0),(4.0,5.0,6.0),(7.0,8.0,9.0)]
        y = [4.0,5.0,6.0]
        bounds = [[0.0,4.0],[3.0,7.0],[6.0,10.0]]
        my_rad = RadialBasis(x,y,bounds,z->norm(z),1)
        est = my_rad((1.0,2.0,3.0))
        @test est ≈ 4.0

        #WITH ADD_POINT, adding singleton
        x = [(1.0,2.0,3.0),(4.0,5.0,6.0),(7.0,8.0,9.0)]
        y = [4.0,5.0,6.0]
        bounds = [[0.0,4.0],[3.0,7.0],[6.0,10.0]]
        my_rad = RadialBasis(x,y,bounds,z->norm(z),1)
        add_point!(my_rad,(9.0,10.0,11.0),10.0)
        est = my_rad((1.0,2.0,3.0))
        @test est ≈ 4.0

        #WITH ADD_POINT, adding more

        x = [(1.0,2.0,3.0),(4.0,5.0,6.0),(7.0,8.0,9.0)]
        y = [4.0,5.0,6.0]
        bounds = [[0.0,4.0],[3.0,7.0],[6.0,10.0]]
        my_rad = RadialBasis(x,y,bounds,z->norm(z),1)
        add_point!(my_rad,[(9.0,10.0,11.0),(12.0,13.0,14.0)],[10.0,11.0])
        est = my_rad((1.0,2.0,3.0))
        @test est ≈ 4.0

    end
end