export mipdecoHeuristic, heuristicKnapsack

function heuristicKnapsack(mc,gc,delta)
    mt = copy(mc)
    pred = 0
    cIndex = reverse(sortperm(abs.(gc)))
    wCount = 0
	for j=1:length(cIndex)
		ind = cIndex[Int(j)]
			  if  gc[ind] > 0 && mc[ind] == 1
				mt[ind] = 0
				pred -= gc[ind]
                wCount += 1
			  elseif  gc[ind] < 0 && mc[ind] == 0
				mt[ind] = 1
				pred += gc[ind]
                wCount += 1
			  end
            if (wCount >= delta)
                break
            end
	  end
    return mt,pred
end

function  mipdecoHeuristic(mc::BitArray{1}, pInv::InverseParam, pMis;
	                out::Int=2, getNeighborhood::Function=x->trues(size(x)),
					delta::Int=32, sigma::Float64=0.25,outFile::String="")

	mc          = copy(mc)
    iter        = 0
	TermFlag	= 0
	alpha       = pInv.alpha
	maxIter     = pInv.maxIter      #  Max. no. iterations.

	## Time statistics
	timeMIP = 0.0

	## evaluate function and derivatives
    Jc,gc,Dc,Fc,Rc,timeF,timeR = jInvObj(mc,pMis,pInv)
    rId = getNeighborhood(mc)
    gc .*= rId
    J0 = Jc;
	His = zeros(maxIter,3)

	### Initialization for trust region algorithm
	outStr = @sprintf("\n %4s\t%08s\t%08s\t%08s\t%08s\t%08s\t%08s\t%06s\n",
					  	"iter", "Fc", "Rc","alpha","Jc/J0","pred","ared","delta")
    if out>=2; print(outStr); end
	if !isempty(outFile)
		f = open(outFile, "w")
		write(f, outStr)
	end

      outStr = @sprintf("%3d\t%3.2e\t%3.2e\t%3.2e\t%3.2e\t%3.2e\t%3.2e\t%3d\n",
		         iter, Fc, Rc,alpha[1],Jc/J0,0,0,delta)
      if out>=2; print(outStr); end
      if !isempty(outFile)
		  write(f, outStr)
	  end

	while (delta > 0)
	  iter  += 1
	  toMIP = @elapsed begin
          mt,pred = heuristicKnapsack(mc,gc,delta)
      end
      pred *=-1
      timeMIP += toMIP

      Jt,gt,Dt,Ft,Rt,toF,toR = jInvObj(mt,pMis,pInv)
      timeF += toF
      timeR += toR
	  ared  = Jc - Jt  	# actual reduction
	  deltaO = delta

      outStr = @sprintf("%3d\t%3.2e\t%3.2e\t%3.2e\t%3.2e\t%3.2e\t%3.2e\t%3d\n",
		         iter, Fc, Rc,alpha[1],Jc/J0,pred,ared,delta)
      if out>=2; print(outStr); end
      if !isempty(outFile)
		  write(f, outStr)
	  end

        if ared > 0
	    mc .= mt
	    (Jc,Dc,Fc,Rc) = (Jt,Dt,Ft,Rt)

	    if ared >= sigma*pred;delta *= 2;end
	      gc = gt
          rId = getNeighborhood(mc)
          gc .*= rId
        else
         delta = floor(.5*delta)
	  end
	  His[iter,:] = [Jc Jt deltaO]
	  if iter >= maxIter
          TermFlag = 1
          break
      end
	end 	#while of trust region
	println("Trust region on entire space: MisFit, Regularizer, Total $(Fc) , $(Rc)	, $(Jc)")
	println("Time to compute Data term, Regularizer, and MIP: $(timeF), $(timeR), $(timeMIP) \n")
	if out>=1
		if TermFlag==0
			println("Trust region iterated $iter times with final objective $(Jc)." )
		else
			println("Trust region reached iteration limit with final objective $(Jc).")
		end
	end
    if !isempty(outFile)
	  close(f)
  end
	return mc,Dc,TermFlag,His
end


Base.@kwdef struct DGPMErlang <: AbstractModel
    # Data
    c0::Vector{Float64}
    y0::Vector{Float64}
    X0::Matrix{Float64}
    y1::Vector{Float64}
    X1::Matrix{Float64}
    # Transformed data
    event::Vector{Bool} = y0 .< c0
    N0::Int = length(y0)
    N1::Int = length(y1)
    # HyperParameters
    a_φ0::Float64 = 2.0
    b_φ0::Float64 = 0.1
    a_λ0::Float64 = 0.1
    b_λ0::Float64 = 0.1
    # Parameters
    φ::Vector{Float64} = rand(Gamma(a_φ0, 1.0 / b_φ0), 1)
    λ::Vector{Float64} = rand(Gamma(a_λ0, 1.0 / b_λ0), 1)
    # Transformed parameters
    sumy::Base.RefValue{Float64} = Ref(0.0)
    sumlogy::Vector{Float64} = zeros(1)
    S::Vector{Float64} = zeros(N1)
    # Skeleton
    skl::Skeleton = Skeleton(; y0, y1, X0, X1)
end

function skeleton(m::DGPMErlang)
    return m.skl
end

function kernel_pdf(m::DGPMErlang, yi::Float64, j::Int)
    kernel = Erlang(ceil(Int, m.φ[j]), 1.0 / m.λ[])
    return pdf(kernel, yi)
end

function kernel_cdf(m::DGPMErlang, yi::Float64, j::Int)
    kernel = Erlang(ceil(Int, m.φ[j]), 1.0 / m.λ[])
    return 1 - cdf(kernel, yi)
end

function update_atoms!(m::DGPMErlang)
    update_suffstats!(m)
    update_λ!(m)
    update_φ!(m)
    update_y!(m)
end

function update_suffstats!(m::DGPMErlang)
    (; y0, sumlogy, sumy, skl) = m
    d = cluster_labels(skl)
    if length(sumlogy) < rmax(skl)
        resize!(sumlogy, rmax(skl))
    end
    for j in 1:rmax(skl)
        sumlogy[j] = 0.0
    end
    sumy[] = 0.0
    for i in 1:length(y0)
        sumy[] += y0[i]
        sumlogy[d[i]] += log(y0[i])
    end
    return nothing
end

function update_λ!(m::DGPMErlang)
    (; λ, φ, a_λ0, b_λ0, sumy, skl) = m
    (; d) = skl
    a_λ1 = a_λ0
    b_λ1 = b_λ0 + sumy[]
     for di in d
        a_λ1 += ceil(φ[di])
    end
    λ[] = rand(Gamma(a_λ1, 1.0 / b_λ1))
    return nothing
end

function update_φ!(m::DGPMErlang)
    (; φ, a_φ0, b_φ0, skl) = m
    pφ = Gamma(a_φ0, 1.0 / b_φ0)
    while length(φ) < rmax(skl)
        push!(φ, rand(pφ))
    end
    for j in 1:rmax(skl)
        φ0 = φ[j]
        d0 = truncated(Normal(φ0, 0.1), 0, Inf)
        φ1 = rand(d0)
        d1 = truncated(Normal(φ1, 0.1), 0, Inf)
        log_ar = 0.0
        log_ar += (logpφ(m, φ1, j) + logpdf(d1, φ0))
        log_ar -= (logpφ(m, φ0, j) + logpdf(d0, φ1))
        if rand() < exp(log_ar)
            (φ[j] = φ1)
        end
    end
    return nothing
end

function logpφ(m::DGPMErlang, φ0, j)
    (; λ, a_φ0, b_φ0, sumlogy, skl) = m
    n = cluster_sizes(skl)
    δ0 = ceil(Int, φ0)
    return (
        n[j] * δ0 * log(λ[]) +
        (δ0 - 1) * sumlogy[j] -
        n[j] * logfactorial(δ0 - 1) +
        (a_φ0 - 1) * log(φ0) -
        b_φ0 * φ0
    )
end

function update_y!(m::DGPMErlang)
    (; y0, c0, event, φ, λ, skl) = m
    d = cluster_labels(skl)
    for i in 1:length(y0)
        if !event[i]
            pdf = Erlang(ceil(Int, φ[d[i]]), 1.0 / λ[])
            tpdf = Truncated(pdf, c0[i], Inf)
            if cdf(pdf, c0[i]) < 0.99
                y0[i] = rand(tpdf)
            end
        end
    end
end

function update_S!(m::DGPMErlang)
    (; S, skl) = m
    (; N1, y1, ϕ1, s, rmax) = skl
    for i in 1:N1
        S[i] = 0.0
        for j in 1:rmax[]
            wj = get_w(ϕ1[i], s[], j)
            S[i] += wj * kernel_cdf(m, y1[i], j)
        end
    end
end

function sample!(m::DGPMErlang; mcmcsize = 4000, burnin = 2000, thin = 1)
    (; S, skl) = m
    (; N1, D1, f, β) = skl
    chainf = [zeros(N1) for _ in 1:(mcmcsize - burnin) ÷ thin]
    chainS = [zeros(N1) for _ in 1:(mcmcsize - burnin) ÷ thin]
    chainβ = [zeros(D1) for _ in 1:(mcmcsize - burnin) ÷ thin]
    chaing = [zeros(Bool, D1) for _ in 1:(mcmcsize - burnin) ÷ thin]
    for iter in 1:mcmcsize
        step!(m)
        update_S!(m)
        if iter > burnin && iszero((iter - burnin) % thin)
            chainf[(iter - burnin) ÷ thin] .= f
            chainS[(iter - burnin) ÷ thin] .= S
            chainβ[(iter - burnin) ÷ thin] .= β
            chaing[(iter - burnin) ÷ thin] .= (β .!= 0.0)
        end
    end
    return chainf, chainS, chainβ, chaing
end


import Pkg; Pkg.activate(@__DIR__); Pkg.instantiate();
#Write Altro folder instead of @__DIR__
using TrajectoryOptimization
using Plots, LinearAlgebra

#number individual systems
agent_number=10

#defining product system dynamics
function dynamics!(ẋ::AbstractVector{T},x::AbstractVector{T},u::AbstractVector{T}) where T

	for k=1:10
		ẋ[3*k-2] = u[2*k-1]*cos(x[3*k]);
		ẋ[3*k-1] = u[2*k-1]*sin(x[3*k]);
	    ẋ[3*k] = u[2*k];
	end
end

#Total number of states
n=3*agent_number
#Total number of control inputs
m=2*agent_number

#defining descrete model
model = Model(dynamics!,n,m)
model_d = rk4(model)
n = model.n;
m = model.m;

T = Float64;

#inital point for product system
x0 = [0.0;2.0;0.0;
0.0;4.0;0.0;
0.0;6.0;0.0;
0.0;8.0;0.0;
2.0;0.0;pi/2;
4.0;0.0;pi/2;
6.0;0.0;pi/2;
8.0;0.0;pi/2;
0.0;0.0;pi/4;
0.0;10.0;-pi/4]

#goal point for product system
xf = [10.0;2.0;0.0;
10.0;4.0;0.0;
10.0;6.0;0.0;
10.0;8.0;0.0;
2.0;10.0;pi/2;
4.0;10.0;pi/2;
6.0;10.0;pi/2;
8.0;10.0;pi/2;
10.0;10.0;pi/4;
10.0;0.0;-pi/4]

#number of points in trajectory(number of time samples)
N = 104
#sampling time
dt = 0.1
#total time
tf = (N-1)*dt

#LQR Cost Matrices (same as default)
Q = 0.01*Diagonal(I,n)
Qf = 1000.0*Diagonal(I,n)
R = 0.01*Diagonal(I,m)
obj = LQRObjective(Q,R,Qf,xf,N)

#creating upperbound and lower bound for control input with bound constraint
constraints = Constraints(N)
u_min = Array{Float64}(undef,2*agent_number)
u_max = Array{Float64}(undef,2*agent_number)
u_c=2;
for k=1:2*agent_number
	u_min[k]=-u_c;
	u_max[k]=u_c;
end
bnd = BoundConstraint(n,m,u_min=u_min,u_max=u_max)
for k = 1:N-1
     constraints[k] += bnd
 end

 #constraint for obstacle
 #this function is defining penalty for collision constraint
function My_rectangle(x)
	M=10000
	ϵ=1.46
		for i=1:10
        	temp=max((x[3*i-2]-5)^2,(x[3*i-1]-5)^2)
			M=min(temp,M)
    	end
    return -(M-ϵ^2);
end

#constraint for obstacle
#this function is defining penalty for collision constraint
function My_constraint(x)
	ϵ=0.64
	M=1000;
	for i=1:9 #agent number -1
		for j=i+1:10  #agent number
		temp=max((x[3*i-2]-x[3*j-2])^2 , (x[3*i-1]-x[3*j-1])^2);
		M=min(temp,M);
		end
	end
	return -(M - ϵ^2)
end

function My_obs(c,x,u)
    c[1]=My_constraint(x);
    c[2]=My_rectangle(x);
	return nothing
end

obs= Constraint{Inequality}(My_obs,n,m,2,:obs)
for k = 1:N-1
     constraints[k] += obs
 end

#goal constraint
goal = goal_constraint(xf);
constraints[N] += goal

prob = Problem(model_d,obj,constraints=constraints,x0=x0,N=N,dt=dt,xf=xf)

# options for solver
max_con_viol = 1.0e-8
verbose=false

opts_ilqr = iLQRSolverOptions{T}(verbose=verbose,
    iterations=300)

opts_al = AugmentedLagrangianSolverOptions{T}(verbose=verbose,
    opts_uncon=opts_ilqr,
    iterations=40,
    cost_tolerance=1.0e-5,
    cost_tolerance_intermediate=1.0e-4,
    constraint_tolerance=max_con_viol,
    penalty_scaling=10.,
    penalty_initial=1.)

opts_pn = ProjectedNewtonSolverOptions{T}(verbose=verbose,
    feasibility_tolerance=max_con_viol,
    solve_type=:feasible)

opts_altro = ALTROSolverOptions{T}(verbose=verbose,
    opts_al=opts_al,
    R_inf=1.0e-8,
    resolve_feasible_problem=false,
    opts_pn=opts_pn,
    projected_newton=true,
    projected_newton_tolerance=1.0e-4);

	solver = AugmentedLagrangianSolver(prob)

	#here we call the solver to solve
    @time begin
    a=solve!(prob, opts_altro); # solve with ALTRO
    end

#parsing the solution to differenet arrays

	x=Array{Array{Float64,1}}(undef,agent_number)
	z=Array{Array{Float64,1}}(undef,agent_number)
	for i=1:agent_number
    	x[i] = [prob.X[k][3*i-2] for k = 1:N]
    	z[i] = [prob.X[k][3*i-1] for k = 1:N]
	end
#plotting
	a=plot()

	rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
	plot!(rectangle(2,2,4,4), opacity=.5)
	plot!(rectangle(2.8,2.8,3.6,3.6), opacity=.3)

	plot!(x[1],z[1],xlabel="x",ylabel="y",label="1",legend=:topleft,color=:blue,width=2,ratio=:equal,title="Multi-Agent Example Trajectories")
	plot!(x[2],z[2],xlabel="x",ylabel="y",label="2",legend=:topleft,width=2,ratio=:equal)
	plot!(x[3],z[3],xlabel="x",ylabel="y",label="3",legend=:topleft,width=2,ratio=:equal)
	plot!(x[4],z[4],xlabel="x",ylabel="y",label="4",legend=:topleft,width=2,ratio=:equal)
	plot!(x[5],z[5],xlabel="x",ylabel="y",label="5",legend=:topleft,width=2,ratio=:equal)
	plot!(x[6],z[6],xlabel="x",ylabel="y",label="6",legend=:topleft,width=2,ratio=:equal)
	plot!(x[7],z[7],xlabel="x",ylabel="y",label="7",legend=:topleft,width=2,ratio=:equal)
	plot!(x[8],z[8],xlabel="x",ylabel="y",label="8",legend=:topleft,width=2,ratio=:equal)
	plot!(x[9],z[9],xlabel="x",ylabel="y",label="9",legend=:topleft,width=2,ratio=:equal)
	plot!(x[10],z[10],xlabel="x",ylabel="y",label="10",legend=:topleft,width=2,ratio=:equal)

	plot!((x[1][1],z[1][1]),marker=:circle,color=:red,label="")
    plot!((x[1][end],z[1][end]),marker=:utriangle,color=:green,label="")

	plot!((x[2][1],z[2][1]),marker=:circle,color=:red,label="")
    plot!((x[2][end],z[2][end]),marker=:utriangle,color=:green,label="")

	plot!((x[3][1],z[3][1]),marker=:circle,color=:red,label="")
    plot!((x[3][end],z[3][end]),marker=:utriangle,color=:green,label="")

	plot!((x[4][1],z[4][1]),marker=:circle,color=:red,label="")
    plot!((x[4][end],z[4][end]),marker=:utriangle,color=:green,label="")

	plot!((x[5][1],z[5][1]),marker=:circle,color=:red,label="")
    plot!((x[5][end],z[5][end]),marker=:utriangle,color=:green,label="")

	plot!((x[6][1],z[6][1]),marker=:circle,color=:red,label="")
    plot!((x[6][end],z[6][end]),marker=:utriangle,color=:green,label="")

	plot!((x[7][1],z[7][1]),marker=:circle,color=:red,label="")
    plot!((x[7][end],z[7][end]),marker=:utriangle,color=:green,label="")

	plot!((x[8][1],z[8][1]),marker=:circle,color=:red,label="")
    plot!((x[8][end],z[8][end]),marker=:utriangle,color=:green,label="")

	plot!((x[9][1],z[9][1]),marker=:circle,color=:red,label="")
    plot!((x[9][end],z[9][end]),marker=:utriangle,color=:green,label="")

	plot!((x[10][1],z[10][1]),marker=:circle,color=:red,label="")
    plot!((x[10][end],z[10][end]),marker=:utriangle,color=:green,label="")

	io = open("nom_tr.txt", "w")
	for i in prob.U
		for k=1:m-1
    		print(io,i[k])
    		print(io," ")
		end
    	println(io,i[m])
	end
	close(io)


"""
    sv = ds_therm(v, M)

`sv` is an `M` by length(`v`) matrix of ±1s where sum(`sv`) = `v`
and the -1s are located in the bottom rows
"""
function ds_therm(v, M)
    sv = -ones(M, length(v))
    for i = 1:length(v)
        sv[1:((M+v[i])÷2), i] .= 1
    end
    return sv
end


module MyITensorsPkg

using ITensors

export norm2

norm2(A::ITensor) = (A*dag(A))[]

end


using Luna
using Test

@testset "Luna" begin
    @testset "Subtracting two points" begin
        p1 = point(3, 2, 1)
        p2 = point(5, 6, 7)
        @test p1 - p2 == vector(-2, -4, -6)
    end

    @testset "Subtracting a vector from a point" begin
        p = point(3, 2, 1)
        v = vector(5, 6, 7)
        @test p - v == point(-2, -4, -6)
    end

    @testset "Subtracting two vectors" begin
        v1 = vector(3, 2, 1)
        v2 = vector(5, 6, 7)
        @test v1 - v2 == vector(-2, -4, -6)
    end

    @testset "Multiplying by a translation matrix" begin
        t = Luna._translate(5, -3, 2) |> Transform
        p = point(-3, 4, 5)
        @test matrix(t) * p == point(2, 1, 7)
    end

    @testset "Individual transformations are applied in sequence" begin
        A = Luna._rotx(π / 2)
        B = Luna._scale(5, 5, 5)
        C = Luna._translate(10, 5, 7)
        p1 = point(1, 0, 1)
        p2 = A * p1
        p3 = B * p2
        p4 = C * p3
        @test p4 == point(15, 0, 7)
    end

    @testset "Chained transformations are applied in reverse order" begin
        A = Luna._rotx(π / 2)
        B = Luna._scale(5, 5, 5)
        C = Luna._translate(10, 5, 7)
        T = C * B * A
        p = point(1, 0, 1)
        @test T * p == point(15, 0, 7)
    end

    @testset "Fluent API transformations can be chained in order" begin
        T = eye() |> rotx(π / 2) |> scale(5, 5, 5) |> translate(10, 5, 7)
        p = point(1, 0, 1)
        @test T * p == point(15, 0, 7)
    end
end


using PyPlot, AcousticWave
# =============generate synthetic data set================================
nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3;
v  = 3000. * ones(nz, nx);
# initialize source term
pos_src = [101 75; 101 125]; f0 = 30.0; ot = [0.0, 0.0];
src = InitSources(pos_src, f0, ot, dt);
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
iz = collect(5:5:201); ix = 20 * ones(Int64, length(iz));
pos_rec = hcat(iz, ix);
shot = MultiStepForward(pos_rec, src, fidMtx, tmax=0.5);
wt = Ricker(f0, dt); ntw = length(wt);
mu = 0.1; lambda = 120.0; tmax = 0.5; maxit = 10;

(dis, obj) =CG_dis(shot, fidMtx, pos_rec, wt, maxit=10, mu=0.001, tmax=0.5)
dis = reshape(dis, nz, nx)
SeisPlot(dis, name="/Users/wenlei/Desktop/CG.pdf", dx=10, dy=10, xlabel="Length", xunits="(m)", ylabel="Depth", yunits="(m)")
dis = zeros(Float64,nz,nx)
dis[pos_src[1,1], pos_src[1,2]] = 1.; dis[pos_src[2,1], pos_src[2,2]] = 1.;
dis = vec(dis)
ntw = 5*ntw;
(w, obj) = CG_w(ntw, shot, fidMtx, pos_rec, dis, maxit=10, mu=0.001, tmax=0.5)
plot(w);
plot(wt);
xlable = "Time (s)"
ylable = "Amplitude"
savefig("/Users/wenlei/Desktop/wavelet.pdf")

dis = fista(shot, wt, mu, lambda, fidMtx, pos_rec, tmax, maxit)

# (dis, w) = alterMini(w, shot, fidMtx, pos_rec, mu, lambda, outer_it=6, fistait=10, CGit=5, tmax=0.5)
# (dis, w) = alternateCG(w, shot, fidMtx, pos_rec, outer_it=6)

# ============================================================================================================
# using PyPlot, AcousticWave
# nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
# v  = 2500. * ones(nz, nx);
# v[101:end, 1:end] = 3000.;
# vmax = maximum(v); vmin = minimum(v);
#
# Nz = nz + 2*ext; Nx = nx + 2*ext
# w = Ricker(f0, dt); dis = randn(nz*nx); it = floor(Int64, length(w)/2);
# spt = dis2spt(w, dis, nz, nx, ext, iflag, dt, it);
# # SeisPlot(reshape(spt.pz, Nz, Nx));
# # SeisPlot(reshape(spt.px, Nz, Nx));
#
# tmp = reshape(spt.pz, Nz, Nx) + reshape(spt.px, Nz, Nx);
# tmp = tmp[ext+1:ext+nz, ext+1:ext+nx];
# norm(w[it]*dis-vec(tmp))
#
# spt1 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, it)
# spt1.vz = randn(length(spt1.vz));
# spt1.vx = randn(length(spt1.vx));
# spt1.pz = randn(length(spt1.pz));
# spt1.px = randn(length(spt1.px));
# dis1 = spt2dis(w, spt1);
# tmp_spt = dot(spt.vz, spt1.vz) + dot(spt.vx, spt1.vx) + dot(spt.pz, spt1.pz) + dot(spt.px, spt1.px)
# tmp_dis = dot(dis, dis1)

# ===========================================================================================================
# using PyPlot, AcousticWave
# nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
# v  = 2500. * ones(nz, nx);
# v[101:end, 1:end] = 3000.;
# vmax = maximum(v); vmin = minimum(v);
# w = Ricker(f0, dt); it = 33;
# spt2 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, it)
# dis  = zeros(nz, nx); dis[101,75] = 1.0; dis[101,125] = 1.0; dis = vec(dis);
# AddSources!(spt2, w, dis)
# spt3 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, it)
# f0 = 30.0; ot = [0.0,0.0]; pos_src = [101 75;101 125]
# src = InitSources(pos_src, f0, ot, dt);
# AddSources!(spt3, src)
# norm(spt2.vx-spt3.vx)
# norm(spt2.vz-spt3.vz)
# norm(spt2.pz-spt3.pz)
# norm(spt2.px-spt3.px)

# =============================================================================================================
using PyPlot, AcousticWave
nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3;
v  = 3000. * ones(nz, nx);
# initialize source term
pos_src = [101 75; 101 125]; f0 = 30.0; ot = [0.0, 0.0];
src = InitSources(pos_src, f0, ot, dt);
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
iz = collect(5:5:201); ix = 20 * ones(Int64, length(iz));
pos_rec = hcat(iz, ix);
w = Ricker(f0, dt);
dis = zeros(nz, nx); dis[pos_src[1,1], pos_src[1,2]] = 1.0;
dis[pos_src[2,1], pos_src[2,2]] = 1.0; dis = vec(dis);
shot = MultiStepForward(pos_rec, w, dis, fidMtx, tmax=0.5);
dis1 = MultiStepAdjoint(w, shot, fidMtx)

# ============================================================================================================
# dot product test
# using PyPlot, AcousticWave
# nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3;
# v  = 3000. * ones(nz, nx);
# # initialize source term
# pos_src = [101 75; 101 125]; f0 = 30.0; ot = [0.0, 0.0];
# src = InitSources(pos_src, f0, ot, dt);
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# iz = collect(5:5:201); ix = 20 * ones(Int64, length(iz));
# pos_rec = hcat(iz, ix);
# w = Ricker(f0, dt);
# dis = zeros(nz, nz);
# for ix = 1: nx
#     for iz = 1 : nz
#         if rand() <= 0.1
#            dis[iz, ix] = randn()
#         end
#     end
# end
# dis[pos_src[2,1], pos_src[2,2]] = 1.0; dis = vec(dis);
# shot = MultiStepForward(pos_rec, w, dis, fidMtx, tmax=0.5);
#
# nt = shot.nt
# shot1 = InitShot(0, 0, pos_rec, 0.0, nt, dt)
# nrec  = size(pos_rec, 1)
# ls    = floor(Int64, length(w)/2)
# for irec = 1 : nrec
#     tmp  = conv(rand(nt), w)
#     shot1.d[:, irec] = tmp[ls+1:end-ls]
# end
# dis1 = MultiStepAdjoint(w, shot1, fidMtx)
#
# tmp  = dot(dis, dis1)
# tmp1 = dot(vec(shot.d), vec(shot1.d))

# ============================================================================================================
# using PyPlot, AcousticWave
# nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
# v  = 2500. * ones(nz, nx);
# v[101:end, 1:end] = 3000.;
# vmax = maximum(v); vmin = minimum(v);
# Nz = nz + 2*ext; Nx = nx + 2*ext
# w = Ricker(f0, dt); dis = randn(nz*nx);
# nt = length(w)
# spts = Array(SnapShot, nt)
# for it = 1 : nt
#     spts[it] = dis2spt(w, dis, nz, nx, ext, iflag, dt, it)
# end
# it = 36
# a = reshape(spts[it].pz, 261, 261);
# b = a[ext+1:ext+nz, ext+1:ext+nx];
# tmp = b ./ reshape(dis, nz, nx)
#
#
# spts1 = Array(SnapShot, nt)
# for it = 1 : nt
#     spts1[it] = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, it)
#     spts1[it].vz = randn(length(spts1[it].vz))
#     spts1[it].vx = randn(length(spts1[it].vx))
#     spts1[it].pz = randn(length(spts1[it].pz))
#     spts1[it].px = randn(length(spts1[it].px))
# end
# w1 = zeros(nt)
# for it = 1 : nt
#     spt2wlet(w1, dis, spts1[it])
# end
#
# tmp = 0.0
# for it = 1 : nt
#     tmp = tmp + dot(spts[it].vz, spts1[it].vz)
#     tmp = tmp + dot(spts[it].vx, spts1[it].vx)
#     tmp = tmp + dot(spts[it].pz, spts1[it].pz)
#     tmp = tmp + dot(spts[it].px, spts1[it].px)
# end
# tmp1 = dot(w, w1)

# ============================================================================================================
# dot product test
# using PyPlot, AcousticWave
# nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3;
# v  = 3000. * ones(nz, nx);
# # initialize source term
# pos_src = [101 75; 101 125]; f0 = 30.0; ot = [0.0, 0.0];
# src = InitSources(pos_src, f0, ot, dt);
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# iz = collect(5:5:201); ix = 20 * ones(Int64, length(iz));
# pos_rec = hcat(iz, ix);
# w = Ricker(f0, dt);
# dis = zeros(nz, nz);
# for ix = 1: nx
#     for iz = 1 : nz
#         if rand() <= 0.1
#            dis[iz, ix] = randn()
#         end
#     end
# end
# dis = vec(dis)
# shot = MultiStepForward(pos_rec, w, dis, fidMtx, tmax=0.5);
#
#
# nt = shot.nt
# shot1 = InitShot(0, 0, pos_rec, 0.0, nt, dt)
# nrec  = size(pos_rec, 1)
# ls    = floor(Int64, length(w)/2)
# for irec = 1 : nrec
#     tmp  = conv(rand(nt), w)
#     shot1.d[:, irec] = tmp[ls+1:end-ls]
# end
# ntw = length(w)
# w1  = MultiStepAdjoint(ntw, dis, shot1, fidMtx)
#
# tmp  = dot(w, w1)
# tmp1 = dot(vec(shot.d), vec(shot1.d))

# ===========================================================================================================
using PyPlot, AcousticWave
# =============generate synthetic data set================================
nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3;
v  = 3000. * ones(nz, nx);
# initialize source term
pos_src = [101 75; 101 125]; f0 = 30.0; ot = [0.0, 0.0];
src = InitSources(pos_src, f0, ot, dt);
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
iz = collect(5:5:201); ix = 20 * ones(Int64, length(iz));
pos_rec = hcat(iz, ix); tmax = 0.5; maxit=20;
w = Ricker(f0, dt)
lambda = power(maxit, fidMtx, w, pos_rec, tmax)


# Mandatory
export polyhedron, loadpolyhedron!
export hrepiscomputed, hrep, vrepiscomputed, vrep
export volume, volume_simplex, unscaled_volume_simplex, surface

"""
    polyhedron(rep::Representation{T})

Creates a polyhedron from the representation `rep` using the default library included in the Polyhedra package.
"""
polyhedron(rep::Representation{T}) where {T} = polyhedron(rep, default_library(FullDim(rep), T))

"""
    hrepiscomputed(p::Polyhedron)

Returns whether the H-representation of this polyhedron has been computed.
"""
function hrepiscomputed end

"""
    hrep(p::Polyhedron)

Returns an H-representation for the polyhedron `p`.
"""
hrep(p::Polyhedron) = error("`hrep` not implemented for `$(eltype(p))`")

"""
    vrepiscomputed(p::Polyhedron)

Returns whether the V-representation of this polyhedron has been computed.
"""
function vrepiscomputed end

"""
    vrep(p::Polyhedron)

Returns a V-representation for the polyhedron `p`.
"""
vrep(p::Polyhedron) = error("`vrep` not implemented for `$(eltype(p))`")

"""
    volume(p::Polyhedron{T}) where {T}

Returns the `fulldim(p)`-dimensional hyper-volume of the polyhedron `p`.
Returns `Inf` or `-one(T)` if it is infinite depending on whether the type `T` has an infinite value.
"""
function volume(p::Polyhedron)
    return sum(unscaled_volume_simplex(Δ) for Δ in triangulation(p)) / factorial(fulldim(p))
end

function unscaled_volume_simplex(s)
    A = Matrix{coefficient_type(s)}(undef, fulldim(s), npoints(s) - 1)
    v = first(points(s))
    for (i, p) in enumerate(Iterators.drop(points(s), 1))
        A[:, i] = p - v
    end
    return abs(det(A))
end
function volume_simplex(s)
    return unscaled_volume_simplex(s) / factorial(fulldim(s))
end

"""
    surface(p::Polyhedron{T}) where {T}

Returns the `fulldim(p)-1`-dimensional hyper-volume of the surface of the polyhedron `p`.
Returns `Inf` or `-one(T)` if it is infinite depending on whether the type `T` has an infinite value.
"""
function surface end

loadpolyhedron!(p::Polyhedron, filename::AbstractString, extension::Type{Val{:ine}}) = error("loadpolyhedron! not implemented for .ine")
loadpolyhedron!(p::Polyhedron, filename::AbstractString, extension::Type{Val{:ext}}) = error("loadpolyhedron! not implemented for .ext") # FIXME ExtFileVRepresentation or just ExtFile

loadpolyhedron!(p::Polyhedron, filename::AbstractString, extension::Symbol) = loadpolyhedron!(p, filename, Val{extension})

function loadpolyhedron!(p::Polyhedron, filename::AbstractString, extension::AbstractString)
    s = something(findfirst(isequal(extension), ["ext", "ine"]), 0)
    if s == 0
        error("Invalid extension $extension, please give 'ext' for V-representation or 'ine' for H-representation")
    end
    loadpolyhedron!(p, filename, [:ext, :ine][s])
end


module EmpiricalRisks

using ArrayViews

import Base.LinAlg: BlasReal
import Base.LinAlg.BLAS: axpy!, gemv!, gemm!

export

    ## common.jl
    no_op,
    shrink,

    ## loss.jl

    Loss,
    UnivariateLoss,
    MultivariateLoss,

    AbsLoss,
    SqrLoss,
    QuantileLoss,
    EpsilonInsLoss,
    HuberLoss,
    HingeLoss,
    SqrHingeLoss,
    SmoothedHingeLoss,
    LogisticLoss,
    SumLoss,
    SumSqrLoss,
    MultiLogisticLoss,

    value,
    value!,
    deriv,
    grad,
    grad!,
    value_and_deriv,
    value_and_grad,
    value_and_grad!,

    ## prediction.jl

    PredictionModel,
    UnivariatePredictionModel,
    MultivariatePredictionModel,
    LinearPred,
    MvLinearPred,
    AffinePred,
    MvAffinePred,

    inputlen,
    inputsize,
    outputlen,
    outputsize,
    paramlen,
    paramsize,
    ninputs,
    isvalidparam,
    predict,
    predict!,

    # risks

    RiskModel,
    SupervisedRiskModel,
    riskmodel,
    value_and_addgrad!,

    # regularizers

    Regularizer,
    ZeroReg,
    SqrL2Reg,
    L1Reg,
    ElasticReg,
    NonNegReg,
    SimplexReg,
    L1BallReg,

    prox, prox!

# source files
include("common.jl")
include("prediction.jl")
include("loss.jl")
include("risks.jl")
include("regularizer.jl")

end # module


using FastRefine
using LazySets
import FastRefine: forward_network, forward_affine_map, ishull

nnet = read_nnet("nnet/test3.nnet")

delta = 0.05

solver = DimTree(delta)

in_hyper = Hyperrectangle(fill(1.0, 3), fill(1.0, 3))
out_hyper = Hyperrectangle(fill(0.0, 2), fill(10.0, 2))
problem = Problem(nnet, in_hyper, out_hyper)

file = open("results/group5.txt", "a")
print(file, "Test Result of Group 5:\n\n")

#solver4

time4 = 0

for i = 1:1
    timed_result =@timed solve(solver, problem)
    print(file, "DimTree - test " * string(i) * " - Time: " * string(timed_result.time) * " s")
    print(file, " - Output: " * string(timed_result.value) * "\n")
    global time4 += timed_result.time
end

print(file, "Average time: " * string(time4/1) * " s\n\n")

close(file)


"""
Implement some ball matrix arithmetic, i.e., verified matrix operations
for (abstract) interval matrices stored in midpoint-radius format.

The 'radius' for a matrix can be either a matrix of the same dimension, or
a scalar which represents a norm.

A type parameter specifies the "norm type" for the radius. It can be `:Abs` for
entrywise radii, or `OpnormL1`, `OpnormLinf` for normwise radii (specifying which norm to use)
"""

using FastRounding

abstract type NormKind end
struct OpnormL1 <: NormKind end
struct OpnormLinf <: NormKind end
struct Abs <: NormKind end

"""
Certified upper bound to ||A||, abs(A), or whatever is used in the specified radius
"""
function normbound(A::AbstractVecOrMat{T}, ::OpnormL1) where T
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(float(real(zero(T))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for j = 1:size(A, 2)
            nrmj::Tsum = 0
            for i = 1:size(A, 1)
                nrmj = nrmj ⊕₊ abs(A[i,j])
            end
            nrm = max(nrm,nrmj)
        end
    end
    return convert(Tnorm, nrm)
end

function normbound(A::AbstractVecOrMat{T}, ::OpnormLinf) where T
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(float(real(zero(T))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for i = 1:size(A, 1)
            nrmi::Tsum = 0
            for j = 1:size(A, 2)
                nrmi = nrmi ⊕₊ abs(A[i,j])
            end
            nrm = max(nrm,nrmi)
        end
    end
    return convert(Tnorm, nrm)
end

function normbound(A::AbstractArray{T}, ::Abs) where T
    return abs.(A)
end

radiustype(MatrixType, ::OpnormLinf) = Float64
radiustype(MatrixType, ::OpnormL1) = Float64
radiustype(MatrixType, ::Abs) = MatrixType

"""
Type for ball arrays. Comes in two variants, one with cached normbound(midpoint)
(handy if you have to compute many products with it) and one without.
"""
abstract type BallArray{N<:NormKind, MatrixType<:Union{AbstractArray,Number}, RadiusType<:Union{AbstractArray, Number}} end

struct BallArrayWithoutCachedNorm{N, MatrixType, RadiusType} <: BallArray{N, MatrixType, RadiusType}
    midpoint::MatrixType
    radius::RadiusType
    function BallArrayWithoutCachedNorm{N, MatrixType, RadiusType}(midpoint, radius=zero(RadiusType)) where {N, MatrixType, RadiusType}
        if !isa(radius, radiustype(MatrixType, N()))
            error("Wrong radius type specified")
        end
        return new(midpoint, radius)
    end
end

struct BallArrayWithCachedNorm{N, MatrixType, RadiusType} <: BallArray{N, MatrixType, RadiusType}
    midpoint::MatrixType
    radius::RadiusType
    norm::RadiusType
    function BallArrayWithCachedNorm{N, MatrixType, RadiusType}(midpoint, radius=zero(RadiusType)) where {N, MatrixType, RadiusType}
        if !isa(radius, radiustype(MatrixType, N()))
            error("Wrong radius type specified")
        end
        return new{N, MatrixType, RadiusType}(midpoint, radius, normbound(midpoint, N()))
    end
end

function BallArray{N}(midpoint, radius) where {N}
    return BallArrayWithoutCachedNorm{N, typeof(midpoint), typeof(radius)}(midpoint, radius)
end

function BallArrayWithCachedNorm{N}(midpoint, radius) where {N}
    return BallArrayWithCachedNorm{N, typeof(midpoint), typeof(radius)}(midpoint, radius)
end

function midpoint_norm(A::BallArrayWithCachedNorm)
    return A.norm
end

function midpoint_norm(A::BallArrayWithoutCachedNorm{N, MT, RT}) where {N, MT, RT}
    return normbound(A, N())
end

"""
Product between ball matrices / vectors.
"""
function Base.:*(A::BallArray, B::BallArray)
    # A*maximum(B, dims=2) TODO
end

# TODO: using this file for tests

function matvec(A, v)
    w = zero(v)
    rows = rowvals(A)
    vals = nonzeros(A)
    m, n = size(A)
    for j = 1:n
       for k in nzrange(A, j)
          @inbounds i = rows[k]
          @inbounds Aij = vals[k]
          @inbounds w[i] +=  Aij * v[j]
       end
    end
end

function matvec_roundup(A, v)
    w = zero(v)
    rows = rowvals(A)
    vals = nonzeros(A)
    m, n = size(A)
    for j = 1:n
       for k in nzrange(A, j)
          @inbounds i = rows[k]
          @inbounds Aij = vals[k]
          @inbounds w[i] = w[i] ⊕₊ Aij ⊗₊ v[j]
       end
    end
end


abstract type Closure end
@inline closure_relation(γ, br, u, r) = exp.(-u .+ γ .+ br) .- γ .- 1.0

struct HypernettedChain <: Closure end
@inline closure_relation(γ, ::HypernettedChain, u, r) = closure_relation(γ, 0.0, u, r)

struct PercusYevick <: Closure end
@inline function closure_relation(γ, ::PercusYevick, u, r)
    return closure_relation(γ, -γ .+ log.(1.0 .+ γ), u, r)
end

struct ModifiedVerlet{R<:Real} <: Closure
    α::R
    β::R
    δ::R
end
ModifiedVerlet() = ModifiedVerlet{Float64}(-0.5, 0.0, 0.8)
ModifiedVerlet(a::T, b::T, c::T) where {T<:Real} = ModifiedVerlet{T}(a, b, c)

function closure_relation(γ, c::ModifiedVerlet, u, r)
    new_br = @. (c.α * γ^2) * (1.0 + c.β * γ)
    @. new_br /= (1.0 + c.δ * abs(γ))

    return closure_relation(γ, new_br, u, r)
end

struct MeanSpherical <: Closure end

function closure_relation(γ, ::MeanSpherical, u, r)
    cr = similar(r)
    less_one = r .< 1.0
    larger_one = r .> 1.0

    @. cr[less_one] = -1.0 - γ[less_one]
    @. cr[larger_one] = -u[larger_one]

    return cr
end

struct SoftMeanSpherical <: Closure end
function _softmeanspherical!(cr, u1, u2, gamma, idxs)
    @. cr[idxs] = exp(-u1[idxs]) * (1.0 + gamma[idxs] - u2[idxs])
    @. cr[idxs] += -gamma[idxs] - 1.0

    return nothing
end

function closure_relation(γ, ::SoftMeanSpherical, u, r)
    cr = similar(r)
    u₁ = similar(u)
    u₂ = similar(u)

    less_one = r .< 1.0
    larger_one = r .>= 1.0

    u₁[less_one] .= u[less_one]
    u₁[larger_one] .= 0.0
    u₂[less_one] .= 0.0
    u₂[larger_one] .= u[larger_one]

    _softmeanspherical!(cr, u₁, u₂, γ, less_one)
    _softmeanspherical!(cr, u₁, u₂, γ, larger_one)

    return cr
end

struct HMSA{T<:Real} <: Closure
    α::T
end

function _fmix!(fm, alpha, r, idxs)
    @. fm[idxs] = 1.0 - exp(-alpha * r[idxs])

    return nothing
end

function _hmsa(cr, u1, u2, gamma, fm, idxs)
    @. cr[idxs] = exp(-u1[idxs])
    @. cr[idxs] *= 1.0 + ((exp((gamma[idxs] - u2[idxs]) * fm[idxs]) - 1.0) / fm[idxs])
    @. cr[idxs] += -gamma[idxs] - 1.0

    return nothing
end

function closure_relation(γ, cls::HMSA, u, r)
    cr = similar(r)
    f_mix = similar(r)
    u₁ = similar(u)
    u₂ = similar(u)

    less_one = r .< 1.0
    larger_one = r .>= 1.0

    u₁[less_one] .= u[less_one]
    u₁[larger_one] .= 0.0
    u₂[less_one] .= 0.0
    u₂[larger_one] .= u[larger_one]

    _fmix!(f_mix, cls.α, r, less_one)
    _fmix!(f_mix, cls.α, r, larger_one)

    _hmsa(cr, u₁, u₂, γ, f_mix, less_one)
    _hmsa(cr, u₁, u₂, γ, f_mix, larger_one)

    return cr
end


@testset "Jordan and Rational canonical form" begin

  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x^6 + rand(Qx, 1:5, -3:3)
  M = companion_matrix(f)
  @test minpoly(Qx, M) == f
  @test charpoly(Qx, M) == f

  ff = factor(f)
  g = first(keys(ff.fac))
  g = divexact(g, leading_coefficient(g))
  J = Hecke.jordan_block(g, 2)
  @test minpoly(Qx, J) == g^2
  @test charpoly(Qx, J) == g^2

  M = identity_matrix(FlintQQ, 6)
  for i = 1:6
    for j = i+1:6
      M[i, j] = rand(FlintQQ, -2:2)
    end
  end
  M1 = transpose(M)
  @test Hecke.issimilar(M, M1)
  MT = Hecke.conjugating_matrix(M, M1)
  @test MT*M1*inv(MT) == M
  J, S = jordan_normal_form(M)
  @test S*M*inv(S) == J
  for i = 1:6
    for j = 1:i-1
      @test iszero(J[i, j])
    end
  end
  for i = 1:5
    @test isone(J[i, i+1]) || iszero(J[i, i+1])
  end
  for i = 1:6
    for j = i+2:6
      @test iszero(J[i, i+2])
    end
  end

  f = x^3+ rand(Qx, 1:2, -2:2)
  g = x^3 + rand(Qx, 1:2, -2:2)
  f1 = f
  f2 = f*g
  C = zero_matrix(FlintQQ, 9, 9)
  Hecke._copy_matrix_into_matrix(C, 1, 1, companion_matrix(f1))
  Hecke._copy_matrix_into_matrix(C, 4, 4, companion_matrix(f2))
  CF, TM = rational_canonical_form(C)
  @test CF == C
  @test TM * C * inv(TM) == CF

end


function gsreg(
        depvar,
        expvars,
        data;
        intercept=nothing,
        outsample=nothing,
        criteria=nothing,
        ttest=nothing,
        vectoroperation=nothing,
        modelavg=nothing,
        residualtest=nothing,
        time=nothing,
        summary=nothing,
        datanames=nothing,
        datatype=nothing,
        orderresults=nothing,
        onmessage=nothing,
        parallel=nothing
    )
    result = GSRegResult(
        depvar,
        expvars,
        data,
        intercept,
        outsample,
        criteria,
        ttest,
        vectoroperation,
        modelavg,
        residualtest,
        time,
        datanames,
        datatype,
        orderresults,
        onmessage,
        parallel
    )
    proc!(result)
    if summary != nothing
        f = open(summary, "w")
        write(f, to_string(result))
        close(f)
    end
    return result
end

function gsreg_single_proc_result!(
        order,
        data,
        results,
        depvar,
        expvars,
        intercept,
        outsample,
        criteria,
        ttest,
        vectoroperation,
        residualtest,
        time,
        datanames,
        datatype,
        header)

    cols = get_selected_cols(order)
    data_cols_num = size(data, 2)

    if intercept
        append!(cols, data_cols_num)
    end

    depvar_data = @view(data[1:end-outsample, 1])
    expvars_data = @view(data[1:end-outsample, cols])

    nobs = size(depvar_data, 1)
    ncoef = size(expvars_data, 2)
    qrf = qr(expvars_data)
    b = qrf \ depvar_data                   # estimate
    ŷ = expvars_data * b                    # predicted values
    er = depvar_data - ŷ                    # in-sample residuals
    er2 = er .^ 2                           # squared errors
    sse = sum(er2)                          # residual sum of squares
    df_e = nobs - ncoef                     # degrees of freedom
    rmse = sqrt(sse / nobs)                 # root mean squared error
    r2 = 1 - var(er) / var(depvar_data)     # model R-squared

    if ttest
        bstd = sqrt.( sum( (UpperTriangular(qrf.R) \ Matrix(1.0LinearAlgebra.I, ncoef, ncoef) ) .^ 2, dims=2) * (sse / df_e) ) # std deviation of coefficients
    end

    if outsample > 0
        depvar_out = @view(data[end-outsample:end, 1])
        expvars_out = @view(data[end-outsample:end, cols])
        erout = depvar_out - expvars_out * b          # out-of-sample residuals
        sseout = sum(erout .^ 2)                      # residual sum of squares
        rmseout = sqrt(sseout / outsample)            # root mean squared error
        results[order, header[:rmseout]] = rmseout
    end

    results[order, header[:index]] = order
    for (index, col) in enumerate(cols)
        header[Symbol(string(expvars[col-1], "_b"))]
        results[order, header[Symbol(string(expvars[col-1], "_b"))]] = datatype(b[index])
        if ttest
            results[order, header[Symbol(string(expvars[col-1], "_bstd"))]] = datatype(bstd[index])
        end
    end

    results[order, header[:nobs]] = nobs
    results[order, header[:ncoef]] = ncoef
    results[order, header[:sse]] = datatype(sse)
    results[order, header[:r2]] = datatype(r2)
    results[order, header[:rmse]] = datatype(rmse)
    results[order, header[:order]] = 0

    if residualtest != nothing && residualtest

        #jbtest
        x = er
        n = length(x)
        m1 = sum(x) / n
        m2 = sum((x .- m1) .^ 2) / n
        m3 = sum((x .- m1) .^ 3) / n
        m4 = sum((x .- m1) .^ 4) / n
        b1 = (m3 / m2 ^ (3 / 2)) ^ 2
        b2 = (m4 / m2 ^ 2)
        statistic = n * b1 / 6 + n * (b2 - 3) ^ 2 / 24
        d = Chisq(2.)
        jbtest = 1 .- cdf(d, statistic)

        #wtest
        regmatw = hcat((ŷ .^ 2), ŷ , ones(size(ŷ, 1)))
        qrfw = qr(regmatw)
        regcoeffw = qrfw \ er2
        residw = er2 - regmatw * regcoeffw
        rsqw = 1 - dot(residw, residw) / dot(er2, er2) # uncentered R^2
        statisticw = n * rsqw
        wtest = ccdf(Chisq(2), statisticw)

        results[order, header[:wtest]] = wtest
        results[order, header[:jbtest]] = jbtest
        if time != nothing
            e = er
            lag = 1
            xmat = expvars_data

            n = size(e,1)
            elag = zeros(Float64,n,lag)
            for ii = 1:lag  # construct lagged residuals
                elag[ii+1:end, ii] = e[1:end-ii]
            end

            offset = lag
            regmatbg = [xmat[offset+1:end,:] elag[offset+1:end,:]]
            qrfbg = qr(regmatbg)
            regcoeffbg = qrfbg \ e[offset+1:end]
            residbg = e[offset+1:end] .- regmatbg * regcoeffbg
            rsqbg = 1 - dot(residbg,residbg) / dot(e[offset+1:end], e[offset+1:end]) # uncentered R^2
            statisticbg = (n - offset) * rsqbg
            bgtest = ccdf(Chisq(lag), statisticbg)
            results[order, header[:bgtest]] = bgtest
        end
    end

    if vectoroperation == false
        if ttest
            for (index, col) in enumerate(cols)
                pos_b = header[Symbol(string(expvars[col-1], "_b"))]
                pos_bstd = header[Symbol(string(expvars[col-1], "_bstd"))]
                pos_t = header[Symbol(string(expvars[col-1], "_t"))]
                results[order, pos_t] = results[order, pos_b] / results[order, pos_bstd]
            end
        end

        if :aic in criteria
            results[order,header[:aic]] = 2 * results[order, header[:ncoef]] + results[order, header[:nobs]] * log(results[order,header[:sse]] / results[order, header[:nobs]])
        end

        if :aicc in criteria
            if :aic in criteria
                results[order, header[:aicc]] = results[order, header[:aic]] + (2(results[order, header[:ncoef]] + 1) * (results[order, header[:ncoef]]+2)) / (results[order, header[:nobs]] - (results[order, header[:ncoef]] + 1 ) - 1)
            else
                aic = 2 * results[order, header[:ncoef]] + results[order, header[:nobs]] * log(results[order, header[:sse]] / results[order, header[:nobs]])
                results[order, header[:aicc]] = aic + (2(results[order, header[:ncoef]] + 1) * (results[order, header[:ncoef]]+2)) / (results[order, header[:nobs]] - (results[order, header[:ncoef]] + 1 ) - 1)
            end
        end

        if :bic in criteria
            results[order, header[:bic]] = results[order, header[:nobs]] * log.(results[order,header[:rmse]]) + ( results[order, header[:ncoef]] - 1 ) * log.(results[order, header[:nobs]]) + results[order, header[:nobs]] + results[order, header[:nobs]] * log(2π)
        end

        if :r2adj in criteria
            results[order,header[:r2adj]] = 1 - (1 - results[order, header[:r2]]) * ((results[order, header[:nobs]] - 1) / (results[order, header[:nobs]] - results[order, header[:ncoef]]))
        end

        results[order, header[:F]] = (results[order, header[:r2]] / (results[order,header[:ncoef]] - 1)) / ((1 - results[order, header[:r2]]) / (results[order, header[:nobs]] - results[order, header[:ncoef]]))
    end
end

function gsreg_proc_result!(num_job, num_jobs, ops_by_worker, opts...)
    for j = 1:ops_by_worker
        order = (j-1) * num_jobs + num_job
        gsreg_single_proc_result!(order, opts...)
    end
end

function proc!(result::GSRegResult)
    expvars_num = size(result.expvars, 1)
    if result.intercept
        expvars_num = expvars_num-1
    end
    num_operations = 2 ^ expvars_num - 1

    result.onmessage("Creating environment")

    pdata = convert(SharedArray, result.data)
    presults = fill!(SharedArray{result.datatype}(num_operations, length(keys(result.header))),NaN)

    result.onmessage("Doing $num_operations regressions")

    if nprocs() == nworkers()
        for order = 1:num_operations
            gsreg_single_proc_result!(order, pdata, presults, result.depvar, result.expvars, result.intercept, result.outsample, result.criteria, result.ttest, result.vectoroperation,  result.residualtest, result.time, result.datanames, result.datatype, result.header)
        end
    else
        num_workers = (result.parallel != nothing) ? result.parallel : nworkers()
        ops_by_worker = div(num_operations, num_workers)
        num_jobs = (num_workers > num_operations) ? num_operations : num_workers
        remainder = num_operations - ops_by_worker * num_jobs
        jobs = []
        for num_job = 1:num_jobs
            push!(jobs, @spawnat num_job+1 gsreg_proc_result!(num_job, num_jobs, ops_by_worker, pdata, presults, result.depvar, result.expvars, result.intercept, result.outsample, result.criteria, result.ttest, result.vectoroperation, result.residualtest, result.time, result.datanames, result.datatype, result.header))
        end
        for job in jobs
            fetch(job)
        end

        if( remainder > 0 )
            for j = 1:remainder
                order = j + ops_by_worker * num_jobs
                gsreg_single_proc_result!(order, pdata, presults, result.depvar, result.expvars, result.intercept, result.outsample, result.criteria, result.ttest, result.vectoroperation, result.residualtest, result.time, result.datanames, result.datatype, result.header)
            end
        end
    end

    result.results = Array(presults)

    presult = nothing
    pdata = nothing

    if result.vectoroperation
        result.onmessage("Calculating Vector operations")
        if result.ttest
            for expvar in result.expvars
                pos_b = result.header[Symbol(string(expvar, "_b"))]
                pos_bstd = result.header[Symbol(string(expvar, "_bstd"))]
                pos_t = result.header[Symbol(string(expvar, "_t"))]
                result.results[:,pos_t] = result.results[:,pos_b] ./ result.results[:,pos_bstd]
            end
        end

        if :aic in result.criteria || :aicc in result.criteria
            result.results[:,result.header[:aic]] = 2 * result.results[:,result.header[:ncoef]] .+ result.results[:,result.header[:nobs]] .* log.(result.results[:,result.header[:sse]] ./ result.results[:,result.header[:nobs]])
        end

        if :aicc in result.criteria
            result.results[:,result.header[:aicc]] = result.results[:,result.header[:aic]] .+ (2(result.results[:,result.header[:ncoef]] .+ 1) .* (result.results[:,result.header[:ncoef]] .+ 2)) ./ (result.results[:,result.header[:nobs]] .- (result.results[:,result.header[:ncoef]] .+ 1 ) .- 1)
        end

        if :bic in result.criteria
            result.results[:,result.header[:bic]] = result.results[:,result.header[:nobs]] .* log.(result.results[:,result.header[:rmse]]) .+ ( result.results[:,result.header[:ncoef]] .- 1 ) .* log.(result.results[:,result.header[:nobs]]) .+ result.results[:,result.header[:nobs]] .+ result.results[:,result.header[:nobs]] .* log(2π)
        end

        if :r2adj in result.criteria
            result.results[:,result.header[:r2adj]] = 1 .- (1 .- result.results[:,result.header[:r2]]) .* ((result.results[:,result.header[:nobs]] .- 1) ./ (result.results[:,result.header[:nobs]] .- result.results[:,result.header[:ncoef]]))
        end

        result.results[:,result.header[:F]] = (result.results[:,result.header[:r2]] ./ (result.results[:,result.header[:ncoef]] .- 1)) ./ ((1 .- result.results[:,result.header[:r2]]) ./ (result.results[:,result.header[:nobs]] .- result.results[:,result.header[:ncoef]]))
    end

    # CP must be computed with vector operations
    if :cp in result.criteria
        result.results[:,result.header[:cp]] = (result.results[:,result.header[:nobs]] .- maximum(result.results[:,result.header[:ncoef]]) .- 2) .* (result.results[:,result.header[:rmse]] ./ minimum(result.results[:,result.header[:rmse]])) .- (result.results[:,result.header[:nobs]] .- 2 .* result.results[:,result.header[:ncoef]])
    end
    len_criteria = length(result.criteria)
    for criteria in result.criteria
        result.results[:, result.header[:order]] += AVAILABLE_CRITERIA[criteria]["index"] * (1 / len_criteria) * ( (result.results[:, result.header[criteria]] .- mean(result.results[:, result.header[criteria]]) ) ./ std(result.results[:, result.header[criteria]]) )
    end
    if result.modelavg
        result.onmessage("Generating model averaging results")
        # usar order para weight
        delta = maximum(result.results[:,result.header[:order]]) .- result.results[:,result.header[:order]]
        w1 = exp.(-delta/2)
        result.results[:,result.header[:weight]] = w1./sum(w1)
        result.average = Array{Float64}(undef, 1, length(keys(result.header)))
        weight_pos = (result.ttest) ? 4 : 2
        for expvar in result.expvars
            obs = result.results[:,result.header[Symbol(string(expvar, "_b"))]]
            if result.ttest
                obs = hcat(obs, result.results[:,result.header[Symbol(string(expvar, "_bstd"))]])
                obs = hcat(obs, result.results[:,result.header[Symbol(string(expvar, "_t"))]])
            end
            obs = hcat(obs, result.results[:,result.header[:weight]])

            #filter NaN values from selection
            obs = obs[findall(x -> !isnan(obs[x,1]), 1:size(obs,1)),:]

            #weight resizing
            obs[:, weight_pos] /= sum(obs[:, weight_pos])

            result.average[result.header[Symbol(string(expvar, "_b"))]] = sum(obs[:, 1] .* obs[:, weight_pos])
            if result.ttest
                result.average[result.header[Symbol(string(expvar, "_bstd"))]] = sum(obs[:, 2] .* obs[:, weight_pos])
                result.average[result.header[Symbol(string(expvar, "_t"))]] = sum(obs[:, 3] .* obs[:, weight_pos])
            end
        end

        for criteria in [:nobs, :r2adj, :F, :order]
            result.average[result.header[criteria]] = sum(result.results[:, result.header[criteria]] .* result.results[:, result.header[:weight]])
        end
    end

    if result.orderresults
        result.onmessage("Sorting results")
        result.results = gsregsortrows(result.results, [result.header[:order]]; rev=true)
        result.bestresult = result.results[1,:]
    else
        result.onmessage("Looking for the best result")
        max_order = result.results[1,result.header[:order]]
        best_result_index = 1
        for i = 1:num_operations
            if result.results[i,result.header[:order]] > max_order
                max_order = result.results[i,result.header[:order]]
                best_result_index = i
            end
        end
        result.bestresult = result.results[best_result_index,:]
    end
end

function to_string(result::GSRegResult)
    out = ""
    out *= @sprintf("\n")
    out *= @sprintf("══════════════════════════════════════════════════════════════════════════════\n")
    out *= @sprintf("                              Best model results                              \n")
    out *= @sprintf("══════════════════════════════════════════════════════════════════════════════\n")
    out *= @sprintf("                                                                              \n")
    out *= @sprintf("                                     Dependent variable: %s                   \n", result.depvar)
    out *= @sprintf("                                     ─────────────────────────────────────────\n")
    out *= @sprintf("                                                                              \n")
    out *= @sprintf(" Selected covariates                 Coef.")
    if result.ttest
        out *= @sprintf("        Std.         t-test")
    end
    out *= @sprintf("\n")
    out *= @sprintf("──────────────────────────────────────────────────────────────────────────────\n")

    cols = get_selected_cols(Int64(result.bestresult[result.header[:index]]))

    data_cols_num = size(result.data, 2)
    if result.intercept
        append!(cols, data_cols_num)
    end

    for pos in cols
        varname = result.expvars[pos-1]
        out *= @sprintf(" %-35s", varname)
        out *= @sprintf(" %-10f", result.bestresult[result.header[Symbol(string(varname, "_b"))]])
        if result.ttest
            out *= @sprintf("   %-10f", result.bestresult[result.header[Symbol(string(varname, "_bstd"))]])
            out *= @sprintf("   %-10f", result.bestresult[result.header[Symbol(string(varname, "_t"))]])
        end
        out *= @sprintf("\n")
    end

    out *= @sprintf("──────────────────────────────────────────────────────────────────────────────\n")
    out *= @sprintf(" Observations                        %-10d\n", result.bestresult[result.header[:nobs]])
    out *= @sprintf(" F-statistic                         %-10f\n", result.bestresult[result.header[:F]])
    for criteria in result.criteria
        if AVAILABLE_CRITERIA[criteria]["verbose_show"]
    out *= @sprintf(" %-30s      %-10f\n", AVAILABLE_CRITERIA[criteria]["verbose_title"], result.bestresult[result.header[criteria]])
        end
    end

    if !result.modelavg
        out *= @sprintf("──────────────────────────────────────────────────────────────────────────────\n")
    end

    if result.modelavg
        out *= @sprintf("\n")
        out *= @sprintf("\n")
        out *= @sprintf("══════════════════════════════════════════════════════════════════════════════\n")
        out *= @sprintf("                            Model averaging results                           \n")
        out *= @sprintf("══════════════════════════════════════════════════════════════════════════════\n")
        out *= @sprintf("                                                                              \n")
        out *= @sprintf("                                     Dependent variable: %s                   \n", result.depvar)
        out *= @sprintf("                                     ─────────────────────────────────────────\n")
        out *= @sprintf("                                                                              \n")
        out *= @sprintf(" Covariates                          Coef.")
        if result.ttest
            out *= @sprintf("        Std.         t-test")
        end
        out *= @sprintf("\n")
        out *= @sprintf("──────────────────────────────────────────────────────────────────────────────\n")

        for varname in result.expvars
            out *= @sprintf(" %-35s", varname)
            out *= @sprintf(" %-10f", result.average[result.header[Symbol(string(varname, "_b"))]])
            if result.ttest
                out *= @sprintf("   %-10f", result.average[result.header[Symbol(string(varname, "_bstd"))]])
                out *= @sprintf("   %-10f", result.average[result.header[Symbol(string(varname, "_t"))]])
            end
            out *= @sprintf("\n")
        end
        out *= @sprintf("\n")
        out *= @sprintf("──────────────────────────────────────────────────────────────────────────────\n")
        out *= @sprintf(" Observations                        %-10d\n", result.average[result.header[:nobs]])
        out *= @sprintf(" F-statistic                         %-10f\n", result.average[result.header[:F]])
        out *= @sprintf(" Combined criteria                   %-10f\n", result.average[result.header[:order]])
        out *= @sprintf("──────────────────────────────────────────────────────────────────────────────\n")

    end
    return out
end

function to_dict(res::GSRegResult)
    result = Dict()
    for (header, index) in res.header
        push!(result, Pair(string(header), res.bestresult[index]))
    end

    dic = Dict(
        "selected_cols" => map( x -> res.datanames[x], get_selected_cols(Int64(res.bestresult[res.header[:index]]))),
        "bestresult" => result
    )

    if res.modelavg
        average = Dict()
        for (header, index) in res.header
            push!(average, Pair(string(header), res.average[index]))
        end
        push!(dic, Pair("avgresults", average))
    end

    dic
end

Base.show(io::IO, result::GSRegResult) = print(io, to_string(result))


module MLJExampleInterface

# Example.jl doesn't actually provide machine learning models, so we
# provide a module below with the same name to furnish us with simple
# constant probabilistic classification. Note that `Example.fit`
# ignores the training features `Xmatrix`.

module Example

function fit(Xmatrix::Matrix, yint::AbstractVector{<:Integer})
    classes = sort(unique(yint))
    counts = [count(==(c), yint) for c in classes]
    Θ = counts / sum(counts)
end

predict(Xnew::Matrix, Θ) = vcat(fill(Θ', size(Xnew, 1))...)

# julia> yint = rand([1,3,4], 100);

# julia> Θ = fit(rand(100, 3), yint)
# 3-element Vector{Float64}:
#  0.35
#  0.23
#  0.42

# julia> predict(rand(5, 3), Θ)
# 5×3 Matrix{Float64}:
#  0.35  0.23  0.42
#  0.35  0.23  0.42
#  0.35  0.23  0.42
#  0.35  0.23  0.42
#  0.35  0.23  0.42

end # of module

### CONTINUATION OF TEMPLATE

import .Example # substitute model-providing package name here (no dot)
import MLJModelInterface
import ScientificTypesBase

const PKG = "Example"          # substitute model-providing package name
const MMI = MLJModelInterface
const STB = ScientificTypesBase

"""
    CoolProbabilisticClassifier()

A cool classifier that predicts `UnivariateFinite` probability
distributions. These are distributions for a finite sample space whose
elements are *labeled*.

"""
MMI.@mlj_model mutable struct CoolProbabilisticClassifier <: MMI.Probabilistic
    dummy_hyperparameter1::Float64 = 1.0::(_ ≥ 0)
    dummy_hyperparameter2::Int = 1::(0 < _ ≤ 1)
    dummy_hyperparameter3
end

function MMI.fit(::CoolProbabilisticClassifier, verbosity, X, y)

    Xmatrix = MMI.matrix(X)

    yint = MMI.int(y)
    decode = MMI.decoder(y[1])                # for decoding int repr.
    classes_seen = decode(sort(unique(yint))) # ordered by int repr.

    Θ = Example.fit(Xmatrix, yint)            # probability vector
    fitresult = (Θ, classes_seen)
    report = (n_classes_seen = length(classes_seen),)
    cache = nothing

    return fitresult, cache, report

end

function MMI.predict(::CoolProbabilisticClassifier, fitresult, Xnew)
    Xmatrix = MMI.matrix(Xnew)

    Θ, classes_seen = fitresult
    prob_matrix = Example.predict(Xmatrix, Θ)

    # `classes_seen` is a categorical vector whose pool actually
    # includes *all* classes. The `UnivariateFinite` constructor
    # automatically assigns zero probability to the unseen classes.

    return MMI.UnivariateFinite(classes_seen, prob_matrix)
end

# for returning user-friendly form of the learned parameters:
function MMI.fitted_params(::CoolProbabilisticClassifier, fitresult)
    Θ, classes_seen = fitresult
    return (raw_probabilities = Θ, classes_seen_in_training = classes_seen)
end

## META DATA

MMI.metadata_pkg(CoolProbabilisticClassifier,
             name="$PKG",
             uuid="7876af07-990d-54b4-ab0e-23690620f79a",
             url="https://github.com/JuliaLang/Example.jl",
             is_pure_julia=true,
             license="MIT",
             is_wrapper=false
)

MMI.metadata_model(CoolProbabilisticClassifier,
               input_scitype = MMI.Table(STB.Continuous),
               target_scitype = AbstractVector{<:STB.Finite},# ie, a classifier
               docstring = "Really cool classifier",         # brief description
               path = "$PKG.CoolProbabilisiticClassifier"
               )

end # module


rand_reward{T<:FloatingPoint}(::Type{T}, dims...) = 1 .- 2*rand(T, dims...)

rand_reward{T<:Integer}(::Type{T}, dims...) = rand(T, dims...)

sum_columns_to_one(a::Array) = a ./ sum(a, 1)

# replace all elements of `a` with 0 where `mask` is `false`
function zero_mask!{T}(a::AbstractArray{T,3}, mask::AbstractArray{Bool,2})
    all(sum(mask, 1) .> 0) ||
        error("All columns of mask must have at least one 'true' value.")
    size(mask) == size(a)[1:2] ||
        error("Mask size must be `states`×`states`.")
    for k = 1:size(a)[3]
        setindex!(sub(a, :, :, k), zero(T), find(!mask))
    end
end

function zero_mask!{T}(a::AbstractArray{T,3}, mask::AbstractArray{Bool,3})
    all(sum(mask, 1) .> 0) ||
        error("All columns of mask must have at least one 'true' value.")
    size(mask) == size(a) ||
        error("Mask size must be `states`×`states`×`actions`.")
    setindex!(a, zero(T), find(!mask))
end

function random{P<:FloatingPoint,R<:Real,N}(
    ::Type{P},
    ::Type{R},
    states,
    actions,
    mask::Nullable{Array{Bool,N}},
)
    transition = rand(P, states, states, actions)
    isnull(mask) ? nothing : zero_mask!(transition, get(mask))
    transition = sum_columns_to_one(transition)
    reward = rand_reward(R, states, actions)
    (transition, reward)
end

random{P,R}(::Type{P}, ::Type{R}, states, actions, mask::Array{Bool}) =
    random(P, R, states, actions, Nullable(mask))

random{P,R}(::Type{P}, ::Type{R}, states, actions) =
    random(P, R, states, actions, Nullable{Array{Bool,1}}())

random{N}(states, actions, mask::Nullable{Array{Bool,N}}) =
    random(Float64, Float64, states, actions, mask)

random(states, actions, mask::Array{Bool}) =
 random(states, actions, Nullable(mask))

random(states, actions) = random(states, actions, Nullable{Array{Bool,1}}())


@testset "Test `-`" begin
    @test typeof(-EngineeringStrain(1:6)) == EngineeringStrain{Int}
    @test typeof(-EngineeringStrain(1.0:6.0)) == EngineeringStrain{Float64}
    @test -EngineeringStrain(1:6) == EngineeringStrain(-1, -2, -3, -4, -5, -6)
    @test -EngineeringStress(1.0:6.0) == EngineeringStress(-1.0, -2, -3, -4, -5, -6)
    @test -(-EngineeringStress(1.0:6.0)) == EngineeringStress(1.0:6.0)
end

@testset "Test `*` and `/`" begin
    a = EngineeringStrain(1:6)
    @test typeof(2a) == EngineeringStrain{Int}
    @test typeof(2.0a) == EngineeringStrain{Float64}
    @test 2a == 2 * a == a * 2 == 2.0a == EngineeringStress(2:2:12)
    @test a / 2 == EngineeringStress(0.5, 1, 1.5, 2, 2.5, 3)
    @test 2a / 2 == a
end

@testset "Test addition and subtraction" begin
    a = EngineeringStrain(1:6)
    b = -EngineeringStrain(1:1:6)
    @test a - b == 2a
    @test b - a == -2a
    @test a + b == EngineeringStrain(zeros(6))
    a = TensorStress(EngineeringStress(1:6))
    b = -TensorStress(EngineeringStress(1:6))
    @test a - b == 2a
    @test b - a == -2a
    @test a + b == TensorStress(zeros(3, 3))
end


using BitCircuits
using Base.Test

or(left, right) = left | right
and(left, right) = left & right

# Demo using the following equivalency:
# p <- ab + ac
# q <- a (b + c)
# p = q

a = Variable(0)
b = Variable(1)
c = Variable(2)

ab = Operation(and, a, b)
ac = Operation(and, a, c)
p = Operation(or, ab, ac)
boc = Operation(or, b, c)
q = Operation(and, a, boc)

@test equal(p, q)
@test evaluate(ab, 0b000011) == true
@test evaluate(ab, 0b000010) == false
@test evaluate(ab, 0b000001) == false
@test evaluate(ab, 0b000000) == false


using Test
using CUSTATEVEC
import CUSTATEVEC: CuStateVec, applyMatrix!, expectation, sample

@testset "CUSTATEVEC" begin
    # build a simple state and compute expectations
    n_q = 2
    @testset for elty in [ComplexF32, ComplexF64]
        H = convert(Matrix{elty}, (1/√2).*[1 1; 1 -1])
        X = convert(Matrix{elty}, [0 1; 1 0])
        Z = convert(Matrix{elty}, [1 0; 0 -1])
        sv = CuStateVec(elty, n_q)
        sv = applyMatrix!(sv, H, false, Int32[0], Int32[])
        exp, res = expectation(sv, Z, Int32[0])
        @test exp ≈ 0.0 atol=1e-6
        exp, res = expectation(sv, X, Int32[0])
        @test exp ≈ 1.0 atol=1e-6
    end
    # build a simple state with controls and compute expectations
    n_q = 2
    @testset for elty in [ComplexF32, ComplexF64]
        H = convert(Matrix{elty}, (1/√2).*[1 1; 1 -1])
        X = convert(Matrix{elty}, [0 1; 1 0])
        Z = convert(Matrix{elty}, [1 0; 0 -1])
        sv = CuStateVec(elty, n_q)
        sv = applyMatrix!(sv, H, false, Int32[0], Int32[])
        sv = applyMatrix!(sv, X, false, Int32[1], Int32[0]) # CNOT
        exp, res = expectation(sv, Z, Int32[0])
        @test exp ≈ 0.0 atol=1e-6
        exp, res = expectation(sv, X, Int32[0])
        @test exp ≈ 0.0 atol=1e-6
    end
    # build a simple state and compute samples
    n_q = 10
    @testset for elty in [ComplexF32, ComplexF64]
        H = convert(Matrix{elty}, (1/√2).*[1 1; 1 -1])
        sv = CuStateVec(elty, n_q)
        for q in 0:n_q-1
            sv = applyMatrix!(sv, H, false, Int32[q], Int32[])
        end
        samples = sample(sv, Int32.(collect(0:n_q-1)), 10)
        # check that this succeeded
        @test length(samples) == 10
    end
end


import AdaptiveRejectionSampling
ars = AdaptiveRejectionSampling

using Test

@testset "Line" begin
    @test ars.Line(2.0, 3) isa ars.Line
    @test ars.intersection(ars.Line(-1.0, 1.0), ars.Line(1.0, -1.0)) == 1.0
    @test_throws AssertionError ars.intersection(ars.Line(-1.0, 1.0), ars.Line(-1.0, -1.0))
end

@testset "Envelop" begin
    @test begin
        l1 = ars.Line(1.0, 1.0)
        l2 = ars.Line(-3.0, 2.0)
        support = (-Inf, Inf)
        ars.Envelop([l1, l2], support)
    end isa ars.Envelop
    @test_throws AssertionError begin
        l1 = ars.Line(1.0, 1.0)
        l2 = ars.Line(-3.0, 2.0)
        ars.Envelop([l2, l1], (-Inf, Inf))
    end
    @test_throws AssertionError begin
        l1 = ars.Line(1.0, 1.0)
        l2 = ars.Line(-3.0, 2.0)
        ars.Envelop([l2, l1], (1.0, -1.0))
    end
    @test begin
        l1 = ars.Line(1.0, 1.0)
        l2 = ars.Line(-3.0, 2.0)
        l3 = ars.Line(-5.0, 5.0)
        e = ars.Envelop([l1, l2, l3], (-Inf, Inf))
        e.cutpoints == [-Inf, 0.25, 1.5, Inf]
    end
    @test begin
        l1 = ars.Line(1.0, 1.0)
        newline = ars.Line(0.0, 0.0)
        l3 = ars.Line(-1.0, 1.0)
        e = ars.Envelop([l1, l3], (-Inf, Inf))
        ars.add_segment!(e, newline)
        e.lines[2] == newline
    end
end

@testset "Objective" begin
    @test ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) isa ars.Objective
    @test begin
        objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi))
        objective.logf(0.0) ≈ 0.3989422804014327
    end
    @test begin
        objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi))
        objective.grad(0.0) ≈ 0.0
    end
    @test begin
        objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi))
        objective.logf(Inf) == 0.0 && objective.logf(Inf) == 0.0
    end
end

@testset "RejectionSampler" begin
    @test ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf), (-1.0, 1.0)) isa ars.RejectionSampler
    @test begin
        sampler = ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf), (-1.0, 1.0))
        ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat
    end
    @test begin
        sampler = ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf))
        ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat
    end
    @test_throws AssertionError begin
        sampler = ars.RejectionSampler(x -> exp(-0.5 * x) / sqrt(2pi), (-Inf, Inf))
        ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat
    end
end


## Include needed packages
using IncrementalInference
using RoMEPlotting

# import getSample to be extended for user factor MultiModalConditional
import IncrementalInference: getSample

## create a new facor type MultiModalConditional
## FIXME, this approach is unnecessary, see `::Mixture` instead.  See Caesar.jl documentation for details.
mutable struct MultiModalConditional <: AbstractRelativeRoots
  x::Vector{Distribution}
  hypo::Categorical
  MultiModalConditional(x::Vector{<:Distribution}, p::Categorical) = new(x, p)
end
function getSample(cf::CalcFactor{<:MultiModalConditional}, N::Int=1)
  d = length(cf.factor.hypo.p)
  p = rand(cf.factor.hypo, N)
  ret = zeros(1,N)
  for i in 1:N
    ret[i] = rand(cf.factor.x[p[i]])
  end
  return (ret, p)
end

function (cf::CalcFactor{<:MultiModalConditional})(meas, x1, x2)
  #
  return meas[1] - (x2[1]-x1[1])
end

## build factor graph and populate
fg = initfg()

N=100

doors = [-20.0 0.0 20.0]
pd = kde!(doors,[2.0])
pd = resample(pd,N);
bws = getBW(pd)[:,1]
doors2 = getPoints(pd);
v1 = addVariable!(fg,:x1,ContinuousScalar,N=N)
f1 = addFactor!(fg,[v1],Prior(pd)) #, samplefnc=getSample

# not initialized
v2 = addVariable!(fg,:x2, ContinuousScalar, N=N)

mmc = MultiModalConditional([Normal(-5,0.5),Normal(5,0.5)],Categorical([0.5,0.5]))
f2 = addFactor!(fg, [:x1; :x2], mmc )

pts = approxConv(fg, :x1x2f1, :x2)

## do some plotting
meas = sampleFactor(f2,2000)
q2 = kde!(meas)
h1 = plotKDE([getBelief(v1), q2],c=["red";"green"],fill=true, xlbl="")
h2 = plotKDE(kde!(pts),fill=true,xlbl="", title="N = 100")

draw(PDF("approxconv.pdf",14cm,10cm),vstack(h1,h2))
# @async run(`evince approxconv.pdf`)


# # Heat Equation
#
# ## Introduction
#
# The heat equation is the "Hello, world!" equation of finite elements.
# Here we solve the equation on a unit square, with a uniform internal source.
# The strong form of the (linear) heat equation is given by
#
# ```math
#  -\nabla \cdot (k \nabla u) = f  \quad x \in \Omega,
# ```
#
# where $u$ is the unknown temperature field, $k$ the heat conductivity,
# $f$ the heat source and $\Omega$ the domain. For simplicity we set
# $k = 1$. We will consider homogeneous Dirichlet boundary conditions such that
# ```math
# u(x) = 0 \quad x \in \partial \Omega,
# ```
# where $\partial \Omega$ denotes the boundary of $\Omega$.
#
# The resulting weak form is given by
# ```math
# \int_{\Omega} \nabla v \cdot \nabla u \ d\Omega = \int_{\Omega} f\cdot v \ d\Omega,
# ```
# where $v$ is a suitable test function.
#-
# ## Commented Program
using HDiscontinuousGalerkin
using Tensors
using SparseArrays

# We start  generating a simple grid with 20x20 quadrilateral elements
# using `generate_grid`. The generator defaults to the unit square,
# so we don't need to specify the corners of the domain.
mesh = rectangle_mesh(TriangleCell, (10,10), Vec{2}((0.0,0.0)), Vec{2}((1.0,1.0)));

# ### Initiate function Spaces
dim = 2
Wh = ScalarFunctionSpace(mesh, ContinuousLagrange{dim,RefTetrahedron,1}(); update_face_data = false)

# Declare variables
u_h = TrialFunction(Wh)

# ### Degrees of freedom
# Next we need to define a `DofHandler`, which will take care of numbering
# and distribution of degrees of freedom for our approximated fields.
# We create the `DofHandler` and then add a single field called `u`.
# Lastly we `close!` the `DofHandler`, it is now that the dofs are distributed
# for all the elements.
dh = DofHandler([u_h], mesh);

# Now that we have distributed all our dofs we can create our tangent matrix,
# using `create_sparsity_pattern`. This function returns a sparse matrix
# with the correct elements stored.
K = create_sparsity_pattern(dh);

# We can inspect the pattern using the `spy` function from `UnicodePlots.jl`.
# By default the stored values are set to $0$, so we first need to
# fill the stored values, e.g. `K.nzval` with something meaningful.
using UnicodePlots
fill!(K.nzval, 1.0);
spy(K; height = 15)

# ### Boundary conditions
dbc = Dirichlet(u_h, dh, "boundary", [0.0])

# ### RHS function
f(x::Vec{dim}) = 2*π^2*sin(π*x[1])*sin(π*x[2])

# ### Assembling the linear system
# Now we have all the pieces needed to assemble the linear system, $K u = f$.
function doassemble(Wh, K::SparseMatrixCSC, dh::DofHandler)
    # We allocate the element stiffness matrix and element force vector
    # just once before looping over all the cells instead of allocating
    # them every time in the loop.
    n_basefuncs = getnbasefunctions(Wh)
    Ke = zeros(n_basefuncs, n_basefuncs)
    fe = zeros(n_basefuncs)
    b = zeros(ndofs(dh))
    cell_dofs = Vector{Int}(undef, ndofs_per_cell(dh))

    # Next we define the global force vector `f` and
    # create an assembler. The assembler
    # is just a thin wrapper around `f` and `K` and some extra storage
    # to make the assembling faster.
    assembler = start_assemble(K, b)

    # It is now time to loop over all the cells in our grid.
    @inbounds for (cellcount, cell) in enumerate(CellIterator(mesh))
        # Always remember to reset the element stiffness matrix and
        # force vector since we reuse them for all elements.
        fill!(Ke, 0)
        fill!(fe, 0)
        reinit!(Wh, cell)

        # It is now time to loop over all the quadrature points in the cell and
        # assemble the contribution to `Ke` and `fe`. The integration weight
        # can be queried from `cellvalues` by `getdetJdV`.
        for q_point in 1:getnquadpoints(Wh)
            dΩ = getdetJdV(Wh, q_point)
            fh = function_value(f, Wh, cell, q_point)
            # For each quadrature point we loop over all the (local) shape functions.
            # We need the value and gradient of the testfunction `v` and also the gradient
            # of the trial function `u`. We get all of these from `cellvalues`.
            for i in 1:n_basefuncs
                v  = shape_value(Wh, q_point, i)
                ∇v = shape_gradient(Wh, q_point, i)
                fe[i] += fh*v * dΩ
                for j in 1:n_basefuncs
                    ∇u = shape_gradient(Wh, q_point, j)
                    Ke[i, j] += (∇v ⋅ ∇u) * dΩ
                end
            end
        end
        # The last step in the element loop is to assemble `Ke` and `fe`
        # into the global `K` and `f` with `assemble!`.
        celldofs!(cell_dofs, dh, cell)
        assemble!(assembler, cell_dofs, fe, Ke)
    end
    return K, b
end
#md nothing # hide

# ### Solution of the system
# The last step is to solve the system. First we call `doassemble`
# to obtain the global stiffness matrix `K` and force vector `f`.
K, b = doassemble(Wh, K, dh);

# To account for the boundary conditions we use the `apply!` function.
# This modifies elements in `K` and `f` respectively, such that
# we can get the correct solution vector `u` by using `\`.
apply!(K,b,dbc);
u = K \ b;

reconstruct!(u_h, u, dh)
#
# # ### Compute errors
u_ex(x::Vec{dim}) = sin(π*x[1])*sin(π*x[2])
Etu_h = errornorm(u_h, u_ex)
Etu_h <= 0.0002
#
### Plot Solution
# TODO: Turn into a recipe
using Makie, CairoMakie
#AbstractPlotting.__init__()
using ColorSchemes
connectivity = getcells_matrix(mesh)
coordinates = get_vertices_matrix(mesh)
color = u
scene = Scene(resolution = (400, 200), colorrange = (0.0, 1.0), colormap = ColorSchemes.viridis.colors)
#poly!(scene, coordinates, connectivity, color = color, strokecolor = (:black, 0.6), strokewidth = 1)
for row in 1:size(connectivity)[1]
	#read coordinates
	points = node(:poly, Point2f0[coordinates[i,:] for i in connectivity[row,:]])
	colors = get(ColorSchemes.viridis, [color[i] for i in connectivity[row,:]])
	poly!(scene, points, strokewidth = 1, strokecolor = (:black,0.6), color = colors, show_axis = true, scale_plot = false)
end
display(scene)


if basename(pwd()) == "aoc"
    cd("2016/8")
end

function part1(filename::AbstractString, screensize::Tuple{Int,Int})
    screen = falses(screensize)
    for instr in eachline(filename)
        A, B = parse.(Int, match(r"(\d+)(?:x| by )(\d+)", instr))
        if startswith(instr, "rect")
            screen[1:B, 1:A] .= true
        elseif startswith(instr, "rotate row")
            screen[A+1, :] = circshift(screen[A+1, :], B)
        else
            screen[:, A+1] = circshift(screen[:, A+1], B)
        end
    end
    screen
end

function prettyprint(A)
    for row in eachrow(A)
        println((b ? '#' : '.' for b in row)...)
    end
    println()
end

# part1("example.txt", (3, 7))
result = part1("input.txt", (6, 50))
count(result)
prettyprint(result)


# Goals: minimize calls to ordered_states (allocates memory)

# needs pomdp for stateindex in pdf(b, s)
# needs list of ordered_states for rand(b)

"""
    DiscreteBelief

A belief specified by a probability vector.

Normalization of `b` is NOT enforced at all times, but the `DiscreteBeleif(pomdp, b)` constructor will warn, and `update(...)` always returns a belief with normalized `b`.

Constructor:

`DiscreteBelief(pomdp, b::Vector{Float64}; check::Bool=true)`

# Fields
    - `pomdp` : the POMDP problem
    - `state_list` : a vector of ordered states
    - `b` : the probability vector
"""
struct DiscreteBelief{P<:POMDP, S}
    pomdp::P
    state_list::Vector{S}       # vector of ordered states
    b::Vector{Float64}
end

function DiscreteBelief(pomdp, b::Vector{Float64}; check::Bool=true)
    if check
        if !isapprox(sum(b), 1.0, atol=0.001)
            @warn """
                  b in DiscreteBelief(pomdp, b) does not sum to 1.

                  To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                  """
        end
        if !all(0.0 <= p <= 1.0 for p in b)
            @warn """
                  b in DiscreteBelief(pomdp, b) contains entries outside [0,1].

                  To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                  """
        end
    end
    return DiscreteBelief(pomdp, ordered_states(pomdp), b)
end

"""
     uniform_belief(pomdp)

Return a DiscreteBelief with equal probability for each state.
"""
function uniform_belief(pomdp)
    state_list = ordered_states(pomdp)
    ns = length(state_list)
    return DiscreteBelief(pomdp, state_list, ones(ns) / ns)
end

pdf(b::DiscreteBelief, s) = b.b[stateindex(b.pomdp, s)]

function Random.rand(rng::Random.AbstractRNG, b::DiscreteBelief)
    i = sample(rng, Weights(b.b))
    return b.state_list[i]
end

function Base.fill!(b::DiscreteBelief, x::Float64)
    fill!(b.b, x)
    return b
end

Base.length(b::DiscreteBelief) = length(b.b)

support(b::DiscreteBelief) = b.state_list

==(b1::DiscreteBelief, b2::DiscreteBelief) = b1.state_list == b2.state_list && b1.b == b2.b

Base.hash(b::DiscreteBelief, h::UInt) = hash(b.b, hash(b.state_list, h))

"""
    DiscreteUpdater

An updater type to update discrete belief using the discrete Bayesian filter.

# Fields
    - pomdp <: POMDP

"""
mutable struct DiscreteUpdater{P<:POMDP} <: Updater
    pomdp::P
end

uniform_belief(up::DiscreteUpdater) = uniform_belief(up.pomdp)

function initialize_belief(bu::DiscreteUpdater, dist::Any)
    state_list = ordered_states(bu.pomdp)
    ns = length(state_list)
    b = zeros(ns)
    belief = DiscreteBelief(bu.pomdp, state_list, b)
    for s in support(dist)
        sidx = stateindex(bu.pomdp, s)
        belief.b[sidx] = pdf(dist, s)
    end
    return belief
end

function update(bu::DiscreteUpdater, b::DiscreteBelief, a, o)
    pomdp = b.pomdp
    state_space = b.state_list
    bp = zeros(length(state_space))

    bp_sum = 0.0   # to normalize the distribution

    for (spi, sp) in enumerate(state_space)

        # po = O(a, sp, o)
        od = observation(pomdp, a, sp)
        po = pdf(od, o)

        if po == 0.0
            continue
        end

        b_sum = 0.0
        for (si, s) in enumerate(state_space)
            td = transition(pomdp, s, a)
            pp = pdf(td, sp)
            b_sum += pp * b.b[si]
        end

        bp[spi] = po * b_sum
        bp_sum += bp[spi]
    end

    if bp_sum == 0.0
        error("""
              Failed discrete belief update: new probabilities sum to zero.

              b = $b
              a = $a
              o = $o

              Failed discrete belief update: new probabilities sum to zero.
              """)
    else
        for i = 1:length(bp); bp[i] /= bp_sum; end
    end

    return DiscreteBelief(pomdp, b.state_list, bp)
end

update(bu::DiscreteUpdater, b::Any, a, o) = update(bu, initialize_belief(bu, b), a, o)


using Revise
using Pkg
Pkg.activate("C:/Users/cmmenzen/.julia/dev/BigMat")
using BigMat
using Test
using Random
using LinearAlgebra
using Optim

@testset "BigMat.jl" begin
    include("MPT.jl")
    include("ComputeTN.jl")
    include("MatrixAlgebra.jl") # inversion
end


function jt02()
#
#  test polynomial suggested by Jenkins and Traub
#
p = poly([(1.0:17)...])
    z = [(1.0:17)...]
    z = [z ones(17)]
    p, PolyZeros(z)
end


using DistStat, LinearAlgebra, Test

type=[Float32,Float64]

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CUDA
    ArrayType = CuArray
else
    ArrayType = Array
end

for T in type
    A=ArrayType{T}(reshape(collect(1:45),5,9))
    A_dist=distribute(A)

    @test isapprox(opnorm(A_dist,1),opnorm(A,1))
    @test isapprox(opnorm(A_dist,2),opnorm(A,2))
    @test isapprox(opnorm(A_dist,Inf),opnorm(A,Inf))

end


export Dense
using Knet.Ops21: mmul
using Knet: atype
using AutoGrad: Param

"""
    Dense(wsize::Integer...; f, dims, dropout, atype, init, binit)
    Dense(weights, bias=nothing; f, dims, dropout)
    (d::Dense)(x)

Generalizes matrix multiplication to possibly more than 2 dims with reshapes:

    w(M...,N...) * x(N...,K...) => y(M...,K...)

where `wsize=size(weights)=(M...,N...)` and `dims=length(N)`. Optionally dropout is applied
to the input and a bias of size `(M...)` is added and an elementwise activation function `f`
is applied to the output. The last `dims` dimensions of `w` has to match the first `dims`
dimensions of `x` and the size of bias (if there is one) has to match the first
`ndims(w)-dims` dimensions of `w`.

Keyword arguments:
* `f=nothing`: apply activation function to output if not nothing
* `dims=1`: number of input dimensions in the weight tensor
* `dropout=0`: apply dropout with this probability to input if non-zero
* `atype=Knet.atype()`: array and element type for parameter initialization
* `init=𝑼(√(6/(fanin+fanout)))`: initialization function for weights
* `binit=zeros`: initialization function for bias, if `nothing` do not use bias

References:
* torch.nn.Linear
* tf.keras.layers.Dense
"""
struct Dense; w; b; f; dims; dropout; end

function Dense(weights, bias=nothing; f=nothing, dims=1, dropout=0)
    @assert ndims(weights) > dims "ndims(weights) must be > dims"
    @assert bias === nothing || size(bias) === size(weights)[1:end-dims] "weights and bias do not match"
    w = (weights isa Param ? weights : Param(weights))
    b = (bias isa Nothing || bias isa Param ? bias : Param(bias))
    Dense(w, b, f, dims, dropout)
end

function Dense(wsize::Integer...; f=nothing, dims=1, dropout=0, atype=atype(), binit=zeros,
               init=𝑼(√(6/(densein(wsize,dims)+denseout(wsize,dims)))))
    @assert length(wsize) > dims "ndims(weights) must be > dims"
    w = Param(convert(atype,init(wsize...)))
    b = binit isa Nothing ? nothing : Param(convert(atype,binit(wsize[1:end-dims]...)))
    Dense(w, b, f, dims, dropout)
end

function (l::Dense)(x)
    if l.dropout != 0; x = dropout(x, l.dropout); end
    y = mmul(l.w, x)
    if l.b !== nothing; y = y .+ l.b; end
    if l.f !== nothing; y = l.f.(y); end
    return y
end

densein(w,d)=prod(w[end-d+1:end])
denseout(w,d)=prod(w[1:end-d])

# function (l::Dense)(x::MaskedArray)
#     (a,m) = (x.array, x.mask)
#     @assert m===nothing || all(size(m,i) == 1 || size(m,i) == size(a,i) for i in 1:ndims(a))
#     if m === nothing
#         return MaskedArray(l(a), nothing)
#     elseif size(m,1) == 1   # avoid mask multiplication if possible
#         b = l(a)
#         if ndims(b) > ndims(m)
#             m = reshape(m, ntuple(i->1, ndims(b)-ndims(m))..., size(m)...)
#         end
#         return MaskedArray(b, m)
#     else
#         return MaskedArray(l(a .* oftype(a,m)), nothing)
#     end
# end

##
# In RNNs where input has shape (X,B,T) we want the result to be (Y,B,T). In CNNs where
# input has shape (X1,X2,C,B) we want the output to be (Y,B). So in general we want to be
# able to specify how many initial dims of x should be replaced. Does it ever get replaced
# with more than one dim? In my vaswani it does.


using Plots
using FLoops
using BenchmarkTools
using Octavian

struct Spins
	N::Int32
	M::Int32
	A::Array{Float32}
end

function Init(N = 100::Int,M = 100::Int)
	S = Spins(N, M, rand(Float32, N+2 ,M+2)*2π .- π)
	S.A[1,:] .= 0.0
	S.A[N + 2, :] .= 0.0
	S.A[:,1] .= 0.0
	S.A[:,M + 2] .= 0.0

	boolarr1 = S.A[2:N + 1,2:M + 1] .== 0.0
	temparr = rand(Float32,N,M)
	boolarr2 = temparr .> 0.7
	boolarr = boolarr1 .& boolarr2
	S.A[2:N + 1,2:M + 1] -= boolarr
	boolarr = boolarr1 .& .~boolarr2
	S.A[2:N + 1,2:M + 1] += boolarr

	return S
end

function EnergyParallel(S::Spins)
	E = zeros(Float32,S.N+2,S.M+2)
	p = [(-1,0),(1,0),(0,-1),(0,1)]
	@floop ThreadedEx(basesize = 1, nestlevel=Val(:inf), simd = true) for i in p
			Epart = matmul(S.A[2:S.N+1,2:S.M+1],
					S.A[2+i[1]:S.N+1+i[1],
						2+i[2]:S.M+1+i[2]])
			@reduce(E += Epart)
	end
	return sum(E)
end

function Magnetisation(S::Spins)
	mx = sum(cos.(S.A))/ *(S.N,S.M)
	my = sum(sin.(S.A))/ *(S.N,S.M)
	return √(mx^2 + my^2)
end

function Advance(S::Spins,max_steps=10000::Int,T=5.0::Float32)
	eps = 0.001

	E = zeros(Float32,max_steps + 1)
	m = zeros(Float32,max_steps + 1)

	Enow = EnergyParallel(S)
	E[1] = Enow
	m[1] = Magnetisation(S)

	step = 1

	#while *(step<max_steps, m[step] > eps, 1 - m[step] > eps)
	for _ in 1:max_steps

		step += 1
		i = rand(2:S.N+1)
		j = rand(2:S.M+1)

		InitVal = S.A[i,j]

		S.A[i,j]= (S.A[i,j] + 2π*(rand()-0.5))%π

		Enext = EnergyParallel(S)
		dE = Enext - Enow

        if dE < 0
			Enow = Enext
		elseif	rand() < exp(-dE/T)
            Enow = Enext
        else
            S.A[i,j] = InitVal
		end

		E[step] = Enow
		#m[step] = Magnetisation(S)
	end
	return E,m,S.A
end

function Run(N=100::Int,M=100::Int,max_steps=10000::Int,T=Float32(5.0)::Float32)
	S = Init(N,M)
	h1 = heatmap(S.A)
	E,m,S=Advance(S,max_steps,T)

	"""l = @layout [a{0.6h};b{0.5h}]
	p1 = plot(m, linecolor = :blue, label="m")
	p2 = plot(E, lc = :red, label="E")
	plot(p1,p2,layout = l)"""

	h2 = heatmap(S)
	savefig(plot(h1,h2), "./Out")
end

@time Run(100,100,10)
println(Threads.nthreads())


using PyCall
@pyimport math
@show math.cos(1) # 0.5403023058681398


"""
    SeisRadonFreqFor(in, nt; <keyword arguments>)

Transform a tau-p gather to a time-offset gather using a frequency domain
forward parabolic or linear Radon operator.

# Arguments
- `in::Array{Float64,2}`: 2D Radon panel, `in[1:ntau,1:np]`, where `ntau` is the
number of intercept times and `np` the number of curvatures or ray parameters.
- `nt::Int`: number of time samples.

# Keyword arguments
- `order="parab"`: `"parab"` for parabolic transform, `"linear"`
for linear transform.
- `dt=0.004`: time sampling interval in seconds.
- `h=collect(0.0:20.0:1000.0)`: offset vector; `h[1:nh]`.
- `href=0.0`: reference offset for parabolic Radon Transform. If the
defautl value `href=0.0` is given, `href` is set to `max(abs(h))`.
- `p=collect(-0.05:0.01:2.2)`: `p[1:np]`. If `order="parab"`, `p`
is a vector of residual moveout ("curvatures") at reference offset `href` in
seconds; if `order=linear`, `p` is a vector of ray parameters in s/m.
- `flow=0.0`: minimum frequency in the data in Hz.
- `fhigh=125.0`: maximum frequency in the data in Hz.

# Output
- `d`: data synthetized via forward Radon modeling, `d[1:nt, 1:nh]`.

# References
*  Hampson, D., 1986, Inverse velocity stacking for multiple elimination:
Canadian Journal of Exploration Geophysics, 22, 44-55.
*  Sacchi, M.D. and Ulrych, T.J., 1995, High-resolution velocity gathers and
offset space reconstruction: Geophysics, 60, 1169-1177.
"""
function SeisRadonFreqFor(m::Array{Tm,2}, nt::Int; order="parab",
                            dt=0.004, href=0.0,
                            h=collect(0.0:20.0:1000.0),
                            p=collect(-0.05:0.01:2.2),
                            flow=0.0, fhigh=125.0)where{Tm<:Real}

    if order=="parab"
        ind = 2
        href == 0 && (href = maximum(abs.(h)))
    elseif order=="linear"
        ind = 1
        href = 1.0
    else
        error("Order should be equal to \"parab\" or \"linear\"")
    end
    ntau = size(m, 1)
    np = length(p)
    np == size(m, 2) || error("lenght(p) must be equal to size(m, 2)")
    nh = length(h)
    nw = 2*nextpow(2,ntau)
    m = cat( m, zeros(nw-ntau, np),dims=1)
    M = fft(m, 1)
    iw_low = round(Int, flow*dt*nw+1)
    iw_high = round(Int, fhigh*dt*nw+1)
    D = zeros(ComplexF64, nw, nh)
    for iw = iw_low:iw_high
        w  = 2.0*pi*(iw-1)/(nw*dt)
        L = zeros(ComplexF64, nh, np)
        for ip = 1:np
            for ih = 1:nh
                phi = w*p[ip]*(h[ih]/href)^ind
      	        L[ih, ip] = exp(-im*phi)
            end
        end
        D[iw, :] = L*M[iw,:]
    end
    for iw = round(Int, nw/2)+2:nw
        D[iw, :] = conj(D[nw-iw+2, :])
    end
    d = real(ifft(D, 1))
    d = d[1:nt, :]
    return d
end