using Singular
using Singular: mul!, addeq!, zero!

function shreq!(v :: Vector{spoly{A}}, u :: Vector{spoly{A}}; step :: Int = 1) where A
  (length(v) > 0 && length(u) > 0) || error("vector dimension equal to zero")
  R = parent(v[1])
  for ix in 1:length(u)-step
    v[ix] = u[ix+step]
  end
end

function _addeq!(l0 :: Vector{spoly{A}},
                 l1 :: Vector{spoly{A}},
                 l2 :: Vector{spoly{A}},
                 I  :: sideal{spoly{A}}) where A
  length(l0) == length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l0)
  R = parent(l0[1])
  temp1 = zero(R)
  temp2 = zero(R)
  temp3 = zero(R)
  l0_ix_unred = zero(R)
  carry_unred = zero(R)
  carry = zero(R)
  for ix in 1:ln
    #w/o inlining -- l0[ix] = carry+l1[ix]+l2[ix]
    zero!(l0_ix_unred)
    addeq!(l0_ix_unred, l1[ix])
    addeq!(l0_ix_unred, l2[ix])
    addeq!(l0_ix_unred, carry)
    l0[ix] = reduce(l0_ix_unred, I)
    #note that xy ∨ xz ∨ yz == xy + xz + yz
    #
    #w/o inlining -- carry = carry*l1[ix]+carry*l2[ix]+l1[ix]*l2[ix]
    mul!(temp1, l1[ix], l2[ix])
    mul!(temp2, carry, l2[ix])
    mul!(temp3, carry, l1[ix])
    zero!(carry_unred)
    addeq!(carry_unred, temp1)
    addeq!(carry_unred, temp2)
    addeq!(carry_unred, temp3)
    zero!(carry)
    carry = reduce(carry_unred, I)
  end
end

function _addeq!(l0 :: Vector{spoly{A}},
                 l1 :: Vector{spoly{A}},
                 l2 :: Vector{spoly{A}}) where A
  length(l0) == length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l0)
  R = parent(l0[1])
  temp1 = zero(R)
  temp2 = zero(R)
  temp3 = zero(R)
  l0_ix_unred = zero(R)
  carry_unred = zero(R)
  carry = zero(R)
  for ix in 1:ln
    #w/o inlining -- l0[ix] = carry+l1[ix]+l2[ix]
    addeq!(l0[ix], l1[ix])
    addeq!(l0[ix], l2[ix])
    addeq!(l0[ix], carry)
    #note that xy ∨ xz ∨ yz == xy + xz + yz
    #
    #w/o inlining -- carry = carry*l1[ix]+carry*l2[ix]+l1[ix]*l2[ix]
    mul!(temp1, l1[ix], l2[ix])
    mul!(temp2, carry, l2[ix])
    mul!(temp3, carry, l1[ix])
    zero!(carry)
    addeq!(carry, temp1)
    addeq!(carry, temp2)
    addeq!(carry, temp3)
  end
end

function subeq!(l0 :: Vector{spoly{A}},
                l1 :: Vector{spoly{A}},
                l2 :: Vector{spoly{A}},
                I  :: sideal{spoly{A}}) where A
  length(l0) == length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l0)
  R = parent(l0[1])
  temp1 = zero(R)
  temp2 = zero(R)
  temp3 = zero(R)
  l0_ix_unred = zero(R)
  carry_unred = one(R)
  carry = one(R)
  for ix in 1:ln
    #w/o inlining -- l0[ix] = carry+l1[ix]+l2[ix]+one(R)
    zero!(l0_ix_unred)
    addeq!(l0_ix_unred, one(R))
    addeq!(l0_ix_unred, l2[ix])
    addeq!(l0_ix_unred, l1[ix])
    addeq!(l0_ix_unred, carry)
    l0[ix] = reduce(l0_ix_unred, I)
    #note that xy ∨ xz ∨ yz == xy + xz + yz
    #
    #w/o inlining -- carry = carry*l1[ix]+carry*l2[ix]+l1[ix]*l2[ix]+l1[ix]+carry
    mul!(temp1, l1[ix], l2[ix])
    mul!(temp2, carry, l2[ix])
    mul!(temp3, carry, l1[ix])
    zero!(carry_unred)
    addeq!(carry_unred, carry)
    addeq!(carry_unred, l1[ix])
    addeq!(carry_unred, temp1)
    addeq!(carry_unred, temp2)
    addeq!(carry_unred, temp3)
    zero!(carry)
    carry = reduce(carry_unred, I)
  end
end

function andeq!(l0 :: Vector{spoly{A}},
                l1 :: Vector{spoly{A}},
                l2 :: Vector{spoly{A}},
                I  :: sideal{spoly{A}}) where A
  length(l0) == length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l0)
  R = parent(l0[1])
  l0_ix_unred = zero(R)
  for ix in 1:ln
    zero!(l0_ix_unred)
    mul!(l0_ix_unred,l1[ix],l2[ix])
    l0[ix] = reduce(l0_ix_unred, I)
  end
end

vars = ["b$ix" for ix in 0:31]
R, b = PolynomialRing(Fp(2), vars)
I = Ideal(R, [b[ix]^2-b[ix] for ix in 1:32])
I = std(I) #compute Gröbner basis of I
(b0,b1,b2,b3,b4,b5,b6,b7,
 b8,b9,b10,b11,b12,b13,b14,b15,
 b16,b17,b18,b19,b20,b21,b22,b23,
 b24,b25,b26,b27,b28,b29,b30,b31) = b
reduce(b0^2+b1^3, I)

_55555555 = [one(R),zero(R),one(R),zero(R),one(R),zero(R),one(R),zero(R),
             one(R),zero(R),one(R),zero(R),one(R),zero(R),one(R),zero(R),
             one(R),zero(R),one(R),zero(R),one(R),zero(R),one(R),zero(R),
             one(R),zero(R),one(R),zero(R),one(R),zero(R),one(R),zero(R)]

_33333333 = [one(R),one(R),zero(R),zero(R),one(R),one(R),zero(R),zero(R),
             one(R),one(R),zero(R),zero(R),one(R),one(R),zero(R),zero(R),
             one(R),one(R),zero(R),zero(R),one(R),one(R),zero(R),zero(R),
             one(R),one(R),zero(R),zero(R),one(R),one(R),zero(R),zero(R)]

_0f0f0f0f = [one(R),one(R),one(R),one(R),zero(R),zero(R),zero(R),zero(R),
             one(R),one(R),one(R),one(R),zero(R),zero(R),zero(R),zero(R),
             one(R),one(R),one(R),one(R),zero(R),zero(R),zero(R),zero(R),
             one(R),one(R),one(R),one(R),zero(R),zero(R),zero(R),zero(R)]

_01010101 = [one(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),
             one(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),
             one(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),
             one(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R)]
#mov    edx,edi
#shr    edx,1
#and    edx,0x55555555
#sub    edi,edx
#mov    eax,edi
#shr    edi,0x2
#and    eax,0x33333333
#and    edi,0x33333333
#add    edi,eax
#mov    eax,edi
#shr    eax,0x4
#add    eax,edi
#and    eax,0x0f0f0f0f

edi1 = [b0,b1,b2,b3,b4,b5,b6,b7,
        b8,b9,b10,b11,b12,b13,b14,b15,
        b16,b17,b18,b19,b20,b21,b22,b23,
        b24,b25,b26,b27,b28,b29,b30,b31]
edx1 = [zero(R) for _ in 1:32]
shreq!(edx1, edi1);
edx2 = [zero(R) for _ in 1:32]
andeq!(edx2, edx1,_55555555, I);
edi2 = [zero(R) for _ in 1:32]
subeq!(edi2, edi1, edx2, I);
eax1 = deepcopy(edi2);
edi3 = [zero(R) for _ in 1:32]
shreq!(edi3, edi2; step=2);
eax2 = [zero(R) for _ in 1:32]
andeq!(eax2, eax1, _33333333, I)
edi4 = [zero(R) for _ in 1:32]
andeq!(edi4, edi3, _33333333, I)
edi5 = [zero(R) for _ in 1:32]
_addeq!(edi5, edi4, eax2, I)
eax3 = deepcopy(edi5)
eax4 = [zero(R) for _ in 1:32]
shreq!(eax4, eax3, step=4);
eax5 = [zero(R) for _ in 1:32]
_addeq!(eax5, eax4, edi5, I);
eax6 = [zero(R) for _ in 1:32]
andeq!(eax6, eax5, _0f0f0f0f, I);

#time for imul

vars1 = ["b0_7_1","b0_7_2","b0_7_4","b0_7_8",
         "b8_15_1","b8_15_2","b8_15_4","b8_15_8",
         "b16_23_1","b16_23_2","b16_23_4","b16_23_8",
         "b24_31_1","b24_31_2","b24_31_4","b24_31_8"]
R1, x = PolynomialRing(Fp(2), vars1)
I1 = Ideal(R1, [x[ix]^2-x[ix] for ix in 1:16])
I1 = std(I1)
(x0,x1,x2,x3,x4,x5,x6,x7,
 x8,x9,x10,x11,x12,x13,x14,x15) = x

(b0_7_1,b0_7_2,b0_7_4,b0_7_8,
 b8_15_1,b8_15_2,b8_15_4,b8_15_8,
 b16_23_1,b16_23_2,b16_23_4,b16_23_8,
 b24_31_1,b24_31_2,b24_31_4,b24_31_8) = x

v1 = [x0,x1,x2,x3,zero(R1),zero(R1),zero(R1),zero(R1),
      x4,x5,x6,x7,zero(R1),zero(R1),zero(R1),zero(R1),
      x8,x9,x10,x11,zero(R1),zero(R1),zero(R1),zero(R1),
      x12,x13,x14,x15,zero(R1),zero(R1),zero(R1),zero(R1)]
v2 = [zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),
      x0,x1,x2,x3,zero(R1),zero(R1),zero(R1),zero(R1),
      x4,x5,x6,x7,zero(R1),zero(R1),zero(R1),zero(R1),
      x8,x9,x10,x11,zero(R1),zero(R1),zero(R1),zero(R1)]
v3 = [zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),
      zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),
      x0,x1,x2,x3,zero(R1),zero(R1),zero(R1),zero(R1),
      x4,x5,x6,x7,zero(R1),zero(R1),zero(R1),zero(R1)]
v4 = [zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),
      zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),
      zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),zero(R1),
      x0,x1,x2,x3,zero(R1),zero(R1),zero(R1),zero(R1)]

#imul   eax,eax,0x01010101
#shr    eax,0x18
#ret

v5 = [zero(R1) for _ in 1:32]
_addeq!(v5,v1,v2,I1)
v6 = [zero(R1) for _ in 1:32]
_addeq!(v6,v5,v3,I1)
v7 = [zero(R1) for _ in 1:32]
_addeq!(v7,v6,v4,I1)
v8 = [zero(R1) for _ in 1:32]
shreq!(v8,v7;step=24)

@assert(all(v8[ix]==zero(R1) for ix in 7:32))

#note that
#
#b0_7_1 = all 1-factor products                 with factors in b0,b1,b2,...,b7
#b0_7_2 = all 2-factor products
#b0_7_1*b0_7_2 = all 3-factor products
#b0_7_4 = all 4-factor products
#b0_7_1*b0_7_4 = all 5-factor products
#b0_7_2*b0_7_4 = all 6-factor products
#b0_7_1*b0_7_2*b0_7_4 = all 7-factor products
#b0_7_8 = all 8-factor products

#hence v8[2] is sum of all 2-factor products    with factors in b0,b1,b2,...,b31
#      v8[3] is sum of all 4-factor products
#      v8[4] is sum of all 8-factor products
#      v8[5] is sum of all 16-factor products
#      v8[6] is sum of all 32-factor products


abstract type AbstractAdvection end

struct UpWind2ndOrder <: AbstractAdvection
end

function (a::UpWind2ndOrder)(rhs, ρC, buvec, uvec, Tc, Tbf, ∇Tc, bf2c, bfn, bcv, f2c, fn, cv, fc, ccenter, floops::Face2CellLoop)
    NBF = nbfaces(floops)
    @inbounds for i=1:NBF
        j = bf2c[i][1]
        rhs[j] -= (ρC * Tbf[i] * (buvec[i] ⋅ bfn[i]))*bcv[i]
    end
    NF = nfaces(floops)
    @inbounds for i=1:NF
        j1, j2 = f2c[i]
        n = fn[i]
        v1, v2 = cv[i]
        udotn = uvec[i] ⋅ n
        if udotn >= 0
            Tf = Tc[j1] + (fc[i] - ccenter[i][1])⋅∇Tc[j1]
        else
            Tf = Tc[j2] + (fc[i] - ccenter[i][2])⋅∇Tc[j2]
        end
        qf = ρC*Tf*udotn
        rhs[j1] -= qf*v1
        rhs[j2] += qf*v2
    end

end

function (a::UpWind2ndOrder)(rhs,p)
    floops = p.mesh.f2cloops
    a(rhs, p.ρC, p.u, p.u, p.Tc, p.Tbf, p.∇Tc, floops.bf2c, floops.bfn, floops.bcv, floops.f2c, floops.fn, floops.cv, floops.fc, floops.ccenter, floops)
end

struct FullUpWind <: AbstractAdvection
end

function (a::FullUpWind)(rhs, ρC, buvec, uvec, Tc, Tbf, bf2c, bfn, bcv, f2c, fn, cv, fc, ccenter, floops::Face2CellLoop)
    NBF = nbfaces(floops)
    @inbounds for i=1:NBF
        j = bf2c[i][1]
        rhs[j] -= (ρC * Tbf[i] * (buvec[i] ⋅ bfn[i]))*bcv[i]
    end
    NF = nfaces(floops)
    @inbounds for i=1:NF
        j1, j2 = f2c[i]
        n = fn[i]
        v1, v2 = cv[i]
        udotn = uvec[i] ⋅ n
        if udotn >= 0
            Tf = Tc[j1]
        else
            Tf = Tc[j2]
        end
        qf = ρC*Tf*udotn
        rhs[j1] -= qf*v1
        rhs[j2] += qf*v2
    end

end

function (a::FullUpWind)(rhs,p)
    floops = p.mesh.f2cloops
    a(rhs, p.ρC, p.u, p.u, p.Tc, p.Tbf, floops.bf2c, floops.bfn, floops.bcv, floops.f2c, floops.fn, floops.cv, floops.fc, floops.ccenter, floops)
end

struct UpWindAdvection{FT,Fbf,VT,UBT,fl}
    f::FT
    fbf::Fbf
    an::FT
    anm1::FT
    u::VT
    ubf::UBT
    floops::fl
end

function UpWindAdvection(f,fbf,u,ubf,mesh)
    an = similar(f)
    anm1 = similar(an)
    fill!(an,zero(eltype(an)))
    fill!(anm1,zero(eltype(an)))
    fl = mesh.f2cloops
    return UpWindAdvection{typeof(f),typeof(fbf),typeof(u),typeof(ubf),typeof(fl)}(f,fbf,an,anm1,u,ubf,fl)
end

function mycopy!(rhs,inp)
    @inbounds for i in linearindices(rhs)
        rhs[i] = inp[i]
    end
end

function (Up::UpWindAdvection)()

    # Copy An to An-1
    mycopy!(Up.anm1,Up.an)

    rhs = Up.an
    fill!(rhs,zero(eltype(rhs)))

    floops = Up.floops
    bf2c = floops.bf2c
    Tbf = Up.fbf
    buvec = Up.ubf
    bfn = floops.bfn
    bcv = floops.bcv

    NBF = nbfaces(floops)
    @inbounds for i=1:NBF
        j = bf2c[i][1]
        rhs[j] += (Tbf[i] * (buvec[i] ⋅ bfn[i]))*bcv[i]
    end

    f2c = floops.f2c
    fn = floops.fn
    cv = floops.cv
    u = Up.u
    Tc = Up.f

    NF = nfaces(floops)
    @inbounds for i=1:NF
        j1, j2 = f2c[i]
        n = fn[i]
        v1, v2 = cv[i]
        uvec = 0.5*(u[j1]+u[j2])
        udotn = uvec ⋅ n
        if udotn >= 0
            Tf = Tc[j1]
        else
            Tf = Tc[j2]
        end
        qf = Tf*udotn
        rhs[j1] += qf*v1
        rhs[j2] -= qf*v2
    end

end

function (Up::UpWindAdvection)(rhs)
    an = Up.an
    anm1 = Up.anm1

    @inbounds for i in linearindices(rhs)
        rhs[i] -= 1.5*an[i] - 0.5*anm1[i]
        #rhs[i] -= an[i]
    end

end


module remesh_smoothing_examples
using FinEtools
using FinEtools.MeshExportModule
using FinEtools.TetRemeshingModule

function remesh1()
    L= 0.3;
    W = 0.3;
    a = 0.15;
    nL=46; nW=46; na=36;

    fens,fes = T4block(a,L,W,nL,nW,na,:a);
    t = deepcopy(connasarray(fes));
    v = deepcopy(fens.xyz);
    tmid = ones(Int, size(t,1));

    desired_ts =a;
    bfes = meshboundary(fes);
    f = connectednodes(bfes);
    bv = zeros(Bool, size(v,1));
    bv[f] = true;

    println("Mesh size: initial = $(size(t,1))")
    t0 = time()

    t, v, tmid = TetRemeshingModule.coarsen(t, v, tmid; bv = bv, desired_ts = desired_ts);

    println("Mesh size: final = $(size(t,1)) [$(time() - t0) sec]")

    fens.xyz = deepcopy(v)
    fes = fromarray!(fes, t)
    setlabel!(fes, tmid)
    geom  =  NodalField(fens.xyz)

    femm  =  FEMMBase(IntegData(fes, SimplexRule(3, 1)))
    V = integratefunction(femm, geom, (x) ->  1.0)
    println("V = $(V) compared to $(L * W * a)")

    # File = "test1.vtk"
    # MeshExportModule.vtkexportmesh(File, t, v, MeshExportModule.T4)
    # @async run(`"paraview.exe" $File`)
end # remesh1

function mesh_smoothing()
    println("""
    Meshing, deforming  and smoothing
    """)

    A = 100. # strip width
    N = 16
    tolerance = A / N / 1.0e5
    fens, fes = T3block(A, A, N, N)

    bnl = connectednodes(meshboundary(fes))
    for ixxxx = 1:length(bnl)
        x, y = fens.xyz[bnl[ixxxx], :]
        fens.xyz[bnl[ixxxx], 1] += A / N * sin(2 * pi * y / A)
        fens.xyz[bnl[ixxxx], 2] += -A / N * sin(2 * pi * x / A)
    end

    File =  "mesh_smoothing_before.vtk"
    vtkexportmesh(File, fens, fes);
    @async run(`"paraview.exe" $File`)
    println("$(fens.xyz[Int(N^2 / 2), :] )")

    fixedv = falses(count(fens))
    fixedv[bnl] = true
    fens = meshsmoothing(fens, fes; fixedv=fixedv, method=:taubin, npass=100)

    println("$(fens.xyz[Int(N^2 / 2), :] )")

    geom = NodalField(fens.xyz)

    File =  "mesh_smoothing_after.vtk"
    vtkexportmesh(File, fes.conn, geom.values, FinEtools.MeshExportModule.T3);
    @async run(`"paraview.exe" $File`)

    println("Done")
    true

end # mesh_smoothing

function allrun()
    println("#####################################################")
    println("# remesh1 ")
    remesh1()
    println("#####################################################")
    println("# mesh_smoothing ")
    mesh_smoothing()
    return true
end # function allrun

end # module remesh_smoothing_examples


function diagnose_velocities_from_streamfunction(ψ, grid)
    ψᶠᶠᶜ = Field(Face, Face,   Center, CPU(), grid)
    uᶠᶜᶜ = Field(Face, Center, Center, CPU(), grid)
    vᶜᶠᶜ = Field(Center, Face, Center, CPU(), grid)

    for (f, grid_face) in enumerate(grid.faces)
        Nx, Ny, Nz = size(grid_face)

        ψᶠᶠᶜ_face = ψᶠᶠᶜ.data.faces[f]
        uᶠᶜᶜ_face = uᶠᶜᶜ.data.faces[f]
        vᶜᶠᶜ_face = vᶜᶠᶜ.data.faces[f]

        for i in 1:Nx+1, j in 1:Ny+1
            ψᶠᶠᶜ_face[i, j, 1] = ψ(grid_face.λᶠᶠᵃ[i, j], grid_face.φᶠᶠᵃ[i, j])
        end

        for i in 1:Nx+1, j in 1:Ny
            uᶠᶜᶜ_face[i, j, 1] = (ψᶠᶠᶜ_face[i, j, 1] - ψᶠᶠᶜ_face[i, j+1, 1]) / grid.faces[f].Δyᶠᶜᵃ[i, j]
        end

        for i in 1:Nx, j in 1:Ny+1
            vᶜᶠᶜ_face[i, j, 1] = (ψᶠᶠᶜ_face[i+1, j, 1] - ψᶠᶠᶜ_face[i, j, 1]) / grid.faces[f].Δxᶜᶠᵃ[i, j]
        end
    end

    return uᶠᶜᶜ, vᶜᶠᶜ, ψᶠᶠᶜ
end

function set_velocities_from_streamfunction!(u, v, ψ, grid)
    uψ, vψ, ψ₀ = diagnose_velocities_from_streamfunction(ψ, grid)
    set!(u, uψ)
    set!(v, vψ)
    return nothing
end

@testset "Cubed sphere circulation" begin

    arch = CPU()
    grid = ConformalCubedSphereGrid(cs32_filepath, Nz=1, z=(-1, 0))

    FT = eltype(grid)
    Nx, Ny, Nz, Nf = size(grid)
    R = grid.faces[1].radius

    u_field = XFaceField(arch, grid)
    v_field = YFaceField(arch, grid)

    ψ(λ, φ) = R * sind(φ)
    set_velocities_from_streamfunction!(u_field, v_field, ψ, grid)

    fill_horizontal_velocity_halos!(u_field, v_field, arch)

    grid_faces = [grid.faces[f] for f in 1:Nf]
    u_faces = [get_face(u_field, f) for f in 1:Nf]
    v_faces = [get_face(v_field, f) for f in 1:Nf]

    circulation(i, j, f) = Γᶠᶠᵃ(i, j, 1, grid_faces[f], u_faces[f], v_faces[f])

    @testset "Face 1-2 boundary" begin
        # Quick test at a single point.
        Γ1 = circulation(33, 16, 1)
        Γ2 = circulation(1,  16, 2)
        @test Γ1 == Γ2

        # Check the entire boundary.
        Γ1 = [circulation(33, j, 1) for j in 1:Ny+1]
        Γ2 = [circulation(1,  j, 2) for j in 1:Ny+1]

        # There should be no truncation errors here.
        @test Γ1 == Γ2
    end

    @testset "Face 1-3 boundary" begin
        Γ1 = circulation(16, 33, 1)
        Γ3 = circulation(1,  16, 3)

        # We expect a truncation error of around 3ϵ here.
        ϵ = eps(maximum(abs, [Γ1, Γ3]))
        @test isapprox(Γ1, Γ3, atol=3ϵ)

        # ffc locations are shifted by one along this boundary (and go in reverse).
        Γ1 = [circulation(i, 33, 1) for i in 1:Nx]
        Γ3 = [circulation(1,  j, 3) for j in Ny+1:-1:2]

        # We expect truncation errors at this boundary.
        ϵ = eps(maximum(abs, [Γ1; Γ3]))
        @test all(isapprox.(Γ1, Γ3, atol=4ϵ))
    end

    @testset "Face 1-5 boundary" begin
        Γ1 = [circulation(1,  j, 1) for j in 1:Ny]
        Γ5 = [circulation(i, 33, 5) for i in Nx+1:-1:2]

        # Hmmm there's more truncation error here.
        ϵ = eps(maximum(abs, [Γ1; Γ5]))
        @test all(isapprox.(Γ1, Γ5, atol=32ϵ))
    end

    @testset "Face 1-6 boundary" begin
        Γ1 = [circulation(i,  1, 1) for i in 1:Nx+1]
        Γ6 = [circulation(i, 33, 6) for i in 1:Nx+1]
        @test Γ1 == Γ6
    end

    @testset "Face 1-2-3 corner" begin
        Γ1 = circulation(33, 33, 1)
        Γ2 = circulation(1,  33, 2)
        Γ3 = circulation(1,  1,  3)
        @test Γ1 == Γ2 == Γ3
    end

    @testset "Face 1-3-5 corner" begin
        Γ1, Γ3, Γ5 = Γ = [circulation(1, 33, f) for f in (1, 3, 5)]
        ϵ = eps(maximum(abs, Γ))
        @test isapprox(Γ1, Γ3, atol=32ϵ)
        @test isapprox(Γ1, Γ5, atol=32ϵ)
        @test isapprox(Γ3, Γ5, atol=32ϵ)
    end

    @testset "Face 1-5-6 corner" begin
        Γ1 = circulation(1,   1, 1)
        Γ5 = circulation(33, 33, 5)
        Γ6 = circulation(1,  33, 6)
        @test Γ1 == Γ5 == Γ6
    end

    @testset "Face 1-2-6 corner" begin
        Γ1 = circulation(33,  1, 1)
        Γ2 = circulation(1,   1, 2)
        Γ6 = circulation(33, 33, 6)
        @test Γ1 == Γ2 == Γ6
    end

    @testset "Face 2-3-4 corner" begin
        Γ2 = circulation(33, 33, 2)
        Γ3 = circulation(33, 1,  3)
        Γ4 = circulation(1,  1,  4)
        Γ = [Γ2, Γ3, Γ4]

        ϵ = eps(maximum(abs, Γ))
        @test isapprox(Γ2, Γ3, atol=32ϵ)
        @test isapprox(Γ2, Γ4, atol=32ϵ)
        @test isapprox(Γ3, Γ4, atol=32ϵ)
    end

    @testset "Face 2-4-6 corner" begin
        Γ2, Γ4, Γ6 = Γ = [circulation(33, 1, f) for f in (2, 4, 6)]
        @show ϵ = eps(maximum(abs, Γ))
        @test isapprox(Γ2, Γ4, atol=32ϵ)
        @test isapprox(Γ2, Γ6, atol=32ϵ)
        @test isapprox(Γ4, Γ6, atol=32ϵ)
    end

end


@testset "Random primitives" begin

@testset "rand(::Type{<:Integer}) statements" begin

function foo(n::Int)
    a = rand(Int8)
    b = rand(Int16, n)
    c = rand(Int32, (n, n))
end

genfoo = genify(foo, Int)

# Check types
trace = simulate(genfoo, (10,))
@test trace[:a] isa Int8
@test trace[:b] isa Vector{Int16}
@test trace[:c] isa Matrix{Int32}

# Test array sizes
@test size(trace[:a]) == ()
@test size(trace[:b]) == (10,)
@test size(trace[:c]) == (10, 10)

# Test sample range
@test all(typemin(Int8) .<= trace[:a] .<= typemax(Int8))
@test all(typemin(Int16) .<= trace[:b] .<= typemax(Int16))
@test all(typemin(Int32) .<= trace[:c] .<= typemax(Int32))

# Test generate with constraints
trace, weight = generate(genfoo, (10,), choicemap(:a => 0))
@test weight ≈ -8*log(2)
trace, weight = generate(genfoo, (10,), choicemap(:b => zeros(Int16, 10)))
@test weight ≈ -16*log(2) * 10
trace, weight = generate(genfoo, (10,), choicemap(:c => zeros(Int16, 10, 10)))
@test weight ≈ -32*log(2) * 10 * 10

end

@testset "rand(::Type{<:AbstractFloat}) statements" begin

function foo(n::Int)
    a = rand(Float16)
    b = rand(Float32, n)
    c = rand(Float64, n, n)
end

genfoo = genify(foo, Int)

# Check types
trace = simulate(genfoo, (10,))
@test trace[:a] isa Float16
@test trace[:b] isa Vector{Float32}
@test trace[:c] isa Matrix{Float64}

# Test array sizes
@test size(trace[:a]) == ()
@test size(trace[:b]) == (10,)
@test size(trace[:c]) == (10, 10)

# Test ranges
@test all(0 .<= trace[:a] .<= 1)
@test all(0 .<= trace[:b] .<= 1)
@test all(0 .<= trace[:c] .<= 1)

# Test generate with constraints
trace, weight = generate(genfoo, (10,), choicemap(:a => 0.5))
@test weight == 0
trace, weight = generate(genfoo, (10,), choicemap(:b => [-1; fill(0, 9)]))
@test weight == -Inf
trace, weight = generate(genfoo, (10,), choicemap(:c => rand(Float64, 10, 10)))
@test weight == 0

end

@testset "rand(::Indexable|Setlike) statements" begin

function foo(dims...)
    a = rand([:h, :e, :l, :l, :o], dims...)
    b = rand("world", dims...)
    c = rand(1:5, dims...)
    d = rand((9, 6, 42), dims...)
    e = rand(Set("hello"), dims...)
    f = rand(Dict(zip("world", "hello")), dims...)
end

genfoo = genify(foo)

# Check types
trace = simulate(genfoo, ())
@test trace[:a] isa Symbol
@test trace[:b] isa Char
@test trace[:c] isa Int
@test trace[:d] isa Int
@test trace[:e] isa Char
@test trace[:f] isa Pair{Char,Char}

# Test array sizes
trace = simulate(genfoo, (10, 10))
@test size(trace[:a]) == (10, 10)
@test size(trace[:b]) == (10, 10)
@test size(trace[:c]) == (10, 10)
@test size(trace[:d]) == (10, 10)
@test size(trace[:e]) == (10, 10)
@test size(trace[:f]) == (10, 10)

# Test ranges
@test all(x in [:h, :e, :l, :l, :o] for x in trace[:a])
@test all(x in "world" for x in trace[:b])
@test all(x in 1:5 for x in trace[:c])
@test all(x in (9, 6, 42) for x in trace[:d])
@test all(x in Set("hello") for x in trace[:e])
@test all(x in Dict(zip("world", "hello")) for x in trace[:f])

# Test generate with constraints
trace, weight = generate(genfoo, (), choicemap(:a => :l))
@test weight ≈ log(2/5)
trace, weight = generate(genfoo, (), choicemap(:b => 'w'))
@test weight ≈ log(1/5)
trace, weight = generate(genfoo, (), choicemap(:c => 3))
@test weight ≈ log(1/5)
trace, weight = generate(genfoo, (), choicemap(:d => 42))
@test weight ≈ log(1/3)
trace, weight = generate(genfoo, (), choicemap(:e => 'l'))
@test weight ≈ log(1/4)
trace, weight = generate(genfoo, (), choicemap(:f => ('o' => 'e')))
@test weight ≈ log(1/5)

end

@testset "rand(::Distributions.Distribution) statements" begin

function foo(dims...)
    a = rand(Dirichlet(5, 2.0))
    b = rand(Categorical(a))
    c = rand(Wishart(b+2, [1.0 0.0; 0.0 1.0]))
    d = rand(Normal(0, 1), dims...)
end

genfoo = genify(foo)

# Check types
trace = simulate(genfoo, (10,))
@test trace[:a] isa Vector{Float64}
@test trace[:b] isa Int
@test trace[:c] isa Matrix{Float64}
@test trace[:d] isa Vector{Float64}

# Test array sizes
@test size(trace[:a]) == (5,)
@test size(trace[:b]) == ()
@test size(trace[:c]) == (2, 2)
@test size(trace[:d]) == (10,)

# Test ranges
@test all(0 .<= trace[:a] .<= 1)
@test all(1 .<= trace[:b] .<= 5)
@test insupport(Wishart, trace[:c])

# Test generate with constraints
trace, weight = generate(genfoo, (), choicemap(:a => ones(5)/5))
@test weight ≈ Distributions.logpdf(Dirichlet(5, 2), ones(5)/5)
trace, weight = generate(genfoo, (), choicemap(:a => ones(5)/5, :b => 5))
@test weight ≈ Distributions.logpdf(Dirichlet(5, 2), ones(5)/5) + log(1/5)
trace, weight = generate(genfoo, (), choicemap(:c => [1 0; 0 1]))
@test weight ≈ Distributions.logpdf(Wishart(trace[:b]+2, [1 0; 0 1]), [1 0; 0 1])
trace, weight = generate(genfoo, (10, 10), choicemap(:d => zeros(10, 10)))
@test weight ≈ Gen.logpdf(normal, 0, 0, 1) * 100

end

@testset "randn() and randexp() statements" begin

function foo(dims...)
    a = randn()
    b = randexp()
    c = randn(Float64, dims)
    d = randexp(dims...)
end

genfoo = genify(foo)

# Check types
trace = simulate(genfoo, (10,))
@test trace[:a] isa Float64
@test trace[:b] isa Float64
@test trace[:c] isa Vector{Float64}
@test trace[:d] isa Vector{Float64}

# Test array sizes
@test size(trace[:a]) == ()
@test size(trace[:b]) == ()
@test size(trace[:c]) == (10,)
@test size(trace[:d]) == (10,)

# Test ranges
@test all(trace[:b] .>= 0)
@test all(trace[:d] .>= 0)

# Test generate with constraints
trace, weight = generate(genfoo, (), choicemap(:a => 1, :b => 1))
@test weight ≈ Gen.logpdf(normal, 0, 1, 1) + Gen.logpdf(exponential, 1, 1)
trace, weight = generate(genfoo, (10,), choicemap(:c => ones(10), :d => ones(10)))
@test weight ≈ 10 * (Gen.logpdf(normal, 0, 1, 1) + Gen.logpdf(exponential, 1, 1))

end

@testset "randperm(), randcycle(), and shuffle() statements" begin

function foo(n::Int)
    for i in 1:100
        a = randperm(n)
        b = randcycle(n)
    end
    c = shuffle([:a, :b, :c, :d, :e])
    d = shuffle(['h', 'e', 'l', 'l', 'o', 'o'])
end

genfoo = genify(foo, Int)

# Check types
trace = simulate(genfoo, (5,))
@test all(trace[:a => i] isa Vector{Int} for i in 1:100)
@test all(trace[:b => i] isa Vector{Int} for i in 1:100)
@test trace[:c] isa Vector{Symbol}
@test trace[:d] isa Vector{Char}

# Test array sizes
@test all(length(trace[:a => i]) == 5 for i in 1:100)
@test all(length(trace[:b => i]) == 5 for i in 1:100)
@test length(trace[:c]) == 5
@test length(trace[:d]) == 6

# Test ranges
@test all(isperm(trace[:a => i]) for i in 1:100)
@test all(Genify.iscycle(trace[:b => i]) for i in 1:100)
@test Set(trace[:c]) == Set([:a, :b, :c, :d, :e])
@test sort(trace[:d]) == ['e', 'h', 'l', 'l', 'o', 'o']

# Test generate with constraints
choices = choicemap((:a => 1, randperm(5)), (:b => 1, randcycle(5)))
trace, weight = generate(genfoo, (5,), choices)
@test weight ≈ log(1/factorial(5)) + log(1/factorial(4))
trace, weight = generate(genfoo, (5,), choicemap((:c,  [:e, :b, :d, :c, :a])))
@test weight == log(1/factorial(5))
trace, weight = generate(genfoo, (5,), choicemap((:d,  ['h', 'l', 'l', 'o', 'e', 'o'])))
@test weight == log((factorial(2) * factorial(2)) / factorial(6))

trace, weight = generate(genfoo, (5,), choicemap((:a => 1,  ones(Int, 5))))
@test weight == -Inf
trace, weight = generate(genfoo, (5,), choicemap((:b => 1,  [1, 2, 3, 4, 5])))
@test weight == -Inf
trace, weight = generate(genfoo, (5,), choicemap((:c,  [:e, :b, :d, :c, :f])))
@test weight == -Inf
trace, weight = generate(genfoo, (5,), choicemap((:d,  ['h', 'l', 'l', 'e', 'o'])))
@test weight == -Inf

end

@testset "sample() statements" begin

function foo()
    a = sample([:h, :e, :l, :l, :o], weights([1, 1, 5, 1, 1]))
    b = sample(collect("world"), (10, 10))
end

genfoo = genify(foo)

# Check types
trace = simulate(genfoo, ())
@test trace[:a] isa Symbol
@test trace[:b] isa Matrix{Char}

# Test array sizes
trace = simulate(genfoo, ())
@test size(trace[:b]) == (10, 10)

# Test ranges
@test trace[:a] in [:h, :e, :l, :l, :o]
@test all(x in "world" for x in trace[:b])

# Test generate with constraints
trace, weight = generate(genfoo, (), choicemap(:a => :l))
@test weight ≈ log(6/9)
trace, weight = generate(genfoo, (), choicemap(:b => fill('w', (10, 10))))
@test weight ≈ log(1/5) * 100

end

@testset "Primitives with RNGs" begin

function foo(n::Int)
    rng = MersenneTwister(0)
    a = rand(rng, Int8)
    b = rand(rng, Float32, n)
    c = rand(rng, "world")
    d = rand(rng, Wishart(5, [1.0 0.0; 0.0 1.0]))
    e = randexp(rng)
    f = randn(rng)
    g = randperm(rng, n)
    h = randcycle(rng, n)
end

genfoo = genify(foo, Int)

# Check types
trace = simulate(genfoo, (10,))
@test trace[:a] isa Int8
@test trace[:b] isa Vector{Float32}
@test trace[:c] isa Char
@test trace[:d] isa Matrix{Float64}
@test trace[:e] isa Float64
@test trace[:f] isa Float64
@test trace[:g] isa Vector{Int}
@test trace[:h] isa Vector{Int}

end

@testset "In-place primitives" begin

function foo(A::Vector{Float64}, B::Vector{Int})
    a = rand!(A)
    b = randn!(A)
    c = randexp!(A)
    d = randperm!(B)
    e = randcycle!(B)
    f = shuffle!(B)
    g = sample!(B, A)
end

genfoo = genify(foo, Vector{Float64}, Vector{Int})

# Check that arrays were modified
A, B = Vector{Float64}(undef, 10), Vector{Int}(undef, 10)
trace = simulate(genfoo, (A, B))
@test A == trace[:g]
@test B == trace[:f]

# Check types
@test trace[:a] isa Vector{Float64}
@test trace[:b] isa Vector{Float64}
@test trace[:c] isa Vector{Float64}
@test trace[:d] isa Vector{Int}
@test trace[:e] isa Vector{Int}
@test trace[:f] isa Vector{Int}
@test trace[:g] isa Vector{Int}
@test trace[] isa Vector{Float64}

# Test ranges
@test all(0 .<= trace[:a] .<= 1)
@test all(-Inf .<= trace[:b] .<= Inf)
@test all(0 .<= trace[:c] .<= Inf)
@test isperm(trace[:d])
@test Genify.iscycle(trace[:e])
@test isperm(trace[:f])
@test all(1 .<= trace[:g] .<= 10)

end

end


# Tom Anderson
# vdiv_filter.jl
# MIT License
# 10/30/2016
# Demo program to show the frequency response of an
# RC lowpass filter and an RLC lowpass filter

include("/Users/toma/tools/julia/rc_graph.jl")

fstart = 1
fstop = 10meg
npts = 2810
freq = logsweep(fstart, fstop, npts)

# RC lowpass filter
# R=100 Ohms, C= 0.01uF
filt1 = db(vdiv(zr(100,freq), zc(.01u,freq)))

# RLC lowpass filter with peaking.
# R=100 Ohms, L=1mH, c=0.01uF
filt2 = db(vdiv(zr(100,freq) + zl(1m,freq), zc(.01u,freq)))

plt1 = FramedPlot()
setattr(plt1, "xlog", true)
setattr(plt1, "xrange", (fstart, fstop))
setattr(plt1.frame1, "draw_grid", true)
setattr(plt1.x1, "label", "Frequency, Hz")
setattr(plt1.y1, "label", "LPF out, dB")

a = Curve(freq, filt1, "color", "blue")
b = Curve(freq, filt2, "color", "black")
add(plt1, a)
add(plt1, b)


# module for estimating pose by extended kalman filter
# initial pose is decided randomly
# global localization problem

using Distributions, LinearAlgebra, StatsBase

include(joinpath(split(@__FILE__, "src")[1], "src/model/map/map.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/common/covariance_ellipse/covariance_ellipse.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/common/state_transition/state_transition.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/common/observation_function/observation_function.jl"))

mutable struct GlobalKf
  belief
  motion_noise_stds
  estimated_pose
  estimated_cov
  map
  dist_dev
  dir_dev

  function GlobalKf(init_pose::Array;
                    motion_noise_stds::Dict=Dict("nn"=>0.20, "no"=>0.001, "on"=>0.11, "oo"=>0.20),
                    env_map=nothing, dist_dev_rate=0.14, dir_dev=0.05)
    self = new()
    self.belief = MvNormal([rand(Uniform(-5.0, 5.0)), rand(Uniform(-5.0, 5.0)), rand(Uniform(-pi, pi))],
                           diagm(0 => [1e+4, 1e+4, 1e+4]))
    self.motion_noise_stds = motion_noise_stds
    self.estimated_pose = self.belief.μ
    self.estimated_cov = self.belief.Σ
    self.map = env_map
    self.dist_dev = dist_dev_rate
    self.dir_dev = dir_dev
    return self
  end
end

function mat_M(speed, yaw_rate, time, stds)
  return diagm(0 => [stds["nn"]^2*abs(speed)/time + stds["no"]^2*abs(yaw_rate)/time,
                     stds["on"]^2*abs(speed)/time + stds["oo"]^2*abs(yaw_rate)/time])
end

function mat_A(speed, yaw_rate, time, theta)
  st, ct = sin(theta), cos(theta)
  stw, ctw = sin(theta + yaw_rate * time), cos(theta + yaw_rate * time)
  return [(stw - st)/yaw_rate  -speed/(yaw_rate^2)*(stw - st) + speed/yaw_rate*time*ctw;
          (-ctw + ct)/yaw_rate -speed/(yaw_rate^2)*(-ctw + ct) + speed/yaw_rate*time*stw;
          0                    time]
end

function mat_F(speed, yaw_rate, time, theta)
  F = diagm(0 => [1.0, 1.0, 1.0])
  F[1, 3] = speed / yaw_rate * (cos(theta + yaw_rate * time) - cos(theta))
  F[2, 3] = speed / yaw_rate * (sin(theta + yaw_rate * time) - sin(theta))
  return F
end

function mat_H(mu_pose, obj_pose)
  obj_x, obj_y = obj_pose[1], obj_pose[2]
  mu_x, mu_y = mu_pose[1], mu_pose[2]
  mu_l = sqrt((mu_x - obj_x)^2 + (mu_y - obj_y)^2)
  return [(mu_x - obj_x)/mu_l (mu_y - obj_y)/mu_l 0.0;
          (obj_y - mu_y)/(mu_l^2) (mu_x - obj_x)/(mu_l^2) -1.0]
end

function mat_Q(dist_dev, dir_dev)
  return [dist_dev^2 0.0;
          0.0 dir_dev^2]
end

function motion_update(self::GlobalKf, speed,
                       yaw_rate, time)
  if abs(yaw_rate) < 1e-5 # to prevent division by zero
    yaw_rate = 1e-5
  end
  M = mat_M(speed, yaw_rate, time, self.motion_noise_stds)
  A = mat_A(speed, yaw_rate, time, self.estimated_pose[3])
  F = mat_F(speed, yaw_rate, time, self.estimated_pose[3])
  self.estimated_cov = F * self.estimated_cov * F' + A * M * A'
  self.estimated_pose = state_transition(speed, yaw_rate, time, self.estimated_pose)
end

function observation_update(self::GlobalKf, observation)
  for obs in observation
    z = obs[1] # [distance, direction]
    id = obs[2]

    H = mat_H(self.estimated_pose, self.map.objects[id].pose)
    estimated_z = observation_function(self.estimated_pose,
                                       self.map.objects[id].pose)

    # observation noise
    Q = mat_Q(estimated_z[1]*self.dist_dev, self.dir_dev)

    # kalman gain
    K = self.estimated_cov * H' * inv(Q + H*self.estimated_cov*H')
    self.estimated_pose += K * (z - estimated_z)
    self.estimated_cov = (Matrix{Float64}(I, 3, 3) - K*H) * self.estimated_cov
  end
end

function draw!(self::GlobalKf)
  draw_covariance_ellipse!(self.estimated_pose, self.estimated_cov, 3)
end


using Turing
using Distributions
using Base.Test

# Test the @assume and @predict macros on a model without conditioning.

@model test begin
  @assume x ~ Bernoulli(1)
  @assume y ~ Bernoulli(x / 2)
  @predict y
  @predict x
end

s = SMC(10000)
p = PG(10,1000)

res = sample(test, s)

@test reduce(&, res[:x]) == 1  #check that x is always 1
@test_approx_eq_eps mean(res[:y]) 0.5 0.1  #check that the mean of y is between 0.4 and 0.6

res = sample(test, p)

@test reduce(&, res[:x]) == 1  #check that x is always 1
@test_approx_eq_eps mean(res[:y]) 0.5 0.1  #check that the mean of y is between 0.4 and 0.6


function d3uftau(t, fft_u, fft_v, fft_du, fft_dv, fft_du2, fft_dv2, fft_du3,
            fft_dv3, A1, A2, matr, conjmatr, sigma, llambda,
            Ntaumm)

    sigma = 1

    u = ifft(exp(1im * t * A1) * fft_u) * matr
    v = ifft(exp(1im * t * A1) * fft_v) * matr

    du = ifft(exp(1im * t * A1) * fft_du) * matr
    dv = ifft(exp(1im * t * A1) * fft_dv) * matr

    du2 = ifft(exp(1im * t * A1) * fft_du2) * matr
    dv2 = ifft(exp(1im * t * A1) * fft_dv2) * matr

    du3 = ifft(exp(1im * t * A1) * fft_du3) * matr
    dv3 = ifft(exp(1im * t * A1) * fft_dv3) * matr

    z = (u + conj(v)) / 2
    dz = (du + conj(dv)) / 2
    dz2 = (du2 + conj(dv2)) / 2
    dz3 = (du3 + conj(dv3)) / 2

    fz1 = 2 * dz3 * dz * conj(dz2) + 2 * dz3 * conj(dz) * dz2 + 2 * conj(dz3) * dz * dz2
    champu = -1im * llambda * A2 * exp(-1im * t * A1) * fft(conjmatr * fz1)
    champv = -1im * llambda * A2 * exp(-1im * t * A1) * fft(conjmatr * conj(fz1))

    return champu, champv

end


# Constants (in atomic units) used in some of the potentials.

"""
Coulomb constant in atomic units.
https://en.wikipedia.org/wiki/Coulomb_constant
"""
const kₑ = 1.0

"""
Vacuum permittivity in atomic units.
https://en.wikipedia.org/wiki/Vacuum_permittivity
"""
const ϵ₀ = 1 / 4π


#find how many different letters appear
#in all the persons from a group
#here that is each string separated by space
function find_yes(group_ans)

   #split string by spaces
   #string is coming with "" in the front, and it gets caught in the split
   person_ans = split(group_ans, " ")

   #find shortest string
   ans_l = [length(i) for i in person_ans if i!=""]
   shortest_string = findmin(ans_l)
   shortest_index = shortest_string[2] + 1 #correct index due to "" being counted as element
   shortest_ans = person_ans[shortest_index]

   n_words = length(person_ans)
   occur_in_all = 0

   #see how many times  the shortest string letters occur in
   #the group answers, if it occurs as many times as the number
   #of total words, then it appears in all words
   for letter in shortest_ans
      rx = Regex(string(letter))
      m = eachmatch(rx,group_ans)
      ocurrences = size(collect(m))[1]
      if ocurrences == n_words - 1 #correct because n_words counts "" as an element
         occur_in_all += 1
      end
   end

   return occur_in_all
end

f = open("input_day6","r")

temp_s = ""
sum_y = 0

#note that I modified my inputs to have an empty line at the end
for lines in readlines(f)
   #if not empty line concatenate de string with space
   if lines != ""
      global temp_s = temp_s*" "*lines
   else
      #println(temp_s)
      n = find_yes(temp_s)
      global sum_y += n
      global temp_s = ""

   end
end

println(sum_y)


# This file is a part of Julia. License is MIT: https://julialang.org/license

ndgrid(v::AbstractVector) = copy(v)

function ndgrid(v1::AbstractVector{T}, v2::AbstractVector{T}) where T
    m, n = length(v1), length(v2)
    v1 = reshape(v1, m, 1)
    v2 = reshape(v2, 1, n)
    (repmat(v1, 1, n), repmat(v2, m, 1))
end

function ndgrid_fill(a, v, s, snext)
    for j = 1:length(a)
        a[j] = v[div(rem(j-1, snext), s)+1]
    end
end

function ndgrid(vs::AbstractVector{T}...) where T
    n = length(vs)
    sz = map(length, vs)
    out = ntuple(i->Array{T}(sz), n)
    s = 1
    for i=1:n
        a = out[i]::Array
        v = vs[i]
        snext = s*size(a,i)
        ndgrid_fill(a, v, s, snext)
        s = snext
    end
    out
end

meshgrid(v::AbstractVector) = meshgrid(v, v)

function meshgrid(vx::AbstractVector{T}, vy::AbstractVector{T}) where T
    m, n = length(vy), length(vx)
    vx = reshape(vx, 1, n)
    vy = reshape(vy, m, 1)
    (repmat(vx, m, 1), repmat(vy, 1, n))
end

function meshgrid(vx::AbstractVector{T}, vy::AbstractVector{T},
                  vz::AbstractVector{T}) where T
    m, n, o = length(vy), length(vx), length(vz)
    vx = reshape(vx, 1, n, 1)
    vy = reshape(vy, m, 1, 1)
    vz = reshape(vz, 1, 1, o)
    om = ones(Int, m)
    on = ones(Int, n)
    oo = ones(Int, o)
    (vx[om, :, oo], vy[:, on, oo], vz[om, on, :])
end


using SparseDiffTools
using ForwardDiff: Dual, jacobian
using SparseArrays, Test
using LinearAlgebra
using BlockBandedMatrices
using BandedMatrices
using StaticArrays

fcalls = 0
function f(dx,x)
    global fcalls += 1
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx[end] = x[end-1] - 2x[end]
    nothing
end

function oopf(x)
    global fcalls += 1
    dx = zero(x)
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx[end] = x[end-1] - 2x[end]
    dx
end

function nsqf(x)#length(dx)<length(x)
    global fcalls +=1
    dx = zero(x)[1:div(length(x),2)]
    for i in 2:length(dx)
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx
end

function nsqf2(x)#length(dx)>length(x)
    global fcalls +=1
    dx = zeros(eltype(x),length(x)*2)
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx
end

function nsqf!(dx,x)
    global fcalls +=1
    for i in 2:length(dx)
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    nothing
end

function nsqf2!(dx,x)
    global fcalls +=1
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    nothing
end

function staticf(x,N=length(x))
    global fcalls += 1
    SVector{N}([i == 1 ? -2x[1]+x[2] : (i == N ? x[N-1]-2x[N] : x[i-1]-2x[i]+x[i+1]) for i in 1:N])
end

function staticnsqf(x,N=div(length(x),2))
    global fcalls += 1
    SVector{N}(vcat([-2x[1]+x[2]],[x[i-1]-2x[i]+x[i+1] for i in 2:N]))
end

function second_derivative_stencil(N)
    A = zeros(N,N)
    for i in 1:N, j in 1:N
        (j-i==-1 || j-i==1) && (A[i,j]=1)
        j-i==0 && (A[i,j]=-2)
    end
    A
end

@info "ended definitions"

x = rand(30)
dx = rand(30)

J = jacobian(f, dx, x)
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
@test fcalls == 3

fcalls = 0
_J1 = similar(_J)
forwarddiff_color_jacobian!(_J1, f, x, colorvec = repeat(1:3,10))
@test _J1 ≈ J
@test fcalls == 1

@info "second passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J)
@test _J1 ≈ J
@test typeof(_J1) == typeof(_J)
@test fcalls == 1

@info "third passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J)
@test _J1 ≈ J
@test fcalls == 1

#oop with in-place Jacobian
fcalls = 0
_oop_jacout = sparse(1.01 .* J) # want to be nonzero to check that the pre-allocated matrix is overwritten properly
forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J)
@test _oop_jacout ≈ J
@test typeof(_oop_jacout) == typeof(_J)
@test fcalls == 1

# BandedMatrix
_oop_jacout = BandedMatrix(-1 => diag(J, -1) .* 1.01, 0 => diag(J, 0) .* 1.01,
                           1 => diag(J, 1) .* 1.01) # check w/BandedMatrix instead of sparse
fcalls = 0
forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec = repeat(1:3,10), sparsity = _J)
@test _oop_jacout ≈ J
@test isa(_oop_jacout, BandedMatrix)
@test fcalls == 1

@info "4th passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = SMatrix{30,30}(_J))
@test _J1 ≈ J
@test fcalls == 1

@info "5"

_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), jac_prototype = SMatrix{30,30}(_J))
@test _J1 ≈ J
_J1 = forwarddiff_color_jacobian(oopf, x, jac_prototype = similar(_J))
@test _J1 ≈ J
_J1 = forwarddiff_color_jacobian(oopf, x)
@test _J1 ≈ J

#Non-square Jacobian
#length(dx)<length(x)
nsqJ = jacobian(nsqf,x)
spnsqJ = sparse(nsqJ)
_nsqJ = forwarddiff_color_jacobian(nsqf, x, dx = nothing)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, colorvec = repeat(1:3,10), sparsity = spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, colorvec = repeat(1:3,10), sparsity = spnsqJ, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, jac_prototype = SMatrix{15,30}(nsqJ))
@test _nsqJ ≈ nsqJ
@test typeof(_nsqJ) == typeof(SMatrix{15,30}(nsqJ))
_nsqJ = forwarddiff_color_jacobian(staticnsqf, SVector{30}(x), jac_prototype = SMatrix{15,30}(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(staticnsqf, SVector{30}(x), jac_prototype = SMatrix{15,30}(nsqJ), colorvec = repeat(1:3,10), sparsity = spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf!, x)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf!, x, colorvec = repeat(1:3,10), sparsity = spnsqJ )
@test _nsqJ ≈ nsqJ

#length(dx)>length(x)
nsqJ = jacobian(nsqf2,x)
spnsqJ = sparse(nsqJ)
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, dx = nothing)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec = repeat(1:3,10), sparsity = spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec = repeat(1:3,10), sparsity = spnsqJ, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype = SMatrix{60,30}(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x, colorvec = repeat(1:3,10), sparsity = spnsqJ )
@test _nsqJ ≈ nsqJ

fcalls = 0
_J1 = similar(_J)
jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1)
forwarddiff_color_jacobian!(_J1, f, x, jac_cache)
@test _J1 ≈ J
@test fcalls == 1

fcalls = 0
_J1 = similar(_J)
_denseJ1 = collect(_J1)
forwarddiff_color_jacobian!(_denseJ1, f, x, colorvec = repeat(1:3,10), sparsity = _J1)
@test _denseJ1 ≈ J
@test fcalls == 1

fcalls = 0
_J1 = similar(_J)
_denseJ1 = collect(_J1)
jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1)
forwarddiff_color_jacobian!(_denseJ1, f, x, jac_cache)
@test _denseJ1 ≈ J
@test fcalls == 1

_Jt = similar(Tridiagonal(J))
forwarddiff_color_jacobian!(_Jt, f, x, colorvec = repeat(1:3,10), sparsity = _Jt)
@test _Jt ≈ J

#https://github.com/JuliaDiff/FiniteDiff.jl/issues/67#issuecomment-516871956
function f(out, x)
	x = reshape(x, 100, 100)
	out = reshape(out, 100, 100)
	for i in 1:100
		for j in 1:100
			out[i, j] = x[i, j] + x[max(i -1, 1), j] + x[min(i+1, size(x, 1)), j] +  x[i, max(j-1, 1)]  + x[i, min(j+1, size(x, 2))]
		end
	end
	return vec(out)
end
x = rand(10000)
J = BandedBlockBandedMatrix(Ones(10000, 10000), fill(100, 100), fill(100, 100), (1, 1), (1, 1))
Jsparse = sparse(J)
colors = matrix_colors(J)
forwarddiff_color_jacobian!(J, f, x, colorvec=colors)
forwarddiff_color_jacobian!(Jsparse, f, x, colorvec=colors)
@test J ≈ Jsparse

# Non vector input
x = rand(2,2)
oopf(x) = x
iipf(fx,x) = (fx.=x)
J = forwarddiff_color_jacobian(oopf,x)
@test J ≈ Matrix(I,4,4)
J = zero(J)
forwarddiff_color_jacobian!(J,iipf,x,dx=similar(x))
@test J ≈ Matrix(I,4,4)

#1x1 SVector test
x = SVector{1}([1.])
f(x) = x
J = forwarddiff_color_jacobian(f,x)
@test J isa SArray
@test J ≈ SMatrix{1,1}([1.])


@testset "Tensor operations" begin
@testset "right propagation of A and B" begin
    nl1 = 2
    nl2 = 3
    na = 4
    nb = 5
    ns = 6

    # Float64 propagation.
    L = reshape(collect(1. : nl1*nl2), (nl1, nl2))
    A = reshape(collect(1. : nl1*ns*na), (nl1, ns, na))
    B = reshape(collect(3. : 2+nl2*ns*nb), (nl2, ns, nb))
    # Do manual contraction first.
    new_L = zeros(ComplexF64, (na, nb))
    for a=1:na, b=1:nb
        for l1=1:nl1, l2=1:nl2, s=1:ns
            new_L[a, b] += L[l1, l2]*A[l1, s, a]*B[l2, s, b]
        end
    end
    @test MPStates.prop_right2(L, A, B) ≈ new_L

    # ComplexF64 propagation.
    L = reshape(complex.(collect(1. : nl1*nl2)), (nl1, nl2))
    A = reshape(complex.(collect(1. : nl1*ns*na), collect(-2. : -3+nl1*ns*na)),
                (nl1, ns, na))
    B = reshape(complex.(collect(3. : 2+nl2*ns*nb), collect(-1. : -2+nl2*ns*nb)),
                (nl2, ns, nb))
    # Do manual contraction first.
    new_L = zeros(ComplexF64, (na, nb))
    for a=1:na, b=1:nb
        for l1=1:nl1, l2=1:nl2, s=1:ns
            new_L[a, b] += L[l1, l2]*A[l1, s, a]*conj(B[l2, s, b])
        end
    end
    @test MPStates.prop_right2(L, A, B) ≈ new_L
end

@testset "right propagation of A, M, and B" begin
    nl1 = 2
    nl2 = 3
    nl3 = 4
    na = 4
    nb = 5
    nc = 3
    ns1 = 6
    ns2 = 2
    L = reshape(complex.(collect(1. : nl1*nl2*nl3)), (nl1, nl2, nl3))
    A = reshape(collect(1. : nl1*ns1*na) + 1im*collect(-2. : -3+nl1*ns1*na),
                (nl1, ns1, na))
    M = reshape(collect(1. : nl2*ns1*ns2*nb) + 1im*collect(0. : -1+nl2*ns1*ns2*nb),
                (nl2, ns1, ns2, nb))
    B = reshape(collect(3. : 2+nl3*ns2*nc) + 1im*collect(-1. : -2+nl3*ns2*nc),
                (nl3, ns2, nc))
    # Do manual contraction first.
    new_L = zeros(ComplexF64, (na, nb, nc))
    for a=1:na, b=1:nb, c=1:nc
        for l1=1:nl1, l2=1:nl2, l3=1:nl3, s1=1:ns1, s2=1:ns2
            new_L[a, b, c] += L[l1, l2, l3]*A[l1, s1, a]*M[l2, s1, s2, b]*conj(B[l3, s2, c])
        end
    end
    @test MPStates.prop_right3(L, A, M, B) ≈ new_L
end

@testset "right propagation of M, and B for DMRG3S" begin
    nl1 = 2
    nl2 = 3
    nl3 = 4
    nr2 = 5
    nr3 = 3
    ns1 = 6
    ns2 = 2
    L = reshape(complex.(collect(1. : nl1*nl2*nl3)), (nl1, nl2, nl3))
    M = reshape(complex.(collect(1. : nl2*ns1*ns2*nr2), collect(0. : -1+nl2*ns1*ns2*nr2)),
                (nl2, ns1, ns2, nr2))
    B = reshape(complex.(collect(3. : 2+nl3*ns2*nr3), collect(-1. : -2+nl3*ns2*nr3)),
                (nl3, ns2, nr3))
    # Do manual contraction first.
    P = zeros(ComplexF64, (nl1, ns1, nr2, nr3))
    for l1=1:nl1, s1=1:ns1, r2=1:nr2, r3=1:nr3
        for l2=1:nl2, l3=1:nl3, s2=1:ns2
            P[l1, s1, r2, r3] += L[l1, l2, l3]*M[l2, s1, s2, r2]*conj(B[l3, s2, r3])
        end
    end
    @test MPStates.prop_right_subexp(L, M, B) ≈ P
end

@testset "right propagation of A, M1, M2, and B" begin
    nl1 = 2
    nl2 = 3
    nl3 = 4
    nl4 = 4
    na = 4
    nb = 5
    nc = 3
    nd = 5
    ns1 = 6
    ns2 = 2
    ns3 = 4
    L = reshape(complex.(collect(1. : nl1*nl2*nl3*nl4)), (nl1, nl2, nl3, nl4))
    A = reshape(collect(1. : nl1*ns1*na) + 1im*collect(-2. : -3+nl1*ns1*na),
                (nl1, ns1, na))
    M1 = reshape(collect(1. : nl2*ns1*ns2*nb) + 1im*collect(0. : -1+nl2*ns1*ns2*nb),
                 (nl2, ns1, ns2, nb))
    M2 = reshape(collect(1. : nl3*ns2*ns3*nc) + 1im*collect(0. : -1+nl3*ns2*ns3*nc),
                 (nl3, ns2, ns3, nc))
    B = reshape(collect(3. : 2+nl4*ns3*nd) + 1im*collect(-1. : -2+nl4*ns3*nd),
                (nl4, ns3, nd))
    # Do manual contraction first.
    new_L = zeros(ComplexF64, (na, nb, nc, nd))
    for a=1:na, b=1:nb, c=1:nc, d=1:nd
        for l1=1:nl1, l2=1:nl2, l3=1:nl3, l4=1:nl4, s1=1:ns1, s2=1:ns2, s3=1:ns3
            new_L[a, b, c, d] += (L[l1, l2, l3, l4]*A[l1, s1, a]*M1[l2, s1, s2, b]
                                  *M2[l3, s2, s3, c]*conj(B[l4, s3, d]))
        end
    end
    @test MPStates.prop_right4(L, A, M1, M2, B) ≈ new_L
end

@testset "left propagation of A and B" begin
    nr1 = 2
    nr2 = 3
    na = 4
    nb = 5
    ns = 6

    # Float64 propagation.
    R = reshape(collect(1. : nr1*nr2), (nr1, nr2))
    A = reshape(collect(1. : nr1*ns*na), (na, ns, nr1))
    B = reshape(collect(3. : 2+nr2*ns*nb), (nb, ns, nr2))
    # Do manual contraction first.
    new_R = zeros(ComplexF64, (na, nb))
    for a=1:na, b=1:nb
        for r1=1:nr1, r2=1:nr2, s=1:ns
            new_R[a, b] += R[r1, r2]*A[a, s, r1]*B[b, s, r2]
        end
    end
    @test MPStates.prop_left2(A, B, R) ≈ new_R

    # ComplexF64 propagation.
    R = reshape(complex.(collect(1. : nr1*nr2)), (nr1, nr2))
    A = reshape(complex.(collect(1. : nr1*ns*na), collect(-2. : -3+nr1*ns*na)),
                (na, ns, nr1))
    B = reshape(complex.(collect(3. : 2+nr2*ns*nb), collect(-1. : -2+nr2*ns*nb)),
                (nb, ns, nr2))
    # Do manual contraction first.
    new_R = zeros(ComplexF64, (na, nb))
    for a=1:na, b=1:nb
        for r1=1:nr1, r2=1:nr2, s=1:ns
            new_R[a, b] += R[r1, r2]*A[a, s, r1]*conj(B[b, s, r2])
        end
    end
    @test MPStates.prop_left2(A, B, R) ≈ new_R
end

@testset "left propagation of A, M, and B" begin
    nr1 = 2
    nr2 = 3
    nr3 = 4
    na = 4
    nb = 5
    nc = 3
    ns1 = 6
    ns2 = 2
    R = reshape(complex.(collect(1. : nr1*nr2*nr3)), (nr1, nr2, nr3))
    A = reshape(collect(1. : nr1*ns1*na) + 1im*collect(-2. : -3+nr1*ns1*na),
                (na, ns1, nr1))
    M = reshape(collect(1. : nr2*ns1*ns2*nb) + 1im*collect(0. : -1+nr2*ns1*ns2*nb),
                (nb, ns1, ns2, nr2))
    B = reshape(collect(3. : 2+nr3*ns2*nc) + 1im*collect(-1. : -2+nr3*ns2*nc),
                (nc, ns2, nr3))
    # Do manual contraction first.
    new_R = zeros(ComplexF64, (na, nb, nc))
    for a=1:na, b=1:nb, c=1:nc
        for r1=1:nr1, r2=1:nr2, r3=1:nr3, s1=1:ns1, s2=1:ns2
            new_R[a, b, c] += R[r1, r2, r3]*A[a, s1, r1]*M[b, s1, s2, r2]*conj(B[c, s2, r3])
        end
    end
    @test MPStates.prop_left3(A, M, B, R) ≈ new_R
end

@testset "left propagation of M, and B for DMRG3S" begin
    nr1 = 2
    nr2 = 3
    nr3 = 4
    nl2 = 5
    nl3 = 3
    ns1 = 6
    ns2 = 2
    R = reshape(complex.(collect(1. : nr1*nr2*nr3)), (nr1, nr2, nr3))
    M = reshape(complex.(collect(1. : nr2*ns1*ns2*nl2), (0. : -1+nr2*ns1*ns2*nl2)),
                (nl2, ns1, ns2, nr2))
    B = reshape(complex.(collect(3. : 2+nr3*ns2*nl3), (-1. : -2+nr3*ns2*nl3)),
                (nl3, ns2, nr3))
    # Do manual contraction first.
    P = zeros(ComplexF64, (ns1, nr1, nl2, nl3))
    for r1=1:nr1, s1=1:ns1, l2=1:nl2, l3=1:nl3
        for r2=1:nr2, r3=1:nr3, s2=1:ns2
            P[s1, r1, l2, l3] += R[r1, r2, r3]*M[l2, s1, s2, r2]*conj(B[l3, s2, r3])
        end
    end
    @test MPStates.prop_left_subexp(M, B, R) ≈ P
end

@testset "left propagation of A, M1, M2, and B" begin
    nr1 = 2
    nr2 = 3
    nr3 = 4
    nr4 = 4
    na = 4
    nb = 5
    nc = 3
    nd = 5
    ns1 = 6
    ns2 = 2
    ns3 = 4
    R = reshape(complex.(collect(1. : nr1*nr2*nr3*nr4)), (nr1, nr2, nr3, nr4))
    A = reshape(collect(1. : nr1*ns1*na) + 1im*collect(-2. : -3+nr1*ns1*na),
                (na, ns1, nr1))
    M1 = reshape(collect(1. : nr2*ns1*ns2*nb) + 1im*collect(0. : -1+nr2*ns1*ns2*nb),
                 (nb, ns1, ns2, nr2))
    M2 = reshape(collect(1. : nr3*ns2*ns3*nc) + 1im*collect(0. : -1+nr3*ns2*ns3*nc),
                 (nc, ns2, ns3, nr3))
    B = reshape(collect(3. : 2+nr4*ns3*nd) + 1im*collect(-1. : -2+nr4*ns3*nd),
                (nd, ns3, nr4))
    # Do manual contraction first.
    new_R = zeros(ComplexF64, (na, nb, nc, nd))
    for a=1:na, b=1:nb, c=1:nc, d=1:nd
        for r1=1:nr1, r2=1:nr2, r3=1:nr3, r4=1:nr4, s1=1:ns1, s2=1:ns2, s3=1:ns3
            new_R[a, b, c, d] += (R[r1, r2, r3, r4]*A[a, s1, r1]*M1[b, s1, s2, r2]
                                  *M2[c, s2, s3, r3]*conj(B[d, s3, r4]))
        end
    end
    @test MPStates.prop_left4(A, M1, M2, B, R) ≈ new_R
end
end # @testset "Tensor operations"


using Plots
using LaTeXStrings

function arch(θ₁,θ₂;radius=1.,Δθ=1.)
  θ₁ = π*θ₁/180
  θ₂ = π*θ₂/180
  Δθ = π*Δθ/180
  l = round(Int,(θ₂-θ₁)/Δθ)
  x = zeros(l)
  y = zeros(l)
  for i in 1:l
    θ = θ₁ + i*Δθ
    x[i] = radius*cos(θ)
    y[i] = radius*sin(θ)
  end
  return x, y
end

plot()

x, y = arch(0,360)
plot(x,y,seriestype=:shape,label="",alpha=0.5)

x, y = arch(0,360,radius=0.95)
plot!(x,y,seriestype=:shape,label="",fillcolor=:white)

x, y = arch(0,360,radius=0.7)
plot!(x,y,seriestype=:shape,label="",alpha=0.5,fillcolor=:red)

x, y = arch(0,360,radius=0.65)
plot!(x,y,seriestype=:shape,label="",fillcolor=:white)

plot!([0,0],[0,1.1],arrow=true,color=:black,linewidth=2,label="")
plot!([0,1.1],[0,0],arrow=true,color=:black,linewidth=2,label="")

x, y = arch(15,16,radius=0.65)
plot!([0,x[1]],[0,y[1]],arrow=true,color=:black,linewidth=1,label="")

x, y = arch(35,36,radius=0.95)
plot!([0,x[1]],[0,y[1]],arrow=true,color=:black,linewidth=1,label="")

plot!(draw_arrow=true)
plot!(showaxis=:no,ticks=nothing,xlim=[-0.1,1.1],ylim=[-0.1,1.1],)
plot!(xlabel="x",ylabel="y",size=(400,400))

annotate!(0.58,-0.07,text(L"\Delta v_1",10))
annotate!(0.88,-0.07,text(L"\Delta v_2",10))

savefig("./velocities.pdf")


# return the i-th bit
@inline bitat(chunk::UInt64, i::Int) = (chunk >>> (64 - i)) & 0x01 == 0x01
@inline bitat(::Type{T}, chunk::T, i::Int) where {T<:Unsigned} = (chunk >>> (sizeof(T) * 8 - i)) & 1 == 1

# assume 1 byte = 8 bits
@inline bitsof(::Type{T}) where {T<:Unsigned} = sizeof(T) * 8

@inline mod64(i) = i & 63

# make a bit mask
@inline lmask(typ::Type{T}, n::Int) where {T<:Unsigned} = typemax(typ) << (bitsof(typ) - n)
@inline rmask(typ::Type{T}, n::Int) where {T<:Unsigned} = typemax(typ) >> (bitsof(typ) - n)


"""
    distance_profile(dist, Q, T; kwargs...)

Calculate the distance profile corresponding to sliding `Q` over `T` and measuring distances between time windows using `dist`

The output is a vector of length `length(T) - length(Q) + 1`.
If inputs are matrices or higher-dimensional arrays, time is considered to be the last axis, and windows on the form `T[:, t:t+lastlength(Q)-1]`, where `lastlength` measures the length along the last axis, are compared.

Keyword arguments are sent to `dist`, which is called like this `dist(Q,T; kwargs...)`. `dist` can be a function or a distance from Distances.jl.

`distance_profile!(D, dist, Q, T; kwargs...)` is an in-place version.
"""
function distance_profile(dist, Q::AbstractArray{S}, T::AbstractArray{S}; kwargs...) where S
    m = lastlength(Q)
    n = lastlength(T)
    n >= m || throw(ArgumentError("Q cannot be longer than T"))
    l = n-m+1
    D = map(1:l) do i
        dist(Q, getwindow(T, m, i); kwargs...)
    end
end

"$(SIGNATURES)"
function distance_profile!(D::AbstractArray, dist, Q::AbstractArray{S}, T::AbstractArray{S}; kwargs...) where S
    m = lastlength(Q)
    n = lastlength(T)
    n >= m || throw(ArgumentError("Q cannot be longer than T"))
    l = n-m+1
    length(D) >= l || throw(ArgumentError("D is too short"))
    map!(D, 1:l) do i
        dist(Q, getwindow(T, m, i); kwargs...)
    end
end

"$(SIGNATURES)"
function distance_profile!(D::AbstractVector{S},::ZEuclidean, QT::AbstractVector{S}, μ, σ, m::Int, i::Int) where S <: Number
    @assert i <= length(D)
    @avx for j = eachindex(D)
        frac = (QT[j] - m*μ[i]*μ[j]) / (m*σ[i]*σ[j])
        D[j] = sqrt(max(2m*(1-frac), 0))
    end
    D[i] = typemax(eltype(D))
    D
end

"$(SIGNATURES)"
distance_profile(
    ::ZEuclidean,
    QT::AbstractVector{S},
    μ::AbstractVector{S},
    σ::AbstractVector{S},
    m::Int,
) where {S<:Number} = distance_profile!(similar(μ), ZEuclidean(), QT, μ, σ, m, 1)

"""
    $(SIGNATURES)

Accepts precomputed sliding mean and std of the input arrays. `QT` is the windowed dot product and `i` is the index into `T` (the longer time series).
"""
function distance_profile!(D::AbstractVector{S},::ZEuclidean, QT::AbstractVector{S}, μA, σA, μT, σT, m::Int, i::Int) where S <: Number
    @assert i <= length(μA)
    @avx for j = eachindex(D,QT,μT,σT)
        frac = (QT[j] - m*μA[i]*μT[j]) / (m*σA[i]*σT[j])
        D[j] = sqrt(max(2m*(1-frac), 0))
    end
    D
end

"$(SIGNATURES)"
distance_profile(::ZEuclidean, QT::AbstractVector{S}, μA, σA, μT, σT, m::Int) where {S<:Number} =
    distance_profile!(similar(μT), ZEuclidean(), QT, μA, σA, μT, σT, m, 1)

"""
    distance_profile(::ZEuclidean, Q, T)

Compute the z-normalized Euclidean distance profile corresponding to sliding `Q` over `T`
"""
function distance_profile!(
    D::AbstractVector{S},
    ::ZEuclidean,
    Q::AbstractVector{S},
    T::AbstractVector{S},
) where {S<:Number}
    m = length(Q)
    μ, σ = sliding_mean_std(T, m)
    QT = window_dot(znorm(Q), T) # TODO: allocates a new znorm(Q) each time
    @avx for j in eachindex(D,QT,σ)
        frac = QT[j] / (m * σ[j])
        D[j] = sqrt(max(2m * (1 - frac), 0))
    end
    D
end

"$(SIGNATURES)"
distance_profile(::ZEuclidean, Q::AbstractArray{S}, T::AbstractArray{S}) where {S} =
    distance_profile!(similar(T, length(T) - length(Q) + 1), ZEuclidean(), Q, T)

"""
The dot product between query Q and all subsequences of the same length as Q in time series T
"""
function window_dot(Q, T)
    n   = length(T)
    m   = length(Q)
    QT  = conv(reverse(Q), T)
    return QT[m:n]
end

"""
    $(SIGNATURES)

return vectors with mean and std of sliding windows of length `m`
"""
function sliding_mean_std(x::AbstractVector{T}, m) where T
    @assert length(x) >= m
    n = length(x)-m+1
    s = ss = zero(float(T))
    μ = Vector{float(T)}(undef, n)
    σ = Vector{float(T)}(undef, n)
    @avx for i = 1:m
        s  += x[i]
        ss += x[i]^2
    end
    μ[1] = s/m
    σ[1] = sqrt(max(ss/m - μ[1]^2, 0))
    @fastmath @inbounds for i = 1:n-1 # fastmath making it more accurate here as well, but not faster
        s -= x[i]
        ss -= x[i]^2
        s += x[i+m]
        ss += x[i+m]^2
        μ[i+1] = s/m
        σ[i+1] = sqrt(max(ss/m - μ[i+1]^2, 0))
    end
    μ,σ
end

@inbounds @fastmath function sliding_mean_std(x::AbstractMatrix{T}, m) where T
    @assert lastlength(x) >= m
    n = lastlength(x)-m+1
    s  = zeros(T, size(x,1))
    ss = zeros(T, size(x,1))
    μ = Matrix{float(T)}(undef, size(x,1), n)
    σ = Matrix{float(T)}(undef, size(x,1), n)
    for i = 1:m
        @avx s  .+= view(x, :, i)
        @avx ss .+= view(x, :, i).^2
    end
    @avx μ[:,1] .= s./m
    @avx σ[:,1] .= sqrt.(max.(ss./m .- view(μ, :, 1).^2, 0))
    for i = 1:n-1 # fastmath making it more accurate here as well, but not faster
        @avx s .-= view(x, :, i)
        @avx ss .-= view(x, :, i).^2
        @avx s .+= view(x, :, i+m)
        @avx ss .+= view(x, :, i+m).^2
        @avx μ[:,i+1] .= s./m
        @avx σ[:,i+1] .= sqrt.(max.(ss./m .- view(μ, :, i+1).^2, 0))
    end
    μ,σ
end

"""
    $(SIGNATURES)

return mean of sliding windows of length `m`. Operates in-place and stores result in first argument
"""
function sliding_mean!(μ,x::AbstractVector{T}, m) where T
    n = length(x)-m+1
    @assert length(x) >= m
    @assert length(μ) >= n

    s = zero(T)
    @avx for i = 1:m
        s += x[i]
    end
    μ[1] = s/m
    @fastmath @inbounds for i = 1:n-1 # fastmath making it more accurate here as well, but not faster
        s -= x[i]
        s += x[i+m]
        μ[i+1] = s/m
    end
    μ
end

"""
$(SIGNATURES)

Returns the sliding-window entropy
"""
function sliding_entropy(x::AbstractVecOrMat{T}, m=actuallastlength(x)) where T
    N = actuallastlength(x)
    n = N-m+1
    e = Vector{T}(undef, n)
    if x isa AbstractVector
        ent = vmap(x->x*log(x), x)
    else
        ent = zeros(T, N)
        @avx for i = 1:N
            for j = 1:size(x,1)
                ent[i] += x[j,i]*log(x[j,i])
            end
        end
    end
    e0 = zero(T)
    @avx for i = 1:m
        e0 += ent[i]
    end
    e[1] = e0
    @fastmath @inbounds for i = 1:n-1
        e0 -= ent[i]
        e0 += ent[i+m]
        e[i+1] = e0
    end
    e .= .-e
end

"""
$(SIGNATURES)

Returns the sliding-window entropy, normalized to have unit sum for each feature channel.
"""
function sliding_entropy_normalized(x::AbstractArray{T}, m=actuallastlength(x)) where T
    sums = sum(x, dims=ndims(x))
    sliding_entropy(x ./ sums, m)
end


function solve(
    i::SpanningTreeInstance{T, Maximise},
    ::GreedyAlgorithm,
) where {T}
    return solve(i, PrimAlgorithm())
end

function solve(i::SpanningTreeInstance{T, Maximise}, ::PrimAlgorithm) where {T}
    remaining_edges = PriorityQueue{Edge{T}, Float64}(Base.Reverse) # Easy retrieval of highest-reward edge.
    node_done = falses(nv(i.graph))

    # Helper methods.
    function edges_around(g, v)
        return (
            (v < x) ? edgetype(g)(v, x) : edgetype(g)(x, v) for
            x in neighbors(g, v)
        )
    end

    # Initialise with the source node (arbitrarily, the first one).
    first_node = first(vertices(i.graph))
    node_done[first_node] = true
    for e in edges_around(i.graph, first_node)
        enqueue!(remaining_edges, e => reward(i, e))
    end

    # Build the spanning tree.
    solution = Edge{T}[]
    while length(solution) < nv(i.graph) - 1 # While there is still an edge to add...
        # Find an admissible edge, starting with those of minimum reward.
        e = nothing
        new_node = nothing
        while isnothing(new_node)
            e = dequeue!(remaining_edges)

            if node_done[src(e)] && !node_done[dst(e)]
                new_node = dst(e)
            end
            if !node_done[src(e)] && node_done[dst(e)]
                new_node = src(e)
            end
        end

        @assert !node_done[new_node]

        # Use this edge in the solution.
        push!(solution, e)
        node_done[new_node] = true

        # Prepare to use the neighbours of this new node, but skip nodes already done.
        for ϵ in edges_around(i.graph, new_node)
            if haskey(remaining_edges, ϵ)
                # No need to update the priority, the edge still has the same reward.
                continue
            end

            # Ensure the edge links a node already in the tree and another one outside.
            if (node_done[src(ϵ)] && !node_done[dst(ϵ)]) ||
               (!node_done[src(ϵ)] && node_done[dst(ϵ)])
                enqueue!(remaining_edges, ϵ => reward(i, ϵ))
            end
        end
    end

    return SpanningTreeSolution(i, solution)
end


using SpectralKit
using Test, DocStringExtensions, StaticArrays, BenchmarkTools, Sobol
using ForwardDiff: derivative, gradient

include("utilities.jl")

include("test_generic_api.jl")
include("test_chebyshev.jl")
include("test_univariate.jl")
include("test_smolyak.jl")


using FractionalDiffEq
using Plots

α = 0.7
β = 1.8
κ = 1
T = 2
m = 21
n = 148

result = solve(α, β, κ, T, m, n, FPDEMatrixDiscrete())

XX, YY = meshgrid(0.05^2/6 .*(0:147), 0:0.05:1)

plotlyjs()

plot(XX, YY, result, st=:surface)


using SIAN

println("Setting up the problem")

ode = @ODEmodel(
  x1'(t) = -(k3 + k7) * x1(t) + k4 * x2(t),
  x2'(t) = k3 * x1(t) - (k4 + a(t) * k5 + b(t) * d(t) * k5) * x2(t) + k6 * x3(t) + k6 * x4(t) + k5 * x2(t) * x3(t) + k5 * x2(t) * x4(t),
  x3'(t) = a(t) * k5 * x2(t) - k6 * x3(t) - k5 * x2(t) * x3(t),
  x4'(t) = b(t) * d(t) * k5 * x2(t) - k6 * x4(t) - k5 * x2(t) * x4(t),
  x5'(t) = k7 * x1(t),
  a'(t) = 0,
  b'(t) = 0,
  d'(t) = 0,
  y1(t) = x5(t),
  y2(t) = a(t),
  y3(t) = b(t),
  y4(t) = d(t)
)

res = identifiability_ode(ode, get_parameters(ode); p=0.99, p_mod=2^29 - 3, nthrds=1)

println(res)


# https://fenicsproject.org/docs/dolfin/latest/python/demos/singular-poisson/demo_singular-poisson.py.html

module singular_poisson

using FenicsPy

@assert has_linear_algebra_backend("PETSc") "DOLFIN has not been configured with PETSc. "

parameters["linear_algebra_backend"] = "PETSc"

# Create mesh and define function space
mesh = UnitSquareMesh(64, 64)
V = FunctionSpace(mesh, "CG", 1)

u = TrialFunction(V)
v = TestFunction(V)

f = Expression("10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
g = Expression("-sin(5*x[0])", degree=2)

a = inner(grad(u), grad(v))*dx
L = f*v*dx + g*v*ds

# Assemble system
A = assemble(a)
b = assemble(L)

# Solution Function
u = FeFunction(V)

# Create Krylov solver
solver = PETScKrylovSolver("cg")
solver.set_operator(A)

# Create vector that spans the null space and normalize
null_vec = u.vector()
V.dofmap().set(null_vec, 1.0)
null_vec *= 1.0/null_vec.norm("l2")

# Create null space basis object and attach to PETSc matrix
null_space = VectorSpaceBasis([null_vec])
as_backend_type(A).set_nullspace(null_space)

null_space.orthogonalize(b)

solver.solve(u.vector(), b)

# Output file
File("singular_poisson/output.pvd") << u

#plot(u)

end  # module singular_poisson


module Observations

using DocStringExtensions
using LinearAlgebra
using Statistics

export Observation

"""
    Observation{FT <: AbstractFloat}

Structure that contains the observations

# Fields

$(TYPEDFIELDS)
"""
struct Observation{FT <: AbstractFloat}
    "vector of observational samples, each of length sample_dim"
    samples::Vector{Vector{FT}}
    "covariance of the observational noise (assumed to be normally
    distributed); `sample_dim x sample_dim` (where `sample_dim` is the number of
    elements in each sample), or a scalar if the sample dim is 1. If not
    supplied, `obs_noise_cov` is set to a diagonal matrix whose non-zero elements
    are the variances of the samples, or to a scalar variance in the case of
    1d samples. `obs_noise_cov` is set to nothing if only a single sample is
    provided."
    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, FT, Nothing}
    "sample mean"
    mean::Union{AbstractVector{FT}, FT}
    "names of the data"
    data_names::Union{AbstractVector{String}, String}
end

# Constructors
function Observation(samples::Vector{Vector{FT}}, data_names::Union{Vector{String}, String}) where {FT <: AbstractFloat}

    N_samples = length(samples)
    # convert to sample_dim x N_samples to determine sample covariance

    if N_samples == 1
        # only one sample - this requires a bit more data massaging
        temp1 = convert(Array, reshape(hcat(samples...)', N_samples, :))
        temp = dropdims(temp1, dims = 1)
        samplemean = vec(mean(temp, dims = 2))
        obs_noise_cov = nothing
    else
        temp = convert(Array, reshape(hcat(samples...)', N_samples, :))
        if size(temp, 2) == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the covariance) are scalars
            samplemean = mean(temp)
            obs_noise_cov = var(temp)
        else
            samplemean = vec(mean(temp, dims = 1))
            obs_noise_cov = cov(temp)
        end
    end
    Observation(samples, obs_noise_cov, samplemean, data_names)
end

function Observation(
    samples::AbstractMatrix{FT},
    data_names::Union{AbstractVector{String}, String},
) where {FT <: AbstractFloat}

    # samples is of size sample_dim x N_samples
    sample_dim, N_samples = size(samples)

    if N_samples == 1
        # Only one sample, so there is no covariance to be computed and the
        # sample mean equals the sample itself
        obs_noise_cov = nothing
        samplemean = vec(samples)
        samples_vec = vec([vec(samples)])

    else
        # convert matrix of samples to a vector of vectors
        samples_vec = [samples[:, i] for i in 1:N_samples]
        if sample_dim == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the variance) are scalars
            samplemean = mean(samples)
            obs_noise_cov = var(samples)
        else
            samplemean = vec(mean(samples, dims = 2))
            obs_noise_cov = cov(samples, dims = 2)
        end
    end
    Observation(samples_vec, obs_noise_cov, samplemean, data_names)
end

function Observation(
    samples::Vector{Vector{FT}},
    obs_noise_cov::AbstractMatrix{FT},
    data_names::Union{AbstractVector{String}, String},
) where {FT <: AbstractFloat}

    N_samples = length(samples)
    sample_dim = length(samples[1])
    # convert to sample_dim x N_samples to determine sample covariance

    if N_samples == 1
        # only one sample - this requires a bit more data massaging
        temp1 = convert(Array, reshape(hcat(samples...)', N_samples, :))
        temp = dropdims(temp1, dims = 1)
        samplemean = vec(mean(temp, dims = 2))
        obsnoisecov = obs_noise_cov
    else
        temp = convert(Array, reshape(hcat(samples...)', N_samples, :))
        if sample_dim == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the covariance) are scalars
            samplemean = mean(temp)
            err = ("When sample_dim is 1, obs_cov_noise must be a 1x1 matrix")
            @assert(size(obs_noise_cov) == (1, 1), err)
            obsnoisecov = obs_noise_cov[1]
        else
            samplemean = vec(mean(temp, dims = 1))
            err = ("obs_cov_noise must be of size sample_dim x sample_dim.
                   \tsample_dim: number of elements per observation sample")
            @assert(size(obs_noise_cov) == (sample_dim, sample_dim), err)
            obsnoisecov = obs_noise_cov
        end
    end

    Observation(samples, obsnoisecov, samplemean, data_names)
end

function Observation(
    samples::AbstractMatrix{FT},
    obs_noise_cov::Union{AbstractMatrix{FT}, Nothing},
    data_names::Union{AbstractVector{String}, String},
) where {FT <: AbstractFloat}

    # samples is of size sample_dim x N_samples
    sample_dim, N_samples = size(samples)
    if N_samples == 1
        # sample mean equals the sample itself
        obsnoisecov = obs_noise_cov
        samplemean = vec(samples)
        samples_vec = vec([vec(samples)])
    else
        # convert matrix of samples to a vector of vectors
        samples_vec = [samples[:, i] for i in 1:N_samples]
        if sample_dim == 1
            # We have 1D samples, so the sample mean and covariance (which in
            # this case is actually the variance) are scalars
            samplemean = mean(samples)
            err = ("When sample_dim is 1, obs_cov_noise must be a 1x1 matrix")
            @assert(size(obs_noise_cov) == (1, 1), err)
            obsnoisecov = obs_noise_cov[1]
        else
            samplemean = vec(mean(samples, dims = 2))
            err = ("obs_cov_noise must be of size sample_dim x sample_dim.
                   \tsample_dim: number of elements per observation sample")
            @assert(size(obs_noise_cov) == (sample_dim, sample_dim), err)
            obsnoisecov = obs_noise_cov
        end
    end

    Observation(samples_vec, obsnoisecov, samplemean, data_names)
end

end # module Observations


using SparseArrays, LinearAlgebra, Luxor, Colors, ProgressMeter

function NonLinSolver(Fm::Matrix{Float64}, Ns, L::Float64, E::Float64, nu::Float64, DT::Float64)
    L1, L2 = [L, 0.], [0., L]
    U, Fext, Fint = zeros(Float64, Ns[6], 1), zeros(Float64, Ns[6], 1), zeros(Float64, Ns[6], 1)
    gps = gaussPoints2D(2)
    function Assembly(U::Matrix{Float64})
        # Update the elements and assemble the system matrices.
        Fint, Krow, Kcol, K = zeros(Float64, Ns[6], 1), Int64[], Int64[], zeros(Float64, Ns[6], Ns[6])
        for i = 1:Ns[5]
            Ce = vec(DOFs[i, :, :]')
            elFint, elK = Quad4(U[Ce, :], i, elCoords[i, :, :], E, nu)
            Fint[Ce] += elFint
            K[Ce, Ce] += elK
        end
        K = sparse(K)
        Kims = K[2*Ns[1]+1:end, 2*Ns[1]+1:end]
        Rims = (Fint-Fext)[2*Ns[1]+1:end, :]
        return Fint, Kims, Rims
    end
    function generateEq(Kims::SparseMatrixCSC, Rims::Matrix{Float64})
        Kii = Kims[1:2*Ns[2], 1:2*Ns[2]]
        Kim = Kims[1:2*Ns[2], 2*Ns[2]+1:2*Ns[8]]
        Kis = Kims[1:2*Ns[2], 2*Ns[8]+1:2*Ns[10]]

        Kmi = Kims[2*Ns[2]+1:2*Ns[8], 1:2*Ns[2]]
        Kmm = Kims[2*Ns[2]+1:2*Ns[8], 2*Ns[2]+1:2*Ns[8]]
        Kms = Kims[2*Ns[2]+1:2*Ns[8], 2*Ns[8]+1:2*Ns[10]]

        Ksi = Kims[2*Ns[8]+1:2*Ns[10], 1:2*Ns[2]]
        Ksm = Kims[2*Ns[8]+1:2*Ns[10], 2*Ns[2]+1:2*Ns[8]]
        Kss = Kims[2*Ns[8]+1:2*Ns[10], 2*Ns[8]+1:2*Ns[10]]

        Ri = Rims[1:2*Ns[2], :]
        Rm = Rims[2*Ns[2]+1:2*Ns[8], :]
        Rs = Rims[2*Ns[8]+1:2*Ns[10], :]

        Kupper = hcat(Kii, Kim + Kis)
        Klower = hcat(Kmi + Ksi, Kmm + Ksm + Kms + Kss)

        K = vcat(Kupper, Klower)
        R = vcat(Ri, Rm + Rs)
        RNorm = norm(R)

        return K, R, RNorm
    end

    function Iterate(TOL::Float64, DT::Float64, Draw::Bool)
        t, tFinal = (0.0, 1.0)
        A(t) = (t/tFinal)*(Fm-I)[1:2, 1:2]

        p1 = Progress(Int64(tFinal/DT))
        while round(t, digits = 2) < tFinal
            t += DT
            Ua, Ub = A(t)*L1, A(t)*L2
            Ucorner = vcat(zeros(2, 1), Ua, Ua + Ub, Ub)
            Uinternal = U[2*Ns[1]+1:2*Ns[7], :]
            Umaster = U[2*Ns[7]+1:2*Ns[9], :]
            Uslave = Umaster + D * Ucorner
            U = vcat(Ucorner, Uinternal, Umaster, Uslave)
            Fint, Kims, Rims = Assembly(U)
            RNorm = Inf
            #@printf("-----t: %.2f-----\n",t)
            #p2 = ProgressThresh(TOL, "Minimizing: ")
            #update!(p2, RNorm)
            while RNorm > TOL
                K, R, RNorm = generateEq(Kims, Rims)
                dU = K \ -R
                Uinternal += dU[1:2*Ns[2], :]    # dUInternal
                Umaster += dU[2*Ns[2]+1:end, :]  # dUmaster
                Uslave = Umaster + D * Ucorner
                U = vcat(Ucorner, Uinternal, Umaster, Uslave)
                Fint, Kims, Rims = Assembly(U)
                #@printf("  | RNorm: %.3f \n", RNorm)
                #update!(p2, RNorm)
            end
            next!(p1)
        end
        if Draw
            deformed = elCoords + U[DOFs]
            maxUx, maxUy = max(maximum(elCoords[:,:,1]), maximum(deformed[:,:,1])), max(maximum(elCoords[:,:,2]), maximum(deformed[:,:,2]))
            scl = 200
            Drawing(Int(round(maxUx*scl*1.4)), Int(round(maxUy*scl*1.4)))
            background("white")
            origin(maxUx*scl*0.2, maxUy*scl*1.2)
            for i = 1:Ns[5]
                undefPoints = [Point(elCoords[i, j, 1]*scl, -elCoords[i, j, 2]*scl) for j in 1:4]
                defPoints = [Point(deformed[i, j, 1]*scl, -deformed[i, j, 2]*scl) for j in 1:4]
                sethue("black")
                poly(undefPoints, :stroke, close = true)
                sethue("blue")
                poly(defPoints, :stroke, close = true)
            end
            finish();
        end
        return U
    end
    return Iterate(1e-6, DT, true)
end