text:
using ITensors, ITensorsGPU
using CuArrays
using Random, Logging, LinearAlgebra, DelimitedFiles

import ITensors: tensors

mutable struct PEPS
    Nx::Int
    Ny::Int
    A_::AbstractMatrix{ITensor}

    PEPS() = new(0, 0, Matrix{ITensor}(),0,0)

    PEPS(Nx::Int, Ny::Int, A::Matrix{ITensor}) = new(Nx, Ny, A)
    function PEPS(sites, lattice::Lattice, Nx::Int, Ny::Int; mindim::Int=1, is_gpu::Bool=false)
        p           = Matrix{ITensor}(undef, Ny, Nx)
        right_links = [ Index(mindim, "Link,c$j,r$i,r") for i in 1:Ny, j in 1:Nx ]
        up_links    = [ Index(mindim, "Link,c$j,r$i,u") for i in 1:Ny, j in 1:Nx ]
        T           = is_gpu ? cuITensor : ITensor
        @inbounds for ii in eachindex(sites)
            col = div(ii-1, Ny) + 1
            row = mod(ii-1, Ny) + 1
            s   = sites[ii]
            if 1 < row < Ny && 1 < col < Nx
                p[row, col] = T(right_links[row, col], up_links[row, col], right_links[row, col-1], up_links[row-1, col], s)
            elseif row == 1 && 1 < col < Nx
                p[row, col] = T(right_links[row, col], up_links[row, col], right_links[row, col-1], s)
            elseif 1 < row < Ny && col == 1
                p[row, col] = T(right_links[row, col], up_links[row, col], up_links[row-1, col], s)
            elseif row == Ny && 1 < col < Nx
                p[row, col] = T(right_links[row, col], right_links[row, col-1], up_links[row-1, col], s)
            elseif 1 < row < Ny && col == Nx
                p[row, col] = T(up_links[row, col], right_links[row, col-1], up_links[row-1, col], s)
            elseif row == Ny && col == 1
                p[row, col] = T(right_links[row, col], up_links[row-1, col], s)
            elseif row == Ny && col == Nx
                p[row, col] = T(right_links[row, col-1], up_links[row-1, col], s)
            elseif row == 1 && col == 1
                p[row, col] = T(right_links[row, col], up_links[row, col], s)
            elseif row == 1 && col == Nx
                p[row, col] = T(up_links[row, col], right_links[row, col-1], s)
            end
            @assert p[row, col] == p[ii]
        end
        new(Nx, Ny, p)
    end
end
Base.eltype(A::PEPS) = eltype(A.A_[1,1])

function cudelt(left::Index, right::Index)
    d_data   = CuArrays.zeros(Float64, ITensors.dim(left), ITensors.dim(right))
    ddi = diagind(d_data, 0)
    d_data[ddi] = 1.0
    delt = cuITensor(vec(d_data), left, right)
    return delt
end

function mydelt(left::Index, right::Index)
    d_data   = zeros(Float64, ITensors.dim(left), ITensors.dim(right))
    ddi = diagind(d_data, 0)
    d_data[ddi] = ones(Float64, length(ddi))
    delt = ITensor(vec(d_data), left, right)
    return delt
end

function checkerboardPEPS(sites, Nx::Int, Ny::Int; mindim::Int=1)
    lattice = square_lattice(Nx, Ny,yperiodic=false)
    A = PEPS(sites, lattice, Nx, Ny, mindim=mindim)
    @inbounds for ii ∈ eachindex(sites)
        row = div(ii-1, Nx) + 1
        col = mod(ii-1, Nx) + 1
        spin_side = isodd(row - 1) ⊻ isodd(col - 1) ? 2 : 1
        si  = findindex(A[ii], "Site")
        lis = findinds(A[ii], "Link")
        ivs = [li(1) for li in lis]
        ivs = vcat(ivs, si(spin_side))
        A[ii][ivs...] = 1.0
    end
    for row in 1:Ny, col in 1:Nx
        A[row, col] += randomITensor(inds(A[row, col]))/10.0
    end
    return A
end

function randomPEPS(sites, Nx::Int, Ny::Int; mindim::Int=1)
    lattice = squareLattice(Nx,Ny,yperiodic=false)
    A = PEPS(sites, lattice, Nx, Ny, mindim=mindim)
    @inbounds for ii ∈ eachindex(sites)
        randn!(A[ii])
        normalize!(A[ii])
    end
    return A
end

is_gpu(A::PEPS)    = all(is_gpu.(A[:,:]))
is_gpu(A::ITensor) = (data(store(A)) isa CuArray)

include("environments.jl")
include("ancillaries.jl")
include("gauge.jl")
include("observables.jl")
include("hamiltonians.jl")

function randomCuPEPS(sites, Nx::Int, Ny::Int; mindim::Int=1)
    lattice = squareLattice(Nx,Ny,yperiodic=false)
    A = PEPS(sites, lattice, Nx, Ny; mindim=mindim, is_gpu=true)
    @inbounds for ii ∈ eachindex(sites)
        randn!(A[ii])
        normalize!(A[ii])
    end
    return A
end

function cuPEPS(A::PEPS)
    cA = copy(A)
    @inbounds for i in 1:Ny, j in 1:Nx
        cA[i, j] = cuITensor(A[i, j])
    end
    return cA
end

function Base.collect(cA::PEPS)
    A = copy(cA)
    @inbounds for i in 1:Ny, j in 1:Nx
        A[i, j] = collect(cA[i, j])
    end
    return A
end

tensors(A::PEPS)   = A.A_
Base.size(A::PEPS) = (A.Ny, A.Nx)

Base.getindex(A::PEPS, i::Integer, j::Integer) = getindex(tensors(A), i, j)::ITensor
Base.getindex(A::PEPS, ::Colon,    j::Integer) = getindex(tensors(A), :, j)::Vector{ITensor}
Base.getindex(A::PEPS, i::Integer, ::Colon)    = getindex(tensors(A), i, :)::Vector{ITensor}
Base.getindex(A::PEPS, ::Colon,    j::UnitRange{Int}) = getindex(tensors(A), :, j)::Matrix{ITensor}
Base.getindex(A::PEPS, i::UnitRange{Int}, ::Colon)    = getindex(tensors(A), i, :)::Matrix{ITensor}
Base.getindex(A::PEPS, ::Colon, ::Colon)       = getindex(tensors(A), :, :)::Matrix{ITensor}
Base.getindex(A::PEPS, i::Integer)             = getindex(tensors(A), i)::ITensor

Base.setindex!(A::PEPS, val::ITensor, i::Integer, j::Integer)       = setindex!(tensors(A), val, i, j)
Base.setindex!(A::PEPS, vals::Vector{ITensor}, ::Colon, j::Integer) = setindex!(tensors(A), vals, :, j)
Base.setindex!(A::PEPS, vals::Vector{ITensor}, i::Integer, ::Colon) = setindex!(tensors(A), vals, i, :)
Base.setindex!(A::PEPS, vals::Matrix{ITensor}, ::Colon, j::UnitRange{Int}) = setindex!(tensors(A), vals, :, j)
Base.setindex!(A::PEPS, vals::Matrix{ITensor}, i::UnitRange{Int}, ::Colon) = setindex!(tensors(A), vals, i, :)

Base.copy(A::PEPS)    = PEPS(A.Nx, A.Ny, copy(tensors(A)))
Base.similar(A::PEPS) = PEPS(A.Nx, A.Ny, similar(tensors(A)))

function Base.show(io::IO, A::PEPS)
  print(io,"PEPS")
  (size(A)[1] > 0 && size(A)[2] > 0) && print(io,"\n")
  @inbounds for i in 1:A.Nx, j in 1:A.Ny
      println(io,"$i $j $(A[i,j])")
  end
end

@enum Op_Type Field=0 Vertical=1 Horizontal=2
struct Operator
    sites::Vector{Pair{Int,Int}}
    ops::Vector{ITensor}
    site_ind::Index
    dir::Op_Type
end

getDirectional(ops::Vector{Operator}, dir::Op_Type) = collect(filter(x->x.dir==dir, ops))

function spinI(s::Index; is_gpu::Bool=false)::ITensor
    I_data      = is_gpu ? CuArrays.zeros(Float64, ITensors.dim(s)*ITensors.dim(s)) : zeros(Float64, ITensors.dim(s), ITensors.dim(s))
    idi         = diagind(reshape(I_data, ITensors.dim(s), ITensors.dim(s)), 0)
    I_data[idi] = is_gpu ? CuArrays.ones(Float64, ITensors.dim(s)) : ones(Float64, ITensors.dim(s))
    I           = is_gpu ? cuITensor( I_data, IndexSet(s, s') ) : ITensor(vec(I_data), IndexSet(s, s'))
    return I
end

function combine(Aorig::ITensor, Anext::ITensor, tags::String)::ITensor
    ci        = commonindex(Aorig, Anext)
    cmb, cmbi = combiner(IndexSet(ci, prime(ci)), tags=tags)
    return cmb
end

function reconnect(combiner_ind::Index, environment::ITensor)
    environment_combiner        = findindex(environment, "Site")
    new_combiner, combined_ind  = combiner(IndexSet(combiner_ind, prime(combiner_ind)), tags="Site")
    combiner_transfer           = δ(combined_ind, environment_combiner)
    #return new_combiner*combiner_transfer
    replaceindex!(new_combiner, combined_ind, environment_combiner)
    return new_combiner
end

function buildN(A::PEPS, L::Environments, R::Environments, IEnvs, row::Int, col::Int, ϕ::ITensor)::ITensor
    Ny, Nx   = size(A)
    N        = spinI(findindex(A[row, col], "Site"); is_gpu=is_gpu(A))
    workingN = N
    if row > 1
        workingN *= IEnvs[:below][row - 1]
    end
    if row < Ny
        workingN *= IEnvs[:above][end - row]
    end
    if col > 1
        ci = commonindex(A[row, col], A[row, col-1])
        workingN *= multiply_side_ident(A[row, col], ci, copy(L.I[row]))
    end
    workingN *= ϕ
    if col < Nx
        ci = commonindex(A[row, col], A[row, col+1])
        workingN *= multiply_side_ident(A[row, col], ci, copy(R.I[row]))
    end
    return workingN
end

function multiply_side_ident(A::ITensor, ci::Index, side_I::ITensor)
    scmb        = findindex(side_I, "Site")
    acmb, acmbi = combiner(IndexSet(ci, ci'), tags="Site")
    replaceindex!(acmb, acmbi, scmb)
    return side_I * acmb
end

function nonWorkRow(A::PEPS, L::Environments, R::Environments, H::Operator, row::Int, col::Int)::ITensor
    Ny, Nx  = size(A)
    op_rows = H.sites
    is_cu   = is_gpu(A)
    ops     = deepcopy(H.ops)
    @inbounds for op_ind in 1:length(ops)
        as = findindex(A[op_rows[op_ind][1][1], col], "Site")
        ops[op_ind] = replaceindex!(ops[op_ind], H.site_ind, as)
        ops[op_ind] = replaceindex!(ops[op_ind], H.site_ind', as')
    end
    op = spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
    op_ind = findfirst( x -> x == row, op_rows)
    AA = A[row, col] * op * dag(A[row, col])'
    if col > 1
        ci = commonindex(A[row, col], A[row, col-1])
        msi = multiply_side_ident(A[row, col], ci, L.I[row])
        AA *= msi
    end
    if col < Nx
        ci = commonindex(A[row, col], A[row, col+1])
        msi = multiply_side_ident(A[row, col], ci, R.I[row])
        AA *= msi
    end
    return AA
end

function sum_rows_in_col(A::PEPS, L::Environments, R::Environments, H::Operator, row::Int, col::Int, low_row::Int, high_row::Int, above::Bool, IA::ITensor, IB::ITensor)::ITensor
    Ny, Nx  = size(A)
    op_rows = H.sites
    is_cu   = is_gpu(A)
    ops     = deepcopy(H.ops)
    start_row_ = row == op_rows[1][1] ? low_row + 1 : high_row
    stop_row_  = row == op_rows[1][1] ? high_row : low_row + 1
    start_row_ = min(start_row_, Ny)
    stop_row_  = min(stop_row_, Ny)
    step_row   = row == op_rows[1][1] ? 1 : -1;
    start_row, stop_row = minmax(start_row_, stop_row_)
    op_row_a = H.sites[1][1]
    op_row_b = H.sites[2][1]
    @inbounds for op_ind in 1:length(ops)
        this_A = A[op_rows[op_ind][1][1], col]
        as = findindex(this_A, "Site")
        ops[op_ind] = replaceindex!(ops[op_ind], H.site_ind, as)
        ops[op_ind] = replaceindex!(ops[op_ind], H.site_ind', as')
    end
    nwrs  = is_cu ? cuITensor(1.0) : ITensor(1.0)
    nwrs_ = is_cu ? cuITensor(1.0) : ITensor(1.0)
    Hterm = ITensor()
    if row == op_row_a
        Hterm = IA
        Hterm *= nonWorkRow(A, L, R, H, op_row_b, col)
        if col > 1
            ci  = commonindex(A[row, col], A[row, col-1])
            msi = multiply_side_ident(A[row, col], ci, copy(L.I[row]))
            Hterm *= msi
        end
        if col < Nx
            ci  = commonindex(A[row, col], A[row, col+1])
            msi = multiply_side_ident(A[row, col], ci, copy(R.I[row]))
            Hterm *= msi
        end
        Hterm *= IB
    else
        Hterm = IB
        Hterm *= nonWorkRow(A, L, R, H, op_row_a, col)
        if col > 1
            ci  = commonindex(A[row, col], A[row, col-1])
            msi = multiply_side_ident(A[row, col], ci, copy(L.I[row]))
            Hterm *= msi
        end
        if col < Nx
            ci  = commonindex(A[row, col], A[row, col+1])
            msi = multiply_side_ident(A[row, col], ci, copy(R.I[row]))
            Hterm *= msi
        end
        Hterm *= IA
    end
    op = spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
    op_ind = findfirst( x -> x[1] == row, op_rows)
    if op_ind > 0
        op = ops[op_ind]
    end
    Hterm = Hterm*op
    return Hterm
end

function buildHIedge( A::PEPS, E::Environments, row::Int, col::Int, side::Symbol, ϕ::ITensor )
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    HI     = is_cu ? cuITensor(1.0) : ITensor(1.0)
    IH     = is_cu ? cuITensor(1.0) : ITensor(1.0)
    next_col = side == :left ? 2 : Nx - 1
    @inbounds for work_row in 1:row-1
        AA = A[work_row, col] * dag(prime(A[work_row, col], "Link"))
        ci = commonindex(A[work_row, col], A[work_row, next_col])
        cmb = findindex(E.I[work_row], "Site")
        acmb, acmbi = combiner(IndexSet(ci, ci'), tags="Site")
        replaceindex!(acmb, acmbi, cmb)
        AA *= acmb
        HI *= AA * E.I[work_row]
        IH *= E.H[work_row] * AA
    end
    ci  = commonindex(A[row, col], A[row, next_col])
    cmb = findindex(E.I[row], "Site")
    acmb, acmbi = combiner(IndexSet(ci, ci'), tags="Site")
    replaceindex!(acmb, acmbi, cmb)
    op = spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
    op = is_cu ? cuITensor(op) : op
    HI *= ϕ
    IH *= ϕ
    HI *= op
    IH *= op
    HI *= E.I[row] * acmb
    IH *= E.H[row] * acmb
    @inbounds for work_row in row+1:Ny
        AA = A[work_row, col] * dag(prime(A[work_row, col], "Link"))
        ci = commonindex(A[work_row, col], A[work_row, next_col])
        cmb = findindex(E.I[work_row], "Site")
        acmb, acmbi = combiner(IndexSet(ci, ci'), tags="Site")
        replaceindex!(acmb, acmbi, cmb)
        AA *= acmb
        HI *= AA * E.I[work_row]
        IH *= E.H[work_row] * AA
    end
    AAinds = IndexSet(prime(ϕ))
    @assert hasinds(inds(IH), AAinds)
    @assert hasinds(AAinds, inds(IH))
    return (IH,)
end

function buildHIs(A::PEPS, L::Environments, R::Environments, row::Int, col::Int, ϕ::ITensor)
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    col == 1  && return buildHIedge(A, R, row, col, :left, ϕ)
    col == Nx && return buildHIedge(A, L, row, col, :right, ϕ)
    HLI_a = is_cu ? cuITensor(1.0) : ITensor(1.0)
    HLI_b = is_cu ? cuITensor(1.0) : ITensor(1.0)
    IHR_a = is_cu ? cuITensor(1.0) : ITensor(1.0)
    IHR_b = is_cu ? cuITensor(1.0) : ITensor(1.0)
    HLI   = is_cu ? cuITensor(1.0) : ITensor(1.0)
    IHR   = is_cu ? cuITensor(1.0) : ITensor(1.0)
    @inbounds for work_row in 1:row-1
        AA = A[work_row, col] * dag(prime(A[work_row, col], "Link"))
        lci = commonindex(A[work_row, col], A[work_row, col-1])
        rci = commonindex(A[work_row, col], A[work_row, col+1])
        lcmb = findindex(L.I[work_row], "Site")
        rcmb = findindex(R.I[work_row], "Site")
        lacmb, lacmbi = combiner(IndexSet(lci, lci'), tags="Site")
        racmb, racmbi = combiner(IndexSet(rci, rci'), tags="Site")
        replaceindex!(lacmb, lacmbi, lcmb)
        replaceindex!(racmb, racmbi, rcmb)
        AA *= lacmb
        AA *= racmb
        HLI_b *= L.H[work_row] * AA * R.I[work_row]
        IHR_b *= L.I[work_row] * AA * R.H[work_row]
    end
    lci = commonindex(A[row, col], A[row, col-1])
    rci = commonindex(A[row, col], A[row, col+1])
    lcmb = findindex(L.I[row], "Site")
    rcmb = findindex(R.I[row], "Site")
    lacmb, lacmbi = combiner(IndexSet(lci, lci'), tags="Site")
    racmb, racmbi = combiner(IndexSet(rci, rci'), tags="Site")
    replaceindex!(lacmb, lacmbi, lcmb)
    replaceindex!(racmb, racmbi, rcmb)
    HLI  *= L.H[row] * lacmb * R.I[row] * racmb
    HLI  *= ϕ
    IHR  *= L.I[row] * lacmb * R.H[row] * racmb
    IHR  *= ϕ
    @inbounds for work_row in reverse(row+1:Ny)
        AA = A[work_row, col] * dag(prime(A[work_row, col], "Link"))
        lci = commonindex(A[work_row, col], A[work_row, col-1])
        rci = commonindex(A[work_row, col], A[work_row, col+1])
        lcmb = findindex(L.I[work_row], "Site")
        rcmb = findindex(R.I[work_row], "Site")
        lacmb, lacmbi = combiner(IndexSet(lci, lci'), tags="Site")
        racmb, racmbi = combiner(IndexSet(rci, rci'), tags="Site")
        replaceindex!(lacmb, lacmbi, lcmb)
        replaceindex!(racmb, racmbi, rcmb)
        AA *= lacmb
        AA *= racmb
        HLI_a *= L.H[work_row] * AA * R.I[work_row]
        IHR_a *= L.I[work_row] * AA * R.H[work_row]
    end
    HLI *= HLI_a
    HLI *= HLI_b
    IHR *= IHR_a
    IHR *= IHR_b
    op   = spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
    HLI *= op
    IHR *= op
    AAinds = IndexSet(prime(ϕ))
    @assert hasinds(inds(IHR), AAinds)
    @assert hasinds(inds(HLI), AAinds)
    @assert hasinds(AAinds, inds(IHR))
    @assert hasinds(AAinds, inds(HLI))
    return (HLI, IHR)
end

function verticalTerms(A::PEPS, L::Environments, R::Environments, AI, AV, H, row::Int, col::Int, ϕ::ITensor)::Vector{ITensor}
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    vTerms = ITensor[]#fill(ITensor(), length(H))
    AAinds = IndexSet(prime(ϕ))
    dummy  = is_cu ? cuITensor(1.0) : ITensor(1.0)
    @inbounds for opcode in 1:length(H)
        thisVert = dummy
        op_row_a = H[opcode].sites[1][1]
        op_row_b = H[opcode].sites[2][1]
        if op_row_b < row || row < op_row_a
            local V, I
            if op_row_a > row
                V = AV[:above][opcode][end - row]
                I = row > 1 ? AI[:below][row - 1] : dummy
            elseif op_row_b < row
                V = AV[:below][opcode][row - op_row_b]
                I = row < Ny ? AI[:above][end - row] : dummy
            end
            thisVert *= V
            if col > 1
                ci = commonindex(A[row, col], A[row, col-1])
                msi = multiply_side_ident(A[row, col], ci, copy(L.I[row]))
                thisVert *= msi
            end
            thisVert *= ϕ
            if col < Nx
                ci = commonindex(A[row, col], A[row, col+1])
                msi = multiply_side_ident(A[row, col], ci, copy(R.I[row]))
                thisVert *= msi
            end
            thisVert *= I
            thisVert *= spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
        elseif row == op_row_a
            low_row  = op_row_a - 1
            high_row = op_row_b
            AIL = low_row > 0 ? AI[:below][low_row] : dummy
            AIH = high_row < Ny ? AI[:above][end - high_row] : dummy
            sA   = findindex(A[op_row_a, col], "Site")
            op_a = replaceindex!(copy(H[opcode].ops[1]), H[opcode].site_ind, sA)
            op_a = replaceindex!(op_a, H[opcode].site_ind', sA')
            sB   = findindex(A[op_row_b, col], "Site")
            op_b = replaceindex!(copy(H[opcode].ops[2]), H[opcode].site_ind, sB)
            op_b = replaceindex!(op_b, H[opcode].site_ind', sB')
            thisVert = AIH
            if col > 1
                ci  = commonindex(A[op_row_b, col], A[op_row_b, col-1])
                msi = multiply_side_ident(A[op_row_b, col], ci, copy(L.I[op_row_b]))
                thisVert *= msi
            end
            thisVert *= A[op_row_b, col] * op_b * dag(A[op_row_b, col])'
            if col < Nx
                ci  = commonindex(A[op_row_b, col], A[op_row_b, col+1])
                msi = multiply_side_ident(A[op_row_b, col], ci, copy(R.I[op_row_b]))
                thisVert *= msi
            end
            if col > 1
                ci  = commonindex(A[op_row_a, col], A[op_row_a, col-1])
                msi = multiply_side_ident(A[op_row_a, col], ci, copy(L.I[op_row_a]))
                thisVert *= msi
            end
            thisVert *= ϕ
            if col < Nx
                ci  = commonindex(A[op_row_a, col], A[op_row_a, col+1])
                msi = multiply_side_ident(A[op_row_a, col], ci, copy(R.I[op_row_a]))
                thisVert *= msi
            end
            thisVert *= AIL
            thisVert *= op_a
        elseif row == op_row_b
            low_row  = op_row_a - 1
            high_row = op_row_b
            AIL = low_row > 0 ? AI[:below][low_row] : dummy
            AIH = high_row < Ny ? AI[:above][end - high_row] : dummy
            thisVert = AIL
            if col > 1
                ci  = commonindex(A[op_row_a, col], A[op_row_a, col-1])
                msi = multiply_side_ident(A[op_row_a, col], ci, copy(L.I[op_row_a]))
                thisVert *= msi
            end
            if col < Nx
                ci  = commonindex(A[op_row_a, col], A[op_row_a, col+1])
                msi = multiply_side_ident(A[op_row_a, col], ci, copy(R.I[op_row_a]))
                thisVert *= msi
            end
            sA = findindex(A[op_row_a, col], "Site")
            op_a = replaceindex!(copy(H[opcode].ops[1]), H[opcode].site_ind, sA)
            op_a = replaceindex!(op_a, H[opcode].site_ind', sA')
            sB = findindex(A[op_row_b, col], "Site")
            op_b = replaceindex!(copy(H[opcode].ops[2]), H[opcode].site_ind, sB)
            op_b = replaceindex!(op_b, H[opcode].site_ind', sB')
            thisVert *= A[op_row_a, col] * op_a * dag(A[op_row_a, col])'
            if col > 1
                ci  = commonindex(A[op_row_b, col], A[op_row_b, col-1])
                msi = multiply_side_ident(A[op_row_b, col], ci, copy(L.I[op_row_b]))
                thisVert *= msi
            end
            thisVert *= ϕ
            if col < Nx
                ci  = commonindex(A[op_row_b, col], A[op_row_b, col+1])
                msi = multiply_side_ident(A[op_row_b, col], ci, copy(R.I[op_row_b]))
                thisVert *= msi
            end
            thisVert *= AIH
            thisVert *= op_b
        end
        @assert hasinds(inds(thisVert), AAinds) "inds of thisVert and AAinds differ!\n$(inds(thisVert))\n$AAinds\n"
        @assert hasinds(AAinds, inds(thisVert)) "inds of thisVert and AAinds differ!\n$(inds(thisVert))\n$AAinds\n"
        if hasinds(AAinds, inds(thisVert)) && hasinds(inds(thisVert), AAinds)
            push!(vTerms, thisVert)
        end
    end
    return vTerms
end

function fieldTerms(A::PEPS, L::Environments, R::Environments, AI, AF, H, row::Int, col::Int, ϕ::ITensor)::Vector{ITensor}
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    fTerms = Vector{ITensor}(undef, length(H))
    dummy  = is_cu ? cuITensor(1.0) : ITensor(1.0)
    AAinds = IndexSet(prime(ϕ))
    @inbounds for opcode in 1:length(H)
        thisField = dummy
        op_row    = H[opcode].sites[1][1]
        if op_row != row
            local F, I
            if op_row > row
                F = AF[:above][opcode][end - row]
                I = row > 1 ? AI[:below][row - 1] : dummy
            else
                F = AF[:below][opcode][row - 1]
                I = row < Ny ? AI[:above][end - row] : dummy
            end
            thisField *= F
            if col > 1
                ci = commonindex(A[row, col], A[row, col-1])
                thisField *= multiply_side_ident(A[row, col], ci, copy(L.I[row]))
            end
            thisField *= ϕ
            if col < Nx
                ci = commonindex(A[row, col], A[row, col+1])
                thisField *= multiply_side_ident(A[row, col], ci, copy(R.I[row]))
            end
            thisField *= I
            thisField *= spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
        else
            low_row  = op_row - 1
            high_row = op_row
            AIL = low_row > 0   ? AI[:below][low_row]        : dummy
            AIH = high_row < Ny ? AI[:above][end - high_row] : dummy
            thisField = AIL
            if col > 1
                ci  = commonindex(A[row, col], A[row, col-1])
                msi = multiply_side_ident(A[row, col], ci, copy(L.I[row]))
                thisField *= msi
            end
            thisField *= ϕ
            if col < Nx
                ci  = commonindex(A[row, col], A[row, col+1])
                msi = multiply_side_ident(A[row, col], ci, copy(R.I[row]))
                thisField *= msi
            end
            thisField *= AIH
            sA = findindex(A[row, col], "Site")
            op = copy(H[opcode].ops[1])
            op = replaceindex!(op, H[opcode].site_ind, sA)
            op = replaceindex!(op, H[opcode].site_ind', sA')
            thisField *= op
        end
        @assert hasinds(inds(thisField), AAinds)
        @assert hasinds(AAinds, inds(thisField))
        fTerms[opcode] = thisField;
    end
    return fTerms
end

function connectLeftTerms(A::PEPS, L::Environments, R::Environments, AI, AL, H, row::Int, col::Int, ϕ::ITensor)::Vector{ITensor}
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    lTerms = Vector{ITensor}(undef, length(H))
    AAinds = IndexSet(prime(ϕ))
    dummy  = is_cu ? cuITensor(1.0) : ITensor(1.0)
    @inbounds for opcode in 1:length(H)
        op_row_b = H[opcode].sites[2][1]
        op_b = copy(H[opcode].ops[2])
        as   = findindex(A[op_row_b, col], "Site")
        op_b = replaceindex!(op_b, H[opcode].site_ind, as)
        op_b = replaceindex!(op_b, H[opcode].site_ind', as')
        thisHori = is_cu ? cuITensor(1.0) : ITensor(1.0)
        if op_row_b != row
            local ancL, I
            if op_row_b > row
                ancL = AL[:above][opcode][end - row]
                I = row > 1 ? AL[:below][opcode][row - 1] : dummy
            else
                ancL = AL[:below][opcode][row - 1]
                I = row < Ny ? AL[:above][opcode][end - row] : dummy
            end
            thisHori = ancL
            thisHori *= L.InProgress[row, opcode]
            thisHori *= ϕ
            if col < Nx
                ci = commonindex(A[row, col], A[row, col+1])
                thisHori *= multiply_side_ident(A[row, col], ci, copy(R.I[row]))
            end
            thisHori *= I
            thisHori *= spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
        else
            low_row = (op_row_b <= row) ? op_row_b - 1 : row - 1;
            high_row = (op_row_b >= row) ? op_row_b + 1 : row + 1;
            if low_row >= 1
                thisHori *= AL[:below][opcode][low_row]
            end
            if high_row <= Ny
                thisHori *= AL[:above][opcode][end - high_row + 1]
            end
            if col < Nx
                ci = commonindex(A[row, col], A[row, col+1])
                thisHori *= multiply_side_ident(A[row, col], ci, R.I[row])
            end
            uih = uniqueinds(thisHori, L.InProgress[row, opcode])
            uil = uniqueinds(L.InProgress[row, opcode], thisHori)
            thisHori *= L.InProgress[row, opcode]
            thisHori *= ϕ
            thisHori *= op_b
        end
        @assert hasinds(inds(thisHori), AAinds)
        @assert hasinds(AAinds, inds(thisHori))
        lTerms[opcode] = thisHori
    end
    return lTerms
end

function connectRightTerms(A::PEPS, L::Environments, R::Environments, AI, AR, H, row::Int, col::Int, ϕ::ITensor)::Vector{ITensor}
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    rTerms = Vector{ITensor}(undef, length(H))
    AAinds = IndexSet(prime(ϕ))
    dummy  = is_cu ? cuITensor(1.0) : ITensor(1.0)
    @inbounds for opcode in 1:length(H)
        op_row_a = H[opcode].sites[1][1]
        op_a = copy(H[opcode].ops[1])
        as   = findindex(A[op_row_a, col], "Site")
        op_a = replaceindex!(op_a, H[opcode].site_ind, as)
        op_a = replaceindex!(op_a, H[opcode].site_ind', as')
        thisHori = is_cu ? cuITensor(1.0) : ITensor(1.0)
        if op_row_a != row
            local ancR, I
            if op_row_a > row
                ancR = AR[:above][opcode][end - row]
                I = row > 1 ? AR[:below][opcode][row - 1] : dummy
            else
                ancR = AR[:below][opcode][row - 1]
                I = row < Ny ? AR[:above][opcode][end - row] : dummy
            end
            thisHori = ancR
            thisHori *= R.InProgress[row, opcode]
            thisHori *= ϕ
            if col > 1
                ci = commonindex(A[row, col], A[row, col-1])
                thisHori *= multiply_side_ident(A[row, col], ci, copy(L.I[row]))
            end
            thisHori *= I
            thisHori *= spinI(findindex(A[row, col], "Site"); is_gpu=is_cu)
        else
            low_row = (op_row_a <= row) ? op_row_a - 1 : row - 1;
            high_row = (op_row_a >= row) ? op_row_a + 1 : row + 1;
            if low_row >= 1
                thisHori *= AR[:below][opcode][low_row]
            end
            if high_row <= Ny
                thisHori *= AR[:above][opcode][end - high_row + 1]
            end
            if col > 1
                ci = commonindex(A[row, col], A[row, col-1])
                thisHori *= multiply_side_ident(A[row, col], ci, L.I[row])
            end
            thisHori *= R.InProgress[row, opcode]
            thisHori *= ϕ
            thisHori *= op_a
        end
        @assert hasinds(inds(thisHori), AAinds)
        @assert hasinds(AAinds, inds(thisHori))
        rTerms[opcode] = thisHori
    end
    return rTerms
end

function buildLocalH(A::PEPS, L::Environments, R::Environments, AncEnvs, H, row::Int, col::Int, ϕ::ITensor; verbose::Bool=false)
    field_H_terms = getDirectional(vcat(H[:, col]...), Field)
    vert_H_terms  = getDirectional(vcat(H[:, col]...), Vertical)
    term_count    = 1 + length(field_H_terms) + length(vert_H_terms)
    Ny, Nx        = size(A)
    @timeit "build Ns" begin
        N   = buildN(A, L, R, AncEnvs[:I], row, col, ϕ)
    end
    den = scalar(collect(N * dag(ϕ)'))
    local left_H_terms, right_H_terms
    if col > 1
        left_H_terms = getDirectional(vcat(H[:, col - 1]...), Horizontal)
        term_count += length(left_H_terms)
    end
    if col < Nx
        right_H_terms = getDirectional(vcat(H[:, col]...), Horizontal)
        term_count += length(right_H_terms)
    end
    if 1 < col < Nx
        term_count += 1
    end
    Hs = Vector{ITensor}(undef, term_count)
    term_counter = 1
    @debug "\t\tBuilding I*H and H*I row $row col $col"
    @timeit "build HIs" begin
        HIs = buildHIs(A, L, R, row, col, ϕ)
        Hs[term_counter] = HIs[1]
        if length(HIs) == 2
            Hs[term_counter+1] = HIs[2]
        end
        term_counter += length(HIs)
        if verbose
            println("--- HI TERMS ---")
            for HI in HIs
                println(scalar(HI * dag(ϕ)'))
            end
        end
    end
    @debug "\t\tBuilding vertical H terms row $row col $col"
    @timeit "build vertical terms" begin
        vTs = verticalTerms(A, L, R, AncEnvs[:I], AncEnvs[:V], vert_H_terms, row, col, ϕ)
        Hs[term_counter:term_counter+length(vTs) - 1] = vTs
        term_counter += length(vTs)
        if verbose
            println( "--- vT TERMS ---")
            for vT in vTs
                println(scalar(vT * dag(ϕ)'))
            end
        end
    end
    @debug "\t\tBuilding field H terms row $row col $col"
    @timeit "build field terms" begin
        fTs = fieldTerms(A, L, R, AncEnvs[:I], AncEnvs[:F], field_H_terms, row, col, ϕ)
        Hs[term_counter:term_counter+length(fTs) - 1] = fTs[:]
        term_counter += length(fTs)
        if verbose
            println( "--- fT TERMS ---")
            for fT in fTs
                println(scalar(fT * dag(ϕ)'))
            end
        end
    end
    if col > 1
        @debug "\t\tBuilding left H terms row $row col $col"
        @timeit "build left terms" begin
            lTs = connectLeftTerms(A, L, R, AncEnvs[:I], AncEnvs[:L], left_H_terms, row, col, ϕ)
            Hs[term_counter:term_counter+length(lTs) - 1] = lTs[:]
            term_counter += length(lTs)
            if verbose
                println( "--- lT TERMS ---")
                for lT in lTs
                    println(scalar(lT * dag(ϕ)'))
                end
            end
        end
        @debug "\t\tBuilt left terms"
    end
    if col < Nx
        @debug "\t\tBuilding right H terms row $row col $col"
        @timeit "build right terms" begin
            rTs = connectRightTerms(A, L, R, AncEnvs[:I], AncEnvs[:R], right_H_terms, row, col, ϕ)
            Hs[term_counter:term_counter+length(rTs) - 1] = rTs[:]
            term_counter += length(rTs)
            if verbose
                println( "--- rT TERMS ---")
                for rT in rTs
                    println(scalar(rT * dag(ϕ)'))
                end
            end
        end
        @debug "\t\tBuilt right terms"
    end
    return Hs, N
end

function intraColumnGauge(A::PEPS, col::Int; kwargs...)::PEPS
    Ny, Nx = size(A)
    @inbounds for row in reverse(2:Ny)
        @debug "\tBeginning intraColumnGauge for col $col row $row"
        cmb_is   = IndexSet(findindex(A[row, col], "Site"))
        if col > 1
            cmb_is = IndexSet(cmb_is, commonindex(A[row, col], A[row, col - 1]))
        end
        if col < Nx
            cmb_is = IndexSet(cmb_is, commonindex(A[row, col], A[row, col + 1]))
        end
        cmb, ci = combiner(cmb_is, tags="CMB")
        Lis     = IndexSet(ci) #cmb_is
        if row < Ny
            Lis = IndexSet(Lis, commonindex(A[row, col], A[row + 1, col]))
        end
        Ac = A[row, col]*cmb
        U, S, V = svd(Ac, Lis; kwargs...)
        A[row, col]   = U*cmb
        A[row-1, col] *= (S*V)
    end
    return A
end

function simpleUpdate(A::PEPS, col::Int, next_col::Int, H; kwargs...)::PEPS
    do_side::Bool = get(kwargs, :do_side, true)
    τ::Float64    = get(kwargs, :tau, -0.1)
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    @inbounds for row in Iterators.reverse(1:Ny)
        if do_side
            hori_col   = next_col < col ? next_col : col
            nhori_col  = next_col < col ? col : next_col
            si_a       = findindex(A[row, col], "Site")
            si_b       = findindex(A[row, next_col], "Site")
            ci         = commonindex(A[row, col], A[row, next_col])
            min_dim    = ITensors.dim(ci)
            Ua, Sa, Va = svd(A[row, col], si_a, ci; mindim=min_dim, kwargs...)
            Ub, Sb, Vb = svd(A[row, next_col], si_b, ci; mindim=min_dim, kwargs...)
            Hab_hori   = is_cu ? cuITensor() : ITensor()
            horiH      = getDirectional(vcat(H[:, hori_col]...), Horizontal)
            horiH      = filter(x->x.sites[1][1] == row, horiH)
            for hH in horiH
                op_a = replaceindex!(copy(hH.ops[1]), hH.site_ind, hori_col == col ? si_a : si_b)
                op_a = replaceindex!(op_a, hH.site_ind', hori_col == col ? si_a' : si_b')
                op_b = replaceindex!(copy(hH.ops[2]), hH.site_ind, hori_col == col ? si_b : si_a)
                op_b = replaceindex!(op_b, hH.site_ind', hori_col == col ? si_b' : si_a')
                Hab_hori = ITensors.dim(Hab_hori) < 2 ? op_a * op_b : Hab_hori + op_a * op_b
            end
            cmb, ci   = combiner(findinds(Hab_hori, ("",0)), tags="hab,Site")
            Hab_hori *= cmb
            Hab_hori *= cmb'
            Hab_mat   = is_cu ? matrix(collect(Hab_hori)) : matrix(Hab_hori)
            expiH_mat = exp(τ*Hab_mat)
            expiH     = is_cu ? cuITensor(vec(expiH_mat), ci, ci') : ITensor(expiH_mat, ci, ci')
            expiH *= cmb
            expiH *= cmb'

            bond  = noprime(expiH * Ua * Ub)
            Uf, Sf, Vf = svd(bond, si_a, commonindex(Ua, Sa); vtags="r,Link,r$row,c$hori_col", mindim=min_dim, kwargs...)
            A[row, col] = Sa * Va * Uf * Sf
            A[row, next_col] = Sb * Vb * Vf
        end
        if row < Ny
            si_a       = findindex(A[row, col], "Site")
            si_b       = findindex(A[row+1, col], "Site")
            ci         = commonindex(A[row, col], A[row+1, col])
            min_dim    = ITensors.dim(ci)
            Ua, Sa, Va = svd(A[row, col], si_a, ci; mindim=min_dim, kwargs...)
            Ub, Sb, Vb = svd(A[row+1, col], si_b, ci; mindim=min_dim, kwargs...)
            Hab_vert   = is_cu ? cuITensor() : ITensor()
            vertH      = getDirectional(vcat(H[:, col]...), Vertical)
            vertH      = filter(x->x.sites[1][1] == row, vertH)
            for vH in vertH
                op_a = replaceindex!(copy(vH.ops[1]), vH.site_ind, si_a)
                op_a = replaceindex!(op_a, vH.site_ind', si_a')
                op_b = replaceindex!(copy(vH.ops[2]), vH.site_ind, si_b)
                op_b = replaceindex!(op_b, vH.site_ind', si_b')
                Hab_vert = ITensors.dim(Hab_vert) < 2 ? op_a * op_b : Hab_vert + op_a * op_b
            end
            cmb, ci   = combiner(findinds(Hab_vert, ("",0)), tags="hab,Site")
            Hab_vert *= cmb
            Hab_vert *= cmb'
            Hab_mat   = is_cu ? matrix(collect(Hab_vert)) : matrix(Hab_vert)
            expiH_mat = exp(τ*Hab_mat)
            expiH     = is_cu ? cuITensor(vec(expiH_mat), ci, ci') : ITensor(expiH_mat, ci, ci')
            expiH *= cmb
            expiH *= cmb'
            bond  = noprime(expiH * Ua * Ub)
            Uf, Sf, Vf = svd(bond, si_a, commonindex(Ua, Sa); vtags="u,Link,r$row,c$col", mindim=min_dim, kwargs...)
            A[row, col] = Sa * Va * Uf * Sf
            A[row+1, col] = Sb * Vb * Vf
        end
    end
    return A
end

function buildAncs(A::PEPS, L::Environments, R::Environments, H, col::Int)
    Ny, Nx = size(A)
    is_cu  = is_gpu(A)
    dummy  = is_cu ? cuITensor(1.0) : ITensor(1.0)
    @debug "\tMaking ancillary identity terms for col $col"
    Ia = makeAncillaryIs(A, L, R, col)
    #Ib = fill(dummy, Ny)
    Ib = Vector{ITensor}(undef, Ny)
    Is = (above=Ia, below=Ib)

    @debug "\tMaking ancillary vertical terms for col $col"
    vH = getDirectional(vcat(H[:, col]...), Vertical)
    Va = makeAncillaryVs(A, L, R, vH, col)
    Vb = [Vector{ITensor}() for ii in 1:length(Va)]
    Vs = (above=Va, below=Vb)

    @debug "\tMaking ancillary field terms for col $col"
    fH = getDirectional(vcat(H[:, col]...), Field)
    Fa = makeAncillaryFs(A, L, R, fH, col)
    Fb = [Vector{ITensor}() for ii in 1:length(Fa)]
    Fs = (above=Fa, below=Fb)

    Ls = (above=Vector{ITensor}(), below=Vector{ITensor}())
    Rs = (above=Vector{ITensor}(), below=Vector{ITensor}())
    if col > 1
        @debug "\tMaking ancillary left terms for col $col"
        lH = getDirectional(vcat(H[:, col-1]...), Horizontal)
        La = makeAncillarySide(A, L, R, lH, col, :left)
        Lb = [Vector{ITensor}() for ii in  1:length(La)]
        Ls = (above=La, below=Lb)
    end
    if col < Nx
        @debug "\tMaking ancillary right terms for col $col"
        rH = getDirectional(vcat(H[:, col]...), Horizontal)
        Ra = makeAncillarySide(A, R, L, rH, col, :right)
        Rb = [Vector{ITensor}() for ii in  1:length(Ra)]
        Rs = (above=Ra, below=Rb)
    end
    Ancs = (I=Is, V=Vs, F=Fs, L=Ls, R=Rs)
    return Ancs
end

function updateAncs(A::PEPS, L::Environments, R::Environments, AncEnvs, H, row::Int, col::Int)
    Ny, Nx = size(A)
    is_cu = is_gpu(A)

    Is, Vs, Fs, Ls, Rs = AncEnvs
    @debug "\tUpdating ancillary identity terms for col $col row $row"
    Ib = updateAncillaryIs(A, Is[:below], L, R, row, col)
    Is = (above=Is[:above], below=Ib)

    @debug "\tUpdating ancillary vertical terms for col $col row $row"
    vH = getDirectional(vcat(H[:, col]...), Vertical)
    Vb = updateAncillaryVs(A, Vs[:below], Ib, L, R, vH, row, col)
    Vs = (above=Vs[:above], below=Vb)

    @debug "\tUpdating ancillary field terms for col $col row $row"
    fH = getDirectional(vcat(H[:, col]...), Field)
    Fb = updateAncillaryFs(A, Fs[:below], Ib, L, R, fH, row, col)
    Fs = (above=Fs[:above], below=Fb)

    if col > 1
        @debug "\tUpdating ancillary left terms for col $col row $row"
        lH = getDirectional(vcat(H[:, col-1]...), Horizontal)
        Lb = updateAncillarySide(A, Ls[:below], Ib, L, R, lH, row, col, :left)
        Ls = (above=Ls[:above], below=Lb)
    end
    if col < Nx
        @debug "\tUpdating ancillary right terms for col $col row $row"
        rH = getDirectional(vcat(H[:, col]...), Horizontal)
        Rb = updateAncillarySide(A, Rs[:below], Ib, R, L, rH, row, col, :right)
        Rs = (above=Rs[:above], below=Rb)
    end
    Ancs = (I=Is, V=Vs, F=Fs, L=Ls, R=Rs)
    return Ancs
end

struct ITensorMap
  A::PEPS
  H::Matrix{Vector{Operator}}
  L::Environments
  R::Environments
  AncEnvs::NamedTuple
  row::Int
  col::Int
end
Base.eltype(M::ITensorMap)  = eltype(M.A)
Base.size(M::ITensorMap)    = ITensors.dim(M.A[M.row, M.col])
function (M::ITensorMap)(v::ITensor)
    Hs, N  = buildLocalH(M.A, M.L, M.R, M.AncEnvs, M.H, M.row, M.col, v)
    localH = sum(Hs)
    return noprime(localH)
end

function optimizeLocalH(A::PEPS, L::Environments, R::Environments, AncEnvs, H, row::Int, col::Int; kwargs...)
    Ny, Nx        = size(A)
    is_cu         = is_gpu(A)
    field_H_terms = getDirectional(vcat(H[:, col]...), Field)
    vert_H_terms  = getDirectional(vcat(H[:, col]...), Vertical)
    @debug "\tBuilding H for col $col row $row"
    @timeit "build H" begin
        Hs, N = buildLocalH(A, L, R, AncEnvs, H, row, col, A[row, col])
    end
    initial_N = real(scalar(collect(N * dag(A[row, col])')))
    @timeit "sum H terms" begin
        localH = sum(Hs)
    end
    initial_E = real(scalar(collect(deepcopy(localH) * dag(A[row, col])')))
    @info "Initial energy at row $row col $col : $(initial_E/(initial_N*Nx*Ny)) and norm : $initial_N"
    @debug "\tBeginning davidson for col $col row $row"
    mapper   = ITensorMap(A, H, L, R, AncEnvs, row, col)
    λ, new_A = davidson(mapper, A[row, col]; maxiter=1, kwargs...)
    new_E    = λ #real(scalar(collect(new_A * localH * dag(new_A)')))
    #new_A    = A[row,col]
    #new_E    = initial_E
    N        = buildN(A, L, R, AncEnvs[:I], row, col, new_A)
    new_N    = real(scalar(collect(N * dag(new_A)')))
    @info "Optimized energy at row $row col $col : $(new_E/(new_N*Nx*Ny)) and norm : $new_N"
    @timeit "restore intraColumnGauge" begin
        if row < Ny
            @debug "\tRestoring intraColumnGauge for col $col row $row"
            cmb_is   = IndexSet(findindex(A[row, col], "Site"))
            if col > 1
                cmb_is = IndexSet(cmb_is, commonindex(A[row, col], A[row, col - 1]))
            end
            if col < Nx
                cmb_is = IndexSet(cmb_is, commonindex(A[row, col], A[row, col + 1]))
            end
            cmb, ci = combiner(cmb_is, tags="CMB")
            Lis     = IndexSet(ci) #cmb_is
            if row > 1
                Lis = IndexSet(Lis, commonindex(A[row, col], A[row - 1, col]))
            end
            old_ci = commonindex(A[row, col], A[row+1, col])
            svdA = new_A*cmb
            Ris = uniqueinds(inds(svdA), Lis)
            U, S, V = svd(svdA, Ris; kwargs...)
            new_ci = commonindex(V, S)
            replaceindex!(V, new_ci, old_ci)
            A[row, col]    = V*cmb
            newU = S*U*A[row+1, col]
            replaceindex!(newU, new_ci, old_ci)
            A[row+1, col] = newU
            if row < Ny - 1
                nI    = spinI(findindex(A[row+1, col], "Site"); is_gpu=is_cu)
                newAA = A[row+1, col] * nI * dag(A[row+1, col])'
                if col > 1
                    ci     = commonindex(A[row+1, col], A[row+1, col-1])
                    newAA *= multiply_side_ident(A[row+1, col], ci, L.I[row+1])
                end
                if col < Nx
                    ci     = commonindex(A[row+1, col], A[row+1, col+1])
                    newAA *= multiply_side_ident(A[row+1, col], ci, R.I[row+1])
                end
                AncEnvs[:I][:above][end - row] = newAA * AncEnvs[:I][:above][end - row - 1]
            end
        else
            A[row, col] = initial_N > 0 ? new_A/√new_N : new_A
        end
    end
    return A, AncEnvs
end

function measureEnergy(A::PEPS, L::Environments, R::Environments, AncEnvs, H, row::Int, col::Int)::Tuple{Float64, Float64}
    Ny, Nx    = size(A)
    Hs, N     = buildLocalH(A, L, R, AncEnvs, H, row, col, A[row, col])
    initial_N = collect(N * dag(A[row, col])')
    localH    = sum(Hs)
    initial_E = collect(localH * dag(A[row, col])')
    return real(scalar(initial_N)), real(scalar(initial_E))/real(scalar(initial_N))
end

function sweepColumn(A::PEPS, L::Environments, R::Environments, H, col::Int; kwargs...)
    Ny, Nx = size(A)
    @debug "Beginning intraColumnGauge for col $col"
    A = intraColumnGauge(A, col; kwargs...)
    if col == div(Nx,2)
        L_s = buildLs(A, H; kwargs...)
        R_s = buildRs(A, H; kwargs...)
        EAncEnvs = buildAncs(A, L_s[col - 1], R_s[col + 1], H, col)
        N, E = measureEnergy(A, L_s[col - 1], R_s[col + 1], EAncEnvs, H, 1, col)
        println("Energy at MID: ", E/(Nx*Ny))
        println("Nx: ", Nx)
    end
    @debug "Beginning buildAncs for col $col"
    AncEnvs = buildAncs(A, L, R, H, col)
    @inbounds for row in 1:Ny
        if row > 1
            @debug "Beginning updateAncs for col $col"
            AncEnvs = updateAncs(A, L, R, AncEnvs, H, row-1, col)
        end
        @debug "Beginning optimizing H for col $col"
        A, AncEnvs = optimizeLocalH(A, L, R, AncEnvs, H, row, col; kwargs...)
    end
    return A
end

function rightwardSweep(A::PEPS, Ls::Vector{Environments}, Rs::Vector{Environments}, H; kwargs...)
    simple_update_cutoff = get(kwargs, :simple_update_cutoff, 4)
    Ny, Nx = size(A)
    dummyI = MPS(Ny, fill(ITensor(1.0), Ny), 0, Ny+1)
    dummyEnv = Environments(dummyI, dummyI, fill(ITensor(), 1, Ny))
    prev_cmb_r = Vector{ITensor}(undef, Ny)
    next_cmb_r = Vector{ITensor}(undef, Ny)
    sweep::Int = get(kwargs, :sweep, 0)
    sweep_width::Int = get(kwargs, :sweep_width, Nx)
    offset    = mod(Nx, 2)
    midpoint  = div(Nx, 2)
    rightmost = sweep_width == Nx ? Nx - 1 : midpoint + div(sweep_width, 2) + offset
    leftmost  = sweep_width == Nx ? 1 : midpoint - div(sweep_width, 2)
    if leftmost > 1
        for row in 1:Ny
            prev_cmb_r[row] = reconnect(commonindex(A[row, leftmost], A[row, leftmost-1]), Ls[leftmost-1].I[row])
        end
    end
    @inbounds for col in leftmost:rightmost
        L = col == 1 ? dummyEnv : Ls[col - 1]
        @debug "Sweeping col $col"
        if sweep >= simple_update_cutoff
            @timeit "sweep" begin
                A = sweepColumn(A, L, Rs[col+1], H, col; kwargs...)
            end
        end
        if sweep < simple_update_cutoff
            # Simple update...
            A = simpleUpdate(A, col, col+1, H; do_side=(col < Nx), kwargs...)
        end
        if sweep >= simple_update_cutoff
            # Gauge
            A = gaugeColumn(A, col, :right; kwargs...)
        end
        if col == 1
            left_H_terms = getDirectional(H[1], Horizontal)
            @timeit "left edge env" begin
                Ls[col] = buildEdgeEnvironment(A, H, left_H_terms, prev_cmb_r, :left, 1; kwargs...)
            end
        else
            @timeit "left next env" begin
                Ls[col] = buildNextEnvironment(A, Ls[col-1], H, prev_cmb_r, next_cmb_r, :left, col; kwargs...)
            end
            prev_cmb_r = deepcopy(next_cmb_r)
        end
    end
    return A, Ls, Rs
end

function leftwardSweep(A::PEPS, Ls::Vector{Environments}, Rs::Vector{Environments}, H; kwargs...)
    simple_update_cutoff = get(kwargs, :simple_update_cutoff, 4)
    Ny, Nx = size(A)
    dummyI = MPS(Ny, fill(ITensor(1.0), Ny), 0, Ny+1)
    dummyEnv = Environments(dummyI, dummyI, fill(ITensor(), 1, Ny))
    prev_cmb_l = Vector{ITensor}(undef, Ny)
    next_cmb_l = Vector{ITensor}(undef, Ny)
    sweep::Int = get(kwargs, :sweep, 0)
    sweep_width::Int = get(kwargs, :sweep_width, Nx)
    offset    = mod(Nx, 2)
    midpoint  = div(Nx, 2)
    rightmost = midpoint + div(sweep_width, 2) + offset
    leftmost  = sweep_width == Nx ? 2 : midpoint - div(sweep_width, 2)
    if rightmost < Nx
        for row in 1:Ny
            prev_cmb_l[row] = reconnect(commonindex(A[row, rightmost], A[row, rightmost+1]), Rs[rightmost+1].I[row])
        end
    end
    @inbounds for col in reverse(leftmost:rightmost)
        R = col == Nx ? dummyEnv : Rs[col + 1]
        @debug "Sweeping col $col"
        if sweep >= simple_update_cutoff
            @timeit "sweep" begin
                A = sweepColumn(A, Ls[col - 1], R, H, col; kwargs...)
            end
        end
        if sweep < simple_update_cutoff
            # Simple update...
            A = simpleUpdate(A, col, col-1, H; do_side=(col>1), kwargs...)
        end
        if sweep >= simple_update_cutoff
            # Gauge
            A = gaugeColumn(A, col, :left; kwargs...)
        end
        if col == Nx
            right_H_terms = getDirectional(H[Nx - 1], Horizontal)
            @timeit "right edge env" begin
                Rs[col] = buildEdgeEnvironment(A, H, right_H_terms, prev_cmb_l, :right, Nx; kwargs...)
            end
        else
            @timeit "right next env" begin
                Rs[col] = buildNextEnvironment(A, Rs[col+1], H, prev_cmb_l, next_cmb_l, :right, col; kwargs...)
            end
            prev_cmb_l = deepcopy(next_cmb_l)
        end
    end
    return A, Ls, Rs
end

function doSweeps(A::PEPS, Ls::Vector{Environments}, Rs::Vector{Environments}, H; mindim::Int=1, maxdim::Int=1, simple_update_cutoff::Int=4, sweep_start::Int=1, sweep_count::Int=10, cutoff::Float64=0., env_maxdim=2maxdim, do_mag::Bool=false, prefix="$(Nx)_$(maxdim)_mag", max_gauge_iter::Int=10)
    for sweep in sweep_start:sweep_count
        if iseven(sweep)
            (A, Ls, Rs), this_time, bytes, gctime, memallocs = @timed rightwardSweep(A, Ls, Rs, H; sweep=sweep, mindim=mindim, maxdim=maxdim, simple_update_cutoff=simple_update_cutoff, overlap_cutoff=0.999, cutoff=cutoff, env_maxdim=env_maxdim, max_gauge_iter=max_gauge_iter)
            println("SWEEP RIGHT $sweep, time $this_time")
        else
            (A, Ls, Rs), this_time, bytes, gctime, memallocs = @timed leftwardSweep(A, Ls, Rs, H; sweep=sweep, mindim=mindim, maxdim=maxdim, simple_update_cutoff=simple_update_cutoff, overlap_cutoff=0.999, cutoff=cutoff, env_maxdim=env_maxdim, max_gauge_iter=max_gauge_iter)
            println("SWEEP LEFT $sweep, time $this_time")
        end
        flush(stdout)
        if sweep == simple_update_cutoff - 1
            for col in reverse(2:Nx)
                A = gaugeColumn(A, col, :left; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim, max_gauge_iter=max_gauge_iter)
            end
            Ls = buildLs(A, H; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim)
            Rs = buildRs(A, H; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim)
        end
        if do_mag
            A_ = copy(A)
            x_mag = zeros(Ny, Nx)
            z_mag = zeros(Ny, Nx)
            v_mag = zeros(Ny, Nx)
            h_mag = zeros(Ny, Nx)
            prev_cmb = Vector{ITensor}(undef, Ny)
            next_cmb = Vector{ITensor}(undef, Ny)
            if iseven(sweep)
                L_s = buildLs(A_, H; mindim=mindim, maxdim=maxdim, env_maxdim=env_maxdim)
                R_s = buildRs(A_, H; mindim=mindim, maxdim=maxdim, env_maxdim=env_maxdim)
                for col in reverse(1:Nx)
                    A_  = intraColumnGauge(A_, col; mindim=mindim, maxdim=maxdim)
                    x_mag[:, col] = measureXmag(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    z_mag[:, col] = measureZmag(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    v_mag[:, col] = measureSmagVertical(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    if col < Nx
                        h_mag[:, col] = measureSmagHorizontal(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    end
                    if col > 1
                        A_  = gaugeColumn(A_, col, :left; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim, overlap_cutoff=0.999, max_gauge_iter=max_gauge_iter)
                        if col < Nx
                            R_s[col] = buildNextEnvironment(A_, R_s[col+1], H, prev_cmb, next_cmb, :right, col; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim)
                            prev_cmb = deepcopy(next_cmb)
                        else
                            right_H_terms = getDirectional(H[Nx-1], Horizontal)
                            R_s[col] = buildEdgeEnvironment(A_, H, right_H_terms, prev_cmb, :right, col; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim)
                        end
                    end
                end
            else
                L_s = buildLs(A_, H; mindim=mindim, maxdim=maxdim, env_maxdim=env_maxdim)
                R_s = buildRs(A_, H; mindim=mindim, maxdim=maxdim, env_maxdim=env_maxdim)
                for col in 1:Nx
                    A_  = intraColumnGauge(A_, col; mindim=mindim, maxdim=maxdim)
                    x_mag[:, col] = measureXmag(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    z_mag[:, col] = measureZmag(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    v_mag[:, col] = measureSmagVertical(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    if col < Nx
                        h_mag[:, col] = measureSmagHorizontal(A_, L_s, R_s, col; mindim=mindim, maxdim=maxdim)
                    end
                    if col < Nx
                        A_  = gaugeColumn(A_, col, :right; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim, overlap_cutoff=0.999, max_gauge_iter=max_gauge_iter)
                        if col > 1
                            L_s[col] = buildNextEnvironment(A_, L_s[col-1], H, prev_cmb, next_cmb, :left, col; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim)
                            prev_cmb = deepcopy(next_cmb)
                        else
                            left_H_terms = getDirectional(H[1], Horizontal)
                            L_s[col] = buildEdgeEnvironment(A_, H, left_H_terms, prev_cmb, :left, col; mindim=maxdim, maxdim=maxdim, cutoff=cutoff, env_maxdim=env_maxdim)
                        end
                    end
                end
            end
            writedlm(prefix*"_$(sweep)_x", x_mag)
            writedlm(prefix*"_$(sweep)_z", z_mag)
            writedlm(prefix*"_$(sweep)_v", v_mag)
            writedlm(prefix*"_$(sweep)_h", h_mag)
        end
    end
    return A
end


text:
"""
    Dx,Dk = dfft(x,k)

Measures that make `fft`, `ifft` 2-norm preserving.
Correct measures for mapping between `x`- and `k`-space.
"""
function dfft(x,k)
    dx = x[2]-x[1]; dk = k[2]-k[1]
    Dx = dx/sqrt(2*pi)
    Dk = length(k)*dk/sqrt(2*pi)
    return Dx, Dk
end

"""
    DX,DK = dfftall(X,K)

Evalutes tuple of measures that make `fft`, `ifft` 2-norm preserving for each
`x` or `k` dimension.
"""
function dfftall(X,K)
    M = length(X)
    DX = zeros(M); DK = zeros(M)
    for i ∈ eachindex(X)
        DX[i],DK[i] = dfft(X[i],K[i])
    end
    return DX,DK
end

"""
    ψ = xspace(ϕ,sim)

Transform from `k`- to `x`-space using transforms packed into `sim`.
"""
function xspace(ϕ,sim)
    @unpack T = sim
    return T.Tkx*ϕ
end

"""
    xspace!(ϕ,sim)

Mutating transform from `k`- to `x`-space using transforms packed into `sim`.
"""
function xspace!(ψ,sim)
    @unpack T = sim
    T.Tkx!*ψ
    return nothing
end

"""
    kspace(ψ,sim)

Transform from `x`- to `k`-space using transforms packed into `sim`.
"""
function kspace(ψ,sim)
    @unpack T = sim
    return T.Txk*ψ
end

"""
    kspace!(ψ,sim)

Mutating transform from `x`- to `k`-space using transforms packed into `sim`.
"""
function kspace!(ψ,sim)
    @unpack T = sim
    T.Txk!*ψ
    return nothing
end

"""
    definetransforms(funcs,args,meas,kwargs)

Build all transforms for simulation.
"""
function definetransforms(funcs,args,meas,kwargs)
    trans = []
    for (fun,arg) in zip(funcs,args)
        push!(trans, fun(arg...,flags=kwargs))
    end
    return meas.*trans
end

"""
    T = makeT(X,K,j)

Build `FFTW` transform library for the array tuples `X`, `K`. Defaults to a measure plan.
`j` is number of scratch fields to initialize for in-place evaluation.
"""
function makeT(X,K,j=1;flags=FFTW.MEASURE)
    FFTW.set_num_threads(Sys.CPU_THREADS)
    dim = length(X)
    N = length.(X)
    DX,DK = dfftall(X,K)
    dμx = prod(DX)
    dμk = prod(DK)
    ψtest = ones(N...) |> complex

    trans = (plan_fft,plan_fft!,plan_ifft,plan_ifft!)
    meas = (dμx,dμx,dμk,dμk)
    # flags = FFTW.MEASURE
    args = ((ψtest,),(ψtest,),(ψtest,),(ψtest,))

    # flags != FFTW.MEASURE && @info "Planning FFTs ..."
    Txk,Txk!,Tkx,Tkx! = definetransforms(trans,args,meas,flags)
    # Mxk,Mxk!,Mkx,Mkx! = makeTMixed(X,K,flags=flags)
    # @info "...Plans created."
    return Transforms{dim,j}(Txk,Txk!,Tkx,Tkx!,crandnpartition(dim,j))
end

"""
    T = makeTMixed(X,K,j)

Build mixed transform library for the array tuples `X`, `K`. Defaults to a measure plan.
"""
function makeTMixed(X,K;flags=FFTW.MEASURE)
    FFTW.set_num_threads(Sys.CPU_THREADS)
    N = length.(X)
    dim = length(X)
    DX,DK = dfftall(X,K)
    ψtest = ones(N...) |> complex

args = []
transxk = []
transxk! = []
transkx = []
transkx! = []
measxk = []
measkx = []

for i = 1:dim
    push!(args, (ψtest,i))
    push!(transxk, plan_fft)
    push!(transxk!, plan_fft!)
    push!(transkx, plan_ifft)
    push!(transkx!, plan_ifft!)
    push!(measxk, DX[i])
    push!(measkx, DK[i])
end

#TODO fails in MKL
Mxk = definetransforms(transxk,args,measxk,flags)
Mxk! = definetransforms(transxk!,args,measxk,flags)
Mkx = definetransforms(transkx,args,measkx,flags)
Mkx! = definetransforms(transkx!,args,measkx,flags)
return Mxk,Mxk!,Mkx,Mkx!
end

"""
    maketransarrays(L,N,j=1;flags=FFTW.MEASURE)

Make all transforms and arrays in one call.
"""
function maketransarrays(L,N,j=1;flags=FFTW.MEASURE)
    X,K,dX,dK = makearrays(L,N)
    T = makeT(X,K,j,flags=flags)
    return X,K,dX,dK,DX,DK,T
end

ext: jl


text:
using Plots
gr()

x = 1:10; y = rand(10) # These are the plotting data
plot(x,y)

annotate!([(0,0,  text("this is #5\r\n asdf",16,:red,:center) ) ])

gui()


text:
using Example
using PackageWithBuildSpecificTestDeps
using Test

@test PackageWithBuildSpecificTestDeps.f(3) == 3


text:
# Copyright Peter Luschny. License is MIT.

module PrimeSwingFactorial
using PrimesIterator
export PSfactorial, Swing

"""
Return the accumulated product of an array.
"""
function ∏(A)
    function prod(a, b)
        n = b - a
        if n < 24
            p = BigInt(1)
            for k in a:b
                p *= A[k]
            end
            return BigInt(p)
        end
        m = div(a + b, 2)
        prod(a, m) * prod(m + 1, b)
    end
    A == [] && return 1
    prod(1, length(A))
end

const SwingOddpart = [1,1,1,3,3,15,5,35,35, 315, 63, 693, 231, 3003, 429, 6435,
   6435,109395,12155,230945,46189,969969,88179,2028117,676039,16900975,1300075,
   35102025,5014575,145422675,9694845,300540195,300540195]

"""
Computes the odd part of the swinging factorial ``n≀``. Cf. A163590.
"""
function swing_oddpart(n::Int)
    n < 33 && return ZZ(SwingOddpart[n+1])

    sqrtn = isqrt(n)
    factors = Primes(div(n,2) + 1, n)
    r = Primes(sqrtn + 1, div(n, 3))
    s = filter(x -> isodd(div(n, x)), r)

    for prime in Primes(3, sqrtn)
        p, q = 1, n
        while true
            q = div(q, prime)
            q == 0 && break
            isodd(q) && (p *= prime)
        end
        p > 1 && push!(s, p)
    end

    return ∏(factors)*∏(s)
end

"""
Computes the swinging factorial (a.k.a. Swing numbers n≀). Cf. A056040.
"""
function Swing(n::Int)
    sh = count_ones(div(n, 2))
    swing_oddpart(n) << sh
end

const FactorialOddPart = [1, 1, 1, 3, 3, 15, 45, 315, 315, 2835, 14175, 155925,
    467775, 6081075, 42567525, 638512875, 638512875, 10854718875, 97692469875,
    1856156927625, 9280784638125, 194896477400625, 2143861251406875,
    49308808782358125, 147926426347074375, 3698160658676859375]

"""
Return the largest odd divisor of ``n!``. Cf. A049606.
"""
function factorial_oddpart(n::Int)
    n < length(FactorialOddPart) && return BigInt(FactorialOddPart[n+1])
    swing_oddpart(n)*(factorial_oddpart(div(n,2))^2)
end

"""
Return the factorial ``n! = 1*2*...*n``, which is the order of the
symmetric group S_n or the number of permutations of n letters. Cf. A000142.
"""
function PSfactorial(n::Int)
    n < 0 && ArgumentError("Argument must be >= 0")
    sh = n - count_ones(n)
    factorial_oddpart(n) << sh
end

#START-TEST-########################################################
using Test

function main()
    @testset "PrimeSwing" begin
        for n in 0:999
            S = PSfactorial(n)
            B = Base.factorial(BigInt(n))
            @test S == B
        end
    end
    @time PSfactorial(1000000)
end

main()

end # module

max_stars_repo_name: NaldoDj/Fast-Factorial-Functions

max_issues_repo_name: NaldoDj/Fast-Factorial-Functions

max_forks_repo_name: NaldoDj/Fast-Factorial-Functions


text:
struct WithCache{T}
  data::T
  cache::UInt

end

function with_cache(x::T) where {T}
  h = hash(x)
  return WithCache{T}(x, h)
end

Base.hash(x::WithCache{T}) where {T} = x.cache

struct FacElemWrap{T, S}
  fac::Dict{WithCache{T}, fmpz}
  hash::UInt
  #parent::FacElemMon{S}
end

@inline function Base.iterate(x::FacElemWrap)
  z, st = iterate(x.fac)
  ((z[1]).data => z[2]), st
end

@inline function Base.iterate(x::FacElemWrap, st)
  v = iterate(x.fac, st)
  if v === nothing
    return nothing
  else
    z, st = v
    return ((z[1]).data => z[2]), st
  end
end

function Base.:(*)(x::FacElemWrap, y::FacElemWrap)
  z = copy(x.fac)
  z = copy(x.fac)
  for (a, v) in y.fac
    add_to_key!(z, a, v, remove_zero = true)
  end
  return FacElemWrap(z)
end

function FacElemWrap(x::Dict{T, fmpz}) where {T}
  z = Dict{WithCache{T}, fmpz}(with_cache(k) => p for (k, p) in x)
  return FacElemWrap{T, Nothing}(z, UInt(0))
end


text:
# TIPS Yield Curve and Inflation Compensation
# https://www.federalreserve.gov/data/tips-yield-curve-and-inflation-compensation.htm
#
# READ
# 1. File is compressed with gzip
# 2. Header information at the top of the csv file
# 3. Missing data represetned as "NA"
#
# CLEAN
# 1. Some dates are completely missing data
# 2. Two model parameters have no value until a certain date
# 3. Using heatmap to see missing data pattern
#
# ANALYSIS
# 1. Correlation analysis
# 2. Time series analysis
#
# VISUALIZATION
# 1. Plot time series line chart for a single variable (BETA3)
# 2. Overlay plot with another line chart for a second variable
# 3. Histogram for a single time series variable
# 4. Margin histogram for two variables
# 5. Heatmap for a correlation matrix
# 6. Heatmap for finding missing data

using CSV, DataFrames, Pipe, CodecZlib, Mmap

# Decompress into a byte array first
bytes = transcode(GzipDecompressor, Mmap.mmap("data/feds200805.csv.gz"))

# this doesn't work
df = DataFrame(CSV.File(bytes))

# Let's examine the decompressed text
str = String(bytes);
split(str, "\n")[1:5]
split(str, "\n")[15:20]

# Skip the header comments
bytes = transcode(GzipDecompressor, Mmap.mmap("data/feds200805.csv.gz"))
df = DataFrame(CSV.File(bytes; header = 19))

# Why do all columns have String types?
describe(df)
show(describe(df), allrows = true)

# Let's read the file again and parse those NA's correctly
bytes = transcode(GzipDecompressor, Mmap.mmap("data/feds200805.csv.gz"))
df = DataFrame(CSV.File(bytes; header = 19, missingstring = "NA"))

# How much data is missing?
describe(df)

# Let's check all rows where BETA0 is missing
filter(:BETA0 => ismissing, df)

# How many mising values in those rows?
df240 = filter(:BETA0 => ismissing, df)
count(ismissing, Matrix(df240[:, 2:127]))
126 * 240

# Let's drop those 240 rows.
df
dropmissing!(df, :BETA0)
# filter!(:BETA0 => !ismissing, df)

# See what else is missing
meta = describe(df)
meta[meta.nmissing .> 0, :]  # see https://github.com/JuliaData/DataFrames.jl/pull/2360

# Avoid the `nothing` problem
meta[something.(meta.nmissing, 0) .> 0, :]

sort(meta[something.(meta.nmissing, 0) .> 0, :], :nmissing; rev = true)

# BETA3 and TAU2 have more missing values
select(df, :BETA3, :TAU2)

# Confirm by plotting BETA3 over time
using Plots

plot(df.Date, df.BETA3;
    label = "BETA3", legend = :bottomleft, xrotation = 45,
    title = "BETA Observations", size = (350,250))

plot!(df.Date, df.BETA2; label = "BETA2")

# More analysis

# Check data distribution using histogram
histogram(df.TIPSPY02; bins = 50, size = (300,300), label = "TIPSPY02")

# StatsPlots has a recipe to plot the correlated distribution
using StatsPlots
@df df marginalhist(:TIPSPY02, :TIPSPY20; bins = 50, size = (300,300))

# Any correlation among the TIPS measures?
using Statistics

displaytable(stdout, names(df); index = true)

let cols = 90:108   # TIPSPY02 to TIPSPY20
    m = Matrix(df[:, cols])
    variable = names(df)[cols]
    vcv = cor(m)
    heatmap(variable, variable, vcv;
        xticks = :all,
        yticks = :all,
        title = "Correlation Matrix (TIPSPY02 - TIPSPY20)",
        size = (600, 600),
        xrotation = 45
    )
end

# Heatmap is a great way to find patterns
# let's try to find the missing values again from beginning
df = DataFrame(CSV.File(bytes; header = 19, missingstring = "NA"))

heatmap(df.Date, names(df)[2:end], Matrix(df[:, 2:end])';
    yticks = :all, size=(700,1500))

# Conclusion about BETA3 and TAU2.
# There must be a model change on that date and BETA3 and TAU3 got introduced.


text:
module Jax

using PythonCall: Py, pynew, pyimport, pycopy!
using ChainRulesCore
using DLPack
using Functors: fmap
using Adapt

using ..PythonCallChainRules: PyAdaptor

const inspect = pynew()
const jax = pynew()
const dlpack = pynew()
const stax = pynew()
const numpy = pynew()

const ispysetup = Ref{Bool}(false)

struct JaxFunctionWrapper
    jaxfn::Py
end

function (wrap::JaxFunctionWrapper)(args...; kwargs...)
    # TODO: handle multiple outputs
    out = (wrap.jaxfn(fmap(x->DLPack.share(x, Py, dlpack.from_dlpack), args)...))
    return DLPack.wrap(out, dlpack.to_dlpack)
end

function ChainRulesCore.rrule(wrap::JaxFunctionWrapper, args...; kwargs...)
    T = typeof(first(args))
    project = ProjectTo(args)
    jax_primal, jax_vjpfun = jax.vjp(wrap.jaxfn, fmap(x->DLPack.share(x, Py, dlpack.from_dlpack), args)...; kwargs...)
    function JaxFunctionWrapper_pullback(Δ)
        cΔ = Adapt.adapt(PyAdaptor{T}, Δ)
        dlΔ = DLPack.share(cΔ, Py, dlpack.from_dlpack)
        tangent_vals = fmap(x->DLPack.wrap(x, dlpack.to_dlpack), jax_vjpfun(dlΔ))
        return (NoTangent(), project(tangent_vals)...)
    end
    return (DLPack.wrap(jax_primal, dlpack.to_dlpack)), JaxFunctionWrapper_pullback
end

function __init__()
    try
        pycopy!(jax, pyimport("jax"))
        pycopy!(numpy, pyimport("numpy"))
        pycopy!(stax, pyimport("jax.example_libraries.stax"))
        pycopy!(inspect, pyimport("inspect"))
        ispysetup[] = true
    catch err
        @warn """PythonCallChainRules.jl has failed to import jax from Python.
                 Please make sure these are installed.
                 methods of this package.
        """
        @debug err
        ispysetup[] = false
    end
end

end


text:
function get_additional_state(obj::FieldAdditionalStateProvider, arg0::FieldSpacecraftState)
    return jcall(obj, "getAdditionalState", Vector{RealFieldElement}, (FieldSpacecraftState,), arg0)
end

function get_name(obj::FieldAdditionalStateProvider)
    return jcall(obj, "getName", JString, ())
end


text:
@testset "data" begin
    datasets = Impute.datasets()

    @testset "Impute.dataset($name)" for name in datasets
        result = Impute.dataset(name)
        if occursin("matrix", name)
            @test isa(result, AbstractDict)
        elseif occursin("table", name)
            @test isa(result, CSV.File)
        end
    end
end


text:
using Test

include("../../src/utils/misc.jl")

function test_getkwargs(arg1, arg2; kwarg1, kwarg2)
end

@testset "utils/misc" begin
    @testset "getkwargs" begin
        @test getkwargs(test_getkwargs) == [:kwarg1, :kwarg2]
    end
end

max_stars_repo_name: charlyalizadeh/OPFdecgen

max_issues_repo_name: charlyalizadeh/OPFdecgen

max_forks_repo_name: charlyalizadeh/OPFdecgen


text:
"""
Basic messages in the Paxos protocol.
"""
module Protocols

using UUIDs

export Cluster,
  prepare,
  propose,
  accept,
  onPrepare,
  onPropose,
  onAccept,
  request,
  onRequest,
  respond,
  Result,
  requestAccepted,
  requestFailed,
  requestObsolete,
  noQuorum

# using .Transports
using ..Ballots
using ..Configurations
using ..Ledgers
using ..Nodes
using ..Transports.Common
using ..Utils

abstract type AbstractCluster end
mutable struct Cluster <: AbstractCluster
  timeout::Int
  configuration::Configuration
  connections::Connections
end

"""
Perform a `gcall` against all members of the cluster, creating connections
as needed (and as available), if not already present
"""
function pcall(cluster::AbstractCluster, msg)
  connections =
    connectAll(cluster.connections, memberAddresses(cluster.configuration), cluster.timeout)
  responses = gcall(values(connections), cluster.timeout, msg)
  if !isquorum(cluster.configuration, responders(responses))
    throw(NoQuorumAvailableException())
  end
  responses
end

# Messages

struct PrepareMessage <: Message
  ballotNumbers::Vector{BallotNumber}
end

# Prepare response
struct VoteMessage <: Message
  memberID::NodeID
  ballots::Vector{Ballot}
end

struct ProposeMessage <: Message
  ballots::Vector{Ballot}
end

# Propose response
struct PromiseMessage <: Message
  memberID::NodeID
  ballotNumbers::Vector{BallotNumber}
end

struct AcceptMessage <: Message
  ballotNumbers::Vector{BallotNumber}
end

# Not required -- accept response
struct AcceptedMessage <: Message
  ballotNumbers::Vector{BallotNumber}
end

struct RequestMessage <: Message
  request::Request
end

@enum Result begin
  # """
  # The request was accepted by a quorum
  # """
  requestAccepted
  # """
  # The request was not accepted, but was instead made
  # obsolete by choice of a different request by the quorum.
  # Obsolete requests can be considered for re-execution
  # """
  requestObsolete
  # """
  # An unexpected error occurr during consensus; applications
  # should re-inspect state before retrying the same request
  # """
  requestFailed # only if there is an exception in the protocol
  # """
  # No quorum available to agree on a request
  # """
  noQuorum
end
struct ResponseMessage <: Message
  requestID::RequestID
  instanceID::InstanceID
  result::Result
end

struct TimeoutMessage <: Message end

PaxosMessage = Union{
  PrepareMessage,
  VoteMessage,
  ProposeMessage,
  PromiseMessage,
  AcceptMessage,
  AcceptedMessage,
  RequestMessage,
  ResponseMessage,
  TimeoutMessage,
}

function responders(messages)
  Set(map(messages) do message
    message.memberID
  end)
end

struct Connection
  sender::Channel{PaxosMessage}
  receiver::Channel{PaxosMessage}
end

# Leader protocol

function prepare(cluster::AbstractCluster, ballotNumbers::Vector{BallotNumber})::Vector{Ballot}
  if isempty(ballotNumbers)
    Vector{Ballot}()
  else
    msg = PrepareMessage(ballotNumbers)
    votes = pcall(cluster, msg)
    hcat(map(votes) do vote::VoteMessage
      vote.ballots
    end...)
  end
end

function propose(cluster::AbstractCluster, ballots::Vector{Ballot})::Vector{BallotNumber}()
  if isempty(ballots)
    Vector{Ballot}()
  else
    msg = ProposeMessage(ballots)
    memberPromises = pcall(cluster, msg)
    promises = Vector{Promise}()
    for memberPromise in memberPromises
      memberID = memberPromise.memberID
      for ballotNumber in memberPromise.ballotNumbers
        push!(promises, Promise(memberID, ballotNumber))
      end
    end
  end
  collect(values(promises))
end

function accept(cluster::AbstractCluster, ballotNumbers::Vector{BallotNumber})::Vector{BallotNumber}
  if isempty(ballotNumbers)
    Vector{BallotNumber}()
  else
    msg = AcceptMessage(ballotNumbers)
    acceptances = pcall(cluster, msg)
    hcat(map(acceptances) do acceptance::Vector{AcceptedMessage}
      acceptance.ballotNumbers
    end...)
  end
end

# Follower protocol

function onPrepare(ledger::Ledger, ballotNumbers::Vector{BallotNumber})
  votes!(ledger, ballotNumbers)
end

function onPropose(ledger::Ledger, ballots::Vector{Ballot})
  promises!(ledger, ballots)
end

function onAccept(ledger::Ledger, ballotNumbers::Vector{BallotNumber})
  accepted!(ledger, ballotNumbers)
end

# Client protocol

function request(
  transport::Transport,
  messenger::Union{Messenger,Nothing},
  addresses,
  timeout,
  request::Request,
)
  return callAny(transport, messenger, addresses, timeout, RequestMessage(request))
end

# Server protocol

function onRequest(cluster::AbstractCluster, command::Command) end

function respond(
  client::Messenger,
  requestID::RequestID,
  instanceID::InstanceID,
  result::Result,
)
  msg = ResponseMessage(requestID, instanceID, result)
  sendMessage(client, msg)
end

end


text:
using NLOptControl
using VehicleModels
using DataFrames
using MAVs
using KNITRO

# TODO
#1) could make a check if the vehicle passes  the goal
# 5) detect a crash
# 7) looking at the second iteration it seems like vehicle position did not get updated, but the vehicle did turn
c=defineCase(;(:mode=>:autoGazebo));

n=initializeAutonomousControl(c);
driveStraight!(n)
for ii=1:n.mpc.max_iter
    println("Running model for the: ",n.r.eval_num," time");
    updateAutoParams!(n,c);                      # update model parameters
    status=autonomousControl!(n);                # rerun optimization
    n.mpc.t0_actual=(n.r.eval_num-1)*n.mpc.tex;  # external so that it can be updated easily in PathFollowing
    simPlant!(n) #TODO make sure we are passing the correct control signals
    if n.r.status==:Optimal || n.r.status==:Suboptimal || n.r.status==:UserLimit
      println("Passing Optimized Signals to 3DOF Vehicle Model");
    elseif n.r.status==:Infeasible
      println("\n FINISH:Pass PREVIOUSLY Optimized Signals to 3DOF Vehicle Model \n"); break;
    else
      println("\n There status is nor Optimal or Infeaible -> create logic for this case!\n"); break;
    end
  if n.r.eval_num==n.mpc.max_iter
    warn(" \n The maximum number of iterations has been reached while closing the loop; consider increasing (max_iteration) \n")
  end
  if ((n.r.dfs_plant[end][:x][end]-c.g.x_ref)^2 + (n.r.dfs_plant[end][:y][end]-c.g.y_ref)^2)^0.5 < 2*n.XF_tol[1]
     println("Goal Attained! \n"); n.mpc.goal_reached=true; break;
  end
end

include("postProcess.jl");


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule ReadStat_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("ReadStat")
JLLWrappers.@generate_main_file("ReadStat", UUID("a4dc8951-f1cc-5499-9034-9ec1c3e64557"))
end  # module ReadStat_jll


text:
#!/usr/bin/env julia

using Luxor, Colors

"""
Work out how many rows and columns we need given the number of cells required.
It favours squarer layouts.
"""
function howmanyrowscolumns(n)
  numberofrows = convert(Int, floor(sqrt(n)))
  numberofcols = convert(Int, ceil(n/numberofrows))
  return numberofrows, numberofcols
end

"""
Find the area of an annular sector.
"""
function areaofsector(innerradius, outerradius, startangle, endangle)
  theta = endangle - startangle
  return (theta/2 * (innerradius + outerradius)^2) - (theta/2 * (innerradius)^2)
end

"""
Find the outerradius of a sector if we already know what its area and innerradius are.
"""

function outerradiusgivenarea(area, innerradius, startangle, endangle)
    theta = endangle - startangle
    return sqrt( (area + (theta/2) * innerradius^2) / (theta/2) ) - innerradius
end

"""
Draw a sector chart. You specify the centerposition, the tilewidth and height, some datavalues,
labels (one for each datavalue), and a dictionary that maps languages to colors. And a title.

The incoming values are rescaled to make nice pictures. But the original values are shown in labels.

    sectorchart(centerpos,
        innerradius,
        tilewidth, tileheight,
        rawdatavalues, labels, colordict, title;
        gap=deg2rad(1))

The gap is applied to the left and right of the sector, to leave some white space.
"""

function sectorchart(centerpos, innerradius, tilewidth, tileheight, rawdatavalues, labels, coldict, title; gap=deg2rad(2))
  gsave()
  translate(centerpos)
  n = length(rawdatavalues)
  theta = 2pi/n
  for i in 1:n
    startangle = (i - 1) * theta + gap
    endangle = startangle + theta - gap
    # convert areas to radius values
    # rawdatavalues are areas, to be converted into sector length
    # if datavalue is area, what would outerradius have to be?
    moddatavalues = map(dv -> outerradiusgivenarea(dv, innerradius, startangle, endangle), rawdatavalues)

    # rescale datavalues
    lowv, highv = extrema(moddatavalues)
    scaleddatavalues = map(dv -> rescale(dv, 0, highv, innerradius, min(tilewidth/2, tileheight/2)), moddatavalues)
    outerradius = scaleddatavalues[i]

    # find the required color for this sector
    sethue(coldict[labels[i]])
    sector(innerradius, outerradius, startangle, endangle, :fill)

    gsave()

    # draw the labels for each sector
    fontsize(6)
    sethue("black")
    textoffset = 5
    refpos = Point(
        (innerradius - textoffset) * cos(startangle + (endangle - startangle)/2),
        (innerradius - textoffset) * sin(startangle + (endangle - startangle)/2))
    rotangle = atan2(refpos.y, refpos.x)
    # subtract half the width of the string to center it
    glyph_x_bearing, glyph_y_bearing, glyph_width,
      glyph_height, glyph_x_advance, glyph_y_advance = textextents(labels[i])
    shiftangle = asin((glyph_width/2)/innerradius)
    textcurve(labels[i], rotangle-shiftangle, innerradius - textoffset)

    # show the original raw data value (not the rescaled value used for plotting)
    # refpos is for text placement
    textoffset = 8
    fontsize(10)
    refpos = Point(
        (outerradius + textoffset) * cos(startangle + (endangle - startangle)/2),
        (outerradius + textoffset) * sin(startangle + (endangle - startangle)/2))
    rotangle = atan2(refpos.y, refpos.x)
    translate(refpos)
    roundedtext = string(round(rawdatavalues[i], 2))
    (rotangle > pi/2 || rotangle < -pi/2) ? textright(roundedtext) : text(roundedtext)
    grestore()
  end
  gsave()
  # label this chart
  fontsize(24)
  fontface("Impact")
  sethue("black")
  textcentered(title, 0, min(tilewidth/2, tileheight/2) - 20)
  grestore()
  grestore()
end

function main()
  fname = "/tmp/sector-chart.pdf"
  width, height = 1064, 1064
  Drawing(width, height, fname)
  origin()
  background("ivory")
  setline(1)

  # get the data
  d = readcsv("docs/examples/benchmarks-julia.csv")
  # fields are language | test | timetaken
  d = d[d[:, 1] .!= "octave" , :]
  # but we don't want "Octave" because there are too many outliers :)

  languages = unique(d[:,1])
  benchmarknames = unique(d[:,2])

  # build a color dictionary for languages
  #          "c",     "fortran",   "go",         "java",       "javascript", "julia",      "lua",   "mathematica",  "matlab",  "python",      "r"]
  cols   =  ["gray20", "darkgreen", "blueviolet", "royalblue2",  "orange3",   "red",       "cyan3", "orangered3",   "gray56",  "chartreuse3", "aquamarine"]
  languagecolors = Dict(languages[i] => cols[i] for i in 1:length(languages))

  # how many charts are we plotting?
  # numberofrows, numberofcolumns = howmanyrowscolumns(length(benchmarknames))

  numberofrows, numberofcolumns = 3, 3
  topbottommargins = 200
  pagetiles = Tiler(width, height - topbottommargins, numberofrows, numberofcolumns, margin=50)

  # draw each chart using the next lot of data, for each benchmark
  for (cpos, counter) in pagetiles
    counter > length(benchmarknames) && continue # not all tiles will be filled
    benchmark = benchmarknames[counter]
    data = d[d[:, 2] .== benchmark, [3]]
    language = d[d[:, 2] .== benchmark, [1]]
    length(data) < 2 && continue # don't bother with tests that have only one value!
    sectorchart(cpos, 60, pagetiles.tilewidth, pagetiles.tileheight, data, language, languagecolors, benchmark)
  end

  # heading
  sethue("black")
  fontface("Impact")
  fontsize(40)
  gsave()
  translate(0, -height/2 + 100)
  textcentred("Julia benchmarks")
  grestore()

  # footnotes
  gsave()
  fontface("Monaco")
  fontsize(6)
  translate(0, height/2 - 50)
  textcentred("plotted $(Dates.format(now(), "e, dd u yyyy HH:MM")) Bigger areas mean slower. Humans are bad at comparing areas. Benchmarks are taken from the http://julialang.org web site.")
  translate(0, 10)
  textcentred("Terms and conditions apply. Objects may appear smaller than they actually are. The value of investments may go up as well as down.")
  grestore()

  # legend
  fontsize(14)
  fontface("Helvetica")
  x, y = width/2-150, height/2 - 350
  for language in languages
    sethue("black")
    textright(language, x, y)
    sethue(languagecolors[language])
    rect(Point(x+5, y), 40, -10, :fill)
    y += 25
  end

  finish()
  println("finished test: output in $(fname)")
  preview()
end

main()


text:
"""
    AbstractGate

Abtract unitary quantum gate. `N` is the number of "wires" the gate acts on.
"""
abstract type AbstractGate end

"""
    Base.adjoint(::AbstractGate)

returns adjoint of AbstractGate objects

# Examples
```jldoctest
julia> X ≈ X
true

julia> Y ≈ Z
false

julia> RxGate(0.3) ≈ RxGate(0.3)
true

julia> RxGate(0.3) ≈ RxGate(0.2)
false
```
"""
Base.adjoint(::AbstractGate) = error("Not implemented")

"""
    num_wires(::AbstractGate)

returns number of wires affected by AbstractGate object.

# Examples
```jldoctest
julia> num_wires(X)
1

julia> num_wires(SwapGate())
2
```
"""
num_wires(::AbstractGate) = error("Not implemented")

"""
    matrix(::AbstractGate)

returns dense matrix representation of [`AbstractGate`](@ref) objects

# Examples
```jldoctest
julia> matrix(X)
2×2 Array{Complex{Float64},2}:
 0.0+0.0im  1.0+0.0im
 1.0+0.0im  0.0+0.0im
```
"""
matrix(::AbstractGate) = error("Not implemented")

"""
    sparse_matrix(::AbstractGate)

returns sparse matrix representation of [`AbstractGate`](@ref) objects

# Examples
```jldoctest
julia> sparse_matrix(X)
2×2 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 2 stored entries:
  [2, 1]  =  1.0+0.0im
  [1, 2]  =  1.0+0.0im
```
"""
sparse_matrix(::AbstractGate) = error("Not implemented")

"""
    adjoint(::AbstractGate)

returns adjoint of [`AbstractGate`](@ref)

# Examples
```jldoctest
julia> adjoint(X)
XGate()

julia> adjoint(RxGate(0.3))
RxGate(-0.3)
```
"""

"""
    Pauli X gate

``X = \\begin{pmatrix} 0 & 1 \\\\ 1 & 0 \\end{pmatrix}``
"""
struct XGate <: AbstractGate end

"""
    Pauli Y gate

``Y = \\begin{pmatrix} 0 & -i \\\\ i & 0 \\end{pmatrix}``
"""
struct YGate <: AbstractGate end

"""
    Pauli Z gate

``Z = \\begin{pmatrix} 1 & 0 \\\\ 0 & -1 \\end{pmatrix}``
"""
struct ZGate <: AbstractGate end

matrix(::XGate) = ComplexF64[0  1 ; 1  0]
matrix(::YGate) = ComplexF64[0 -im; im 0]
matrix(::ZGate) = ComplexF64[1  0 ; 0 -1]

sparse_matrix(::XGate) = sparse(matrix(XGate()))
sparse_matrix(::YGate) = sparse(matrix(YGate()))
sparse_matrix(::ZGate) = sparse(matrix(ZGate()))

LinearAlgebra.ishermitian(::XGate) = true
LinearAlgebra.ishermitian(::YGate) = true
LinearAlgebra.ishermitian(::ZGate) = true

# Pauli matrices are Hermitian
Base.adjoint(::XGate) = X
Base.adjoint(::YGate) = Y
Base.adjoint(::ZGate) = Z

# corresponding instances
X = XGate()
Y = YGate()
Z = ZGate()

# wires
num_wires(::XGate)::Int = 1
num_wires(::YGate)::Int = 1
num_wires(::ZGate)::Int = 1

"""
    Hadamard gate

``H = \\frac{1}{\\sqrt{2}} \\begin{pmatrix} 1 & 1 \\\\ 1 & 1 \\end{pmatrix}``
"""
struct HadamardGate <: AbstractGate end

sparse_matrix(::HadamardGate) = sparse([1,1,2,2],[1,2,1,2], ComplexF64[1, 1, 1, -1] / sqrt(2))
matrix(::HadamardGate) = ComplexF64[1  1; 1 -1] / sqrt(2)

LinearAlgebra.ishermitian(::HadamardGate) = true
# Hadamard gate is Hermitian
Base.adjoint(H::HadamardGate) = H

# wires
num_wires(::HadamardGate)::Int = 1

"""
    S gate

``S = \\begin{pmatrix} 1 & 0 \\\\ 0 & i \\end{pmatrix}``
"""
struct SGate <: AbstractGate end

"""
    T gate

``T = \\begin{pmatrix} 1 & 0 \\\\ 0 & e^{\\frac{iπ}{4}} \\end{pmatrix}``
"""
struct TGate <: AbstractGate end

"""
    S^{†} gate

``S^{†} = \\begin{pmatrix} 1 & 0 \\\\ 0 & -i \\end{pmatrix}``
"""
struct SdagGate <: AbstractGate end

"""
    T^{†} gate

``T^{†} = \\begin{pmatrix} 1 & 0 \\\\ 0 & e^{-\\frac{iπ}{4}} \\end{pmatrix}``
"""
struct TdagGate <: AbstractGate end

matrix(::SGate)    = ComplexF64[1 0; 0 im]
matrix(::SdagGate) = ComplexF64[1 0; 0 -im]
matrix(::TGate)    = ComplexF64[1 0; 0 Base.exp(im * π / 4)]
matrix(::TdagGate) = ComplexF64[1 0; 0 Base.exp(-im * π / 4)]

sparse_matrix(::SGate)    = sparse(matrix(SGate()))
sparse_matrix(::TGate)    = sparse(matrix(TGate()))
sparse_matrix(::SdagGate) = sparse(matrix(SdagGate()))
sparse_matrix(::TdagGate) = sparse(matrix(TdagGate()))

LinearAlgebra.ishermitian(::SGate) = false
LinearAlgebra.ishermitian(::TGate) = false

LinearAlgebra.ishermitian(::SdagGate) = false
LinearAlgebra.ishermitian(::TdagGate) = false

Base.adjoint(::SGate) = SdagGate()
Base.adjoint(::TGate) = TdagGate()

Base.adjoint(::SdagGate) = SGate()
Base.adjoint(::TdagGate) = TGate()

# wires
num_wires(::SGate)::Int = 1
num_wires(::TGate)::Int = 1
num_wires(::SdagGate)::Int = 1
num_wires(::TdagGate)::Int = 1

"""
    Rotation-X gate

``R_{x}(\\theta) = \\begin{pmatrix} \\cos(\\frac{\\theta}{2}) & -i\\sin(\\frac{\\theta}{2}) \\\\ -i\\sin(\\frac{\\theta}{2}) & \\cos(\\frac{\\theta}{2}) \\end{pmatrix}``
"""
struct RxGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    θ::Vector{Float64}

    function RxGate(θ::Real)
        new([θ])
    end
end

function matrix(g::RxGate)
    c = cos(g.θ[] / 2.0)
    s = sin(g.θ[] / 2.0)
    ComplexF64[c -im * s; -im * s c]
end

sparse_matrix(g::RxGate) = sparse(matrix(g))

LinearAlgebra.ishermitian(g::RxGate) = abs(sin(g.θ[] / 2)) < 4 * eps()

# wires
num_wires(::RxGate)::Int = 1

"""
    Rotation-Y gate

``R_{y}(\\theta) = \\begin{pmatrix} \\cos(\\frac{\\theta}{2}) & -\\sin(\\frac{\\theta}{2}) \\\\ \\sin(\\frac{\\theta}{2}) & \\cos(\\frac{\\theta}{2}) \\end{pmatrix}``
"""
struct RyGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    θ::Vector{Float64}

    function RyGate(θ::Real)
        new([θ])
    end
end

function matrix(g::RyGate)
    c = cos(g.θ[] / 2.0)
    s = sin(g.θ[] / 2.0)
    ComplexF64[c -s; s c]
end

sparse_matrix(g::RyGate) = sparse(matrix(g))

LinearAlgebra.ishermitian(g::RyGate) = abs(sin(g.θ[] / 2)) < 4 * eps()

# wires
num_wires(::RyGate)::Int = 1

"""
    Rotation-Z gate

``R_{z}(\\theta) = \\begin{pmatrix} e^{\\frac{-i\\theta}{2}} & 0 \\\\ 0 & e^{\\frac{i\\theta}{2}} \\end{pmatrix}``
"""
struct RzGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    θ::Vector{Float64}
    function RzGate(θ::Real)
        new([θ])
    end
end

function matrix(g::RzGate)
    eθ = exp(im * g.θ[] / 2)
    ComplexF64[conj(eθ) 0; 0 eθ]
end

sparse_matrix(g::RzGate) = sparse(matrix(g))

LinearAlgebra.ishermitian(g::RzGate) = abs(sin(g.θ[] / 2)) < 4 * eps()

Base.adjoint(g::RxGate) = RxGate(-g.θ[])
Base.adjoint(g::RyGate) = RyGate(-g.θ[])
Base.adjoint(g::RzGate) = RzGate(-g.θ[])

# wires
num_wires(::RzGate)::Int = 1

"""
    General rotation operator gate: rotation by angle `θ` around unit vector `n`.

``R_{\\vec v}(\\theta) = \\cos(\\frac{\\theta}{2})I - i\\sin(\\frac{\\theta}{2})\\vec v \\sigma, \\\\ \\sigma = [X, Y, Z]``
"""
struct RotationGate <: AbstractGate
    nθ::Vector{Float64}

    function RotationGate(nθ::Vector{<:Real})
        length(nθ) == 3 || error("Rotation axis vector must have length 3.")
        new(nθ)
    end

    function RotationGate(θ::Real, n::Vector{<:Real})
        length(n) == 3 || error("Rotation axis vector must have length 3.")
        norm(n) ≈ 1 || error("Norm of rotation axis vector must be 1.")
        new(n * θ)
    end
end

function matrix(g::RotationGate)
    θ = norm(g.nθ)
    if θ == 0
        return Matrix{ComplexF64}(I, 2, 2)
    end
    n = g.nθ / θ
    return cos(θ/2)*I - im*sin(θ/2)*pauli_vector(n...)
end

sparse_matrix(g::RotationGate) = sparse(matrix(g))

LinearAlgebra.ishermitian(g::RotationGate) = abs(sin(norm(g.nθ) / 2)) < 8 * eps()

Base.adjoint(g::RotationGate) = RotationGate(-g.nθ)

# wires
num_wires(::RotationGate)::Int = 1

"""
    Phase shift gate

``P(\\phi) = \\begin{pmatrix} 1 & 0 \\\\ 0 & e^{i\\phi} \\end{pmatrix}``
"""
struct PhaseShiftGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    ϕ::Vector{Float64}

    function PhaseShiftGate(ϕ::Real)
        new([ϕ])
    end
end

matrix(g::PhaseShiftGate) = ComplexF64[1 0; 0 Base.exp(im * g.ϕ[])]

sparse_matrix(g::PhaseShiftGate) = sparse(matrix(g))

LinearAlgebra.ishermitian(g::PhaseShiftGate) = abs(sin(g.ϕ[])) < 4 * eps()

Base.adjoint(g::PhaseShiftGate) = PhaseShiftGate(-g.ϕ[])

# wires
num_wires(::PhaseShiftGate)::Int = 1

"""
    Swap gate

``SWAP = \\begin{pmatrix} 1 & 0 & 0 & 0 \\\\ 0 & 0 & 1 & 0 \\\\ 0 & 1 & 0 & 0 \\\\ 0 & 0 & 0 & 1 \\end{pmatrix}``
"""
struct SwapGate <: AbstractGate end

matrix(::SwapGate,) = ComplexF64[1 0 0 0; 0 0 1 0; 0 1 0 0; 0 0 0 1]

sparse_matrix(::SwapGate,) = sparse([1, 3, 2, 4], [1, 2, 3, 4], ComplexF64[1, 1, 1, 1])

# swap gate is Hermitian
LinearAlgebra.ishermitian(::SwapGate) = true

Base.adjoint(s::SwapGate) = s

# wires
num_wires(::SwapGate)::Int = 2

"""
    Entanglement-XX gate

``G_{x}(\\theta) = e^{-i \\theta X \\otimes X / 2}``

Reference:\n
    B. Kraus and J. I. Cirac\n
    Optimal creation of entanglement using a two-qubit gate\n
    Phys. Rev. A 63, 062309 (2001)
"""
struct EntanglementXXGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    θ::Vector{Float64}
    function EntanglementXXGate(θ::Real)
        new([θ])
    end
end

function matrix(g::EntanglementXXGate)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    ComplexF64[c  0  0  -im*s; 0  c  -im*s  0; 0  -im*s  c  0; -im*s  0  0  c]
end

function sparse_matrix(g::EntanglementXXGate)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    sparse([1, 4, 2, 3, 2, 3, 1, 4], [1, 1, 2, 2, 3, 3, 4, 4], ComplexF64[c, -im*s, c, -im*s, -im*s, c, -im*s, c])
end

LinearAlgebra.ishermitian(g::EntanglementXXGate) = abs(sin(g.θ[] / 2)) < 4 * eps()

# wires
num_wires(::EntanglementXXGate)::Int = 2

"""
    Entanglement-YY gate

``G_{y}(\\theta) = e^{-i \\theta Y \\otimes Y / 2}``

Reference:\n
    B. Kraus and J. I. Cirac\n
    Optimal creation of entanglement using a two-qubit gate\n
    Phys. Rev. A 63, 062309 (2001)
"""
struct EntanglementYYGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    θ::Vector{Float64}
    function EntanglementYYGate(θ::Real)
        new([θ])
    end
end

function matrix(g::EntanglementYYGate)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    ComplexF64[c  0  0  im*s; 0  c  -im*s  0; 0  -im*s  c  0; im*s  0  0  c]
end

function sparse_matrix(g::EntanglementYYGate)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    sparse([1, 4, 2, 3, 2, 3, 1, 4], [1, 1, 2, 2, 3, 3, 4, 4], ComplexF64[c, im*s, c, -im*s, -im*s, c, im*s, c])
end

LinearAlgebra.ishermitian(g::EntanglementYYGate) = abs(sin(g.θ[] / 2)) < 4 * eps()

# wires
num_wires(::EntanglementYYGate)::Int = 2

"""
    Entanglement-ZZ gate

``G_{z}(\\theta) = e^{-i \\theta Z \\otimes Z / 2}``

Reference:\n
    B. Kraus and J. I. Cirac\n
    Optimal creation of entanglement using a two-qubit gate\n
    Phys. Rev. A 63, 062309 (2001)
"""
struct EntanglementZZGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    θ::Vector{Float64}
    function EntanglementZZGate(θ::Real)
        new([θ])
    end
end

function matrix(g::EntanglementZZGate)
    eθ = exp(im * g.θ[] / 2)
    ComplexF64[conj(eθ) 0 0 0; 0 eθ 0 0; 0 0 eθ 0; 0 0 0 conj(eθ)]
end

function sparse_matrix(g::EntanglementZZGate)
    eθ = exp(im * g.θ[] / 2)
    sparse([1, 2, 3, 4], [1, 2, 3, 4], ComplexF64[conj(eθ), eθ, eθ, conj(eθ)])
end

LinearAlgebra.ishermitian(g::EntanglementZZGate) = abs(sin(g.θ[] / 2)) < 4 * eps()

# wires
num_wires(::EntanglementZZGate)::Int = 2

Base.adjoint(g::EntanglementXXGate) = EntanglementXXGate(-g.θ[])
Base.adjoint(g::EntanglementYYGate) = EntanglementYYGate(-g.θ[])
Base.adjoint(g::EntanglementZZGate) = EntanglementZZGate(-g.θ[])

matrix(::UniformScaling{Bool}) = ComplexF64[1 0; 0 1]
sparse_matrix(::UniformScaling{Bool}) = sparse((1.0+0.0im)*I, 2, 2)

"""
    kron(m::Union{UniformScaling{Bool},G}...) where {G<:AbstractGate}

returns sparse matrix representation of kronecker product of AbstractGate objects
"""
Base.kron(m::Union{UniformScaling{Bool},G}...) where {G<:AbstractGate} = Base.kron(sparse_matrix.(m)...)

"""
    ControlledGate{G}(U::G, M::Int) where {G<:AbstractGate}

returns ControlledGate object. Controls [`AbstractGate`](@ref) of type `G` with `M` control wires
"""
struct ControlledGate{G} <: AbstractGate
    "target gate"
    U::G
    "number of control wires"
    M::Int
end

# wires
num_wires(g::ControlledGate)::Int = g.M + num_wires(g.U)

"""
    target_wires(g::ControlledGate)

return target wires for a [`ControlledGate`](@ref)

# Examples
```jldocttest
# controlled gate with 1 target wire
julia> cnot = ControlledGate(X, 1);
julia> target_wires(cnot)
1
```
"""
target_wires(g::ControlledGate) = num_wires(g.U)

"""
    control_wires(g::ControlledGate)

return control wires for a [`ControlledGate`](@ref)

# Examples
```jldocttest
# controlled gate with 2 control wires
julia> ccnot = ControlledGate(X, 2);
julia> target_wires(ccnot)
2
```
"""
control_wires(g::ControlledGate) = g.M

function matrix(g::ControlledGate{G}) where {G <:AbstractGate}
    N = num_wires(g)
    Umat = matrix(g.U)
    CU = Matrix{ComplexF64}(I, 2^N, 2^N)
    # Note: target qubit(s) corresponds to fastest varying index
    CU[end-size(Umat, 1)+1:end, end-size(Umat, 2)+1:end] = Umat
    return CU
end

function sparse_matrix(g::ControlledGate{G}) where {G <:AbstractGate}
    N = num_wires(g)
    Umat = sparse_matrix(g.U)
    CU = sparse(one(ComplexF64)*I, 2^N, 2^N)
    # Note: target qubit(s) corresponds to fastest varying index
    CU[end-size(Umat, 1)+1:end, end-size(Umat, 2)+1:end] = Umat
    return dropzeros!(CU)
end

LinearAlgebra.ishermitian(g::ControlledGate) = LinearAlgebra.ishermitian(g.U)

Base.adjoint(g::ControlledGate) = ControlledGate(Base.adjoint(g.U), g.M)

controlled_not() = ControlledGate(X, 1)

isunitary(m::AbstractMatrix) = (m * Base.adjoint(m) ≈ I)

"""
    MatrixGate

general gate constructed from an unitary matrix
"""
struct MatrixGate <: AbstractGate
    matrix::SparseMatrixCSC{ComplexF64,Int}
    function MatrixGate(m::AbstractMatrix)
        @assert size(m, 1) == size(m, 2)
        isunitary(m) || error("Quantum gate must be unitary")
        return new(sparse(m))
    end
end

# wires
num_wires(g::MatrixGate)::Int = Int(log(2, size(g.matrix, 1)))

matrix(g::MatrixGate) = Matrix(g.matrix)

sparse_matrix(g::MatrixGate) = g.matrix

Base.adjoint(g::MatrixGate) = MatrixGate(Base.adjoint(g.matrix))

LinearAlgebra.ishermitian(g::MatrixGate) = LinearAlgebra.ishermitian(g.matrix)

"""
    Base.isapprox(g1::G, g2::G) where {G <: AbstractGate}

returns true if input AbstractGates are of same type with same paremeters

# Examples
```jldoctest
julia> X ≈ X
true

julia> X ≈ Y
false

julia> RxGate(0.1) ≈ RxGate(0.2)
false
```
"""
function Base.isapprox(g1::G, g2::G) where {G <: AbstractGate}
    for name in fieldnames(G)
        if !(getfield(g1, name) ≈ getfield(g2, name))
            return false
        end
    end
    return true
end

# handle different gate types or dimensions
Base.isapprox(::AbstractGate, ::AbstractGate) = false

# Define "addition" and "scaling" of parametric gates (required for gradient accumulation)

Base.:+(g1::RxGate, g2::RxGate) = RxGate(g1.θ[] + g2.θ[])
Base.:+(g1::RyGate, g2::RyGate) = RyGate(g1.θ[] + g2.θ[])
Base.:+(g1::RzGate, g2::RzGate) = RzGate(g1.θ[] + g2.θ[])

Base.:+(g1::RotationGate, g2::RotationGate) = RotationGate(g1.nθ + g2.nθ)

Base.:+(g1::PhaseShiftGate, g2::PhaseShiftGate) = PhaseShiftGate(g1.ϕ[] + g2.ϕ[])

Base.:+(g1::EntanglementXXGate, g2::EntanglementXXGate) = EntanglementXXGate(g1.θ[] + g2.θ[])
Base.:+(g1::EntanglementYYGate, g2::EntanglementYYGate) = EntanglementYYGate(g1.θ[] + g2.θ[])
Base.:+(g1::EntanglementZZGate, g2::EntanglementZZGate) = EntanglementZZGate(g1.θ[] + g2.θ[])

Base.:*(α::Real, g::RxGate) = RxGate(α * g.θ[])
Base.:*(α::Real, g::RyGate) = RyGate(α * g.θ[])
Base.:*(α::Real, g::RzGate) = RzGate(α * g.θ[])

Base.:*(α::Real, g::RotationGate) = RotationGate(α * g.nθ)

Base.:*(α::Real, g::PhaseShiftGate) = PhaseShiftGate(α * g.ϕ[])

Base.:*(α::Real, g::EntanglementXXGate) = EntanglementXXGate(α * g.θ[])
Base.:*(α::Real, g::EntanglementYYGate) = EntanglementYYGate(α * g.θ[])
Base.:*(α::Real, g::EntanglementZZGate) = EntanglementZZGate(α * g.θ[])