package ln

import (
	"math"
	"math/rand"
)

struct Sphere
    Center::Vector
    Radius::Float64
    Box::Box
end

type Sphere struct {
	Center Vector
	Radius float64
	Box    Box
}

func NewSphere(center Vector, radius float64) *Sphere {
	min := Vector{center.X - radius, center.Y - radius, center.Z - radius}
	max := Vector{center.X + radius, center.Y + radius, center.Z + radius}
	box := Box{min, max}
	return &Sphere{center, radius, box}
}

func (s *Sphere) Compile() {
}

func (s *Sphere) BoundingBox() Box {
	return s.Box
}

func (s *Sphere) Contains(v Vector, f float64) bool {
	return v.Sub(s.Center).Length() <= s.Radius+f
}

func (s *Sphere) Intersect(r Ray) Hit {
	radius := s.Radius
	to := r.Origin.Sub(s.Center)
	b := to.Dot(r.Direction)
	c := to.Dot(to) - radius*radius
	d := b*b - c
	if d > 0 {
		d = math.Sqrt(d)
		t1 := -b - d
		if t1 > 1e-2 {
			return Hit{s, t1}
		}
		t2 := -b + d
		if t2 > 1e-2 {
			return Hit{s, t2}
		}
	}
	return NoHit
}

func (s *Sphere) Paths4() Paths {
	var paths Paths
	var seen []Vector
	var radii []float64
	for i := 0; i < 140; i++ {
		var v Vector
		var m float64
		for {
			v = RandomUnitVector()
			m = rand.Float64()*0.25 + 0.05
			ok := true
			for i, other := range seen {
				threshold := m + radii[i] + 0.02
				if other.Sub(v).Length() < threshold {
					ok = false
					break
				}
			}
			if ok {
				seen = append(seen, v)
				radii = append(radii, m)
				break
			}
		}
		p := v.Cross(RandomUnitVector()).Normalize()
		q := p.Cross(v).Normalize()
		n := rand.Intn(4) + 1
		for k := 0; k < n; k++ {
			var path Path
			for j := 0; j <= 360; j += 5 {
				a := Radians(float64(j))
				x := v
				x = x.Add(p.MulScalar(math.Cos(a) * m))
				x = x.Add(q.MulScalar(math.Sin(a) * m))
				x = x.Normalize()
				x = x.MulScalar(s.Radius).Add(s.Center)
				path = append(path, x)
			}
			paths = append(paths, path)
			m *= 0.75
		}
	}
	return paths
}

func (s *Sphere) Paths3() Paths {
	var paths Paths
	for i := 0; i < 20000; i++ {
		v := RandomUnitVector()
		v = v.MulScalar(s.Radius).Add(s.Center)
		paths = append(paths, Path{v, v})
	}
	return paths
}

func (s *Sphere) Paths2() Paths {
	var equator Path
	for lng := 0; lng <= 360; lng++ {
		v := LatLngToXYZ(0, float64(lng), s.Radius)
		equator = append(equator, v)
	}
	var paths Paths
	for i := 0; i < 100; i++ {
		m := Identity()
		for j := 0; j < 3; j++ {
			v := RandomUnitVector()
			m = m.Rotate(v, rand.Float64()*2*math.Pi)
		}
		m = m.Translate(s.Center)
		paths = append(paths, equator.Transform(m))
	}
	return paths
}

func (s *Sphere) Paths() Paths {
	var paths Paths
	n := 10
	o := 10
	for lat := -90 + o; lat <= 90-o; lat += n {
		var path Path
		for lng := 0; lng <= 360; lng++ {
			v := LatLngToXYZ(float64(lat), float64(lng), s.Radius).Add(s.Center)
			path = append(path, v)
		}
		paths = append(paths, path)
	}
	for lng := 0; lng <= 360; lng += n {
		var path Path
		for lat := -90 + o; lat <= 90-o; lat++ {
			v := LatLngToXYZ(float64(lat), float64(lng), s.Radius).Add(s.Center)
			path = append(path, v)
		}
		paths = append(paths, path)
	}
	return paths
}

func LatLngToXYZ(lat, lng, radius float64) Vector {
	lat, lng = Radians(lat), Radians(lng)
	x := radius * math.Cos(lat) * math.Cos(lng)
	y := radius * math.Cos(lat) * math.Sin(lng)
	z := radius * math.Sin(lat)
	return Vector{x, y, z}
}

type OutlineSphere struct {
	Sphere
	Eye Vector
	Up  Vector
}

func NewOutlineSphere(eye, up, center Vector, radius float64) *OutlineSphere {
	sphere := NewSphere(center, radius)
	return &OutlineSphere{*sphere, eye, up}
}

func (s *OutlineSphere) Paths() Paths {
	var path Path
	center := s.Sphere.Center
	radius := s.Sphere.Radius

	hyp := center.Sub(s.Eye).Length()
	opp := radius
	theta := math.Asin(opp / hyp)
	adj := opp / math.Tan(theta)
	d := math.Cos(theta) * adj
	r := math.Sin(theta) * adj

	w := center.Sub(s.Eye).Normalize()
	u := w.Cross(s.Up).Normalize()
	v := w.Cross(u).Normalize()
	c := s.Eye.Add(w.MulScalar(d))
	for i := 0; i <= 360; i++ {
		a := Radians(float64(i))
		p := c
		p = p.Add(u.MulScalar(math.Cos(a) * r))
		p = p.Add(v.MulScalar(math.Sin(a) * r))
		path = append(path, p)
	}
	return Paths{path}
}


# distritbutions
@rule typeof(+)(:in1, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_in2::UnivariateNormalDistributionsFamily) = begin
    min2, vin2 = mean_var(m_in2)
    mout, vout = mean_var(m_out)
    return NormalMeanVariance(mout - min2, vout + vin2)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::MultivariateNormalDistributionsFamily, m_in2::MultivariateNormalDistributionsFamily) = begin
    min2, vin2 = mean_cov(m_in2)
    mout, vout = mean_cov(m_out)
    return MvNormalMeanCovariance(mout - min2, vout + vin2)
end

# PointMass
@rule typeof(+)(:in1, Marginalisation) (m_out::PointMass, m_in2::PointMass) = PointMass(mean(m_out) - mean(m_in2))

@rule typeof(+)(:in1, Marginalisation) (m_out::PointMass, m_in2::NormalMeanPrecision) = begin
    min2, win2 = mean_precision(m_in2)
    return NormalMeanPrecision(mean(m_out) - min2, win2)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::NormalMeanPrecision, m_in2::PointMass) = begin
    mout, wout = mean_precision(m_out)
    return NormalMeanPrecision(mout - mean(m_in2), wout)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::PointMass, m_in2::MvNormalMeanPrecision) = begin
    min2, win2 = mean_precision(m_in2)
    return MvNormalMeanPrecision(mean(m_out) - min2, win2)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::MvNormalMeanPrecision, m_in2::PointMass) = begin
    mout, wout = mean_precision(m_out)
    return MvNormalMeanPrecision(mout - mean(m_in2), wout)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::PointMass, m_in2::UnivariateNormalDistributionsFamily) = begin
    min2, vin2 = mean_var(m_in2)
    return NormalMeanVariance(mean(m_out) - min2, vin2)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_in2::PointMass) = begin
    mout, vout = mean_var(m_out)
    return NormalMeanVariance(mout - mean(m_in2), vout)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::PointMass, m_in2::MultivariateNormalDistributionsFamily) = begin
    min2, vin2 = mean_cov(m_in2)
    return MvNormalMeanCovariance(mean(m_out) - min2, vin2)
end

@rule typeof(+)(:in1, Marginalisation) (m_out::MultivariateNormalDistributionsFamily, m_in2::PointMass) = begin
    mout, vout = mean_cov(m_out)
    return MvNormalMeanCovariance(mout - mean(m_in2), vout)
end


include("common.jl")
using LinearAlgebra: rank, eigvals

function balanced_ρ(p, grav, mesh, elm)
  zdg = mesh.zdg
  Nq, K = size(zdg)
  𝟙 = ones(Bool, Nq)

  pdg = reshape(mesh.scatter * p, Nq, K)
  Q = elm.Q

  # Form the RHS vector
  pdg_diff = similar(zdg)
  Fρdg = similar(zdg, Nq * K, Nq * K)
  for e in 1:K
    p̄ = (pdg[:, e] .+ pdg[:, e]') / 2
    pdg_diff[:, e] = ((Q - Q') .* p̄) * 𝟙
  end
  # BCs hack
  pdg_diff[1,1] -= p[1,1]
  pdg_diff[end,end] += p[end,end]
  p_diff = mesh.scatter' * pdg_diff[:]
  @show p_diff

  # Form the ρsystem
  T = eltype(zdg)
  Fρdg = zeros(T, Nq * K, Nq * K)
  for e in 1:K
    rng = (e - 1) * Nq .+ (1: Nq)
    Δϕ = grav * (zdg[:, e]' .- zdg[:, e]) / 2
    A = (Q - Q') .* Δϕ
    Fρdg[rng, rng] .= ((Diagonal(A * 𝟙) + A) / 2 )
  end
  Fρ = mesh.scatter' * Fρdg * mesh.scatter
  ρ_bal = -Fρ \ p_diff
  return ρ_bal
end

let
  T = Float64
  K = 10
  N = 5
  #
  elm = element_operators(N)
  mesh = create_mesh(K, elm, T(0), T(30e4))
  zcg = mesh.zcg
  #
  Nq = N + 1
  #
  params = Params{T}()
  p = similar(zcg)
  ρ = similar(zcg)
  for (i, z) in enumerate(zcg)
    (p[i], ρ[i]) = decaying_temperature_profile(params, z)
  end
  #
  ρdg = reshape(mesh.scatter * ρ, Nq, K)
  pdg = reshape(mesh.scatter * p, Nq, K)
  #
  vdg = fill!(similar(ρdg), 0)
  #
  Q = elm.Q
  zdg = mesh.zdg
  grav = params.grav
  #
  pdg_diff_0 = fill!(similar(pdg), 0)
  ρdg_diff_0 = fill!(similar(ρdg), 0)
  for e in 1:K
    for i in 1:Nq, j in 1:Nq
      p_i, p_j = pdg[i, e], pdg[j, e]
      ρ_i, ρ_j = ρdg[i, e], ρdg[j, e]
      ϕ_i, ϕ_j = grav * zdg[i, e], grav * zdg[j, e]
      # flux terms
      ρ̄ = (ρ_i + ρ_j) / 2
      Δϕ = (ϕ_j - ϕ_i) / 2
      p̄ = (p_i + p_j) / 2
      # flux differences
      pdg_diff_0[i, e] += (Q[i, j] - Q[j, i]) * p̄
      ρdg_diff_0[i, e] += (Q[i, j] - Q[j, i]) * ρ̄ * Δϕ
    end
  end
  # BCs hack
  pdg_diff_0[1,1] -= p[1,1]
  pdg_diff_0[end,end] += p[end,end]
  #
  # @show extrema(mesh.scatter'*(pdg_diff_0 + ρdg_diff_0)[:])
  #
  # Remove the inner for loops
  pdg_diff_1 = fill!(similar(pdg), 0)
  ρdg_diff_1 = fill!(similar(ρdg), 0)
  𝟙 = ones(Bool, Nq)
  for e in 1:K
    p̄ = (pdg[:, e] .+ pdg[:, e]') / 2
    ρ̄ = (ρdg[:, e] .+ ρdg[:, e]') / 2
    Δϕ = grav * (zdg[:, e]' .- zdg[:, e]) / 2
    pdg_diff_1[:, e] = ((Q - Q') .* p̄) * 𝟙
    ρdg_diff_1[:, e] = ((Q - Q') .* ρ̄ .* Δϕ) * 𝟙
  end
  # BCs hack
  pdg_diff_1[1,1] -= p[1,1]
  pdg_diff_1[end,end] += p[end,end]
  @assert all(pdg_diff_0 .≈ pdg_diff_1)
  @assert all(ρdg_diff_0 .≈ ρdg_diff_1)
  #
  # Build the ρ matrix
  Fρdg = zeros(T, Nq * K, Nq * K)
  for e in 1:K
    rng = (e - 1) * Nq .+ (1: Nq)
    Δϕ = grav * (zdg[:, e]' .- zdg[:, e]) / 2
    A = (Q - Q') .* Δϕ
    Fρdg[rng, rng] .= ((Diagonal(A * 𝟙) + A) / 2 )
  end
  ρdg_diff_2 = reshape(Fρdg * ρdg[:], Nq, K)
  @assert all(ρdg_diff_0 .≈ ρdg_diff_2)
  #
  Fρ = mesh.scatter' * Fρdg * mesh.scatter
  ρ_bal = -Fρ \ (mesh.scatter' * pdg_diff_0[:])
  ρdg_bal = reshape(mesh.scatter * ρ_bal, Nq, K)
  display(ρdg)
  display(ρdg_bal)
  display(ρdg_bal - ρdg)
  ρ_bal2 = balanced_ρ(p, grav, mesh, elm)
  @show extrema(ρ_bal2 - ρ_bal)
  nothing
end


immutable LocalLevelFit{T}
    residual_std::T
    locallevel_innovation_std::T
    regression_coefs::Vector{T}
    transform_params::Vector{T}
    y::Vector{T}
    y_residuals::Vector{T}
    locallevel::Vector{T}
    locallevel_std::Vector{T}
    locallevel_innovations::Vector{T}
end #LocalLevelFit

#TODO
#function Base.show(llf::LocalLevelFit)
#end #show

function fit{T}(y::Vector{T}, X::Matrix{T}=zeros(length(y),0))

    n   = length(y) # Number of observations
    @assert n == size(X, 1)
    notmissing = !isnan(y)

    X   = X'
    m   = size(X,1) # Number of regressors

    nₓ  = 1 + m
    Q = zeros(nₓ, nₓ)
    C = [ones(1, n); X] .* notmissing'

    yclean = y .* notmissing

    function ll(params::Vector{T})
        Q[1] = exp(2params[2])
        return likelihood(yclean, C, exp(2params[1]), Q)
    end #ll

    params = optimize(ll, [0., 0.]).minimum

    #println(exp(params))
    σϵ², σμ² = exp(2params)
    y_residuals, x, Pₓ, x_residuals = smooth(yclean, C, σϵ², σμ²)

    return LocalLevelFit(sqrt(σϵ²), sqrt(σμ²), x[2:end,end], zeros(0), y, y_residuals, vec(x[1,:]), vec(Pₓ[1,1,:]), vec(x_residuals[1,:]))

end #fit

function fit{T}(y::Vector{T}, f::Function, p₀::Vector{T}, X::Matrix{T})

    n   = length(y) # Number of observations
    m   = size(f(p₀, X), 2) # Number of regressors
    @assert n == size(X, 1)

    notmissing = !isnan(y)

    nₓ  = 1 + m
    Q = zeros(nₓ, nₓ)
    yclean = y .* notmissing
    c = ones(1, n)

    function ll(params::Vector{T})
        C = [c; f(params[3:end], X)'] .* notmissing'
        Q[1] = exp(2params[2])
        return likelihood(yclean, C, exp(2params[1]), Q)
    end #ll

    params = optimize(ll, [0.; 0.; p₀]).minimum

    σϵ², σμ² = exp(2params[1:2])
    p = params[3:end]
    fX = f(p, X)'
    C = [ones(1, n); fX] .* notmissing'
    y_residuals, x, Pₓ, x_residuals = smooth(yclean, C, σϵ², σμ²)

    return LocalLevelFit(sqrt(σϵ²), sqrt(σμ²), x[2:end,end], p, y, y_residuals, vec(x[1,:]), vec(Pₓ[1,1,:]), vec(x_residuals[1,:]))

end #fit


struct Biped{T}
    n::Int
    m::Int
    d::Int

    g::T
    μ::T

    # torso
    l1
    d1
    m1
    J1

    # leg 1
        # thigh
    l2
    d2
    m2
    J2
        # calf
    l3
    d3
    m3
    J3

    # leg 2
        # thigh
    l4
    d4
    m4
    J4
        # calf
    l5
    d5
    m5
    J5

    # joint limits
    qL
    qU

    # torque limits
    uL
    uU

    nq
    nu
    nc
    nf
    nb
    ns

    idx_u
    idx_λ
    idx_b
    idx_ψ
    idx_η
    idx_s
end

# Dimensions
nq = 2 + 5                # configuration dimension
nu = 4                    # control dimension
nc = 2                    # number of contact points
nf = 2                    # number of parameters for friction cone
nb = nc * nf
ns = 1

# World parameters
μ = 1.0      # coefficient of friction
g = 9.81     # gravity

# Model parameters
	# from Robust sliding-mode control applied to a 5-link biped robot
m_torso = 14.79
m_thigh = 5.28
m_leg = 2.23

J_torso = 3.30e-2
J_thigh = 3.30e-2
J_leg = 3.30e-2

l_torso = 0.486
l_thigh = 0.302
l_leg = 0.332

d_torso = 0.282
d_thigh = 0.236
d_leg = 0.189

n = 2 * nq
m = nu + nc + nb + nc + nb + ns
d = 0

idx_u = (1:nu)
idx_λ = nu .+ (1:nc)
idx_b = nu + nc .+ (1:nb)
idx_ψ = nu + nc + nb .+ (1:nc)
idx_η = nu + nc + nb + nc .+ (1:nb)
idx_s = nu + nc + nb + nc + nb .+ (1:ns)

qL = -Inf * ones(nq)
qU = Inf * ones(nq)

uL = -Inf * ones(nu)
uU = Inf * ones(nu)

function kinematics_1(model::Biped, q; body = :torso, mode = :ee)
	x = q[1]
	z = q[2]

	if body == :torso
		l = model.l1
		d = model.d1
		θ = q[3]
	elseif body == :thigh_1
		l = model.l2
		d = model.d2
		θ = q[4]
	elseif body == :thigh_2
		l = model.l4
		d = model.d4
		θ = q[6]
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return [x + l * sin(θ); z - l * cos(θ)]
	elseif mode == :com
		return [x + d * sin(θ); z - d * cos(θ)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_1(model::Biped, q; body = :torso, mode = :ee)
	# jac = [1.0 0.0 0 0.0 0.0 0.0 0.0;
	# 	   0.0 1.0 0 0.0 0.0 0.0 0.0]
	jac = zeros(eltype(q), 2, model.nq)
	jac[1, 1] = 1.0
	jac[2, 2] = 1.0
	if body == :torso
		r = mode == :ee ? model.l1 : model.d1
		θ = q[3]
		jac[1, 3] = r * cos(θ)
		jac[2, 3] = r * sin(θ)
	elseif body == :thigh_1
		r = mode == :ee ? model.l2 : model.d2
		θ = q[4]
		jac[1, 4] = r * cos(θ)
		jac[2, 4] = r * sin(θ)
	elseif body == :thigh_2
		r = mode == :ee ? model.l4 : model.d4
		θ = q[6]
		jac[1, 6] = r * cos(θ)
		jac[2, 6] = r * sin(θ)
	else
		@error "incorrect body specification"
	end

	return jac
end

function kinematics_2(model::Biped, q; body = :leg_1, mode = :ee)

	if body == :leg_1
		p = kinematics_1(model, q, body = :thigh_1, mode = :ee)

		θa = q[4]
		θb = q[5]

		lb = model.l3
		db = model.d3
	elseif body == :leg_2
		p = kinematics_1(model, q, body = :thigh_2, mode = :ee)

		θa = q[6]
		θb = q[7]

		lb = model.l5
		db = model.d5
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return p + [lb * sin(θa + θb); -1.0 * lb * cos(θa + θb)]
	elseif mode == :com
		return p + [db * sin(θa + θb); -1.0 * db * cos(θa + θb)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_2(model::Biped, q; body = :leg_1, mode = :ee)

	if body == :leg_1
		jac = jacobian_1(model, q, body = :thigh_1, mode = :ee)

		θa = q[4]
		θb = q[5]

		r = mode == :ee ? model.l3 : model.d3

		jac[1, 4] += r * cos(θa + θb)
		jac[1, 5] += r * cos(θa + θb)

		jac[2, 4] += r * sin(θa + θb)
		jac[2, 5] += r * sin(θa + θb)
	elseif body == :leg_2
		jac = jacobian_1(model, q, body = :thigh_2, mode = :ee)

		θa = q[6]
		θb = q[7]

		r = mode == :ee ? model.l5 : model.d5

		jac[1, 6] += r * cos(θa + θb)
		jac[1, 7] += r * cos(θa + θb)

		jac[2, 6] += r * sin(θa + θb)
		jac[2, 7] += r * sin(θa + θb)
	else
		@error "incorrect body specification"
	end

	return jac
end

# # test kinematics
# q = rand(nq)
#
# norm(kinematics_1(model, q, body = :torso, mode = :ee) - [q[1] + model.l1 * sin(q[3]); q[2] - model.l1 * cos(q[3])])
# norm(kinematics_1(model, q, body = :torso, mode = :com) - [q[1] + model.d1 * sin(q[3]); q[2] - model.d1 * cos(q[3])])
# norm(kinematics_1(model, q, body = :thigh_1, mode = :ee) - [q[1] + model.l2 * sin(q[4]); q[2] - model.l2 * cos(q[4])])
# norm(kinematics_1(model, q, body = :thigh_1, mode = :com) - [q[1] + model.d2 * sin(q[4]); q[2] - model.d2 * cos(q[4])])
# norm(kinematics_1(model, q, body = :thigh_2, mode = :ee) - [q[1] + model.l4 * sin(q[6]); q[2] - model.l4 * cos(q[6])])
# norm(kinematics_1(model, q, body = :thigh_2, mode = :com) - [q[1] + model.d4 * sin(q[6]); q[2] - model.d4 * cos(q[6])])
#
# norm(kinematics_2(model, q, body = :leg_1, mode = :ee) - [q[1] + model.l2 * sin(q[4]) + model.l3 * sin(q[4] + q[5]); q[2] - model.l2 * cos(q[4]) - model.l3 * cos(q[4] + q[5])])
# norm(kinematics_2(model, q, body = :leg_1, mode = :com) - [q[1] + model.l2 * sin(q[4]) + model.d3 * sin(q[4] + q[5]); q[2] - model.l2 * cos(q[4]) - model.d3 * cos(q[4] + q[5])])
# norm(kinematics_2(model, q, body = :leg_2, mode = :ee) - [q[1] + model.l4 * sin(q[6]) + model.l5 * sin(q[6] + q[7]); q[2] - model.l4 * cos(q[6]) - model.l5 * cos(q[6] + q[7])])
# norm(kinematics_2(model, q, body = :leg_2, mode = :com) - [q[1] + model.l4 * sin(q[6]) + model.d5 * sin(q[6] + q[7]); q[2] - model.l4 * cos(q[6]) - model.d5 * cos(q[6] + q[7])])
#
# k1(z) = kinematics_1(model, z, body = :torso, mode = :ee)
# norm(ForwardDiff.jacobian(k1,q) - jacobian_1(model, q, body = :torso, mode = :ee))
#
# k1(z) = kinematics_1(model, z, body = :torso, mode = :com)
# norm(ForwardDiff.jacobian(k1,q) - jacobian_1(model, q, body = :torso, mode = :com))
#
# k1(z) = kinematics_1(model, z, body = :thigh_1, mode = :ee)
# norm(ForwardDiff.jacobian(k1,q) - jacobian_1(model, q, body = :thigh_1, mode = :ee))
#
# k1(z) = kinematics_1(model, z, body = :thigh_1, mode = :com)
# norm(ForwardDiff.jacobian(k1,q) - jacobian_1(model, q, body = :thigh_1, mode = :com))
#
# k1(z) = kinematics_1(model, z, body = :thigh_2, mode = :ee)
# norm(ForwardDiff.jacobian(k1,q) - jacobian_1(model, q, body = :thigh_2, mode = :ee))
#
# k1(z) = kinematics_1(model, z, body = :thigh_2, mode = :com)
# norm(ForwardDiff.jacobian(k1,q) - jacobian_1(model, q, body = :thigh_2, mode = :com))
#
# k2(z) = kinematics_2(model, z, body = :leg_1, mode = :ee)
# norm(ForwardDiff.jacobian(k2,q) - jacobian_2(model, q, body = :leg_1, mode = :ee))
#
# k2(z) = kinematics_2(model, z, body = :leg_1, mode = :com)
# norm(ForwardDiff.jacobian(k2,q) - jacobian_2(model, q, body = :leg_1, mode = :com))
#
# k2(z) = kinematics_2(model, z, body = :leg_2, mode = :ee)
# norm(ForwardDiff.jacobian(k2,q) - jacobian_2(model, q, body = :leg_2, mode = :ee))
#
# k2(z) = kinematics_2(model, z, body = :leg_2, mode = :com)
# norm(ForwardDiff.jacobian(k2,q) - jacobian_2(model, q, body = :leg_2, mode = :com))

# Lagrangian

function lagrangian(model::Biped, q, q̇)
	L = 0.0

	# torso
	p_torso = kinematics_1(model, q, body = :torso, mode = :com)
	J_torso = jacobian_1(model, q, body = :torso, mode = :com)
	v_torso = J_torso * q̇

	L += 0.5 * model.m1 * v_torso' * v_torso
	L += 0.5 * model.J1 * q̇[3]^2.0
	L -= model.m1 * model.g * p_torso[2]

	# thigh 1
	p_thigh_1 = kinematics_1(model, q, body = :thigh_1, mode = :com)
	J_thigh_1 = jacobian_1(model, q, body = :thigh_1, mode = :com)
	v_thigh_1 = J_thigh_1 * q̇

	L += 0.5 * model.m2 * v_thigh_1' * v_thigh_1
	L += 0.5 * model.J2 * q̇[4]^2.0
	L -= model.m2 * model.g * p_thigh_1[2]

	# leg 1
	p_leg_1 = kinematics_2(model, q, body = :leg_1, mode = :com)
	J_leg_1 = jacobian_2(model, q, body = :leg_1, mode = :com)
	v_leg_1 = J_leg_1 * q̇

	L += 0.5 * model.m3 * v_leg_1' * v_leg_1
	L += 0.5 * model.J3 * q̇[5]^2.0
	L -= model.m3 * model.g * p_leg_1[2]

	# thigh 2
	p_thigh_2 = kinematics_1(model, q, body = :thigh_2, mode = :com)
	J_thigh_2 = jacobian_1(model, q, body = :thigh_2, mode = :com)
	v_thigh_2 = J_thigh_2 * q̇

	L += 0.5 * model.m4 * v_thigh_2' * v_thigh_2
	L += 0.5 * model.J4 * q̇[6]^2.0
	L -= model.m4 * model.g * p_thigh_2[2]

	# leg 2
	p_leg_2 = kinematics_2(model, q, body = :leg_2, mode = :com)
	J_leg_2 = jacobian_2(model, q, body = :leg_2, mode = :com)
	v_leg_2 = J_leg_2 * q̇

	L += 0.5 * model.m5 * v_leg_2' * v_leg_2
	L += 0.5 * model.J5 * q̇[7]^2.0
	L -= model.m5 * model.g * p_leg_2[2]

	return L
end

function dLdq(model::Biped, q, q̇)
	Lq(x) = lagrangian(model, x, q̇)
	ForwardDiff.gradient(Lq, q)
end

function dLdq̇(model::Biped, q, q̇)
	Lq̇(x) = lagrangian(model, q, x)
	ForwardDiff.gradient(Lq̇, q̇)
end

# q̇ = rand(nq)
# lagrangian(model, q, q̇)
# dLdq(model, q, q̇)
# dLdq̇(model, q, q̇)

# Methods
function M_func(model::Biped, q)
	M = Diagonal([0.0, 0.0, model.J1, model.J2, model.J3, model.J4, model.J5])

	# torso
	J_torso = jacobian_1(model, q, body = :torso, mode = :com)
	M += model.m1 * J_torso' * J_torso

	# thigh 1
	J_thigh_1 = jacobian_1(model, q, body = :thigh_1, mode = :com)
	M += model.m2 * J_thigh_1' * J_thigh_1

	# leg 1
	J_leg_1 = jacobian_2(model, q, body = :leg_1, mode = :com)
	M += model.m3 * J_leg_1' * J_leg_1

	# thigh 2
	J_thigh_2 = jacobian_1(model, q, body = :thigh_2, mode = :com)
	M += model.m4 * J_thigh_2' * J_thigh_2

	# leg 2
	J_leg_2 = jacobian_2(model, q, body = :leg_2, mode = :com)
	M += model.m5 * J_leg_2' * J_leg_2

	return M
end

# eigen(M_func(model, q))
#
# tmp_q(z) = dLdq̇(model, z, q̇)
# tmp_q̇(z) = dLdq̇(model, q, z)
# norm(ForwardDiff.jacobian(tmp_q̇,q̇) - M_func(model, q))

function C_func(model::Biped, q, q̇)
	tmp_q(z) = dLdq̇(model, z, q̇)
	tmp_q̇(z) = dLdq̇(model, q, z)

	ForwardDiff.jacobian(tmp_q, q) * q̇ - dLdq(model, q, q̇)
end

function ϕ_func(model::Biped, q)
	p_leg_1 = kinematics_2(model, q, body = :leg_1, mode = :ee)
	p_leg_2 = kinematics_2(model, q, body = :leg_2, mode = :ee)

	@SVector [p_leg_1[2], p_leg_2[2]]
end

function B_func(model::Biped, q)
	@SMatrix [0.0 0.0 0.0 1.0 0.0 0.0 0.0;
			  0.0 0.0 0.0 0.0 1.0 0.0 0.0;
			  0.0 0.0 0.0 0.0 0.0 1.0 0.0;
			  0.0 0.0 0.0 0.0 0.0 0.0 1.0]
end

function N_func(model::Biped, q)
	J_leg_1 = jacobian_2(model, q, body = :leg_1, mode = :ee)
	J_leg_2 = jacobian_2(model, q, body = :leg_2, mode = :ee)

	return [J_leg_1[2:2,:];
			J_leg_2[2:2,:]]
end

function P_func(model::Biped, q)
	J_leg_1 = jacobian_2(model, q, body = :leg_1, mode = :ee)
	J_leg_2 = jacobian_2(model, q, body = :leg_2, mode = :ee)
	map = [1.0; -1.0]

	return [map * J_leg_1[1:1, :];
			map * J_leg_2[1:1, :]]
end

function friction_cone(model::Biped, u)
	λ = u[model.idx_λ]
	b = u[model.idx_b]
	return @SVector [model.μ * λ[1] - sum(b[1:2]),
					 model.μ * λ[2] - sum(b[3:4])]
end

function maximum_dissipation(model::Biped, x⁺, u, h)
	q3 = x⁺[model.nq .+ (1:model.nq)]
	q2 = x⁺[1:model.nq]
	ψ = u[model.idx_ψ]
	ψ_stack = [ψ[1] * ones(2); ψ[2] * ones(2)]
	η = u[model.idx_η]
	return P_func(model, q3) * (q3 - q2) / h + ψ_stack - η
end

function fd(model::Biped, x⁺, x, u, w, h, t)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2⁺ = view(x⁺, 1:model.nq)
	q2⁻ = view(x, model.nq .+ (1:model.nq))
	q1 = view(x, 1:model.nq)
	u_ctrl = view(u, model.idx_u)
	λ = view(u, model.idx_λ)
	b = view(u, model.idx_b)

    [q2⁺ - q2⁻;
    ((1.0 / h) * (M_func(model, q1) * (SVector{7}(q2⁺) - SVector{7}(q1))
    - M_func(model, q2⁺) * (SVector{7}(q3) - SVector{7}(q2⁺)))
    + transpose(B_func(model, q3)) * SVector{4}(u_ctrl)
    + transpose(N_func(model, q3)) * SVector{2}(λ)
    + transpose(P_func(model, q3)) * SVector{4}(b)
    - h * C_func(model, q3, (q3 - q2⁺) / h))]
end

model = Biped(n, m, d,
			  g, μ,
			  l_torso, d_torso, m_torso, J_torso,
			  l_thigh, d_thigh, m_thigh, J_thigh,
			  l_leg, d_leg, m_leg, J_leg,
			  l_thigh, d_thigh, m_thigh, J_thigh,
			  l_leg, d_leg, m_leg, J_leg,
			  qL, qU,
			  uL, uU,
			  nq,
			  nu,
			  nc,
			  nf,
			  nb,
			  ns,
			  idx_u,
			  idx_λ,
			  idx_b,
			  idx_ψ,
			  idx_η,
			  idx_s)

function initial_configuration(model::Biped, θ)
	q0 = zeros(model.nq)
	q0[3] = pi

	q0[4] = θ
	q0[5] = -1.0 * q0[4]
	q0[2] = model.l2*cos(q0[4]) + model.l3 * cos(q0[4] + q0[5])
	q0[6] = -1.0 * acos(q0[2] / (model.l4 + model.l5))
	q0[1] = -1.0 * model.l2 * sin(θ)

	return q0
end

function loop_configurations(model::Biped, θ)
	q1 = zeros(model.nq)
	q1[3] = pi

	q1[4] = θ
	q1[5] = -1.0 * q1[4]
	q1[2] = model.l2*cos(q1[4]) + model.l3 * cos(q1[4] + q1[5])
	q1[6] = -1.0 * acos(q1[2] / (model.l4 + model.l5))
	q1[1] = -1.0 * model.l2 * sin(θ)

	qT = zeros(model.nq)
	qT[3] = pi

	qT[6] = θ
	qT[7] = -1.0 * qT[6]
	qT[2] = model.l4*cos(qT[6]) + model.l5 * cos(qT[6] + qT[7])
	qT[4] = -1.0 * acos(qT[2] / (model.l2 + model.l3))
	qT[1] = sqrt((model.l2 + model.l3)^2.0 - qT[2]^2.0)

	return q1, qT
end

# visualization
function visualize!(vis, model::Biped, q;
      r = 0.1, Δt = 0.1)

	torso = Cylinder(Point3f0(0.0, 0.0, 0.0), Point3f0(0.0, 0.0, model.l1),
		convert(Float32, 0.025))
	setobject!(vis["torso"], torso,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	thigh_1 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l2),
		convert(Float32, 0.025))
	setobject!(vis["thigh1"], thigh_1,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	leg_1 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l3),
		convert(Float32, 0.025))
	setobject!(vis["leg1"], leg_1,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	thigh_2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l4),
		convert(Float32, 0.025))
	setobject!(vis["thigh2"], thigh_2,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	leg_2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l5),
		convert(Float32, 0.025))
	setobject!(vis["leg2"], leg_2,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q)
	for t = 1:T
		MeshCat.atframe(anim, t) do
			p = [q[t][1]; 0.0; q[t][2]]

			k_torso = kinematics_1(model, q[t], body = :torso, mode = :ee)
			p_torso = [k_torso[1], 0.0, k_torso[2]]

			k_thigh_1 = kinematics_1(model, q[t], body = :thigh_1, mode = :ee)
			p_thigh_1 = [k_thigh_1[1], 0.0, k_thigh_1[2]]

			k_leg_1 = kinematics_2(model, q[t], body = :leg_1, mode = :ee)
			p_leg_1 = [k_leg_1[1], 0.0, k_leg_1[2]]

			k_thigh_2 = kinematics_1(model, q[t], body = :thigh_2, mode = :ee)
			p_thigh_2 = [k_thigh_2[1], 0.0, k_thigh_2[2]]

			k_leg_2 = kinematics_2(model, q[t], body = :leg_2, mode = :ee)
			p_leg_2 = [k_leg_2[1], 0.0, k_leg_2[2]]

			settransform!(vis["thigh1"], cable_transform(p, p_thigh_1))
			settransform!(vis["leg1"], cable_transform(p_thigh_1, p_leg_1))
			settransform!(vis["thigh2"], cable_transform(p, p_thigh_2))
			settransform!(vis["leg2"], cable_transform(p_thigh_2, p_leg_2))
			settransform!(vis["torso"], cable_transform(p_torso,p))
		end
	end

	MeshCat.setanimation!(vis, anim)
end


using SparseArrays
using OffsetArrays
using Colors
function slurp_file(filename)
    coords, folds = split(read(filename, String), "\n\n")
    coords = [parse.(Int, t) for t in split.(split(coords, "\n"), ",")]
    folds = split(folds, "\n")
    Set(coords), folds
end

coords, folds = slurp_file("13.txt")

function fold(grid, fold_line, idx)
    result = Set{Vector{Int}}()
    for p in grid
        p[idx] = p[idx] > fold_line ? fold_line - (p[idx] - fold_line) : p[idx]
        push!(result, p)
    end
    result
end

function solve(filename)
    coords, folds = slurp_file(filename)
    for f in folds
        axis, line = match(r"fold along ([xy])=(\d+)", f)
        line = parse(Int, line)
        idx = axis == "x" ? 1 : 2
        @show axis, line
        coords = fold(coords, line, idx)
    end
    coords
end

function draw(coords)
    result = zeros(50, 50)
    result = OffsetArray(result, 0:49, 0:49)
    for idx in coords
        y, x = idx
        result[x,y] = 20
    end
    Gray.(result)
end

soln = solve("13.txt")
max(last.(soln)...)

soln

draw(soln)


"""
ap_test_cross(re, factors; cov_method = "time iid", lag = 0, test_method = "Shanken")

CROSS-SECTIONAL ASSET PRICING TESTS

Asset pricing models predict a restriction on the error terms of the second-stage cross-sectional regressions.
The asset pricing model is valid if the error terms of cross-sectional regressions are jointly zero-mean.

Two-Pass Regression Estimation Procedure
- First-stage: time-series regressions of excess returns on factors to obtain estimates of betas for each assets
- Second-stage: cross-sectional regressions (without intercept) of time-series average returns on estiamted betas to obtain risk premium for each assets

INPUTS
`re': T x N matrix of excess returns, where T is the number of periods for each test assets and N is the number of test assets
`factors': T x K matrix of factors. Factors must be excess returns for time-series tests
`cov_method': method for calculating the variance-covariance matrix of the error terms:
            - time iid: error terms are iid across time but cross-sectionally correlated
            - iid: error terms are iiid both across time and assets
            - NW: Newey-West standard errors with m lag
            - HH: Hansen-Hodrick standard errors with m lags
`lag': number of lags used in the Newey-West or Hansen-Hodrick standard errors
`test_method': method of the time-series asset pricing tests:
            - OLS: OLS
            - GLS: GLS
            - Shanken: OLS with Shanken Correction (1992)
            - GLS Shanken: GLS with Shanken Correction (1992)
            - GMM: GMM exploiting the moment conditions from the first-stage time-series regressions and second-stage cross-sectional regressions. The model is overidentified
            - GLS GMM: GLS GMM exploiting the moment conditions from the first-stage time-series regressions and second-stage cross-sectional regressions. The model is overidentified

OUTPUTS
`Ere' : N x 1 vector of expected excess returns for N test assets
`lambda': K x 1 vector of risk premiums for K factors
`alpha': N x 1 vector of alphas for N test assets

`test_statistics': Wald-statistics of the asset pricing tests
`pvalue': pvalue of the test statistics. The test statistics has a Chi-square distribution.

"""

function ap_test_cross(re, factors, cov_method = "time iid", lag = 0, test_method = "Shanken")
    (T, K) = size(factors);
    N = size(re, 2);
    cov_f = cov(factors); # coviance matrix for the factors K x K

    X = [ones(T,1) factors];
    (~, beta, resid, cov_resid) = OLSFn(re, X, cov_method, lag); # First-stage time-series regressions

    Ere = mean(re, dims = 1)'; # N x 1 vector of sample average excess returns for each test assets

    # Methods of cross-sectional tests
    if test_method == "OLS" # OLS
        lambda = (beta' * beta) \ (beta' * Ere); # OLS second stage regressions - Risk premiums K x 1
        alpha = Ere- beta * lambda; # Alaphas N x 1
        Q = I - beta * ((beta' * beta) \ beta');
        cov_alpha = 1.0 ./ T * Q * cov_resid * Q;
        test_statistics = alpha' * (cov_alpha \ alpha);
        pvalue = 1.0 .- cdf.(Chisq(N - K), test_statistics);

    elseif test_method == "GLS" # GLS
        lambda = (beta' * (cov_resid \ beta)) \ beta' * (cov_resid \ Ere); # GLS second stage regressions - Risk premiums K x 1
        alpha = Ere- beta * lambda; # Alaphas N x 1
        test_statistics = T * alpha' * (cov_resid \ alpha);
        pvalue = 1.0 .- cdf.(Chisq(N - K), test_statistics);

    elseif test_method == "Shanken" # OLS Shanken Correction (1992)
        lambda = beta' * beta \ (beta' * Ere); # OLS second stage regressions - Risk premiums K x 1
        alpha = Ere- beta * lambda; # Alaphas N x 1
        Q = I - beta * ((beta' * beta) \ beta');
        cov_alpha = 1.0 ./ T * Q * cov_resid * Q .* (1.0 .+ lambda' * (cov_f \ lambda));
        test_statistics = alpha' * (cov_alpha \ alpha);
        pvalue = 1.0 .- cdf.(Chisq(N - K), test_statistics);

    elseif test_method == "GLS Shanken" # GLS Shanken Correction (1992)
        lambda = (beta' * (cov_resid \ beta)) \ beta' * (cov_resid \ Ere); # GLS second stage regressions - Risk premiums K x 1
        alpha = Ere - beta * lambda; # Alaphas N x 1
        cov_alpha = (1.0 ./ T) * (cov_resid - beta * (beta' * (cov_resid \ beta) \ beta')) .* (1.0 .+ lambda' * (cov_f \ lambda));
        test_statistics = T * (1.0 .+ lambda' * (cov_f \ lambda)) * alpha' * (cov_resid \ alpha);
        pvalue = 1.0 .- cdf.(Chisq(N - K), test_statistics);

    elseif test_method == "GMM" || test_method == "GLS GMM" # OLS GMM and GLS GMM
        lambda = beta' * beta \ (beta' * Ere); # OLS second stage regressions - Risk premiums K x 1
        alpha = Ere- beta * lambda; # Alaphas N x 1
        h = [ones(T, 1) factors];
        f = kron(ones(1, K+1), resid) .* kron(h, ones(1, N)); # Time-series moment conditions generalized to k factors
        f = [f re .- kron(ones(T, 1), (beta * lambda)')]; # Adding cross-sectional moment conditions gernerlized to k factors - p.225 in Cochrane (2000)
        d = - kron(h' * h / T, I(N)); # Jacobian of the time-seires GMM moment conditions
        d = [d zeros((1 + K) * N, K); kron([0 lambda'],I(N)) -beta]; # Adding the Jacobian of the cross-sectional GMM moment conditions

        # Weigthing matrix
        if test_method == "GMM"
            a = [I zeros((K + 1) * N, N); zeros(K, (K + 1) * N) beta']; # weighting matrix for GMM
        else
            a = [I zeros((K + 1) * N, N); zeros(K, (K + 1) * N) beta' * (cov_resid \ I)]; # weighting matrix for GLS GMM
        end

        invad = (a * d) \ I;
        premom = I(size(d * invad * a)[1]) - d * invad * a; #  for the sampling distributionof the moments - p.190 eqn (147) in Cochrane (2000)
        Ef2 = CovMoment(f, cov_method, lag); # spectral density matrix
        varb2 = (1.0 ./ T) * invad * a * Ef2 * a' * invad; # variance of the coefficients - p.190 eqn (146) in Cochrane (2000)
        varmom = (1.0 ./ T) * premom * Ef2 * premom'; # asymptotic covariance matrix of the moments - p.190 eqn (147) in Cochrane (2000)

        test_statistics = alpha' * pinv(varmom[N * (K + 1) + 1:end, N * (K + 1) + 1:end]) * alpha; # test the moment conditions that the the error terms(alphas) in the second-stage cross-sectional regressions are jointly zero
        pvalue = 1.0 .- cdf.(Chisq(N - K), test_statistics);
    end

    return(Ere, lambda, alpha, test_statistics, pvalue)
end


using ThreeBodyProblem
using LinearAlgebra
using Test

sys = earth_moon()
p = [sys.μ₁, sys.μ₂, sys.d]
rv = [0.5, 0, 0, 0, 0.5, 0]
rv_dim = [1e4, 0, 0, 0, 1, 0]
μ = sys.μ
t = 0

rvdot = zeros(6)
@test CR3BPdynamics!(rvdot, rv, sys, t) == nothing
@test rvdot == CR3BPdynamics(rv, sys, t)

rvdot = zeros(6)
@test CR3BPdynamics!(rvdot, rv_dim, p, t) == nothing
@test rvdot == CR3BPdynamics(rv_dim, p, t)

rvdot = zeros(6)
@test CR3BPdynamics!(rvdot, rv_dim, sys.μ, t) == nothing
@test rvdot == CR3BPdynamics(rv_dim, sys.μ, t)

rvdot = zeros(6)
@test CR3BPinert!(rvdot, rv, sys, t) == nothing
@test rvdot == CR3BPinert(rv, sys, t)

rvdot = zeros(6)
@test CR3BPinert!(rvdot, rv_dim, p, t) == nothing
@test rvdot == CR3BPinert(rv_dim, p, t)

wdot = zeros(42)
w = [reshape(I(6),36,1); rv]
@test CR3BPstm!(wdot, w, sys, t) == nothing
@test wdot == CR3BPstm(w, sys, t)

rvdot = zeros(6)
@test R2BPdynamics!(rvdot, rv_dim, sys.prim, t) == nothing
@test rvdot == R2BPdynamics(rv_dim, sys.prim, t)

rvdot = zeros(6)
@test R2BPdynamics!(rvdot, rv_dim, sys.μ₁, t) == nothing
@test rvdot == R2BPdynamics(rv_dim, sys.μ₁, t)

rvdot = zeros(6)
n = 1e-3
@test CWdynamics!(rvdot, rv, n, t) == nothing
@test rvdot == CWdynamics(rv, n, t)

bisys = earth_moon_sun()
rvdot = zeros(6)
@test BCPdynamics!(rvdot, rv, bisys, t) == nothing
@test rvdot == BCPdynamics(rv, bisys, t)


export costfunction, LHD

# "Create tools to compute cost function"

# Type to hold cost functions
abstract type costfunction end

# Likelihood function
struct LHD <:costfunction
	A::Array{Float64,2}
	dψ::Array{Float64,2}
	b::Array{Float64,1}
	λ::Float64
	δ::Float64
	nb::Int64
	Ne::Int64

	# Cache variables
	dS::Array{Float64,1}
	logdS::Array{Float64,1}
	dψ_dS::Array{Float64,2}

	# Store gradient and hessian
	J::Array{Float64,1}
	G::Array{Float64,1}
	H::Array{Float64,2}
end

## Create constructor for likelihood function

function LHD(A::Array{Float64,2}, dψ::Array{Float64,2}, λ::Float64, δ::Float64)
	# Number of basis = p+2
	nb = size(A,1)
	@assert size(A,1)==size(A,2) "A must be a square matrix"

	Ne = size(dψ, 1)

	# Add L-2 regularization
	A += (λ/Ne)*I

	b = zeros(nb)
	b .= sum(A, dims=2)[:,1]
	rmul!(b, δ)

	return LHD(A, dψ, b, λ, δ, nb, Ne, zeros(Ne), zeros(Ne), zeros(Ne,nb),zeros(1), zeros(nb), zeros(nb,nb))
end

function Base.show(io::IO, L::LHD)
	if L.λ==0.0 && L.δ==0.0
	println(io,"Cost function with Ne = $(L.Ne) samples")
	elseif L.λ==0.0 && L.δ!=0.0
	println(io,"Cost function with δ = $(L.δ) - regularized log barrier and Ne = $(L.Ne) samples")
	elseif L.λ!=0.0 && L.δ==0.0
	println(io,"Cost function with λ = $(L.λ) L-2 regularization and Ne = $(L.Ne) samples")
	else
	println(io,"Cost function with λ = $(L.λ) L-2 regularization,  δ = $(L.δ) - regularized log barrier and Ne = $(L.Ne) samples")
	end
end

# Function to compute the cost, gradient and loss
function (lhd::LHD)(x::Array{Float64,1}, inplace::Bool; noutput::Int64=3)
	lhd.dS .= lhd.dψ*x
	lhd.dS .+= lhd.δ*sum(lhd.dψ, dims = 2)[:,1]

	@inbounds for i=1:lhd.Ne
		lhd.logdS[i] = log(lhd.dS[i])
	end

	if inplace
		# Compute gradient
		@inbounds for i=1:lhd.Ne
			lhd.dψ_dS[i,:] .= lhd.dψ[i,:]/lhd.dS[i]
		end

		# Return objective
		lhd.J .= 0.5*dot(x, lhd.A*x) - sum(lhd.logdS,dims = 1)[1]*(1/lhd.Ne) + dot(x,lhd.b)

		# Return gradient
		lhd.G .= lhd.A*x - sum(lhd.dψ_dS, dims = 1)[1,:]*(1/lhd.Ne) + lhd.b

		# Return hessian
		lhd.H .= lhd.A + BLAS.gemm('T', 'N', 1/lhd.Ne, lhd.dψ_dS, lhd.dψ_dS)
	else
		if noutput == 1
			# Return objective
			J = 0.5*dot(x, lhd.A*x) - sum(lhd.logdS,dims = 1)[1]*(1/lhd.Ne) + dot(x,lhd.b)
		return J

		elseif noutput == 2
			# Compute gradient
			@inbounds for i=1:lhd.Ne
				lhd.dψ_dS[i,:] .= lhd.dψ[i,:]/lhd.dS[i]
			end

			# Return objective
			J = 0.5*dot(x, lhd.A*x) - sum(lhd.logdS,dims = 1)[1]*(1/lhd.Ne) + dot(x,lhd.b)

			# Return gradient
			G = lhd.A*x - sum(lhd.dψ_dS, dims = 1)[1,:]*(1/lhd.Ne) + lhd.b
			return J,G

		elseif noutput == 3
			# Compute gradient
			@inbounds for i=1:lhd.Ne
				lhd.dψ_dS[i,:] .= lhd.dψ[i,:]/lhd.dS[i]
			end

			# Return objective
			J = 0.5*dot(x, lhd.A*x) - sum(lhd.logdS,dims = 1)[1]*(1/lhd.Ne) + dot(x,lhd.b)

			# Return gradient
			G = lhd.A*x - sum(lhd.dψ_dS, dims = 1)[1,:]*(1/lhd.Ne) + lhd.b

			# Return hessian
			H = lhd.A + BLAS.gemm('T', 'N', 1/lhd.Ne, lhd.dψ_dS, lhd.dψ_dS)

			return J,G,H
		end
	end
end


using CellwiseEditingDifferentiation
using Distributions
using Lora

ncells = 20

λ = Hyperparameter(:λ)

μ = Data(:μ)
coverage = Data(:coverage)
edited = Data(:edited)
rate = Data(:rate)

function wlogtarget(w::Float64, s::Vector)
  v = s[2]*(1-s[2])*inv_logit(w)
  a, b = beta_pars_from_mv(s[2], v)
  mapreduce(x -> logpdf(Beta(a, b), x), +, s[7:(6+ncells)])+vpcprior(s[2], v, a, b, s[1])+logdvdw(w, s[2])
end

w = BasicContUnvParameter(:w, logtarget=wlogtarget, nkeys=6+ncells)

psetpdf = Array(Function, ncells)
for i in 1:ncells
  psetpdf[i] = function (p::Float64, s::Vector)
    a, b = beta_pars_from_mv(s[2], s[2]*(1-s[2])*inv_logit(s[6]))
    Beta(a+s[4][i], b+s[3][i]-s[4][i])
  end
end

p = Array(BasicContUnvParameter, ncells)
for i in 1:ncells
  p[i] = BasicContUnvParameter(Symbol('p', i), setpdf=psetpdf[i], nkeys=6+ncells)
end

model = GenericModel([λ, μ, coverage, edited, rate, w, p...], isindexed=false)

init = Dict(
  :λ => 10.,
  :μ => 0.5,
  :coverage => fill(10.0, 20),
  :edited => rand([1., 2., 3., 8., 9.], 20),
  :w => 0.
)
init[:rate] = float(init[:edited]./init[:coverage])
for i in 1:ncells
  init[Symbol('p', i)] =
    rand(psetpdf[i](NaN, [[init[k] for k in [:λ, :μ, :coverage, :edited, :rate, :w]]; fill(NaN, ncells)]))
end

wjob = BasicMCJob(
  model,
  RAM(),
  BasicMCRange(nsteps=100, burnin=10),
  init;
  pindex=6,
  resetpstate=false,
  verbose=false
)

mcrange = BasicMCRange(nsteps=1100, burnin=100)

outopts = fill(Dict(:monitor => [:value]), ncells+1)

gibbsjob = GibbsJob(model, Dict(:w => wjob), mcrange, init)

@time run(gibbsjob)

output(gibbsjob)

chains = Dict(gibbsjob)

v = map(x -> init[:μ]*(1-init[:μ])*inv_logit(x), chains[:w].value)

for i in 1:ncells
  writecsv(string("pchain", i, ".csv"), chains[Symbol('p', i)].value)
end
writecsv("vchain.csv", v)


using PyPlot
import Base.LAPACK.gbtrf!
import Base.LAPACK.gbtrs!

# Writing a simple test for LAPACK solver to see if I can get the syntax correct:
bandwidth = [1,2,4,8,16,32,64,128,256,512,1024]
#ntime = [64,256,1024,4096,16384,65536,262144]
ntime = [16384]
y = 0.5
dt = zeros(Float64,length(bandwidth))
for n in ntime
count = 1
for band in bandwidth
# First, solve it directly:
#  mat = zeros(n,n)
  ab = zeros(3*band+1,n)
  for i=1:n
    for j=i-band:i+band
      if j >=1 && j<=n
#        mat[i,j] = y^abs(i-j)
        ab[2*band+1+(j-i),i]=y^abs(i-j)
      end
    end
  end
#  println(mat)
  b = randn(n)
#  x = \(mat,b)

# Now use banded matrix:
  b0 = copy(b)
  tic()
  ab0,ipiv = gbtrf!(band,band,n,ab)
  x0 = gbtrs!('N',band,band,n,ab0,ipiv,b0)
  dt[count] = toq()
#  println(band," ",sum(abs(b - *(mat,x)))," ",sum(abs(x-x0))," ",dt[count])
  println(band," ",dt[count])
  count +=1
end
loglog(bandwidth,dt,label="Time elapsed")
loglog(bandwidth,0.1.*(bandwidth./16).^2)
xlabel("J")
ylabel("Time [sec]")
legend(loc="upper left")
end


include("k-means.jl")

function gaussian_weight{T<:Number}(x_i::Vector{T}, x_j::Vector{T}; sigma=1)
    exp(-norm(x_i - x_j)^2 / (2sigma^2))
end

function gaussian_weight(normval::Float64; sigma=1)
    exp(-normval^2 / (2sigma^2))
end

"""
Do spectral clustering based on Ng's algorithm (Ng, Jordan, Weiss, & others, 2002).

    @inputs the input data matrix, each column is a vector in feature space.
    @cluster_num the desired number of clusters

    returns an array, whose elements are pairs of array index and the cluster id
"""
function spectral_clustering(inputs::Matrix{Float64}, cluster_num::Int;
    neighbor_num::Int=10, sigma=1)

    W, D = build_matrix_normal_knn(inputs, k=neighbor_num, sigma=sigma)
    L_sym = Symmetric(D^(-1/2) * W * D^(-1/2))

    eigenvals, eigenvecs = eig(L_sym)

    # normalize the eigen vectors
    N = length(eigenvals)
    U = eigenvecs[:, (N - cluster_num + 1):N]
    for row = 1:size(U)[1]
        U[row, :] /= norm(U[row, :])
    end

    centroids, nearest_centroids = kmeans(cluster_num, transpose(U))

    return nearest_centroids
end

using NearestNeighbors
"""
Build a Adjacency Matrix using Normal Function as similarity function
    and only k-nearest-neighborhood have edges connected.

    @inputs input data as a matrix, each column is an input sample
    @k only the nearest k samples have an edge connected with the given point
    @sigma the parameter of the gaussian function

    return the pair of the adjacency matrix and the diagnoal matrix
"""
function build_matrix_normal_knn(inputs::Matrix{Float64}; k::Int=10, sigma=1)
    kdtree = KDTree(inputs)
    N = size(inputs)[2]
    adj = zeros(N, N)

    for i = 1:N
        idxs, dists = knn(kdtree, inputs[:, i], k, true)
        for j in 1:length(idxs)
            weight = idxs[j] == i ? 0 : gaussian_weight(dists[j], sigma=sigma)
            adj[i, idxs[j]] = weight
        end
    end

    adj = (adj + transpose(adj)) / 2.0

    degree = diagm(collect(sum(adj, 1)))

    return adj, degree
end


# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    _SOCtoPSDaff{T}(f::MOI.VectorAffineFunction{T}, g::MOI.ScalarAffineFunction{T})

Builds a VectorAffineFunction representing the upper (or lower) triangular part of the matrix
[ f[1]     f[2:end]' ]
[ f[2:end] g * I     ]
"""
function _SOCtoPSDaff(
    ::Type{T},
    f::Union{MOI.AbstractVectorFunction,AbstractVector{T}},
    g::Union{MOI.AbstractScalarFunction,T},
) where {T}
    F = MOIU.promote_operation(vcat, T, typeof(g), T)
    f_scalars = MOIU.eachscalar(f)
    dim = length(f_scalars)
    n = div(dim * (dim + 1), 2)
    h = MOIU.zero_with_output_dimension(F, n)
    MOIU.operate_output_index!(+, T, MOIU.trimap(1, 1), h, f_scalars[1])
    for i in 2:dim
        MOIU.operate_output_index!(+, T, MOIU.trimap(1, i), h, f_scalars[i])
        MOIU.operate_output_index!(+, T, MOIU.trimap(i, i), h, g)
    end
    return h
end

"""
The `SOCtoPSDBridge` transforms the second order cone constraint ``\\lVert x \\rVert \\le t`` into the semidefinite cone constraints
```math
\\begin{pmatrix}
  t & x^\\top\\\\
  x & tI
\\end{pmatrix} \\succeq 0
```
Indeed by the Schur Complement, it is positive definite iff
```math
\\begin{align*}
  tI & \\succ 0\\\\
  t - x^\\top (tI)^{-1} x & \\succ 0
\\end{align*}
```
which is equivalent to
```math
\\begin{align*}
  t & > 0\\\\
  t^2 & > x^\\top x
\\end{align*}
```

!!! warning
    This bridge is not added by default by [`MOI.Bridges.full_bridge_optimizer`](@ref)
    as bridging second order cone constraints to semidefinite constraints can be
    achieved by the [`SOCtoRSOCBridge`](@ref) followed by the [`RSOCtoPSDBridge`](@ref)
    while creating a smaller semidefinite constraint.
"""
struct SOCtoPSDBridge{T,F,G} <: SetMapBridge{
    T,
    MOI.PositiveSemidefiniteConeTriangle,
    MOI.SecondOrderCone,
    F,
    G,
}
    constraint::MOI.ConstraintIndex{F,MOI.PositiveSemidefiniteConeTriangle}
end

const SOCtoPSD{T,OT<:MOI.ModelLike} =
    SingleBridgeOptimizer{SOCtoPSDBridge{T},OT}

function concrete_bridge_type(
    ::Type{<:SOCtoPSDBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.SecondOrderCone},
) where {T}
    F = MOIU.promote_operation(vcat, T, MOIU.scalar_type(G), T)
    return SOCtoPSDBridge{T,F,G}
end

function MOIB.map_set(::Type{<:SOCtoPSDBridge}, set::MOI.SecondOrderCone)
    return MOI.PositiveSemidefiniteConeTriangle(MOI.dimension(set))
end

function MOIB.inverse_map_set(
    ::Type{<:SOCtoPSDBridge},
    set::MOI.PositiveSemidefiniteConeTriangle,
)
    return MOI.SecondOrderCone(MOI.side_dimension(set))
end

function MOIB.map_function(::Type{<:SOCtoPSDBridge{T}}, func) where {T}
    return _SOCtoPSDaff(T, func, MOIU.eachscalar(func)[1])
end

function MOIB.inverse_map_function(::Type{<:SOCtoPSDBridge}, func)
    scalars = MOIU.eachscalar(func)
    dim = MOIU.side_dimension_for_vectorized_dimension(length(scalars))
    return scalars[MOIU.trimap.(1, 1:dim)]
end

function MOIB.adjoint_map_function(::Type{<:SOCtoPSDBridge{T}}, func) where {T}
    scalars = MOIU.eachscalar(func)
    dim = MOIU.side_dimension_for_vectorized_dimension(length(scalars))
    tdual = sum(i -> func[MOIU.trimap(i, i)], 1:dim)
    return MOIU.operate(vcat, T, tdual, func[MOIU.trimap.(2:dim, 1)] * 2)
end

function MOIB.inverse_adjoint_map_function(
    ::Type{<:SOCtoPSDBridge{T}},
    func,
) where {T}
    # func is (t, x) such that x'x ≤ t^2 and we need to find a PSD matrix
    # [a   x'/2]
    # [x/2 Q   ]
    # such that a + tr(Q) == t
    # By choosing a = t/2 and Q = t/(2x'x) * xx', we get
    # a + tr(Q) = t/2 + t/(2x'x) * tr(xx') = t/2 + t/(2x'x) * tr(x'x) = t
    # Moreover, by the Schur complement, the matrix is PSD iff
    # xx'/(2t) ⪯ t/(2x'x) * xx'
    # 1/(2t) ≤ t/(2x'x)
    # x'x ≤ t^2
    # which is the SOC inequality
    t = func[1]
    x = func[2:end]
    return MOIB.inverse_adjoint_map_function(RSOCtoPSDBridge{T}, [t / 2; t; x])
end

"""
The `RSOCtoPSDBridge` transforms the second order cone constraint
``\\lVert x \\rVert \\le 2tu`` with ``u \\ge 0`` into the semidefinite cone
constraints
```math
\\begin{pmatrix}
  t & x^\\top\\\\
  x & 2uI
\\end{pmatrix} \\succeq 0
```

Indeed by the Schur Complement, it is positive definite iff
```math
\\begin{align*}
  uI & \\succ 0\\\\
  t - x^\\top (2uI)^{-1} x & \\succ 0
\\end{align*}
```
which is equivalent to
```math
\\begin{align*}
  u & > 0\\\\
  2tu & > x^\\top x
\\end{align*}
```
"""
struct RSOCtoPSDBridge{T,F,G} <: SetMapBridge{
    T,
    MOI.PositiveSemidefiniteConeTriangle,
    MOI.RotatedSecondOrderCone,
    F,
    G,
}
    constraint::MOI.ConstraintIndex{F,MOI.PositiveSemidefiniteConeTriangle}
end

const RSOCtoPSD{T,OT<:MOI.ModelLike} =
    SingleBridgeOptimizer{RSOCtoPSDBridge{T},OT}

function concrete_bridge_type(
    ::Type{<:RSOCtoPSDBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.RotatedSecondOrderCone},
) where {T}
    S = MOIU.scalar_type(G)
    H = MOIU.promote_operation(*, T, T, S)
    F = MOIU.promote_operation(vcat, T, S, H, T)
    return RSOCtoPSDBridge{T,F,G}
end

function MOIB.map_set(
    ::Type{<:RSOCtoPSDBridge},
    set::MOI.RotatedSecondOrderCone,
)
    return MOI.PositiveSemidefiniteConeTriangle(MOI.dimension(set) - 1)
end

function MOIB.inverse_map_set(
    ::Type{<:RSOCtoPSDBridge},
    set::MOI.PositiveSemidefiniteConeTriangle,
)
    return MOI.RotatedSecondOrderCone(MOI.side_dimension(set) + 1)
end

function MOIB.map_function(::Type{<:RSOCtoPSDBridge{T}}, func) where {T}
    scalars = MOIU.eachscalar(func)
    if length(scalars) < 3
        error(
            "Unable to bridge RotatedSecondOrderCone to PSD because the ",
            "dimension is too small: got $(length(scalars)), expected >= 3.",
        )
    end
    h = MOIU.operate!(*, T, scalars[2], convert(T, 2))
    return _SOCtoPSDaff(T, scalars[[1; 3:length(scalars)]], h)
end

function MOIB.inverse_map_function(::Type{<:RSOCtoPSDBridge{T}}, func) where {T}
    scalars = MOIU.eachscalar(func)
    dim = MOIU.side_dimension_for_vectorized_dimension(length(scalars))
    t = scalars[1]
    # It is (2u*I)[1,1] so it needs to be divided by 2 to get u
    u = MOIU.operate!(/, T, scalars[3], convert(T, 2))
    return MOIU.operate(
        vcat,
        T,
        t,
        u,
        scalars[[MOIU.trimap(1, i) for i in 2:dim]],
    )
end

function MOIB.adjoint_map_function(::Type{<:RSOCtoPSDBridge{T}}, func) where {T}
    scalars = MOIU.eachscalar(func)
    dim = MOIU.side_dimension_for_vectorized_dimension(length(scalars))
    udual = sum(i -> func[MOIU.trimap(i, i)], 2:dim)
    return MOIU.operate(
        vcat,
        T,
        func[1],
        2udual,
        func[MOIU.trimap.(2:dim, 1)] * 2,
    )
end

function MOIB.inverse_adjoint_map_function(
    ::Type{<:RSOCtoPSDBridge{T}},
    func,
) where {T}
    # func is (t, u, x) such that x'x ≤ 2tu and we need to find a PSD matrix
    # [t   x'/2]
    # [x/2 Q   ]
    # such that 2tr(Q) == u
    # By choosing Q = u/(2x'x) * xx', we get
    # 2tr(Q) = u/(x'x) * tr(xx') = u/(x'x) * tr(x'x) = u
    # Moreover, by the Schur complement, the matrix is PSD iff
    # xx'/(4t) ⪯ u/(2x'x) * xx'
    # 1/(4t) ≤ u/(2x'x)
    # x'x ≤ 2tu
    # which is the RSOC inequality
    t = func[1]
    u = func[2]
    x = func[3:end]
    Q = (x * x') * (u / (2 * (x' * x)))
    M = [
        t x'/2
        x/2 Q
    ]
    return [M[i, j] for j in 1:size(M, 2) for i in 1:j]
end


using IsoRank
using Base.Test

for k = 1:2:100
    G = sprand(k,k,0.4); G = G'+G; G = Int.(G - Diagonal(G) .> 0)
    R, res, L = isorank(G,G,maxiter=25,tol=1e-9,details=true,verbose=false)
    @test norm(L*vec(R) - res[1].*vec(R)) < 0.1
    b = rand(k,k);
    R, res, L = isorank(G,G,0.5,b,maxiter=25,tol=1e-9,details=true,verbose=false)
    @test norm(L*vec(R) - res[1].*vec(R)) < 0.1
end

function test1()
    G = sparse([1,2,3,1,3,2],[2,1,1,3,2,3],1,3,3)
    R, res, L = isorank(G,G,maxiter=25,tol=1e-9,details=true,verbose=false)
    norm(L*vec(R) .- res[1].*vec(R)) < 1e-2
end
@test test1()

function test2()
    G = sparse([1,2],[2,1],1,2,2)
    R,res,L = isorank(G, G, 0.85, ones(2,2), maxiter=25, tol=1e-9, details=true,verbose=false)
    vec(R) == [0.25,0.25,0.25,0.25] && norm(L*vec(R) - res[1].*vec(R)) == 0.0
end
@test test2()

function test3()
    G = sparse([1,2,3,1,3,2],[2,1,1,3,2,3],1,3,3)
    R,res,L = isorank(G, G, 0.85, ones(size(G,1),size(G,1)), maxiter=25, tol=1e-9, details=true,verbose=false)
    sum(vec(R) .== ones(9)./9)==0 && norm(L*vec(R) - res[1].*vec(R)) < 1e-2
end
@test test3()

function test4()
    diamond = sparse(vec([1 1 2 2 3]), vec([2 3 3 4 4]), 1.0,4,4)
    x = IsoRank.pagerank(diamond,verbose=false)
    isapprox(x, vec([0.128414 0.182991 0.260762 0.427833]), atol=0.01)
end
@test test4()

function test5()
    G1 = sprand(10,10,0.4); G1 = G1'+G1; G1 = Int.(G1 - Diagonal(G1) .> 0);
    G2 = sprand(15,15,0.4); G2 = G2'+G2; G2 = Int.(G2 - Diagonal(G2) .> 0);
    R, res, L = isorank(G1,G2,1.0,maxiter=25,tol=1e-9,details=true,verbose=false);
    A = Float64.(kron(G2,G1)); S = 1.0 ./ (A' * ones(Float64,size(A,2))); scale!(A,S);
    @test norm(A*vec(R) - res[1].*vec(R)) < 0.001
    els,evs = eigs(A);
    @test norm(A*evs[:,1] - L*evs[:,1]) < 0.001
end
test5()


# TODO: DESIGN PROPOSALS

"""
torus(x, a, b)
# Description
- Takes x ∈ ℝ and outputs torus(x) ∈ [a, b] in a periodic way.
- If a particle is moving to the right then it will pop from b to the point a
# Arguments: x, a, b
- `x`: (scalar). Current location of particle
- `a`: (scalar). left endpoint of interval
- `b`: (scalar). right endpoint of interval
# Output
-  `y`: (scalar). a value in the interval [a,b]
"""
torus(x::Number, a::Number, b::Number) = (((x-a)/(b-a))%1 - 0.5 * (sign((x-a)/(b-a)) - 1) )*(b-a) + a

"""
torus(x, a, b)
# Description
- Takes x ∈ ℝⁿ and outputs torus(x) ∈ ∏[aⁿ, bⁿ] in a periodic way.
- If a particle is moving to the right then it will pop from one part of the box to the oher
# Arguments: x, a, b
- `x`: (array). Current location of particle
- `a`: (array). left endpoint of tensor product interval
- `b`: (array). right endpoint of tensor product interval
# Output
-  `y`: (array). a value in the interval ∏[aⁿ, bⁿ]
"""
function torus(x::AbstractArray, a::AbstractArray, b::AbstractArray)
    N = length(x)
    y = zeros(N)
    for i in 1:N
        y[i] = torus(x[i], a[i], b[i])
    end
    return y
end

"""
closure_proprosal(covariance = Σ; left_bounds = [], right_bounds = []))
# Description
- Constructs a proposal for the Monte Carlo method.
# Arguments
- `covariance`: (vector) proposal parameter
# Keyword Arguments
- `left_bounds`: (array), left bounds for parameters
- `right_bounds`: (array), right bounds for parameters
# Output:
- `proposal`: (function), a function that outputs the proposal parameter
"""
function closure_proposal(Σ; left_bounds = [], right_bounds = [])
    perturbation = MvNormal(Σ)
    function proposal(C)
        proposal_C = C + rand(perturbation)
        # limit ranges for the parameters
        if isempty(left_bounds)
            return proposal_C
        else
            return torus(proposal_𝑪, left_bounds, right_bounds)
        end
        return proposal_C
    end
    return proposal
end


using HardTestProblems
using Test

@testset "PMM" begin
    for i in 1:6
        for (uldim, lldim) in zip([2,5,10], [3,5,10])
            # objective function and problem configuration
            F, f, conf = PMM_get_problem(i, uldim = uldim, lldim = lldim)
            cond1 = length(conf[:xmin]) == uldim == length(conf[:xmax])
            cond2 = length(conf[:ymin]) == lldim == length(conf[:ymax])
            cond3 = length(conf[:xbest]) == uldim && length(conf[:ybest]) == lldim
            cond4 = conf[:leader_optimum] ≈ 0 && conf[:follower_optimum] ≈ 0
            !cond4 && display(conf)
            @test cond1 && cond2 && cond3 && cond4
        end
    end
end

@testset "SMD" begin
    for i in 1:12
        for (uldim, lldim) in zip([2,5,10], [3,5,10])
            # objective function and problem configuration
            F, f, conf = SMD_get_problem(i, uldim = uldim, lldim = lldim)
            cond1 = length(conf[:xmin]) == uldim == length(conf[:xmax])
            cond2 = length(conf[:ymin]) == lldim == length(conf[:ymax])
            cond3 = length(conf[:xbest]) == uldim && length(conf[:ybest]) == lldim
            @test cond1 && cond2 && cond3
        end
    end
end

@testset "CEC2020" begin
    for fnum in 1:10
        for D in [5,10,15,20,30,50,100]
            f, conf = get_cec2020_problem(fnum)
            F =  f(rand(77, D))
            fmin = conf[:minimum]
            @test sum(F .< fmin) == 0
        end
    end
end

@testset "RW-MOP-2021" begin
    for i in 1:50
        # objective function and problem configuration
        f, conf = get_RW_MOP_problem(i)

        _, conf_tmp = get_RW_MOP_problem(conf[:function])
        # check get_RW_MOP_problem
        @test conf[:function] == conf_tmp[:function]

        # check dimension
        @test length(conf[:xmin]) == length(conf[:xmax]) == conf[:n]

        # check objectives
        @test length(conf[:nadir]) == conf[:fn]

        fx, gx, hx = f(conf[:xmin])

        # constraint violation sum
        vio = sum(max.(0.0, gx .<= 0)) + sum(max.(0.0, abs.(hx) .- 1e-2 .<= 0))

        # infeasible or feasible and not worst than nadir point.
        @test vio > 0 || (vio == 0 && sum(fx .> conf[:nadir]) == 0)
    end

end


"""
    δ_extrema(x)
Compute both the difference between the maximum and the minimum of `x`.
"""
function δ_extrema(x)
    δ_min, δ_max = extrema(x)
    return δ_max - δ_min
end


struct KroneckerMap{T, As<:LinearMapTuple} <: LinearMap{T}
    maps::As
    function KroneckerMap{T}(maps::LinearMapTuple) where {T}
        for TA in Base.Generator(eltype, maps)
            promote_type(T, TA) == T ||
                error("eltype $TA cannot be promoted to $T in KroneckerMap constructor")
        end
        return new{T,typeof(maps)}(maps)
    end
end

"""
    kron(A::LinearMap, B::LinearMap)::KroneckerMap
    kron(A, B, Cs...)::KroneckerMap

Construct a (lazy) representation of the Kronecker product `A⊗B`. One of the two factors
can be an `AbstractMatrix`, which is then promoted to a `LinearMap` automatically.

To avoid fallback to the generic [`Base.kron`](@ref) in the multi-map case,
there must be a `LinearMap` object among the first 8 arguments in usage like
`kron(A, B, Cs...)`.

For convenience, one can also use `A ⊗ B` or `⊗(A, B, Cs...)` (typed as `\\otimes+TAB`) to construct the
`KroneckerMap`, even when all arguments are of `AbstractMatrix` type.

If `A`, `B`, `C` and `D` are linear maps of such size that one can form the matrix
products `A*C` and `B*D`, then the mixed-product property `(A⊗B)*(C⊗D) = (A*C)⊗(B*D)`
holds. Upon vector multiplication, this rule is checked for applicability.

# Examples
```jldoctest; setup=(using LinearAlgebra, SparseArrays, LinearMaps)
julia> J = LinearMap(I, 2) # 2×2 identity map
2×2 LinearMaps.UniformScalingMap{Bool} with scaling factor: true

julia> E = spdiagm(-1 => trues(1)); D = E + E' - 2I;

julia> Δ = kron(D, J) + kron(J, D); # discrete 2D-Laplace operator

julia> Matrix(Δ)
4×4 Array{Int64,2}:
 -4   1   1   0
  1  -4   0   1
  1   0  -4   1
  0   1   1  -4
```
"""
Base.kron(A::LinearMap, B::LinearMap) =
    KroneckerMap{promote_type(eltype(A), eltype(B))}((A, B))
Base.kron(A::LinearMap, B::KroneckerMap) =
    KroneckerMap{promote_type(eltype(A), eltype(B))}(tuple(A, B.maps...))
Base.kron(A::KroneckerMap, B::LinearMap) =
    KroneckerMap{promote_type(eltype(A), eltype(B))}(tuple(A.maps..., B))
Base.kron(A::KroneckerMap, B::KroneckerMap) =
    KroneckerMap{promote_type(eltype(A), eltype(B))}(tuple(A.maps..., B.maps...))
# hoist out scalings
Base.kron(A::ScaledMap, B::LinearMap) = A.λ * kron(A.lmap, B)
Base.kron(A::LinearMap, B::ScaledMap) = kron(A, B.lmap) * B.λ
Base.kron(A::ScaledMap, B::ScaledMap) = (A.λ * B.λ) * kron(A.lmap, B.lmap)
# reduce UniformScalingMaps
Base.kron(A::UniformScalingMap, B::UniformScalingMap) = UniformScalingMap(A.λ * B.λ, A.M * B.M)
# disambiguation
Base.kron(A::ScaledMap, B::KroneckerMap) = A.λ * kron(A.lmap, B)
Base.kron(A::KroneckerMap, B::ScaledMap) = kron(A, B.lmap) * B.λ
# generic definitions
Base.kron(A::LinearMap, B::LinearMap, C::LinearMap, Ds::LinearMap...) =
    kron(kron(A, B), C, Ds...)
Base.kron(A::AbstractVecOrMat, B::LinearMap) = kron(LinearMap(A), B)
Base.kron(A::LinearMap, B::AbstractVecOrMat) = kron(A, LinearMap(B))
# promote AbstractMatrix arguments to LinearMaps, then take LinearMap-Kronecker product
for k in 3:8 # is 8 sufficient?
    Is = ntuple(n->:($(Symbol(:A, n))::AbstractVecOrMat), Val(k-1))
    # yields (:A1, :A2, :A3, ..., :A(k-1))
    L = :($(Symbol(:A, k))::LinearMap)
    # yields :Ak::LinearMap
    mapargs = ntuple(n -> :(LinearMap($(Symbol(:A, n)))), Val(k-1))
    # yields (:LinearMap(A1), :LinearMap(A2), ..., :LinearMap(A(k-1)))

    @eval Base.kron($(Is...), $L, As::MapOrVecOrMat...) =
        kron($(mapargs...), $(Symbol(:A, k)), convert_to_lmaps(As...)...)
end

struct KronPower{p}
    function KronPower(p::Integer)
        p > 1 || throw(ArgumentError("the Kronecker power is only defined for exponents larger than 1, got $k"))
        return new{p}()
    end
end

"""
    ⊗(k::Integer)

Construct a lazy representation of the `k`-th Kronecker power
`A^⊗(k) = A ⊗ A ⊗ ... ⊗ A`, where `A` can be an `AbstractMatrix` or a `LinearMap`.
"""
⊗(k::Integer) = KronPower(k)

⊗(A, B, Cs...) = kron(convert_to_lmaps(A, B, Cs...)...)

Base.:(^)(A::MapOrMatrix, ::KronPower{p}) where {p} =
    kron(ntuple(n -> convert_to_lmaps_(A), Val(p))...)

Base.size(A::KroneckerMap) = map(*, size.(A.maps)...)

LinearAlgebra.issymmetric(A::KroneckerMap) = all(issymmetric, A.maps)
LinearAlgebra.ishermitian(A::KroneckerMap) = all(ishermitian, A.maps)

LinearAlgebra.adjoint(A::KroneckerMap) = KroneckerMap{eltype(A)}(map(adjoint, A.maps))
LinearAlgebra.transpose(A::KroneckerMap) = KroneckerMap{eltype(A)}(map(transpose, A.maps))

Base.:(==)(A::KroneckerMap, B::KroneckerMap) = (eltype(A) == eltype(B) && A.maps == B.maps)

#################
# multiplication helper functions
#################

@inline function _kronmul!(y, B, x, A, T)
    ma, na = size(A)
    mb, nb = size(B)
    X = reshape(x, (nb, na))
    Y = reshape(y, (mb, ma))
    if B isa UniformScalingMap
        _unsafe_mul!(Y, X, transpose(A))
        lmul!(B.λ, y)
    else
        temp = similar(Y, (ma, nb))
        _unsafe_mul!(temp, A, copy(transpose(X)))
        _unsafe_mul!(Y, B, transpose(temp))
    end
    return y
end
@inline function _kronmul!(y, B, x, A::UniformScalingMap, _)
    ma, na = size(A)
    mb, nb = size(B)
    iszero(A.λ) && return fill!(y, zero(eltype(y)))
    X = reshape(x, (nb, na))
    Y = reshape(y, (mb, ma))
    _unsafe_mul!(Y, B, X)
    !isone(A.λ) && rmul!(y, A.λ)
    return y
end
@inline function _kronmul!(y, B, x, A::VecOrMatMap, _)
    ma, na = size(A)
    mb, nb = size(B)
    X = reshape(x, (nb, na))
    Y = reshape(y, (mb, ma))
    At = transpose(A.lmap)
    if B isa UniformScalingMap
        # the following is (perhaps due to the reshape?) faster than
        # _unsafe_mul!(Y, B * X, At)
        _unsafe_mul!(Y, X, At)
        lmul!(B.λ, y)
    elseif nb*ma <= mb*na
        _unsafe_mul!(Y, B, X * At)
    else
        _unsafe_mul!(Y, Matrix(B * X), At)
    end
    return y
end
const VectorMap{T} = WrappedMap{T,<:AbstractVector}
const AdjOrTransVectorMap{T} = WrappedMap{T,<:LinearAlgebra.AdjOrTransAbsVec}
@inline _kronmul!(y, B::AdjOrTransVectorMap, x, a::VectorMap, _) = mul!(y, a.lmap, B.lmap * x)

#################
# multiplication with vectors
#################

const KroneckerMap2{T} = KroneckerMap{T, <:Tuple{LinearMap, LinearMap}}

function _unsafe_mul!(y::AbstractVecOrMat, L::KroneckerMap2, x::AbstractVector)
    require_one_based_indexing(y)
    A, B = L.maps
    _kronmul!(y, B, x, A, eltype(L))
    return y
end
function _unsafe_mul!(y::AbstractVecOrMat, L::KroneckerMap, x::AbstractVector)
    require_one_based_indexing(y)
    A = first(L.maps)
    B = kron(Base.tail(L.maps)...)
    _kronmul!(y, B, x, A, eltype(L))
    return y
end
# mixed-product rule, prefer the right if possible
# (A₁ ⊗ A₂ ⊗ ... ⊗ Aᵣ) * (B₁ ⊗ B₂ ⊗ ... ⊗ Bᵣ) = (A₁B₁) ⊗ (A₂B₂) ⊗ ... ⊗ (AᵣBᵣ)
function _unsafe_mul!(y::AbstractVecOrMat,
                        L::CompositeMap{<:Any,<:Tuple{KroneckerMap,KroneckerMap}},
                        x::AbstractVector)
    require_one_based_indexing(y)
    B, A = L.maps
    if length(A.maps) == length(B.maps) && all(_iscompatible, zip(A.maps, B.maps))
        _unsafe_mul!(y, kron(map(*, A.maps, B.maps)...), x)
    else
        _unsafe_mul!(y, LinearMap(A)*B, x)
    end
    return y
end
# mixed-product rule, prefer the right if possible
# (A₁⊗B₁) * (A₂⊗B₂) * ... * (Aᵣ⊗Bᵣ) = (A₁*A₂*...*Aᵣ) ⊗ (B₁*B₂*...*Bᵣ)
function _unsafe_mul!(y::AbstractVecOrMat,
                        L::CompositeMap{T, <:Tuple{Vararg{KroneckerMap2}}},
                        x::AbstractVector) where {T}
    require_one_based_indexing(y)
    As = map(AB -> AB.maps[1], L.maps)
    Bs = map(AB -> AB.maps[2], L.maps)
    As1, As2 = Base.front(As), Base.tail(As)
    Bs1, Bs2 = Base.front(Bs), Base.tail(Bs)
    apply = all(_iscompatible, zip(As1, As2)) && all(_iscompatible, zip(Bs1, Bs2))
    if apply
        _unsafe_mul!(y, kron(prod(As), prod(Bs)), x)
    else
        _unsafe_mul!(y, CompositeMap{T}(map(LinearMap, L.maps)), x)
    end
    return y
end

###############
# KroneckerSumMap
###############
struct KroneckerSumMap{T, As<:Tuple{LinearMap, LinearMap}} <: LinearMap{T}
    maps::As
    function KroneckerSumMap{T}(maps::Tuple{LinearMap,LinearMap}) where {T}
        A1, A2 = maps
        (size(A1, 1) == size(A1, 2) && size(A2, 1) == size(A2, 2)) ||
            throw(ArgumentError("operators need to be square in Kronecker sums"))
        for TA in Base.Generator(eltype, maps)
            promote_type(T, TA) == T ||
                error("eltype $TA cannot be promoted to $T in KroneckerSumMap constructor")
        end
        return new{T, typeof(maps)}(maps)
    end
end

"""
    kronsum(A, B)::KroneckerSumMap
    kronsum(A, B, Cs...)::KroneckerSumMap

Construct a (lazy) representation of the Kronecker sum `A⊕B = A ⊗ Ib + Ia ⊗ B`
of two square linear maps of type `LinearMap` or `AbstractMatrix`. Here, `Ia`
and `Ib` are identity operators of the size of `A` and `B`, respectively.
Arguments of type `AbstractMatrix` are automatically promoted to `LinearMap`.

For convenience, one can also use `A ⊕ B` or `⊕(A, B, Cs...)` (typed as
`\\oplus+TAB`) to construct the `KroneckerSumMap`.

# Examples
```jldoctest; setup=(using LinearAlgebra, SparseArrays, LinearMaps)
julia> J = LinearMap(I, 2) # 2×2 identity map
2×2 LinearMaps.UniformScalingMap{Bool} with scaling factor: true

julia> E = spdiagm(-1 => trues(1)); D = LinearMap(E + E' - 2I);

julia> Δ₁ = kron(D, J) + kron(J, D); # discrete 2D-Laplace operator, Kronecker sum

julia> Δ₂ = kronsum(D, D);

julia> Δ₃ = D^⊕(2);

julia> Matrix(Δ₁) == Matrix(Δ₂) == Matrix(Δ₃)
true
```
"""
kronsum(A::MapOrMatrix, B::MapOrMatrix) =
    KroneckerSumMap{promote_type(eltype(A), eltype(B))}(convert_to_lmaps(A, B))
kronsum(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...) =
    kronsum(A, kronsum(B, C, Ds...))

struct KronSumPower{p}
    function KronSumPower(p::Integer)
        p > 1 || throw(ArgumentError("the Kronecker sum power is only defined for exponents larger than 1, got $k"))
        return new{p}()
    end
end

"""
    ⊕(k::Integer)

Construct a lazy representation of the `k`-th Kronecker sum power `A^⊕(k) = A ⊕ A ⊕ ... ⊕ A`,
where `A` can be a square `AbstractMatrix` or a `LinearMap`.
"""
⊕(k::Integer) = KronSumPower(k)

⊕(a, b, c...) = kronsum(a, b, c...)

Base.:(^)(A::MapOrMatrix, ::KronSumPower{p}) where {p} =
    kronsum(ntuple(n->convert_to_lmaps_(A), Val(p))...)

Base.size(A::KroneckerSumMap, i) = prod(size.(A.maps, i))
Base.size(A::KroneckerSumMap) = (size(A, 1), size(A, 2))

LinearAlgebra.issymmetric(A::KroneckerSumMap) = all(issymmetric, A.maps)
LinearAlgebra.ishermitian(A::KroneckerSumMap) = all(ishermitian, A.maps)

LinearAlgebra.adjoint(A::KroneckerSumMap) =
    KroneckerSumMap{eltype(A)}(map(adjoint, A.maps))
LinearAlgebra.transpose(A::KroneckerSumMap) =
    KroneckerSumMap{eltype(A)}(map(transpose, A.maps))

Base.:(==)(A::KroneckerSumMap, B::KroneckerSumMap) =
    (eltype(A) == eltype(B) && A.maps == B.maps)

function _unsafe_mul!(y::AbstractVecOrMat, L::KroneckerSumMap, x::AbstractVector)
    A, B = L.maps
    ma, na = size(A)
    mb, nb = size(B)
    X = reshape(x, (nb, na))
    Y = reshape(y, (nb, na))
    _unsafe_mul!(Y, X, transpose(A))
    _unsafe_mul!(Y, B, X, true, true)
    return y
end


## Label propagation algorithm for detecting communities ##

function _count_neighbor_labels(g::Graph, label::Vector{Int}, n::Integer)
    count = Dict{Int,Int}()
    for i = neighbors(g, n)
        l = label[i]
        count[l] = get(count, l, 0) + 1
    end
    return count
end

"""
    label_propagation(g::Graph)

A simple implementation of the label propagation algorithm.

# References

Raghavan, Usha Nandini, Réka Albert, and Soundar Kumara.
“Near Linear Time Algorithm to Detect Community Structures in Large-Scale Networks.”
Physical Review E 76, no. 3 (2007).
"""
function label_propagation(g::Graph)
    label = collect(nodes(g))
    finished = false
    while !finished
        order = shuffle(nodes(g))
        # Assign new labels
        for i = order
            count = _count_neighbor_labels(g, label, i)
            maxlabel, maxcount = 0, 0
            for l = keys(count)
                if count[l] > maxcount || count[l] == maxcount && order[l] < order[maxlabel]
                    maxlabel = l
                    maxcount = count[l]
                end
            end
            label[i] = maxlabel
        end
        # Check if labels don't need to be changed
        finished = true
        for i = nodes(g)
            count = _count_neighbor_labels(g, label, i)
            lcount = get(count, label[i], 0)
            if any(c > lcount for c = values(count))
                finished = false
                break
            end
        end
    end
    return relabel(label)
end


#=

    This is a port of the R2 model ClinicalTrial.cs

    with the datafiles
    - clinical_trial_control.csv  (size 1000)
    - clinical_trial_treatment.csv (size 1000)

    length :1000 mean(control): 0.513 mean(treatment): 0.51

    Output from the R2 model ClinicalTrial.cs (for isEffective)
    ```
    Mean: 0.113
    Variance: 0.100331
    Number of accepted samples = 318
    0:00:14:048.362.900
    ```

    Output from this Turing.jl model:

Summary Statistics
     parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
         Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

    isEffective    0.0659    0.2481     0.0025    0.0140   246.7994    1.0080       25.5512
    probControl    0.5152    0.2910     0.0029    0.0184   182.4067    0.9999       18.8846
  probTreatment    0.5278    0.2789     0.0028    0.0180   173.0546    1.0028       17.9164
        probAll    0.4985    0.0829     0.0008    0.0050   223.0815    1.0112       23.0957

Quantiles
     parameters      2.5%     25.0%     50.0%     75.0%     97.5%
         Symbol   Float64   Float64   Float64   Float64   Float64

    isEffective    0.0000    0.0000    0.0000    0.0000    1.0000
    probControl    0.0225    0.2196    0.5287    0.7653    0.9847
  probTreatment    0.0196    0.3207    0.5222    0.8098    0.9720
        probAll    0.1196    0.5019    0.5118    0.5190    0.5348

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function clinical_trial(control,treatment)
    n = length(control)

    isEffective   ~ flip(0.5)
    probControl   ~ Beta(1,1)
    probTreatment ~ Beta(1,1)
    probAll       ~ Beta(1,1)

    if isEffective
        control   ~ filldist(Bernoulli(probControl),n)
        treatment ~ filldist(Bernoulli(probTreatment),n)
    else
        control   ~ filldist(Bernoulli(probAll),n)
        treatment ~ filldist(Bernoulli(probAll),n)

    end
end

control = parse.(Int,split(readline("clinical_trial_control.csv"),","))
treatment = parse.(Int,split(readline("clinical_trial_treatment.csv"),","))
println("len(control):$(length(control)) len(treatment):$(length(treatment)) mean(control):$(mean(control)) mean(treatment): $(mean(treatment))")
# println("treatment:$treatment")

model = clinical_trial(control, treatment)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns[[:isEffective,:probControl,:probTreatment,:probAll]])


"""
    ind,rse,IND = PartitionTreeMatrixDhillon_QuadTree(matrix,jmax)

Perform the quadtree partition of matrix.

### Input Arguments
* `matrix`: The matrix to be partitioned
* `jmax`: Specify the max level of partitioning the matrix

### Output Arguments
* `ind`: Vector with the same length as length(matrix[:]), indicate the order of matrix after partition.
* `rse`: Region start and end. Each pair represents the start and end of a region in the matrix, corresponding to one node in the quadtree
* `IND`: IND[i] indicate the order of matrix after iteration i.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""
function PartitionTreeMatrixDhillon_QuadTree(matrix,jmax)

# 0.Preliminary

# Constants
rows,cols = size(matrix);
N = rows*cols
ind = Matrix{Int64}(reshape(1:N,rows,cols))
IND = zeros(rows,cols,jmax);
IND[:,:,1] = ind;

# region starts & ends
rse = Vector{Vector{Int64}}(jmax)
rse[1] = [1,N];

# 1. Partition the graph to yield rse, ind and jmax

j=1;
regioncount = 0;

while j<=jmax-1
    # add a column for j+1, if necessary

    rr = 1#for tracking the child regions

    rse[j+1] = zeros(Int64,4*length(rse[j]))

    for r = 1:2:(length(rse[j]) - 1)
        rs_rows,rs_cols = ind2sub((rows,cols),rse[j][r])
        re_rows,re_cols = ind2sub((rows,cols),rse[j][r+1])

        n_rows = re_rows-rs_rows+1
        n_cols = re_cols-rs_cols+1

        indrs = ind[rs_rows:re_rows,rs_cols:re_cols]

        if n_rows > 1 && n_cols > 1

            # compute the left and right singular vectors
            u_rows,v_cols = second_largest_singular_vectors(reshape(matrix[indrs[:]],n_rows,n_cols))

            # partition the current region
            pm,_ = partition_fiedler(SparseMatrixCSC{Float64,Int}(ones(1,1)),v = vcat(u_rows, v_cols))
            pm_rows = pm[1:n_rows]
            pm_cols = pm[n_rows+1:end]

            # determine the number of points in child region 1
            n1_rows = sum(pm_rows .> 0);
            n1_cols = sum(pm_cols .> 0);

            # update the indexing
            ind[rs_rows:re_rows,rs_cols:re_cols] = indrs[vcat(find(pm_rows .> 0),find(pm_rows.<=0)),vcat(find(pm_cols.>0),find(pm_cols.<=0))]

            if !(n1_rows==0 || n1_rows == n_rows) && !(n1_cols==0 ||n1_cols==n_cols)
                child_rows = [rs_rows, rs_rows+n1_rows-1, rs_rows+n1_rows, re_rows, rs_rows, rs_rows+n1_rows-1, rs_rows+n1_rows, re_rows];
                child_cols = [rs_cols, rs_cols+n1_cols-1, rs_cols, rs_cols+n1_cols-1, rs_cols+n1_cols, re_cols, rs_cols+n1_cols, re_cols];
                for k = 1:8
                    rse[j+1][rr+k-1] = sub2ind((rows,cols),child_rows[k],child_cols[k])
                end
                rr = rr + 8

            elseif (n1_rows == 0 || n1_rows == n_rows) && !(n1_cols == 0 || n1_cols == n_cols)
                child_rows = [rs_rows, re_rows, rs_rows, re_rows];
                child_cols = [rs_cols, rs_cols+n1_cols-1, rs_cols+n1_cols, re_cols];
                for k = 1:4
                    rse[j+1][rr+k-1] = sub2ind((rows,cols),child_rows[k],child_cols[k])
                end
                rr = rr + 4

            elseif (n1_cols == 0 || n1_cols == n_cols) && !(n1_rows == 0 || n1_rows == n_rows)
                child_rows = [rs_rows, rs_rows+n1_rows-1, rs_rows+n1_rows, re_rows];
                child_cols = [rs_cols, re_cols, rs_cols, re_cols];
                for k = 1:4
                    rse[j+1][rr+k-1] = sub2ind((rows,cols),child_rows[k],child_cols[k])
                end
                rr = rr + 4
            else
                rse[j+1][rr:rr+1] = rse[j][r:r+1]
                rr = rr + 2
            end

        elseif n_rows == 1 && n_cols > 1 #no bisection on rows
            pm_cols,_ = partition_fiedler(SparseMatrixCSC{Float64,Int}(ones(1,1)),v = matrix[indrs[:]])
            n1_cols = sum(pm_cols .> 0);
            ind[rs_rows:re_rows,rs_cols:re_cols] = indrs[:,vcat(find(pm_cols.>0),find(pm_cols.<=0))]

            child_rows = [rs_rows, re_rows, rs_rows, re_rows];
            child_cols = [rs_cols, rs_cols+n1_cols-1, rs_cols+n1_cols, re_cols];
            for k = 1:4
                rse[j+1][rr+k-1] = sub2ind((rows,cols),child_rows[k],child_cols[k])
            end
            rr = rr + 4

        elseif n_rows > 1 && n_cols == 1 #no bisection on cols
            pm_rows,_ = partitio_fiedler(SparseMatrixCSC{Float64,Int}(ones(1,1)),v = matrix[indrs[:]])
            n1_rows = sum(pm_rows .> 0);
            ind[rs_rows:re_rows,rs_cols:re_cols] = indrs[vcat(find(pm_rows .> 0),find(pm_rows.<=0)),:]

            child_rows = [rs_rows, rs_rows+n1_rows-1, rs_rows+n1_rows, re_rows];
            child_cols = [rs_cols, re_cols, rs_cols, re_cols];
            for k = 1:4
                rse[j+1][rr+k-1] = sub2ind((rows,cols),child_rows[k],child_cols[k])
            end
            rr = rr + 4

        else
            rse[j+1][rr:rr+1] = rse[j][r:r+1]
            rr = rr + 2
        end

    end
    j = j+1
    rse[j] = rse[j][rse[j].!=0]
    IND[:,:,j] = ind
end
return ind,rse,IND
end

"""
    u,v = second_largest_singular_vectors(A::Matrix{Float64})

Compute the second largest left and right singular vectors of matrix.

### Input Arguments
* `A`: The matrix of which the singular vectors will be computed.

### Output Arguments
* `u`: Second largest left singular vector.
* `v`: Second largest right singular vector.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""
function second_largest_singular_vectors(A::Matrix{Float64})
#Compute the second largest left and right singular vectors of An = D1^-0.5*A*D2^-0.5

rows,cols = size(A)

# compute D1 and D2 and make sure there are no zeros
D1 = sum(A,2)
D1[D1 .== 0] = maximum([0.01, minimum(D1[D1 .> 0])/10])
D2 = sum(A,1)
D2[D2 .== 0] = maximum([0.01, minimum(D2[D2 .> 0])/10])

u,_,v = svd(Diagonal(D1[:].^(-0.5))*A*Diagonal(D2[:].^(-0.5)))

# extract the 2nd singular vectors and multiply by D_i^-0.5
u = Diagonal(D1[:].^-0.5)*u[:,2]
v = Diagonal(D2[:].^-0.5)*v[:,2]
return u,v
end

"""
    BBmatrix, BBmatrix_ind, tag, dmatrix = GHWT_Matrix2D_Analysis_BestBasis(matrix::Matrix{Float64},ind::Matrix{Int64} ,rse::Vector{Array{Int64,1}})

Compute the expanding coeffcients of all levels and search for the best-basis given the partition quadtree of matrix.

### Input Arguments
* `matrix`: The matrix which will be analyzed.
* `ind`: Vector with the same length as length(matrix[:]), indicate the order of matrix after partition.
* `rse`: Region start and end. Each pair represents the start and end of a region in the matrix, corresponding to one node in the quadtree

### Output Arguments
* `BBmatrix`: Coefficients of the best-basis.
* `BBmatrix_ind`: Location of the best-basis in dmatrix.
* `tag`: Tag information of the expanding coefficients.
* `dmatrix`: The expanding coefficients on all levels.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""

function GHWT_Matrix2D_Analysis_BestBasis(matrix::Matrix{Float64},ind::Matrix{Int64} ,rse::Vector{Array{Int64,1}})
    # cost functional
    costfun = function(x) norm(x[:],1) end

    # constants
    rows,cols = size(matrix);
    #tol = 10^4*eps()

    # 1. Recursively partition the rows and columns of the matrix in a coupled manner
    #ind,rse,_ = PartitionTreeMatrixDhillon_QuadTree_ys(matrix)
    jmax = length(rse)

    # 2. Analyze the matrix and perform the best basis search
    dmatrix = zeros(rows,cols,jmax)
    dmatrix[:,:,jmax] = matrix[ind]
    tag = zeros(Int64,rows,cols,jmax)

    for r = 1:2:length(rse[jmax])
        if rse[jmax][r] != rse[jmax][r+1]
            rs_rows,rs_cols = ind2sub((rows,cols),rse[jmax][r])
            re_rows,re_cols = ind2sub((rows,cols),rse[jmax][r+1])

            tag[rs_rows:re_rows,rs_cols:re_cols,jmax] = reshape(0:(re_rows-rs_rows+1)*(re_cols-rs_cols+1)-1,re_rows-rs_rows+1,re_cols-rs_cols+1)
            average = sum(dmatrix[rs_rows:re_rows,rs_cols:re_cols,jmax])/sqrt((re_rows-rs_rows+1)*(re_cols-rs_cols+1))
            dmatrix[rs_rows:re_rows,rs_cols:re_cols,jmax] = 0
            dmatrix[rs_rows,rs_cols,jmax] = average
        end
    end
    #best basis
    BBmatrix = dmatrix[:,:,jmax]
    BBmatrix_ind = jmax*ones(Int64,rows,cols)

    for j = jmax-1:-1:1
        rr = 1;
        for r = 1:2:length(rse[j])
            rs_rows,rs_cols = ind2sub((rows,cols),rse[j][r])
            re_rows,re_cols = ind2sub((rows,cols),rse[j][r+1])

            # the row indices being considered
            indrs_rows = rs_rows:re_rows

            # the column indices being considered
            indrs_cols = rs_cols:re_cols

            n_rows = re_rows-rs_rows+1
            n_cols = re_cols-rs_cols+1

            #copy of children
            if rse[j][r+1] == rse[j+1][rr+1]
                dmatrix[indrs_rows,indrs_cols,j] = dmatrix[indrs_rows,indrs_cols,j+1]
                tag[indrs_rows,indrs_cols,j] = tag[indrs_rows,indrs_cols,j+1]
                rr = rr+2

                #copy of two childern
            elseif rse[j][r+1] == rse[j+1][rr+3]
                a,b = zeros(Int64,4), zeros(Int64,4)
                for k = 1:4
                    a[k],b[k] = ind2sub((rows,cols),rse[j+1][rr+k-1])
                end

                region1 = dmatrix[a[1]:a[2],b[1]:b[2],j+1]
                region2 = dmatrix[a[3]:a[4],b[3]:b[4],j+1]
                tag1 = tag[a[1]:a[2],b[1]:b[2],j+1]
                tag2 = tag[a[3]:a[4],b[3]:b[4],j+1]
                region_vec, tag_vec = twoblockcombine(region1[:],region2[:],tag1[:],tag2[:]);
                tag[indrs_rows,indrs_cols,j] = reshape(tag_vec,n_rows,n_cols);
                dmatrix[indrs_rows,indrs_cols,j] = reshape(region_vec,n_rows,n_cols)
                rr = rr + 4

            elseif rse[j][r+1] == rse[j+1][rr+7]
                a,b = zeros(Int64,8), zeros(Int64,8)
                for k = 1:8
                    a[k],b[k] = ind2sub((rows,cols),rse[j+1][rr+k-1])
                end

                region11 = dmatrix[a[1]:a[2],b[1]:b[2],j+1]
                region21 = dmatrix[a[3]:a[4],b[3]:b[4],j+1]
                tag11 = tag[a[1]:a[2],b[1]:b[2],j+1]
                tag21 = tag[a[3]:a[4],b[3]:b[4],j+1]
                region12 = dmatrix[a[5]:a[6],b[5]:b[6],j+1]
                region22 = dmatrix[a[7]:a[8],b[7]:b[8],j+1]
                tag12 = tag[a[5]:a[6],b[5]:b[6],j+1]
                tag22 = tag[a[7]:a[8],b[7]:b[8],j+1]

                region_vec_1, tag_vec_1 = twoblockcombine(region11[:],region21[:],tag11[:],tag21[:])
                region_vec_2, tag_vec_2 = twoblockcombine(region12[:],region22[:],tag12[:],tag22[:])
                region_vec, tag_vec = twoblockcombine(region_vec_1,region_vec_2,tag_vec_1,tag_vec_2)

                tag[indrs_rows,indrs_cols,j] = reshape(tag_vec, n_rows,n_cols)
                dmatrix[indrs_rows,indrs_cols,j] = reshape(region_vec,n_rows,n_cols)
                rr = rr + 8

            end
            #Obtaining BBmatrix
            if costfun(BBmatrix[indrs_rows,indrs_cols]) > costfun(dmatrix[indrs_rows,indrs_cols,j])
                BBmatrix[indrs_rows,indrs_cols] = dmatrix[indrs_rows,indrs_cols,j]
                BBmatrix_ind[indrs_rows,indrs_cols] = j;

            end
        end
    end
    return BBmatrix, BBmatrix_ind, tag, dmatrix
end

"""
    region, tag = twoblockcombine(region1::Vector{Float64}, region2::Vector{Float64}, tag1::Vector{Int64}, tag2::Vector{Int64})

Compute the coefficients of coarser basis-vectors after combining finer basis-vectors supported on two sub-regions.

### Input Arguments
* `region1`: The coeffcient of finer basis-vectors supported on region 1.
* `tag1`: The tag information of finer basis-vectors supported on region 1.
* `region2`: The coeffcient of finer basis-vectors supported on region 2.
* `tag2`: The tag information of finer basis-vectors supported on region 2.

### Output Arguments
* `region`: The coeffcient of coarser basis-vectors supported on region 1 and region 2.
* `tag`: The tag information of coarser basis-vectors supported on region 1 and region 2.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""

function twoblockcombine(region1::Vector{Float64}, region2::Vector{Float64}, tag1::Vector{Int64}, tag2::Vector{Int64})
    # the number of points in the first subregion
    n1 = length(region1)
    # the number of points in the second subregion
    n2 = length(region2)

    n = n1+n2
    region = zeros(n)
    tag = zeros(Int64,n)

    ### SCALING COEFFICIENT (n > 1)
    region[1] = ( sqrt(n1)*region1[1] + sqrt(n2)*region2[1] )/sqrt(n)
    tag[1] = 0

    ### HAAR-LIKE COEFFICIENT
    region[2] = ( sqrt(n2)*region1[1] - sqrt(n1)*region2[1] )/sqrt(n);
    tag[2] = 1

    ### WALSH-LIKE COEFFICIENTS
    # sweep through the coefficients in subregion 1 and subregion 2

    # the index of the new coefficient(s) to be created on level j
    parent = 3

    # the index of the current coefficients in subregions 1 and 2
    child1 = 2
    child2 = 2

    while parent <= n
        # no matching coefficient (use subregion 1)
        if child1 <= n1 && (child2 == n2+1 || tag1[child1] < tag2[child2])
            region[parent] = region1[child1]
            tag[parent] = 2*tag1[child1]
            child1 = child1+1;
            parent = parent+1;

            # no matching coefficient (use subregion 2)
        elseif child2 <=n2 && (child1 == n1+1 || tag2[child2] < tag1[child1])
            region[parent] = region2[child2]
            tag[parent] = 2*tag2[child2]
            child2 = child2+1
            parent = parent+1

            # matching coefficients
        else
            region[parent] = ( region1[child1] + region2[child2] )/sqrt(2)
            tag[parent] = 2*tag[child1]

            region[parent+1] = ( region1[child1] - region2[child2] )/sqrt(2)
            tag[parent+1] = 2*tag[child1]+1

            child1 = child1+1
            child2 = child2+1
            parent = parent+2
        end
    end
    return region, tag
end

"""
    matrix_r = GHWT_Matrix2D_Analysis_Synthesis(BBmatrix::Matrix{Float64},BBmatrix_ind::Matrix{Int64},tag::Array{Int64,3} ,rse::Vector{Array{Int64,1}})

Synthesize the matrix given coefficients of selected basis-vectors.

### Input Arguments
* `BBmatrix`: Coefficients of the best-basis.
* `BBmatrix_ind`: Location of the best-basis in dmatrix.
* `tag`: Tag information of the expanding coefficients.
* `rse`: Region start and end. Each pair represents the start and end of a region in the matrix, corresponding to one node in the quadtree.

### Output Arguments
* `matrix_r`: The synthesized matrix.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""
function GHWT_Matrix2D_Analysis_Synthesis(BBmatrix::Matrix{Float64},BBmatrix_ind::Matrix{Int64},tag::Array{Int64,3} ,rse::Vector{Array{Int64,1}})
    costfun = function(x) norm(x[:],1) end
    # constants
    rows,cols = size(BBmatrix);
    #tol = 10^4*eps()
    jmax = length(rse)

    cmatrix = zeros(rows,cols,jmax);
    for i  = 1:rows
        for j = 1:cols
            cmatrix[i,j,BBmatrix_ind[i,j]] = BBmatrix[i,j];
        end
    end

    for j = 1:1:jmax-1
        rr = 1
        for r = 1:2:length(rse[j])
            rs_rows,rs_cols = ind2sub((rows,cols),rse[j][r])
            re_rows,re_cols = ind2sub((rows,cols),rse[j][r+1])

            if rse[j][r+1] == rse[j+1][rr+1]
                if countnz(cmatrix[rs_rows:re_rows,rs_cols:re_cols,j+1]) == 0 && countnz(cmatrix[rs_rows:re_rows,rs_cols:re_cols,j]) > 0
                    cmatrix[rs_rows:re_rows,rs_cols:re_cols,j+1] = cmatrix[rs_rows:re_rows,rs_cols:re_cols,j]
                end
                rr = rr + 2

            elseif rse[j][r+1] == rse[j+1][rr+3]
                if countnz(cmatrix[rs_rows:re_rows,rs_cols:re_cols,j+1]) == 0 && countnz(cmatrix[rs_rows:re_rows,rs_cols:re_cols,j]) > 0
                    a,b = zeros(Int64,4), zeros(Int64,4)
                    for k = 1:4
                        a[k],b[k] = ind2sub((rows,cols),rse[j+1][rr+k-1])
                    end

                    region = cmatrix[rs_rows:re_rows,rs_cols:re_cols,j]
                    tag1 = tag[a[1]:a[2],b[1]:b[2],j+1]
                    tag2 = tag[a[3]:a[4],b[3]:b[4],j+1]
                    region1_vec, region2_vec = twoblocksplit(region[:],tag1[:],tag2[:]);
                    cmatrix[a[1]:a[2],b[1]:b[2],j+1] = reshape(region1_vec,a[2]-a[1]+1,b[2]-b[1]+1)
                    cmatrix[a[3]:a[4],b[3]:b[4],j+1] = reshape(region2_vec,a[4]-a[3]+1,b[4]-b[3]+1)
                end
                rr = rr + 4
            else
                if countnz(cmatrix[rs_rows:re_rows,rs_cols:re_cols,j+1]) == 0 && countnz(cmatrix[rs_rows:re_rows,rs_cols:re_cols,j]) > 0
                    a,b = zeros(Int64,8), zeros(Int64,8)
                    for k = 1:8
                        a[k],b[k] = ind2sub((rows,cols),rse[j+1][rr+k-1])
                    end
                    tag11 = tag[a[1]:a[2],b[1]:b[2],j+1]
                    tag21 = tag[a[3]:a[4],b[3]:b[4],j+1]
                    tag12 = tag[a[5]:a[6],b[5]:b[6],j+1]
                    tag22 = tag[a[7]:a[8],b[7]:b[8],j+1]
                    tag11 = tag11[:]
                    tag21 = tag21[:]
                    tag12 = tag12[:]
                    tag22 = tag22[:]
                    tag1 = twoblockcombine_tag(tag11,tag21)
                    tag2 = twoblockcombine_tag(tag12,tag22)
                    region = cmatrix[rs_rows:re_rows,rs_cols:re_cols,j]
                    region1_vec,region2_vec = twoblocksplit(region[:], tag1,tag2)
                    region11,region21 = twoblocksplit(region1_vec,tag11,tag21)
                    region12,region22 = twoblocksplit(region2_vec,tag12,tag22)
                    cmatrix[a[1]:a[2],b[1]:b[2],j+1] = reshape(region11,a[2]-a[1]+1,b[2]-b[1]+1)
                    cmatrix[a[3]:a[4],b[3]:b[4],j+1] = reshape(region21,a[4]-a[3]+1,b[4]-b[3]+1)
                    cmatrix[a[5]:a[6],b[5]:b[6],j+1] = reshape(region12,a[6]-a[5]+1,b[6]-b[5]+1)
                    cmatrix[a[7]:a[8],b[7]:b[8],j+1] = reshape(region22,a[8]-a[7]+1,b[8]-b[7]+1)
                end
                rr = rr + 8
            end
        end
    end

    matrix_r = cmatrix[:,:,end]
    for r = 1:2:length(rse[jmax])
        if rse[jmax][r] != rse[jmax][r+1]
            rs_rows,rs_cols = ind2sub((rows,cols),rse[j][r])
            re_rows,re_cols = ind2sub((rows,cols),rse[j][r+1])

            matrix_r[rs_rows:re_rows,rs_cols:re_cols] = matrix_r[rs_rows,rs_cols]/sqrt((re_rows-rs_rows+1)*(re_cols-rs_cols+1))
        end
    end
    return matrix_r
end

"""
    tag = twoblockcombine_tag(tag1::Vector{Int64},tag2::Vector{Int64})

Compute the tag information of coarser basis-vectors after combining finer basis-vectors supported on two sub-regions.

### Input Arguments
* `tag1`: The tag information of finer basis-vectors supported on region 1.
* `tag2`: The tag information of finer basis-vectors supported on region 2.

### Output Arguments
* `tag`: The tag information of coarser basis-vectors supported on region 1 and region 2.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""
function twoblockcombine_tag(tag1::Vector{Int64},tag2::Vector{Int64})
    # the number of points in the first subregion
    n1 = length(tag1)
    # the number of points in the second subregion
    n2 = length(tag2)

    n = n1+n2
    tag = zeros(Int64,n)

    ### SCALING COEFFICIENT (n > 1)
    tag[1] = 0

    ### HAAR-LIKE COEFFICIENT
    tag[2] = 1

    ### WALSH-LIKE COEFFICIENTS
    # sweep through the coefficients in subregion 1 and subregion 2

    # the index of the new coefficient(s) to be created on level j
    parent = 3

    # the index of the current coefficients in subregions 1 and 2
    child1 = 2
    child2 = 2

    while parent <= n
        # no matching coefficient (use subregion 1)
        if child1 <= n1 && (child2 == n2+1 || tag1[child1] < tag2[child2])
            tag[parent] = 2*tag1[child1]
            child1 = child1+1
            parent = parent+1

            # no matching coefficient (use subregion 2)
        elseif child2 <=n2 && (child1 == n1+1 || tag2[child2] < tag1[child1])
            tag[parent] = 2*tag2[child2]
            child2 = child2+1
            parent = parent+1

            # matching coefficients
        else
            tag[parent] = 2*tag[child1]
            tag[parent+1] = 2*tag[child1]+1

            child1 = child1+1
            child2 = child2+1
            parent = parent+2
        end
    end
    return tag
end

"""
    region1,region2 = twoblocksplit(region::Matrix{Float64}, tag1::Vector{Int64}, tag2::Vector{Int64})

Compute the coefficients of finer basis-vectors supported on two sub-regions given the coeffcients of coarser basis-vectors supported on both regions.

### Input Arguments
* `region`: The coeffcient of coarser basis-vectors supported on region 1 and region 2.
* `tag1`: The tag information of finer basis-vectors supported on region 1.
* `tag2`: The tag information of finer basis-vectors supported on region 2.

### Output Arguments
* `region1`: The coeffcient of finer basis-vectors supported on region 1.
* `region2`: The coeffcient of finer basis-vectors supported on region 2.

Copyright 2018 The Regents of the University of California

Implemented by Yiqun Shao (Adviser: Dr. Naoki Saito)
"""
function twoblocksplit(region::Matrix{Float64}, tag1::Vector{Int64}, tag2::Vector{Int64})
    # the number of points in the first subregion
    n1 = length(tag1)
    # the number of points in the second subregion
    n2 = length(tag2)

    n = n1+n2
    region1 = zeros(n1)
    region2 = zeros(n2);

    ### SCALING COEFFICIENT (n > 1)
    region1[1] = ( sqrt(n1)*region[1] + sqrt(n2)*region[2] )/sqrt(n);

    ### HAAR-LIKE COEFFICIENT
    region2[1] = ( sqrt(n2)*region[1] - sqrt(n1)*region[2] )/sqrt(n);

    ### WALSH-LIKE COEFFICIENTS
    # sweep through the coefficients in subregion 1 and subregion 2

    # the index of the new coefficient(s) to be created on level j
    parent = 3

    # the index of the current coefficients in subregions 1 and 2
    child1 = 2
    child2 = 2

    while parent <= n
        # no matching coefficient (use subregion 1)
        if child1 <= n1 && (child2 == n2+1 || tag1[child1] < tag2[child2])
            region1[child1] = region[parent]
            child1 = child1+1
            parent = parent+1

            # no matching coefficient (use subregion 2)
        elseif child2 <=n2 && (child1 == n1+1 || tag2[child2] < tag1[child1])
            region2[child2] = region[parent]
            child2 = child2+1
            parent = parent+1

            # matching coefficients
        else
            region1[child1] = ( region[parent] + region[parent+1] )/sqrt(2)
            region2[child2] = ( region[parent] - region[parent+1] )/sqrt(2)
            child1 = child1+1
            child2 = child2+1
            parent = parent+2
        end

    end
    return region1,region2
end