text:
using Arpack
using StaticArrays
using LinearAlgebra

using NPZ

"""
    idx::Int
    orb_list::Vector{Int}
"""
struct Cluster
    idx::Int
    orb_list::Vector{Int}
end
"""
	length(c::Cluster)

Return number of orbitals in `Cluster`
"""
function Base.length(c::Cluster)
    return length(c.orb_list)
end
"""
	dim_tot(c::Cluster)

Return dimension of hilbert space spanned by number of orbitals in `Cluster`.
This is all sectors
"""
function dim_tot(c::Cluster)
    return 2^(2*length(c))
end
"""
	dim_tot(c::Cluster, na, nb)

Return dimension of hilbert space spanned by number of orbitals in `Cluster`
with `na` and `nb` number of alpha/beta electrons.
"""
function dim_tot(c::Cluster, na, nb)
    return binomial(length(c), na)*binomial(length(c), na)
end
"""
	display(c::Cluster)
"""
function Base.display(c::Cluster)
    @printf("IDX%03i:DIM%04i:" ,c.idx,dim_tot(c))
    for si in c.orb_list
        @printf("%03i|", si)
    end
    @printf("\n")
end
function Base.isless(ci::Cluster, cj::Cluster)
    return Base.isless(ci.idx, cj.idx)
end
function Base.isequal(ci::Cluster, cj::Cluster)
    return Base.isequal(ci.idx, cj.idx)
end
######################################################################################################

"""
    cluster::Cluster                            # Cluster to which basis belongs
    basis::Dict{Tuple,Matrix{Float64}}          # Basis vectors (nα, nβ)=>[I,s]
These basis coefficients map local slater determinants to local vectors
`(nα, nβ):
V[αstring*βstring, cluster_state]`
"""
struct ClusterBasis
    cluster::Cluster
    basis::Dict{Tuple,Matrix{Float64}}
end
ClusterBasis(ci::Cluster) = ClusterBasis(ci, Dict{Tuple,Matrix{Float64}}())

Base.iterate(cb::ClusterBasis, state=1) = iterate(cb.basis, state)
Base.length(cb::ClusterBasis) = length(cb.basis)
Base.getindex(cb::ClusterBasis,i) = cb.basis[i]
Base.setindex!(cb::ClusterBasis,val,key) = cb.basis[key] = val
Base.haskey(cb::ClusterBasis,key) = haskey(cb.basis, key)
function Base.display(cb::ClusterBasis)
    @printf(" ClusterBasis for Cluster: %4i\n",cb.cluster.idx)
    norb = length(cb.cluster)
    sum_total_dim = 0
    sum_dim = 0
    for (sector, vecs) in cb.basis
        dim = size(vecs,2)
        total_dim = binomial(norb,sector[1]) * binomial(norb,sector[2])
        sum_dim += dim
        sum_total_dim += total_dim

        @printf("   FockSector = (%2iα, %2iβ): Total Dim = %5i: Dim = %4i\n", sector[1],sector[2],total_dim, dim)
    end

    @printf("   -----------------------------\n")
    @printf("   Total Dim = %5i: Dim = %4i\n", sum_total_dim, sum_dim)
end

######################################################################################################
struct ClusterOps
    cluster::Cluster
    data::Dict{String,Dict{Tuple{Tuple{Int16,Int16},Tuple{Int16,Int16}},Array}}
end
Base.iterate(i::ClusterOps, state=1) = iterate(i.data, state)
Base.length(i::ClusterOps) = length(co.data)
Base.getindex(co::ClusterOps,i) = co.data[i]
Base.setindex!(co::ClusterOps,val,key) = co.data[key] = val
Base.haskey(co::ClusterOps,key) = haskey(co.data, key)
Base.keys(co::ClusterOps) = keys(co.data)
function Base.display(co::ClusterOps)
    @printf(" ClusterOps for Cluster: %4i\n",co.cluster.idx)
    norb = length(co.cluster)
    for (op,sectors) in co
        print("   Operator: \n", op)
        print("     Sectors: \n")
        for sector in sectors
            println(sector)
        end
    end
end
function ClusterOps(ci::Cluster)
    dic1 = Dict{Tuple{Tuple{Int16,Int16},Tuple{Int16,Int16}},Array{Float64}}()
    dic2 = Dict{String,typeof(dic1)}()
    return ClusterOps(ci, dic2)
end

"""
    adjoint(co::ClusterOps; verbose=0)

Take ClusterOps, `co`, and return a new ClusterOps'
"""

"""
Defines a single cluster's subspace for Tucker. Each focksector is allowed to have a distinct cluster state range
for the subspace.

    cluster::Cluster
    data::OrderedDict{Tuple{UInt8,UInt8}, UnitRange{Int}}
"""
struct ClusterSubspace
    cluster::Cluster
    data::OrderedDict{Tuple{UInt8,UInt8}, UnitRange{Int}}
end
function ClusterSubspace(cluster::Cluster)
    return ClusterSubspace(cluster,OrderedDict{Tuple{UInt8,UInt8}, UnitRange{Int}}())
end
Base.haskey(css::ClusterSubspace, i) = return haskey(css.data,i)
Base.setindex!(tss::ClusterSubspace, i, j) = tss.data[j] = i
Base.getindex(tss::ClusterSubspace, i) = return tss.data[i]
function Base.display(tss::ClusterSubspace)
    @printf(" Subspace for Cluster: %4i : ", tss.cluster.idx)
    display(tss.cluster)
    for (fock,range) in tss.data
        @printf("  %10s   Range: %4i → %-4i Dim %4i\n",Int.(fock), first(range), last(range), length(range))
    end
end

"""
    get_ortho_compliment(tss::ClusterSubspace, cb::ClusterBasis)

For a given `ClusterSubspace`, `tss`, return the subspace remaining
"""
function get_ortho_compliment(tss::ClusterSubspace, cb::ClusterBasis)
#={{{=#
    data = OrderedDict{Tuple{UInt8,UInt8}, UnitRange{Int}}()
    for (fock,basis) in cb

        if haskey(tss.data,fock)
            first(tss.data[fock]) == 1 || error(" p-space doesn't include ground state?")
            newrange = last(tss[fock])+1:size(cb[fock],2)
            if length(newrange) > 0
                data[fock] = newrange
            end
        else
            newrange = 1:size(cb[fock],2)
            if length(newrange) > 0
                data[fock] = newrange
            end
        end
    end

    return ClusterSubspace(tss.cluster, data)
#=}}}=#
end

function Base.:+(a::Tuple{T,T}, b::Tuple{T,T}) where T<:Integer
    return (a[1]+b[1], a[2]+b[2])
end
function Base.:-(a::Tuple{T,T}, b::Tuple{T,T}) where T<:Integer
    return (a[1]-b[1], a[2]-b[2])
end

"""
    compute_cluster_eigenbasis(ints::InCoreInts, clusters::Vector{Cluster};
        init_fspace=nothing, delta_elec=nothing, verbose=0, max_roots=10,
        rdm1a=nothing, rdm1b=nothing)

Return a Vector of `ClusterBasis` for each `Cluster`
- `ints::InCoreInts`: In-core integrals
- `clusters::Vector{Cluster}`: Clusters
- `verbose::Int`: Print level
- `init_fspace`: list of pairs of (nα,nβ) for each cluster for defining reference space
                 for selecting out only certain fock sectors
- `delta_elec`: number of electrons different from reference (init_fspace)
- `max_roots::Int`: Maximum number of vectors for each focksector basis
- `rdm1a`: background density matrix for embedding local hamiltonian (alpha)
- `rdm1b`: background density matrix for embedding local hamiltonian (beta)
"""
function compute_cluster_eigenbasis(ints::InCoreInts, clusters::Vector{Cluster};
                init_fspace=nothing, delta_elec=nothing, verbose=0, max_roots=10,
                rdm1a=nothing, rdm1b=nothing)
#={{{=#
    # initialize output
    cluster_bases = Vector{ClusterBasis}()

    for ci in clusters
        verbose == 0 || display(ci)

        #
        # Get subset of integrals living on cluster, ci
        ints_i = subset(ints, ci.orb_list, rdm1a, rdm1b)

        if (rdm1a != nothing && init_fspace == nothing)
            error(" Cant embed withing init_fspace")
        end

        if all( (rdm1a,rdm1b,init_fspace) .!= nothing)
            #
            # Verify that density matrix provided is consistent with reference fock sectors
            occs = diag(rdm1a)
            occs[ci.orb_list] .= 0
            na_embed = sum(occs)
            occs = diag(rdm1b)
            occs[ci.orb_list] .= 0
            nb_embed = sum(occs)
            verbose == 0 || @printf(" Number of embedded electrons a,b: %f %f", na_embed, nb_embed)
        end

        delta_e_i = ()
        if all( (delta_elec,init_fspace) .!= nothing)
            delta_e_i = (init_fspace[ci.idx][1], init_fspace[ci.idx][2], delta_elec)
        end

        #
        # Get list of Fock-space sectors for current cluster
        #
        sectors = FermiCG.possible_focksectors(ci, delta_elec=delta_e_i)

        #
        # Loop over sectors and do FCI for each
        basis_i = ClusterBasis(ci)
        for sec in sectors

            #
            # prepare for FCI calculation for give sector of Fock space
            problem = FermiCG.StringCI.FCIProblem(length(ci), sec[1], sec[2])
            verbose == 0 || display(problem)
            nr = min(max_roots, problem.dim)

            #
            # Build full Hamiltonian matrix in cluster's Slater Det basis
            Hmat = FermiCG.StringCI.build_H_matrix(ints_i, problem)
            if problem.dim < 1000
                F = eigen(Hmat)

                e = F.values[1:nr]
                basis_i[sec] = F.vectors[:,1:nr]
                #display(e)
            else
                e,v = Arpack.eigs(Hmat, nev = nr, which=:SR)
                e = real(e)[1:nr]
                basis_i[sec] = v[:,1:nr]
            end
            if verbose > 0
                state=1
                for ei in e
                    @printf("   State %4i Energy: %12.8f %12.8f\n",state,ei, ei+ints.h0)
                    state += 1
                end
            end
        end
        push!(cluster_bases,basis_i)
    end
    return cluster_bases
end
#=}}}=#

"""
    compute_cluster_ops(cluster_bases::Vector{ClusterBasis})
"""
function compute_cluster_ops(cluster_bases::Vector{ClusterBasis},ints)

    clusters = Vector{Cluster}()
    for ci in cluster_bases
        push!(clusters, ci.cluster)
    end

    cluster_ops = Vector{ClusterOps}()
    for ci in clusters
        push!(cluster_ops, ClusterOps(ci))
    end

    for ci in clusters
        cb = cluster_bases[ci.idx]

        cluster_ops[ci.idx]["H"] = FermiCG.tdm_H(cb, subset(ints, ci.orb_list), verbose=0)

        cluster_ops[ci.idx]["A"], cluster_ops[ci.idx]["a"] = FermiCG.tdm_A(cb,"alpha")
        cluster_ops[ci.idx]["B"], cluster_ops[ci.idx]["b"] = FermiCG.tdm_A(cb,"beta")
        cluster_ops[ci.idx]["AA"], cluster_ops[ci.idx]["aa"] = FermiCG.tdm_AA(cb,"alpha")
        cluster_ops[ci.idx]["BB"], cluster_ops[ci.idx]["bb"] = FermiCG.tdm_AA(cb,"beta")
        cluster_ops[ci.idx]["Aa"] = FermiCG.tdm_Aa(cb,"alpha")
        cluster_ops[ci.idx]["Bb"] = FermiCG.tdm_Aa(cb,"beta")
        cluster_ops[ci.idx]["Ab"], cluster_ops[ci.idx]["Ba"] = FermiCG.tdm_Ab(cb)
        # remove BA and ba account for these terms
        cluster_ops[ci.idx]["AB"], cluster_ops[ci.idx]["ba"], cluster_ops[ci.idx]["BA"], cluster_ops[ci.idx]["ab"] = FermiCG.tdm_AB(cb)
        cluster_ops[ci.idx]["AAa"], cluster_ops[ci.idx]["Aaa"] = FermiCG.tdm_AAa(cb,"alpha")
        cluster_ops[ci.idx]["BBb"], cluster_ops[ci.idx]["Bbb"] = FermiCG.tdm_AAa(cb,"beta")
        cluster_ops[ci.idx]["ABa"], cluster_ops[ci.idx]["Aba"] = FermiCG.tdm_ABa(cb,"alpha")
        cluster_ops[ci.idx]["ABb"], cluster_ops[ci.idx]["Bba"] = FermiCG.tdm_ABa(cb,"beta")
        #cluster_ops[ci.idx]["ABa"], cluster_ops[ci.idx]["Aba"], cluster_ops[ci.idx]["BAa"], cluster_ops[ci.idx]["Aab"] = FermiCG.tdm_ABa(cb,"alpha")
        #cluster_ops[ci.idx]["ABb"], cluster_ops[ci.idx]["Bba"], cluster_ops[ci.idx]["BAb"], cluster_ops[ci.idx]["Bab"] = FermiCG.tdm_ABa(cb,"beta")

        # spin operators

        #
        # S+
        op = Dict{Tuple,Array}()
        for (fock,mat) in cluster_ops[ci.idx]["Ab"]
            dims = size(mat)
            op[fock] = zeros(dims[2:4]...)
            for j in 1:dims[4]
                for i in 1:dims[3]
                    for p in 1:dims[2]
                        op[fock][p,i,j] = mat[p,p,i,j]
                    end
                end
            end
        end
        cluster_ops[ci.idx]["S+"] = op

        #
        # S-
        op = Dict{Tuple,Array}()
        for (fock,mat) in cluster_ops[ci.idx]["Ba"]
            dims = size(mat)
            op[fock] = zeros(dims[2:4]...)
            for j in 1:dims[4]
                for i in 1:dims[3]
                    for p in 1:dims[2]
                        op[fock][p,i,j] = mat[p,p,i,j]
                    end
                end
            end
        end
        cluster_ops[ci.idx]["S-"] = op

        #
        # Sz
        op = Dict{Tuple,Array}()
        #
        # loop over fock-space transitions
        for (fock,basis) in cb
            focktrans = (fock,fock)

            sz = (fock[1] - fock[2]) / 2.0
            op[focktrans] = sz*Matrix(1.0I, size(cb[fock],2), size(cb[fock],2))
            op[focktrans] = reshape(op[focktrans],1,size(op[focktrans],1),size(op[focktrans],2))

        end
        cluster_ops[ci.idx]["Sz"] = op

        #
        # S2
        cluster_ops[ci.idx]["S2"] = FermiCG.tdm_S2(cb, subset(ints, ci.orb_list), verbose=0)

        to_delete = [
                     #"AAa",
                     #"Aaa",
                     #"BBb",
                     #"Bbb",
                     #
                     #"ABa",
                     #"Aba",
                     ##"BAa",
                     ##"Aab",
                     #
                     #"ABb",
                     #"Bba",
                     ##"BAb",
                     ##"Bab",
                     #"Aa",
                     #"Bb",
                     #"Ab",
                     #"Ba",
                     #"AB",
                     #"ba",
                     #"BA",
                     #"ab",
                     #"AA",
                     #"BB",
                     #"aa",
                     #"bb"
                     ]
        for op in to_delete
            for (ftran,array) in cluster_ops[ci.idx][op]
                cluster_ops[ci.idx][op][ftran] .*= 0
            end
        end

        # Compute single excitation operator
        tmp = Dict{Tuple,Array}()
        for (fock,basis) in cb
            tmp[(fock,fock)] = (cluster_ops[ci.idx]["Aa"][(fock,fock)] + cluster_ops[ci.idx]["Bb"][(fock,fock)])
        end
        cluster_ops[ci.idx]["E1"] = tmp

        #
        # reshape data into 3index quantities: e.g., (pqr, I, J)
        for opstring in keys(cluster_ops[ci.idx])
            opstring != "H" || continue
            opstring != "S2" || continue
            for ftrans in keys(cluster_ops[ci.idx][opstring])
                data = cluster_ops[ci.idx][opstring][ftrans]
                dim1 = prod(size(data)[1:(length(size(data))-2)])
                dim2 = size(data)[length(size(data))-1]
                dim3 = size(data)[length(size(data))-0]
                cluster_ops[ci.idx][opstring][ftrans] = copy(reshape(data, (dim1,dim2,dim3)))
            end
        end
    end
    return cluster_ops
end

"""
    possible_focksectors(c::Cluster, delta_elec=nothing)

Get list of possible fock spaces accessible to the cluster

- `delta_elec::Vector{Int}`: (nα, nβ, Δ) restricts fock spaces to: (nα,nβ) ± Δ electron transitions
"""
function possible_focksectors(c::Cluster; delta_elec::Tuple=())
    ref_a = nothing
    ref_b = nothing
    delta = nothing
    if length(delta_elec) != 0
        length(delta_elec) == 3 || throw(DimensionMismatch)
        ref_a = delta_elec[1]
        ref_b = delta_elec[2]
        delta = delta_elec[3]
    end

    no = length(c)

    fsectors::Vector{Tuple} = []
    for na in 0:no
        for nb in 0:no
            if length(delta_elec) != 0
                if abs(na-ref_a)+abs(nb-ref_b) > delta
                    continue
                end
            end
            push!(fsectors,(na,nb))
        end
    end
    return fsectors
end

"""
    tdm_H(cb::ClusterBasis; verbose=0)

Compute local Hamiltonian `<s|H|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space.

Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_H(cb::ClusterBasis, ints; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(cb.cluster)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    verbose == 0 || display(cb.cluster)
    for (fock,basis) in cb
        focktrans = (fock,fock)
        problem = StringCI.FCIProblem(norbs, fock[1], fock[2])
        verbose == 0 || display(problem)
        Hmap = StringCI.get_map(ints, problem)

        dicti[focktrans] = cb[fock]' * Matrix((Hmap * cb[fock]))

        if verbose > 0
            for e in 1:size(cb[fock],2)
                @printf(" %4i %12.8f\n", e, dicti[focktrans][e,e])
            end
        end
    end
    return dicti
#=}}}=#
end
"""
"""
function tdm_S2(cb::ClusterBasis, ints; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(cb.cluster)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    verbose == 0 || display(cb.cluster)
    for (fock,basis) in cb
        focktrans = (fock,fock)
        problem = StringCI.FCIProblem(norbs, fock[1], fock[2])
        verbose == 0 || display(problem)

        dicti[focktrans] = cb[fock]' * (StringCI.build_S2_matrix(problem) * cb[fock])

        if verbose > 0
            for e in 1:size(cb[fock],2)
                @printf(" %4i %12.8f\n", e, dicti[focktrans][e,e])
            end
        end
    end
    return dicti
#=}}}=#
end

"""
    tdm_A(cb::ClusterBasis; verbose=0)

Compute `<s|p'|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space.

Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_A(cb::ClusterBasis, spin_case; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    dicti_adj = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in 0:norbs
        for nb in 0:norbs
            fockbra = ()
            if spin_case == "alpha"
                fockbra = (na+1,nb)
            elseif spin_case == "beta"
                fockbra = (na,nb+1)
            else
                throw(DomainError(spin_case))
            end
            fockket = (na,nb)
            focktrans = (fockbra,fockket)
            focktrans_adj = (fockket,fockbra)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                dicti[focktrans] = FermiCG.StringCI.compute_annihilation(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket, spin_case)
                # adjoint
                basis_bra = cb[fockket]
                basis_ket = cb[fockbra]
                dicti_adj[focktrans_adj] =  permutedims(dicti[focktrans], [1,3,2])
            end
        end
    end
    return dicti, dicti_adj
#=}}}=#
end

"""
    tdm_AA(cb::ClusterBasis; verbose=0)

Compute `<s|p'q'|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space.

Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_AA(cb::ClusterBasis, spin_case; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    dicti_adj = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in 0:norbs
        for nb in 0:norbs
            fockbra = ()
            if spin_case == "alpha"
                fockbra = (na+2,nb)
            elseif spin_case == "beta"
                fockbra = (na,nb+2)
            else
                throw(DomainError(spin_case))
            end

            fockket = (na,nb)
            focktrans = (fockbra,fockket)
            focktrans_adj = (fockket,fockbra)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                #println()
                #println("::::::::::::::: ", fockbra, fockket)
                #println(":: ", size(basis_bra), size(basis_ket))
                dicti[focktrans] = FermiCG.StringCI.compute_AA(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket, spin_case)
                # adjoint
                basis_bra = cb[fockket]
                basis_ket = cb[fockbra]
                dicti_adj[focktrans_adj] =  permutedims(dicti[focktrans], [2,1,4,3])
            end
        end
    end
    return dicti, dicti_adj
#=}}}=#
end

"""
    tdm_Aa(cb::ClusterBasis, spin_case; verbose=0)

Compute `<s|p'q|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space.
- `spin_case`: alpha or beta
Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_Aa(cb::ClusterBasis, spin_case; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in 0:norbs
        for nb in 0:norbs
            fockbra = (na,nb)

            fockket = (na,nb)
            focktrans = (fockbra,fockket)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                dicti[focktrans] = FermiCG.StringCI.compute_Aa(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket, spin_case)
                #dicti[focktrans] = reshape(dicti[focktrans],(norbs*norbs, size(dicti[focktrans],3), size(dicti[focktrans],4)))
            end
        end
    end
    return dicti
#=}}}=#
end

"""
    tdm_Ab(cb::ClusterBasis; verbose=0)

Compute `<s|p'q|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space, where
`p'` is alpha and `q` is beta.

Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_Ab(cb::ClusterBasis; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    dicti_adj = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in -1:norbs+1
        for nb in -1:norbs+1
            fockbra = (na+1,nb-1)

            fockket = (na,nb)
            focktrans = (fockbra,fockket)
            focktrans_adj = (fockket,fockbra)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                dicti[focktrans] = FermiCG.StringCI.compute_Ab(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket)

                # adjoint
                basis_bra = cb[fockket]
                basis_ket = cb[fockbra]
                #dicti[focktrans] = FermiCG.StringCI.compute_Ba(norbs, fockket[1], fockket[2], fockket[1], fockket[2], basis_bra, basis_ket)
                dicti_adj[focktrans_adj] =  permutedims(dicti[focktrans], [2,1,4,3])
            end
        end
    end
    return dicti, dicti_adj
#=}}}=#
end

"""
    tdm_AB(cb::ClusterBasis; verbose=0)

Compute `<s|p'q'|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space, where
`p'` is alpha and `q'` is beta.

Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_AB(cb::ClusterBasis; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    dicti_adj = Dict{Tuple,Array}()
    dictj = Dict{Tuple,Array}()
    dictj_adj = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in -2:norbs+2
        for nb in -2:norbs+2
            fockbra = (na+1,nb+1)

            fockket = (na,nb)
            focktrans = (fockbra,fockket)
            focktrans_adj = (fockket,fockbra)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                dicti[focktrans] = FermiCG.StringCI.compute_AB(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket)
                dictj[focktrans] = -permutedims(dicti[focktrans], [2,1,3,4])

                # adjoint
                dicti_adj[focktrans_adj] =  permutedims(dicti[focktrans], [2,1,4,3])
                dictj_adj[focktrans_adj] =  permutedims(dictj[focktrans], [2,1,4,3])
            end
        end
    end
    return dicti, dicti_adj, dictj, dictj_adj
#=}}}=#
end

"""
    tdm_AAa(cb::ClusterBasis, spin_case; verbose=0)

Compute `<s|p'q'r|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space.
- `spin_case`: alpha or beta
Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_AAa(cb::ClusterBasis, spin_case; verbose=0)
#={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    dicti_adj = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in 0:norbs
        for nb in 0:norbs
            fockbra = ()
            if spin_case == "alpha"
                fockbra = (na+1,nb)
            elseif spin_case == "beta"
                fockbra = (na,nb+1)
            else
                throw(DomainError(spin_case))
            end

            fockket = (na,nb)
            focktrans = (fockbra,fockket)
            focktrans_adj = (fockket,fockbra)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                dicti[focktrans] = FermiCG.StringCI.compute_AAa(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket, spin_case)
                # adjoint
                basis_bra = cb[fockket]
                basis_ket = cb[fockbra]
                dicti_adj[focktrans_adj] =  permutedims(dicti[focktrans], [3,2,1,5,4])
            end
        end
    end
    return dicti, dicti_adj
#=}}}=#
end

"""
    tdm_ABa(cb::ClusterBasis, spin_case; verbose=0)

Compute `<s|p'q'r|t>` between all cluster states, `s` and `t`
from accessible sectors of a cluster's fock space.
- `spin_case`: alpha or beta
Returns `Dict[((na,nb),(na,nb))] => Array`
"""
function tdm_ABa(cb::ClusterBasis, spin_case; verbose=0)
    #={{{=#
    verbose == 0 || println("")
    verbose == 0 || display(ci)
    norbs = length(cb.cluster)

    dicti = Dict{Tuple,Array}()
    dicti_adj = Dict{Tuple,Array}()
    dictj = Dict{Tuple,Array}()
    dictj_adj = Dict{Tuple,Array}()
    #
    # loop over fock-space transitions
    for na in -2:norbs+2
        for nb in -2:norbs+2
            fockbra = ()
            if spin_case == "alpha"
                fockbra = (na,nb+1)
            elseif spin_case == "beta"
                fockbra = (na+1,nb)
            else
                throw(DomainError(spin_case))
            end

            fockket = (na,nb)
            focktrans = (fockbra,fockket)
            focktrans_adj = (fockket,fockbra)

            if haskey(cb, fockbra) && haskey(cb, fockket)
                basis_bra = cb[fockbra]
                basis_ket = cb[fockket]
                if spin_case == "alpha"
                    dicti[focktrans]     = FermiCG.StringCI.compute_ABa(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket)
                    dictj[focktrans] = -permutedims(dicti[focktrans], [2,1,3,4,5])
                elseif spin_case == "beta"
                    dicti[focktrans]     = FermiCG.StringCI.compute_ABb(norbs, fockbra[1], fockbra[2], fockket[1], fockket[2], basis_bra, basis_ket)
                    dictj[focktrans] = -permutedims(dicti[focktrans], [2,1,3,4,5])
                else
                    error("Wrong spin_case: ",spin_case)
                end

                # adjoint
                basis_bra = cb[fockket]
                basis_ket = cb[fockbra]

                dicti_adj[focktrans_adj] =  permutedims(dicti[focktrans], [3,2,1,5,4])
                dictj_adj[focktrans_adj] =  permutedims(dictj[focktrans], [3,2,1,5,4])
            end
        end
    end
    return dicti, dicti_adj, dictj, dictj_adj
    #=}}}=#
end

function add_cmf_operators!(ops::Vector{ClusterOps}, bases::Vector{ClusterBasis}, ints, Da, Db; verbose=0)

    n_clusters = length(bases)
    for ci_idx in 1:n_clusters
        cb = bases[ci_idx]
        ci = cb.cluster
        verbose == 0 || println()
        verbose == 0 || display(ci)
        norbs = length(cb.cluster)

        ints_i = form_casci_ints(ints, ci, Da, Db)

        dicti = Dict{Tuple,Array}()

        #
        # loop over fock-space transitions
        verbose == 0 || display(cb.cluster)
        for (fock,basis) in cb
            focktrans = (fock,fock)
            problem = StringCI.FCIProblem(norbs, fock[1], fock[2])
            verbose == 0 || display(problem)
            Hmap = StringCI.get_map(ints_i, problem)

            dicti[focktrans] = cb[fock]' * Matrix((Hmap * cb[fock]))

            if verbose > 0
                for e in 1:size(cb[fock],2)
                    @printf(" %4i %12.8f\n", e, dicti[focktrans][e,e])
                end
            end
        end
        ops[ci.idx]["Hcmf"] = dicti
    end
    return
end

"""
	form_schmidt_basis
thresh_orb      :   threshold for determining how many bath orbitals to include
thresh_schmidt  :   threshold for determining how many singular vectors to include for cluster basis

Returns new basis for the cluster
"""
function form_schmidt_basis(ints::InCoreInts, ci::Cluster, Da, Db;
        thresh_schmidt=1e-3, thresh_orb=1e-8, thresh_ci=1e-6,do_embedding=true,
        eig_nr=1, eig_max_cycles=200)

    println()
    println("------------------------------------------------------------")
    @printf("Form Embedded Schmidt-style basis for Cluster %4i\n",ci.idx)
    D = Da + Db

    # Form the exchange matrix
    K = zeros(size(ints.h1))
    @tensor begin
	K[q,r]  = ints.h2[p,q,r,s] * D[p,s]
    end

    no = size(ints.h1,1)
    ci_no = length(ci.orb_list)

    na_tot = Int(round(tr(Da)))
    nb_tot = Int(round(tr(Db)))
    println(" Number of electrons in full system:")
    @printf("  α: %12.8f  β:%12.8f \n ",na_tot,nb_tot)

    active = ci.orb_list

    backgr = Vector{Int}()
    for i in 1:no
    	if !(i in active)
    	    append!(backgr,i)
        end
    end

    println("active",active)
    println("backgr",backgr)

    K2 = zeros((ci_no,no-ci_no))

    for (pi,p) in enumerate(active)
    	for (qi,q) in enumerate(backgr)
	    K2[pi,qi] = K[p,q]
	end
    end

    println("The exchange matrix:")
    display(K2)
    F = svd(K2,full=true)

    @printf("\nSing. Val.\n")
    nkeep = 0
    for si in F.S
    	@printf("%16.12f\n",si)
        if si > thresh_orb
            nkeep += 1
        end
    end

    C = zeros(size(ints.h1))
    for (pi,p) in enumerate(active)
    	for (qi,q) in enumerate(active)
	    if pi==qi
	        C[p,qi] = 1
	    end
	end
    end

    #display(C)

    #display(F.Vt)
    #display(F.U)
    for (pi,p) in enumerate(backgr)
    	for (qi,q) in enumerate(backgr)
	    C[p,qi+length(active)] = F.Vt[qi,pi]
	end
    end
    #display(C)

    Cfrag = C[:,1:ci_no]
    Cbath = C[:,ci_no+1:ci_no+nkeep]
    Cenvt = C[:,ci_no+nkeep+1:end]

    @printf("Cfrag\n")
    display(Cfrag)
    @printf("\n NElec: %12.8f\n",(tr(Cfrag'*(Da+Db)*Cfrag)))
    @printf("Cbath\n")
    display(Cbath)
    @printf("\n NElec: %12.8f\n",(tr(Cbath'*(Da+Db)*Cbath)))
    @printf("Cenv\n")
    display(Cenvt)
    @printf("\n NElec: %12.8f\n",(tr(Cenvt'*(Da+Db)*Cenvt)))

    K2 = C'* K * C
    Da2 = C'* Da * C
    Db2 = C'* Db * C

    na = tr(Da2[1:ci_no+nkeep,1:ci_no+nkeep])
    nb = tr(Db2[1:ci_no+nkeep,1:ci_no+nkeep])

    println(" Number of electrons in Fragment+Bath system:")
    @printf("  α: %12.8f  β:%12.8f \n ",na,nb)

    denvt_a = Cenvt*Cenvt'*Da*Cenvt*Cenvt'
    denvt_b = Cenvt*Cenvt'*Db*Cenvt*Cenvt'
    #println(denvt_a)

    na_env = tr(denvt_a)
    nb_env = tr(denvt_b)

    println(" Number of electrons in Environment system:")
    @printf("  α: %12.8f  β:%12.8f \n ",na_env,nb_env)

    na_envt = Int(round(tr(Cenvt'*Da*Cenvt)))
    nb_envt = Int(round(tr(Cenvt'*Db*Cenvt)))

    println(" Number of electrons in Environment system:")
    @printf("  α: %12.8f  β:%12.8f \n ",na_envt,nb_envt)
    #display(Da)

    # rotate integrals to current subspace basis
    denvt_a = C'*denvt_a*C
    denvt_b = C'*denvt_b*C
    ints2 = orbital_rotation(ints, C)

    #avoid very zero numbers in diagonalization
    denvt_a[abs.(denvt_a) .< 1e-15] .= 0
    denvt_b[abs.(denvt_b) .< 1e-15] .= 0
    #display(denvt_a)
    #display(denvt_a)

    # find closest idempotent density for the environment
    if do_embedding
        if size(Cenvt,2)>0

	    #eigenvalue
	    EIG = eigen(denvt_a)
	    U = EIG.vectors
	    n = EIG.values

	    U = U[:, sortperm(n,rev=true)]
	    n = n[sortperm(n,rev=true)]
	    #println(n)
	    #display(U)

            for i in 1:nkeep
                @assert(n[i]>1e-14)
	    end

            denvt_a = U[:,1:na_envt] * U[:,1:na_envt]'

	    EIG = eigen(denvt_b)
	    U = EIG.vectors
	    n = EIG.values

	    U = U[:, sortperm(n,rev=true)]
	    n = n[sortperm(n,rev=true)]

            for i in 1:nkeep
                @assert(n[i]>1e-14)
	    end

            denvt_b = U[:,1:nb_envt] * U[:,1:nb_envt]'

	end
    #form ints in the cluster
    no_range = collect(1:size(Cfrag,2)+size(Cbath,2))

    #ints_f = subset(ints2,collect(1:size(Cfrag,2)+size(Cbath,2)), denvt_a, denvt_b)

    ints_f = FermiCG.form_1rdm_dressed_ints(ints2,no_range,denvt_a,denvt_b)

    else
        denvt_a *= 0
        denvt_b *= 0
        ints_f = form_casci_eff_ints(ints2,collect(1:size(Cfrag,2)+size(Cbath,2)), denvt_a, denvt_b)
    end

    println(" Number of electrons in Environment system:")
    @printf("  α: %12.8f  β:%12.8f \n ",tr(denvt_a),tr(denvt_b))

    na_actv = na_tot - na_envt
    nb_actv = nb_tot - nb_envt
    println(" Number of electrons in Fragment+Bath system:")
    @printf("  α: %12.8f  β:%12.8f \n ",na_actv,nb_actv)

    norb2 = size(ints_f.h1,1)
    problem = FermiCG.StringCI.FCIProblem(norb2, na_actv, nb_actv)
    Hmap = StringCI.get_map(ints_f, problem)
    v0 = svd(rand(problem.dim,eig_nr)).U
    davidson = FermiCG.Davidson(Hmap,v0=v0,max_iter=200, max_ss_vecs=20, nroots=eig_nr, tol=1e-8)
    #FermiCG.solve(davidson)
    @printf(" Now iterate: \n")
    flush(stdout)
    #@time FermiCG.iteration(davidson, Adiag=Adiag, iprint=2)
    @time e,v = FermiCG.solve(davidson);
    e = real(e)[1]
    v = v[:,1]

    basis = FermiCG.StringCI.svd_state(v,problem,length(active),nkeep,thresh_schmidt)
    return basis
end

"""
    compute_cluster_est_basis(ints::InCoreInts, clusters::Vector{Cluster};
        init_fspace=nothing, delta_elec=nothing, verbose=0, max_roots=10,
        rdm1a=nothing, rdm1b=nothing)

Return a Vector of `ClusterBasis` for each `Cluster`  using the Embedded Schmidt Truncation
- `ints::InCoreInts`: In-core integrals
- `clusters::Vector{Cluster}`: Clusters
- `Da`: background density matrix for embedding local hamiltonian (alpha)
- `Db`: background density matrix for embedding local hamiltonian (beta)
- `init_fspace`: list of pairs of (nα,nβ) for each cluster for defining reference space
                 for selecting out only certain fock sectors
- `thresh_schmidt`: the threshold for the EST
- `thresh_orb`: threshold for the orbital
- `thresh_ci`: threshold for the ci problem
"""
function compute_cluster_est_basis(ints::InCoreInts, clusters::Vector{Cluster},Da,Db;
        thresh_schmidt=1e-3, thresh_orb=1e-8, thresh_ci=1e-6,
	do_embedding=true,verbose=0,init_fspace=nothing,delta_elec=nothing,
	est_nr=1, est_max_cycles=200, est_thresh=1e-6)
#={{{=#
    # initialize output
    cluster_bases = Vector{ClusterBasis}()

    for ci in clusters
        verbose == 0 || display(ci)

        # Obtain the schmidt basis
        basis = FermiCG.form_schmidt_basis(ints, ci, Da, Db,thresh_schmidt=thresh_schmidt,
					   eig_nr=est_nr, eig_max_cycles=est_max_cycles, thresh_ci=est_thresh)

        delta_e_i = ()
        if all( (delta_elec,init_fspace) .!= nothing)
            delta_e_i = (init_fspace[ci.idx][1], init_fspace[ci.idx][2], delta_elec)
        end

        #
        # Get list of Fock-space sectors for current cluster
        #
        sectors = FermiCG.possible_focksectors(ci, delta_elec=delta_e_i)

        #
        # Loop over sectors and do FCI for each
        basis_i = ClusterBasis(ci)

        #for (key, value) in basis
        #    basis_i[key] = value
	#    println(key)
	#    display(value)
        #end

        for sec in sectors
            if sec in keys(basis)
                basis_i[sec] = basis[sec]
		#display(basis[sec])
		#st = "fock_"*string(ci.idx)*"_"*string(sec)
		#npzwrite(st, Matrix(basis[sec]))
            else
	    	#println(sec)
            end
        end
        push!(cluster_bases,basis_i)
    end
    return cluster_bases
end
#=}}}=#

"""
    rotate!(cb::ClusterBasis, U::Dict{Tuple,Matrix{T}}) where {T}

Rotate `cb` by unitary matrices in `U`
"""
function rotate!(cb::ClusterBasis,U::Dict{Tuple,Matrix{T}}) where {T}
#={{{=#
    for (fspace,mat) in U
        cb[fspace] .= cb[fspace] * mat
    end
end
#=}}}=#

"""
    rotate!(ops::ClusterOps, U::Dict{Tuple,Matrix{T}}) where {T}

Rotate `ops` by unitary matrices in `U`
"""
function rotate!(ops::ClusterOps,U::Dict{Tuple,Matrix{T}}) where T
#={{{=#
    for (op,fspace_deltas) in ops
        #println(" Rotate ", op)
        for (fspace_delta,tdm) in fspace_deltas
            fspace_l = fspace_delta[1]
            fspace_r = fspace_delta[2]

            if haskey(U, fspace_l)==true && haskey(U, fspace_r)==true
                Ul = U[fspace_l]
                Ur = U[fspace_r]
                if length(size(tdm)) == 2
                    @tensoropt tmp[q,s] := Ul[p,q] * Ur[r,s] * tdm[p,r]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 3
                    @tensoropt tmp[p,s,t] := Ul[q,s] * Ur[r,t] * tdm[p,q,r]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 4
                    @tensoropt tmp[p,q,t,u] := Ul[r,t] * Ur[s,u] * tdm[p,q,r,s]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 5
                    @tensoropt  tmp[p,q,r,u,v] := Ul[s,u] * Ur[t,v] * tdm[p,q,r,s,t]
                    ops[op][fspace_delta] = tmp
                else
                    error("Wrong dimension")
                end

            elseif haskey(U, fspace_l)==true && haskey(U, fspace_r)==false
                Ul = U[fspace_l]
                if length(size(tdm)) == 2
                    @tensoropt tmp[q,r] := Ul[p,q] * tdm[p,r]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 3
                    @tensoropt tmp[p,s,r] := Ul[q,s] * tdm[p,q,r]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 4
                    @tensoropt tmp[p,q,t,s] := Ul[r,t] * tdm[p,q,r,s]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 5
                    @tensoropt tmp[p,q,r,u,t] := Ul[s,u] * tdm[p,q,r,s,t]
                    ops[op][fspace_delta] = tmp
                else
                    error("Wrong dimension")
                end

            elseif haskey(U, fspace_l)==false && haskey(U, fspace_r)==true
                Ur = U[fspace_r]
                if length(size(tdm)) == 2
                    @tensoropt tmp[p,s] := Ur[r,s] * tdm[p,r]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 3
                    @tensoropt tmp[p,q,t] := Ur[r,t] * tdm[p,q,r]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 4
                    @tensoropt tmp[p,q,r,u] := Ur[s,u] * tdm[p,q,r,s]
                    ops[op][fspace_delta] = tmp
                elseif length(size(tdm)) == 5
                    @tensoropt tmp[p,q,r,s,v] := Ur[t,v] * tdm[p,q,r,s,t]
                    ops[op][fspace_delta] = tmp
                else
                    error("Wrong dimension")
                end
            end
        end
    end
end
#=}}}=#

function check_basis_orthogonality(basis::ClusterBasis; thresh=1e-12)
    for (fspace,mat) in basis
        if check_orthogonality(mat,thresh=thresh) == false
            println(" Cluster:", basis.cluster)
            println(" Fockspace:", fspace)
        end
    end
end

function check_orthogonality(mat; thresh=1e-12)
    Id = mat' * mat
    if maximum(abs.(I-Id)) > thresh
        @warn("problem with orthogonality ", maximum(abs.(I-Id)))
        return false
    end
    return true
end

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
function __init__cifar10()
    DEPNAME = "CIFAR10"

    register(DataDep(
        DEPNAME,
        """
        Dataset: The CIFAR-10 dataset
        Authors: Alex Krizhevsky, Vinod Nair, Geoffrey Hinton
        Website: https://www.cs.toronto.edu/~kriz/cifar.html
        Reference: https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf

        [Krizhevsky, 2009]
            Alex Krizhevsky.
            "Learning Multiple Layers of Features from Tiny Images",
            Tech Report, 2009.

        The CIFAR-10 dataset is a labeled subsets of the 80
        million tiny images dataset. It consists of 60000
        32x32 colour images in 10 classes, with 6000 images
        per class.

        The compressed archive file that contains the
        complete dataset is available for download at the
        offical website linked above; specifically the binary
        version for C programs. Note that using the data
        responsibly and respecting copyright remains your
        responsibility. The authors of CIFAR-10 aren't really
        explicit about any terms of use, so please read the
        website to make sure you want to download the
        dataset.
        """,
        "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz",
        "c4a38c50a1bc5f3a1c5537f2155ab9d68f9f25eb1ed8d9ddda3db29a59bca1dd",
        post_fetch_method = file -> (run(BinDeps.unpack_cmd(file, dirname(file), ".gz", ".tar")); rm(file))
    ))
end

"""
    CIFAR10(; Tx=Float32, split=:train, dir=nothing)
    CIFAR10([Tx, split])

The CIFAR10 dataset is a labeled subsets of the 80
million tiny images dataset. It consists of 60000
32x32 colour images in 10 classes, with 6000 images
per class.

# Arguments

$ARGUMENTS_SUPERVISED_ARRAY
- `split`: selects the data partition. Can take the values `:train:` or `:test`.

# Fields

$FIELDS_SUPERVISED_ARRAY
- `split`.

# Methods

$METHODS_SUPERVISED_ARRAY
- [`convert2image`](@ref) converts features to `RGB` images.

# Examples

```julia-repl
julia> using MLDatasets: CIFAR10

julia> dataset = CIFAR10()
CIFAR10:
  metadata    =>    Dict{String, Any} with 2 entries
  split       =>    :train
  features    =>    32×32×3×50000 Array{Float32, 4}
  targets     =>    50000-element Vector{Int64}

julia> dataset[1:5].targets
5-element Vector{Int64}:
 6
 9
 9
 4
 1

julia> X, y = dataset[];

julia> dataset = CIFAR10(Tx=Float64, split=:test)
CIFAR10:
  metadata    =>    Dict{String, Any} with 2 entries
  split       =>    :test
  features    =>    32×32×3×10000 Array{Float64, 4}
  targets     =>    10000-element Vector{Int64}

julia> dataset.metadata
Dict{String, Any} with 2 entries:
  "n_observations" => 10000
  "class_names"    => ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
```
"""
struct CIFAR10 <: SupervisedDataset
    metadata::Dict{String, Any}
    split::Symbol
    features::Array{<:Any, 4}
    targets::Vector{Int}
end

CIFAR10(; split=:train, Tx=Float32, dir=nothing) = CIFAR10(Tx, split; dir)

function CIFAR10(Tx::Type, split::Symbol=:train; dir=nothing)
    DEPNAME = "CIFAR10"
    NCHUNKS = 5
    TESTSET_FILENAME = joinpath("cifar-10-batches-bin", "test_batch.bin")

    filename_for_chunk(file_index::Int) =
        joinpath("cifar-10-batches-bin", "data_batch_$(file_index).bin")

    @assert split ∈ (:train, :test)

    if split == :train
        # placeholders for the chunks
        Xs = Vector{Array{UInt8,4}}(undef, NCHUNKS)
        Ys = Vector{Vector{Int}}(undef, NCHUNKS)
        # loop over all 5 trainingset files (i.e. chunks)
        for file_index in 1:NCHUNKS
            file_name = filename_for_chunk(file_index)
            file_path = datafile(DEPNAME, file_name, dir)
            # load all the data from each file and append it to
            # the placeholders X and Y
            X, Y = CIFAR10Reader.readdata(file_path)
            Xs[file_index] = X
            Ys[file_index] = Y

            #TODO define a lazy version that reads a signle image only when asked
            # file_index = ceil(Int, index / Reader.CHUNK_SIZE)
            # file_name = filename_for_chunk(file_index)
            # file_path = datafile(DEPNAME, file_name, dir)
            ## once we know the file we just need to compute the approriate
            ## offset of the image realtive to that file.
            # sub_index = ((index - 1) % Reader.CHUNK_SIZE) + 1
            # image, label = CIFAR10Reader.readdata(file_path, sub_index)
        end
        # cat all the placeholders into one image array
        # and one label array. (good enough)
        images = cat(Xs..., dims=4)::Array{UInt8,4}
        labels = vcat(Ys...)::Vector{Int}
        # optionally transform the image array before returning
        features, targets = bytes_to_type(Tx, images), labels
    else
        file_path = datafile(DEPNAME, TESTSET_FILENAME, dir)
        # simply read the complete content of the testset file
        images, labels = CIFAR10Reader.readdata(file_path)
        # optionally transform the image array before returning
        features, targets = bytes_to_type(Tx, images), labels
    end

    metadata = Dict{String, Any}()
    metadata["class_names"] = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
    metadata["n_observations"] = size(features)[end]
    return CIFAR10(metadata, split, features, targets)
end

convert2image(::Type{<:CIFAR10}, x::AbstractArray{<:Integer}) =
    convert2image(CIFAR10, reinterpret(N0f8, convert(Array{UInt8}, x)))

function convert2image(::Type{<:CIFAR10}, x::AbstractArray{T,N}) where {T,N}
    @assert N == 3 || N == 4
    x = permutedims(x, (3, 2, 1, 4:N...))
    return  ImageCore.colorview(RGB, x)
end

# DEPRECATED INTERFACE, REMOVE IN v0.7 (or 0.6.x)
function Base.getproperty(::Type{CIFAR10}, s::Symbol)
    if s == :traintensor
        @warn "CIFAR10.traintensor() is deprecated, use `CIFAR10(split=:train).features` instead." maxlog=2
        traintensor(T::Type=N0f8; kws...) = traintensor(T, :; kws...)
        traintensor(i; kws...) = traintensor(N0f8, i; kws...)
        function traintensor(T::Type, i; dir=nothing)
            CIFAR10(; split=:train, Tx=T, dir)[i][1]
        end
        return traintensor
    elseif s == :testtensor
        @warn "CIFAR10.testtensor() is deprecated, use `CIFAR10(split=:test).features` instead."  maxlog=2
        testtensor(T::Type=N0f8; kws...) = testtensor(T, :; kws...)
        testtensor(i; kws...) = testtensor(N0f8, i; kws...)
        function testtensor(T::Type, i; dir=nothing)
            CIFAR10(; split=:test, Tx=T, dir)[i][1]
        end
        return testtensor
    elseif s == :trainlabels
        @warn "CIFAR10.trainlabels() is deprecated, use `CIFAR10(split=:train).targets` instead."  maxlog=2
        trainlabels(; kws...) = trainlabels(:; kws...)
        function trainlabels(i; dir=nothing)
            CIFAR10(; split=:train, dir)[i][2]
        end
        return trainlabels
    elseif s == :testlabels
        @warn "CIFAR10.testlabels() is deprecated, use `CIFAR10(split=:test).targets` instead." maxlog=2
        testlabels(; kws...) = testlabels(:; kws...)
        function testlabels(i; dir=nothing)
            CIFAR10(; split=:test, dir)[i][2]
        end
        return testlabels
    elseif s == :traindata
        @warn "CIFAR10.traindata() is deprecated, use `CIFAR10(split=:train)[]` instead." maxlog=2
        traindata(T::Type=N0f8; kws...) = traindata(T, :; kws...)
        traindata(i; kws...) = traindata(N0f8, i; kws...)
        function traindata(T::Type, i; dir=nothing)
            CIFAR10(; split=:train, Tx=T, dir)[i]
        end
        return traindata
    elseif s == :testdata
        @warn "CIFAR10.testdata() is deprecated, use `CIFAR10(split=:test)[]` instead."  maxlog=2
        testdata(T::Type=N0f8; kws...) = testdata(T, :; kws...)
        testdata(i; kws...) = testdata(N0f8, i; kws...)
        function testdata(T::Type, i; dir=nothing)
            CIFAR10(; split=:test, Tx=T, dir)[i]
        end
        return testdata
    elseif s == :convert2image
        @warn "CIFAR10.convert2image(x) is deprecated, use `convert2image(CIFAR10, x)` instead"
        return x -> convert2image(CIFAR10, x)
    else
        return getfield(CIFAR10, s)
    end
end

ext: jl


text:
"""
    load(file::String, vari::String; poly = Array{Float64}([]), start_date::Tuple, end_date::Tuple, data_units::String = "")

Returns a ClimGrid type with the data in **file** of variable **vari** inside the polygon **poly**. Metadata is built-in the ClimGrid type, from the netCDF attributes.

Inside the ClimgGrid type, the data is stored into an AxisArray data type, with time, longitude/x and latitude/y dimensions.

The polygon provided should be in the -180, +180 longitude format. If the polygon crosses the International Date Line, the polygon should be splitted in multiple parts (i.e. multi-polygons).

Options for data_units are for precipitation : "mm", which converts the usual "kg m-2 s-1" unit found in netCDF files. For temperature : "Celsius", which converts the usual "Kelvin" unit.

Temporal subsetting can be done by providing start_date and end-date Tuples of length 1 (year), length 3 (year, month, day) or 6 (hour, minute, second).

**Note:** load uses [CF conventions](http://cfconventions.org/). If you are unable to read the netCDF file with load, the user will need to read it with low-level functions available in [NetCDF.jl package](https://github.com/JuliaGeo/NetCDF.jl) or [NCDatasets.jl](https://github.com/Alexander-Barth/NCDatasets.jl) or re-create standartized netCDF files.
"""
function load(file::String, vari::String; poly = ([]), start_date::Tuple=(Inf,), end_date::Tuple=(Inf,), data_units::String = "")

  # TODO this file is a complete mess, but it works. Clean it up!

  # Get attributes
  ds = NCDatasets.Dataset(file)
  attribs_dataset = ds.attrib
  attribs = Dict(attribs_dataset)

  # The following attributes should be set for netCDF files that follows CF conventions.
  # project, institute, model, experiment, frequency
  project = ClimateTools.project_id(attribs_dataset)
  institute = ClimateTools.institute_id(attribs_dataset)
  model = ClimateTools.model_id(attribs_dataset)
  experiment = ClimateTools.experiment_id(attribs_dataset)
  frequency = ClimateTools.frequency_var(attribs_dataset)
  runsim = ClimateTools.runsim_id(attribs_dataset)
  grid_mapping = ClimateTools.get_mapping(keys(ds))

  # Get dimensions names
  latname, latstatus = getdim_lat(ds)
  lonname, lonstatus = getdim_lon(ds)

  dataunits = ds[vari].attrib["units"]
  latunits = ds[latname].attrib["units"]
  lonunits = ds[lonname].attrib["units"]
  caltype = ""
    try
        caltype = ds["time"].attrib["calendar"]
    catch
        caltype = "standard"
    end

  # Create dict with latname and lonname
  dimension_dict = Dict(["lon" => lonname, "lat" => latname])

  # lat_raw = NetCDF.ncread(file, latname)
  lat_raw = nomissing(ds[latname][:], NaN)
  # lon_raw = NetCDF.ncread(file, lonname)
  lon_raw = nomissing(ds[lonname][:], NaN)

  # Get variable attributes
  varattrib = Dict(ds[vari].attrib)

  if latstatus # means we don't have a "regular" grid
      # Get names of grid
      latgrid_name = latgridname(ds)
      longrid_name = longridname(ds)

      # latgrid = NetCDF.ncread(file, latgrid_name)
      latgrid = nomissing(ds[latgrid_name][:], NaN)
      # longrid = NetCDF.ncread(file, longrid_name)
      longrid = nomissing(ds[longrid_name][:], NaN)

      # Ensure we have a grid
      if ndims(latgrid) == 1 && ndims(longrid) == 1
          longrid, latgrid = ndgrid(lon_raw, lat_raw)
          map_attrib = Dict(["grid_mapping" => "Regular_longitude_latitude"])
      end

      map_attrib = build_grid_mapping(ds, grid_mapping)
      varattrib["grid_mapping"] = grid_mapping

  else # if no grid provided, create one
      longrid, latgrid = ndgrid(lon_raw, lat_raw)
      map_attrib = Dict(["grid_mapping" => grid_mapping])
      varattrib["grid_mapping"] = grid_mapping
  end

  # =====================
  # TIME
  # Get time resolution
  timeV = ds["time"][:]
  if frequency == "N/A" || !ClimateTools.@isdefined frequency
      try
          try
              frequency = string(diff(timeV)[2])
          catch
              frequency = string(diff(timeV)[1])
          end
      catch
          frequency = "N/A"
      end
  end

  timeattrib = Dict(ds["time"].attrib)
  T = typeof(timeV[1])
  idxtimebeg, idxtimeend = timeindex(timeV, start_date, end_date, T)
  timeV = timeV[idxtimebeg:idxtimeend]

  # ==================
  # Spatial shift if grid is 0-360.
  rotatedgrid = false
  if sum(longrid .> 180) >= 1
      rotatedgrid = true

      # Shift 360 degrees grid to -180, +180 degrees
      longrid_flip = ClimateTools.shiftgrid_180_west_east(longrid)

      # Shift 360 degrees vector to -180, +180 degrees
      lon_raw_flip = ClimateTools.shiftvector_180_west_east(lon_raw)

  else
      longrid_flip = longrid # grid is already "flipped" by design
  end

  # ===================
  # GET DATA
  data_pointer = ds[vari]
  if !isempty(poly)

    # Test to see if the polygon crosses the meridian
    meridian = ClimateTools.meridian_check(poly)

    # Build mask based on provided polygon
    msk = inpolygrid(longrid_flip, latgrid, poly)

    if sum(isnan.(msk)) == length(msk) # no grid point inside polygon
        throw(error("No grid points found inside the provided polygon"))
    end

    if rotatedgrid
        # Regrid msk to original grid if the grid have been rotated to get proper index for data extraction
        msk = ClimateTools.shiftarray_east_west(msk, longrid_flip)
    end

    #Extract data based on mask
    data_ext = ClimateTools.extractdata(data_pointer, msk, idxtimebeg, idxtimeend)
    data_ext = nomissing(data_ext, NaN)

    #new mask (e.g. representing the region of the polygon)
    begin
        I = Base.findall(!isnan, msk)
        idlon, idlat = (getindex.(I, 1), getindex.(I, 2))
      end
    minXgrid = minimum(idlon)
    maxXgrid = maximum(idlon)
    minYgrid = minimum(idlat)
    maxYgrid = maximum(idlat)
    msk = msk[minXgrid:maxXgrid, minYgrid:maxYgrid]
    data_mask = applymask(data_ext, msk) # needed when polygon is not rectangular

    if rotatedgrid
        # Regrid msk to shifted grid if the grid have been rotated to get final mask
        #shiftarray_west_east(msk, longrid)
    end

    # Get lon_raw and lat_raw for such region
    lon_raw = lon_raw[minXgrid:maxXgrid]
    lat_raw = lat_raw[minYgrid:maxYgrid]

    if map_attrib["grid_mapping"] == "Regular_longitude_latitude"

        lon_raw = ClimateTools.shiftvector_180_west_east(lon_raw)
    end

    # Idem for longrid and latgrid
    if meridian
        longrid = ClimateTools.shiftgrid_180_east_west(longrid)#grideast, gridwest)
        longrid = longrid[minXgrid:maxXgrid, minYgrid:maxYgrid]
        latgrid = latgrid[minXgrid:maxXgrid, minYgrid:maxYgrid]

        # Re-translate to -180, 180 if 0, 360

        if rotatedgrid

            longrid_flip = ClimateTools.shiftgrid_180_west_east(longrid)

            data = permute_west_east(data_mask, longrid)#idxwest, idxeast)
            msk = ClimateTools.permute_west_east(msk, longrid)

            # TODO Try to trim padding when meridian is crossed and model was on a 0-360 coords
            # idlon, idlat = findn(.!isnan.(msk))
            # minXgrid = minimum(idlon)
            # maxXgrid = maximum(idlon)
            # minYgrid = minimum(idlat)
            # maxYgrid = maximum(idlat)
            #
            # msk = msk[minXgrid:maxXgrid, minYgrid:maxYgrid]
            # data = data[:, minXgrid:maxXgrid, minYgrid:maxYgrid]
            # lon_raw_flip = lon_raw[minXgrid:maxXgrid]
            #
            # longrid_flip = longrid[minXgrid:maxXgrid, minYgrid:maxYgrid]
            # longrid_flip = ClimateTools.shiftgrid_180_west_east(longrid_flip)
            # latgrid = latgrid[minXgrid:maxXgrid, minYgrid:maxYgrid]

            #     data = applymask(data, msk)

        else
            data = data_mask
        end

    else

        if rotatedgrid # flip in original grid
            longrid = ClimateTools.shiftgrid_180_east_west(longrid) #grideast, gridwest)
        end

        longrid = longrid[minXgrid:maxXgrid, minYgrid:maxYgrid]
        latgrid = latgrid[minXgrid:maxXgrid, minYgrid:maxYgrid]

        if rotatedgrid

            longrid_flip = shiftgrid_180_west_east(longrid)

            lon_raw_flip = shiftvector_180_east_west(lon_raw)

            data = permute_west_east(data_mask, longrid)#idxwest, idxeast)
            msk = ClimateTools.permute_west_east(msk, longrid)
        else
            data = data_mask
        end

    end

  elseif isempty(poly) # no polygon clipping

      msk = Array{Float64}(ones((size(data_pointer, 1), size(data_pointer, 2))))
      data_ext = ClimateTools.extractdata(data_pointer, msk, idxtimebeg, idxtimeend)
      # replace_missing!(data_ext)
      data_ext = nomissing(data_ext, NaN)
      # data_ext = convert(data_ext, Float32)

    if rotatedgrid

        # Flip data "west-east"
        data = ClimateTools.permute_west_east(data_ext, longrid)

    else
        data = data_ext
    end

  end

  if rotatedgrid
      longrid .= longrid_flip
      lon_raw .= lon_raw_flip
  end

  # Convert units of optional argument data_units is provided
  if data_units == "Celsius" && (vari == "tas" || vari == "tasmax" || vari == "tasmin") && dataunits == "K"
    data .-= 273.15
    dataunits = "°C"
    varattrib["units"] = "Celsius"
    # @warn "Using Celsius can be problematic for arithmetic operations. Best practice is to keep Kelvin and only convert to Celsius at the end with the overloaded ClimateTools.uconvert function."
  end

  if data_units == "mm" && vari == "pr" && (dataunits == "kg m-2 s-1" || dataunits == "mm s-1")

    factor = timeresolution(ds["time"])
    # factor = pr_timefactor(rez)
    data .*= factor.value
    dataunits = "mm"
    varattrib["standard_name"] = "precipitation"
    varattrib["units"] = "mm"
  end

  # Create AxisArray from variable "data"
  if ndims(data) == 3
    # Convert data to AxisArray
    dataOut = AxisArray(data, Axis{Symbol(lonname)}(lon_raw), Axis{Symbol(latname)}(lat_raw), Axis{:time}(timeV))
  elseif ndims(data) == 4 # this imply a 3D field (height component)
    # Get level vector
    plev = ds["plev"][:]
    # Convert data to AxisArray
    dataOut = AxisArray(data, Axis{Symbol(lonname)}(lon_raw), Axis{Symbol(latname)}(lat_raw), Axis{:plev}(plev), Axis{:time}(timeV))
  else
    throw(error("load takes only 3D and 4D variables for the moment"))
  end

  C = ClimGrid(dataOut, longrid=longrid, latgrid=latgrid, msk=msk, grid_mapping=map_attrib, dimension_dict=dimension_dict, timeattrib=timeattrib, model=model, frequency=frequency, experiment=experiment, run=runsim, project=project, institute=institute, filename=file, dataunits=dataunits, latunits=latunits, lonunits=lonunits, variable=vari, typeofvar=vari, typeofcal=caltype, varattribs=varattrib, globalattribs=attribs)

  close(ds)
  # NetCDF.close(ncfile)

  return C

end

"""
    load(files::Array{String,1}, vari::String; poly = ([]), start_date::Date = Date(-4000), end_date::Date = Date(-4000), data_units::String = "")

Loads and merge the files contained in the arrar files.
"""
function load(files::Array{String,1}, vari::String; poly=([]), start_date::Tuple=(Inf,), end_date::Tuple=(Inf,), data_units::String="")

    nfiles = length(files)
    C = Array{ClimGrid}(undef, nfiles) # initialize # TODO better initialization
    datesort = Array{Any}(undef, nfiles)
    Cout = []

    p = Progress(nfiles*2, 3, "Loading files: ")

    for ifile = 1:nfiles

        C[ifile] = load(files[ifile], vari, poly = poly, start_date=start_date, end_date=end_date, data_units=data_units)
        datesort[ifile] = get_timevec(C[ifile])[1]

        next!(p)
    end

    # Sort files based on timevector to ensure that merge results in amonotone increase in time
    idx = sortperm(datesort)
    C = C[idx]

    for imod = 1:nfiles
        if imod == 1
            Cout = C[imod]
        else
            Cout = merge(Cout, C[imod])
        end
        next!(p)
    end

    return Cout
end

"""
    load2D(file::String, vari::String; poly=[], data_units::String="")

Returns a 2D array. Should be used for *fixed* data, such as orography.
"""
function load2D(file::String, vari::String; poly=[], data_units::String="")
    # Get attributes

    ds = NCDatasets.Dataset(file)
    attribs_dataset = ds.attrib
    attribs = Dict(attribs_dataset)

    project = ClimateTools.project_id(attribs_dataset)
    institute = ClimateTools.institute_id(attribs_dataset)
    model = ClimateTools.model_id(attribs_dataset)
    experiment = ClimateTools.experiment_id(attribs_dataset)
    frequency = ClimateTools.frequency_var(attribs_dataset)
    runsim = ClimateTools.runsim_id(attribs_dataset)
    grid_mapping = ClimateTools.get_mapping(keys(ds))

    # Get dimensions names
    latname, latstatus = getdim_lat(ds)
    lonname, lonstatus = getdim_lon(ds)

    dataunits = ds[vari].attrib["units"]
    latunits = ds[latname].attrib["units"]
    lonunits = ds[lonname].attrib["units"]

    # Create dict with latname and lonname
    dimension_dict = Dict(["lon" => lonname, "lat" => latname])

    # lat_raw = NetCDF.ncread(file, latname)
    lat_raw = nomissing(ds[latname][:], NaN)
    # lon_raw = NetCDF.ncread(file, lonname)
    lon_raw = nomissing(ds[lonname][:], NaN)

    # Get variable attributes
    varattrib = Dict(ds[vari].attrib)

    if latstatus # means we don't have a "regular" grid
        # Get names of grid
        latgrid_name = latgridname(ds)
        longrid_name = longridname(ds)

        # latgrid = NetCDF.ncread(file, latgrid_name)
        latgrid = nomissing(ds[latgrid_name][:], NaN)
        # longrid = NetCDF.ncread(file, longrid_name)
        longrid = nomissing(ds[longrid_name][:], NaN)

        map_attrib = build_grid_mapping(ds, grid_mapping)
        varattrib["grid_mapping"] = grid_mapping

    else # if no grid provided, create one
        longrid, latgrid = ndgrid(lon_raw, lat_raw)
        map_attrib = Dict(["grid_mapping" => grid_mapping])
        varattrib["grid_mapping"] = grid_mapping
    end

    # ==================
    # Spatial shift if grid is 0-360.
    rotatedgrid = false
    if sum(longrid .> 180) >= 1
        rotatedgrid = true

        # Shift 360 degrees grid to -180, +180 degrees
        longrid_flip = ClimateTools.shiftgrid_180_west_east(longrid)

        # Shift 360 degrees vector to -180, +180 degrees
        lon_raw_flip = ClimateTools.shiftvector_180_west_east(lon_raw)

    else
        longrid_flip = longrid # grid is already "flipped" by design
    end

    # ===================
    # GET DATA
    # data = ds[variable]
    data_pointer = ds[vari]
    if !isempty(poly)

      # Test to see if the polygon crosses the meridian
      meridian = ClimateTools.meridian_check(poly)

      # Build mask based on provided polygon
      msk = inpolygrid(longrid_flip, latgrid, poly)

      if sum(isnan.(msk)) == length(msk) # no grid point insode polygon
          throw(error("No grid points found inside the provided polygon"))
      end

      if rotatedgrid
          # Regrid msk to original grid if the grid have been rotated to get proper index for data extraction
          msk = ClimateTools.shiftarray_east_west(msk, longrid_flip)
      end

      #Extract data based on mask
      data_ext = ClimateTools.extractdata2D(data_pointer, msk)
      # replace_missing!(data_ext)
      data_ext = nomissing(data_ext, NaN)
      # data_ext = convert(data_ext, Float32)

      begin
        I = Base.findall(!isnan, msk)
        idlon, idlat = (getindex.(I, 1), getindex.(I, 2))
      end

      minXgrid = minimum(idlon)
      maxXgrid = maximum(idlon)
      minYgrid = minimum(idlat)
      maxYgrid = maximum(idlat)
      msk = msk[minXgrid:maxXgrid, minYgrid:maxYgrid]
      data_mask = applymask(data_ext, msk) # needed when polygon is not rectangular

      if rotatedgrid
          # Regrid msk to shifted grid if the grid have been rotated to get final mask
          #shiftarray_west_east(msk, longrid)
      end

      # Get lon_raw and lat_raw for such region
      lon_raw = lon_raw[minXgrid:maxXgrid]
      lat_raw = lat_raw[minYgrid:maxYgrid]

      if map_attrib["grid_mapping"] == "Regular_longitude_latitude"

          lon_raw = ClimateTools.shiftvector_180_west_east(lon_raw)

      end

      # Idem for longrid and latgrid
      if meridian
          longrid = ClimateTools.shiftgrid_180_east_west(longrid)#grideast, gridwest)
          longrid = longrid[minXgrid:maxXgrid, minYgrid:maxYgrid]
          latgrid = latgrid[minXgrid:maxXgrid, minYgrid:maxYgrid]

          # Re-translate to -180, 180 if 0, 360

          if rotatedgrid

              longrid_flip = ClimateTools.shiftgrid_180_west_east(longrid)

              data = permute_west_east(data_mask, longrid)#idxwest, idxeast)
              msk = ClimateTools.permute_west_east(msk, longrid)

              # TODO Try to trim padding when meridian is crossed and model was on a 0-360 coords
              # idlon, idlat = findn(.!isnan.(msk))
              # minXgrid = minimum(idlon)
              # maxXgrid = maximum(idlon)
              # minYgrid = minimum(idlat)
              # maxYgrid = maximum(idlat)
              #
              # msk = msk[minXgrid:maxXgrid, minYgrid:maxYgrid]
              # data = data[:, minXgrid:maxXgrid, minYgrid:maxYgrid]
              # lon_raw_flip = lon_raw[minXgrid:maxXgrid]
              #
              # longrid_flip = longrid[minXgrid:maxXgrid, minYgrid:maxYgrid]
              # longrid_flip = ClimateTools.shiftgrid_180_west_east(longrid_flip)
              # latgrid = latgrid[minXgrid:maxXgrid, minYgrid:maxYgrid]

              #     data = applymask(data, msk)

          else
            data = data_mask
          end

      else

          if rotatedgrid # flip in original grid
              longrid = shiftgrid_180_east_west(longrid) #grideast, gridwest)
          end

          longrid = longrid[minXgrid:maxXgrid, minYgrid:maxYgrid]
          latgrid = latgrid[minXgrid:maxXgrid, minYgrid:maxYgrid]

          if rotatedgrid

              longrid_flip = shiftgrid_180_west_east(longrid)

              lon_raw_flip = shiftvector_180_east_west(lon_raw)

              data = permute_west_east(data_mask, longrid)#idxwest, idxeast)
              msk = ClimateTools.permute_west_east(msk, longrid)

          else
            data = data_mask
          end
      end

    elseif isempty(poly) # no polygon clipping
        msk = Array{Float64}(ones((size(data_pointer, 1), size(data_pointer, 2))))
        data_ext = extractdata2D(data_pointer, msk)
        # replace_missing!(data_ext)
        data_ext = nomissing(data_ext, NaN)
        # data_ext = convert(data_ext, Float32)

      if rotatedgrid
          # Flip data "west-east"
          data = permute_west_east(data_ext, longrid)
      else
        data = data_ext
      end

    end

    if rotatedgrid
        longrid = longrid_flip
        lon_raw = lon_raw_flip
    end

    # Convert data to AxisArray
    dataOut = AxisArray(data, Axis{Symbol(lonname)}(lon_raw), Axis{Symbol(latname)}(lat_raw))

    C = ClimGrid(dataOut, longrid=longrid, latgrid=latgrid, msk=msk, grid_mapping=map_attrib, dimension_dict=dimension_dict, model=model, frequency=frequency, experiment=experiment, run=runsim, project=project, institute=institute, filename=file, dataunits=dataunits, latunits=latunits, lonunits=lonunits, variable=vari, typeofvar=vari, typeofcal="fixed", varattribs=varattrib, globalattribs=attribs)

    close(ds)

    return C

end

model_id(attrib::NCDatasets.Attributes) = get(attrib,"model_id", get(attrib, "parent_source_id", get(attrib,"model","N/A")))

experiment_id(attrib::NCDatasets.Attributes) = get(attrib,"experiment_id",get(attrib,"experiment","N/A"))

project_id(attrib::NCDatasets.Attributes) = get(attrib,"project_id", get(attrib, "mip_era", get(attrib,"project","N/A")))

institute_id(attrib::NCDatasets.Attributes) = get(attrib,"institute_id",get(attrib, "institution_id", get(attrib,"institute","N/A")))

frequency_var(attrib::NCDatasets.Attributes) = get(attrib,"frequency","N/A")

runsim_id(attrib::NCDatasets.Attributes) = get(attrib, "parent_experiment_rip", get(attrib,"driving_model_ensemble_member","N/A"))

"""
    getdim_lat(ds::NCDatasets.Dataset)

Returns the name of the "latitude" dimension and the status related to a regular grid. The latitude dimension is usually "latitude", "lat", "y", "yc", "rlat".
"""
function getdim_lat(ds::NCDatasets.Dataset)

    if sum(keys(ds.dim) .== "rlat") == 1
        return "rlat", true
    elseif sum(keys(ds.dim) .== "lat") == 1
        return "lat", false
    elseif sum(keys(ds.dim) .== "latitude") == 1
        return "latitude", false
    elseif sum(keys(ds.dim) .== "y") == 1
        return "y", true
    elseif sum(keys(ds.dim) .== "yc") == 1
        return "yc", true
    else
        error("Manually verify x/lat dimension name")
    end

end

"""
    getdim_lon(ds::NCDatasets.Dataset)

Returns the name of the "longitude" dimension and the status related to a regular grid. The longitude dimension is usually "longitue", "lon", "x", "xc", "rlon".
"""
function getdim_lon(ds::NCDatasets.Dataset)

    if sum(keys(ds.dim) .== "rlon") == 1
        return "rlon", true
    elseif sum(keys(ds.dim) .== "lon") == 1
        return "lon", false
    elseif sum(keys(ds.dim) .== "longitude") == 1
        return "longitude", false
    elseif sum(keys(ds.dim) .== "x") == 1
        return "x", false
    elseif sum(keys(ds.dim) .== "xc") == 1
        return "xc", false
    else
        error("Manually verify x/lat dimension name")
    end

end

"""
    latgridname(ds::NCDatasets.Dataset)

Returns the name of the latitude grid when datasets is not on a rectangular grid.
"""
function latgridname(ds::NCDatasets.Dataset)

    if in("lat", keys(ds))
        return "lat"
    elseif in("latitude", keys(ds))
        return "latitude"
    else
        error("Variable name is not supported. File an issue on https://github.com/Balinus/ClimateTools.jl/issues")
    end
end

"""
    longridname(ds::NCDatasets.Dataset)

Returns the name of the longitude grid when datasets is not on a rectangular grid.
"""
function longridname(ds::NCDatasets.Dataset)

    if in("lon", keys(ds))
        return "lon"
    elseif in("longitude", keys(ds))
        return "longitude"
    else
        error("Variable name is not supported. File an issue on https://github.com/Balinus/ClimateTools.jl/issues")
    end
end

"""
    extractdata(data, msk, idxtimebeg, idxtimeend)

Returns the data contained in netCDF file, using the appropriate mask and time index. Used internally by `load`.

"""
function extractdata(data, msk, idxtimebeg, idxtimeend)

    # idlon, idlat = findn(.!isnan.(msk))
    begin
        I = Base.findall(!isnan, msk)
        idlon, idlat = (getindex.(I, 1), getindex.(I, 2))
    end
    minXgrid = minimum(idlon)
    maxXgrid = maximum(idlon)
    minYgrid = minimum(idlat)
    maxYgrid = maximum(idlat)

    if ndims(data) == 3
        dataout = data[minXgrid:maxXgrid, minYgrid:maxYgrid, idxtimebeg:idxtimeend]
        # Permute dims
        # data = permutedims(data, [3, 1, 2])
    elseif ndims(data) == 4
        dataout = data[minXgrid:maxXgrid, minYgrid:maxYgrid, :, idxtimebeg:idxtimeend]
        # Permute dims
        # data = permutedims(data, [4, 1, 2, 3])
    end

    return dataout
end

function extractdata2D(data, msk)

    # idlon, idlat = findn(.!isnan.(msk))
    begin
        I = Base.findall(!isnan, msk)
        idlon, idlat = (getindex.(I, 1), getindex.(I, 2))
    end
    minXgrid = minimum(idlon)
    maxXgrid = maximum(idlon)
    minYgrid = minimum(idlat)
    maxYgrid = maximum(idlat)

    data = data[minXgrid:maxXgrid, minYgrid:maxYgrid]

    return data

end

"""
    get_mapping(ds::Array{String,1})

Returns the grid_mapping of Dataset *ds*
"""
function get_mapping(K::Array{String,1})

    if in("rotated_pole", K)
        return "rotated_pole"

    elseif in("lambert_conformal_conic", K)
        return "lambert_conformal_conic"

    elseif in("rotated_latitude_longitude", K)
        return "rotated_latitude_longitude"

    elseif in("rotated_mercator", K)
        return "rotated_mercator"

    elseif in("crs", K)
        return "crs"

    elseif in("polar_stereographic", K)
        return "polar_stereographic"

    else
        return "Regular_longitude_latitude"
    end

end

function build_grid_mapping(ds::NCDatasets.Dataset, grid_mapping::String)

    if ClimateTools.@isdefined grid_mapping
        # map_dim = varattrib[grid_mapping]
        map_attrib = Dict(ds[grid_mapping].attrib)
        map_attrib["grid_mapping"] = grid_mapping
    else
        error("File an issue on https://github.com/Balinus/ClimateTools.jl/issues to get the grid supported")
    end

    return map_attrib

end


text: println("Building BDisposal..")


text:
module aoc_18

export ×, ⨸

const newplus = '⨸'
const newtimes = '×'

const inferior_times = '⊕'

const opdict = Dict('+' => newplus, '*' => newtimes)

function read_input(filename, times = newtimes)
    lines = readlines(filename)
    replace.(replace.(lines,('+' => newplus,)),('*' => times,))
end

⨸(x,y) = x + y
×(x,y) = x * y
⊕(x,y) = x * y

part1() = sum(eval.(Meta.parse.( read_input("input.txt") )))
part2() = sum(eval.(Meta.parse.( read_input("input.txt", inferior_times) )))

end


text:
# f(x,y)=x'-y'

FXtransposeMYtranspose(x::Float64,y::Float64)=Fxmy(x,y)
FXtransposeMYtranspose(x::Float64,y)=Fxmy(x,y')
FXtransposeMYtranspose(x,y::Float64)=Fxmy(x',y)
FXtransposeMYtranspose(x,y)=Fxmy(x',y')
FXtransposeMYtranspose_inplace(value,auxvalue,x,y)=Fxmy_inplace(value,auxvalue,x',y')

function DXtransposeMYtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::Array,y::Array)
    if derivativeIDX==1
        if size(x)==(1,1)
            axpy!(1.0,[sum(grad_c)],grad_n)
        else
            axpy!(1.0,grad_c',grad_n)
        end
    elseif derivativeIDX==2
        if size(y)==(1,1)
            axpy!(-1.0,[sum(grad_c)],grad_n)
        else
            axpy!(-1.0,grad_c',grad_n)
        end
    end
end

if PROC=="GPU"

    function FXtransposeMYtranspose_inplace(value,auxvalue,x::CudaArray,y::CudaArray)
        if size(x)==(1,1)
            FXMYtranspose_inplace(value,auxvalue,x::CudaArray,y::CudaArray)
        elseif size(y)==(1,1)
            FXtransposeMY_inplace(value,auxvalue,x::CudaArray,y::CudaArray)
        else
            CUBLAS.geam!('T','T',1.0,x,-1.0,y,value)
        end
    end

    function DXtransposeMYtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray,y::CudaArray)
        if derivativeIDX==1
            DXtransposeMY(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray,y::CudaArray)
        elseif derivativeIDX==2
            DXMYtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray,y::CudaArray)
        end
    end

end

Derivative[FXtransposeMYtranspose]=DXtransposeMYtranspose
Inplace[FXtransposeMYtranspose]=FXtransposeMYtranspose_inplace

export FXtransposeMYtranspose

max_stars_repo_name: davidbarber/AutoDiff

max_issues_repo_name: davidbarber/AutoDiff

max_forks_repo_name: davidbarber/AutoDiff


text:
# syntax: proto3
using ProtoBuf
import ProtoBuf.meta

mutable struct TensorProto <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function TensorProto(; kwargs...)
        obj = new(meta(TensorProto), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
        end
        obj
    end
end # mutable struct TensorProto (has cyclic type dependency)
const __meta_TensorProto = Ref{ProtoMeta}()
function meta(::Type{TensorProto})
    ProtoBuf.metalock() do
        if !isassigned(__meta_TensorProto)
            __meta_TensorProto[] = target = ProtoMeta(TensorProto)
            fnum = Int[1,2,3,4,13,5,6,7,8,9,10,11,12,14,15,16,17]
            pack = Symbol[:half_val,:float_val,:double_val,:int_val,:scomplex_val,:int64_val,:bool_val,:dcomplex_val,:uint32_val,:uint64_val]
            allflds = Pair{Symbol,Union{Type,String}}[:dtype => Int32, :tensor_shape => TensorShapeProto, :version_number => Int32, :tensor_content => Vector{UInt8}, :half_val => Base.Vector{Int32}, :float_val => Base.Vector{Float32}, :double_val => Base.Vector{Float64}, :int_val => Base.Vector{Int32}, :string_val => Base.Vector{Vector{UInt8}}, :scomplex_val => Base.Vector{Float32}, :int64_val => Base.Vector{Int64}, :bool_val => Base.Vector{Bool}, :dcomplex_val => Base.Vector{Float64}, :resource_handle_val => Base.Vector{ResourceHandleProto}, :variant_val => "Base.Vector{VariantTensorDataProto}", :uint32_val => Base.Vector{UInt32}, :uint64_val => Base.Vector{UInt64}]
            meta(target, TensorProto, allflds, ProtoBuf.DEF_REQ, fnum, ProtoBuf.DEF_VAL, pack, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES)
        end
        __meta_TensorProto[]
    end
end
function Base.getproperty(obj::TensorProto, name::Symbol)
    if name === :dtype
        return (obj.__protobuf_jl_internal_values[name])::Int32
    elseif name === :tensor_shape
        return (obj.__protobuf_jl_internal_values[name])::TensorShapeProto
    elseif name === :version_number
        return (obj.__protobuf_jl_internal_values[name])::Int32
    elseif name === :tensor_content
        return (obj.__protobuf_jl_internal_values[name])::Vector{UInt8}
    elseif name === :half_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Int32}
    elseif name === :float_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Float32}
    elseif name === :double_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Float64}
    elseif name === :int_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Int32}
    elseif name === :string_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Vector{UInt8}}
    elseif name === :scomplex_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Float32}
    elseif name === :int64_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Int64}
    elseif name === :bool_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Bool}
    elseif name === :dcomplex_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{Float64}
    elseif name === :resource_handle_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{ResourceHandleProto}
    elseif name === :variant_val
        return (obj.__protobuf_jl_internal_values[name])::Any
    elseif name === :uint32_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{UInt32}
    elseif name === :uint64_val
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{UInt64}
    else
        getfield(obj, name)
    end
end

mutable struct VariantTensorDataProto <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function VariantTensorDataProto(; kwargs...)
        obj = new(meta(VariantTensorDataProto), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
        end
        obj
    end
end # mutable struct VariantTensorDataProto (has cyclic type dependency)
const __meta_VariantTensorDataProto = Ref{ProtoMeta}()
function meta(::Type{VariantTensorDataProto})
    ProtoBuf.metalock() do
        if !isassigned(__meta_VariantTensorDataProto)
            __meta_VariantTensorDataProto[] = target = ProtoMeta(VariantTensorDataProto)
            allflds = Pair{Symbol,Union{Type,String}}[:type_name => AbstractString, :metadata => Vector{UInt8}, :tensors => Base.Vector{TensorProto}]
            meta(target, VariantTensorDataProto, allflds, ProtoBuf.DEF_REQ, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES)
        end
        __meta_VariantTensorDataProto[]
    end
end
function Base.getproperty(obj::VariantTensorDataProto, name::Symbol)
    if name === :type_name
        return (obj.__protobuf_jl_internal_values[name])::AbstractString
    elseif name === :metadata
        return (obj.__protobuf_jl_internal_values[name])::Vector{UInt8}
    elseif name === :tensors
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{TensorProto}
    else
        getfield(obj, name)
    end
end

export TensorProto, VariantTensorDataProto, TensorProto, VariantTensorDataProto


text:
#############################################################################
#Copyright (c) 2016-2019 Charles Le Losq
#
#The MIT License (MIT)
#
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the #Software without restriction, including without limitation the rights to use, copy, #modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, #and to permit persons to whom the Software is furnished to do so, subject to the #following conditions:
#
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, #INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# functions.jl contains several mathematic functions
#
#
#############################################################################
"""
	poly(p::Vector{Float64},x::Array{Float64})

This function just allows to build a polynomial curve.

Inputs
------

	p: Vector{Float64}
		polynomial parameters. For a linear curve, p = [1.0,1.0], for a second order polynomial, p = [1.0,1.0,1.0], etc.;
	x: Array{Float64}
		x values for calculation.

Outputs
-------

	y: Array{Float64}, containing the result of calculation.
"""
function poly(p::Vector{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(p)[1])
    for i = 1:size(p)[1]
        segments[:,i] = p[i].*x[:].^(i-1)
    end
    return sum(segments, dims=2)
end

"""
	gaussiennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64};style::String = "None")

gaussiennes, written in the plural french form there, is a function that allows to build gaussian peaks. The gaussian function used there is:

    y = amplitude x exp(-ln(2) x [(x-centre)/hwhm]^2 )

You can enter the amplitude, centre and half-width at half-maximum (hwhm) values as arrays of float 64 (even containing one float value), without specifying style. hwhm is proportional to the standard deviation sigma:

	hwhm= sqrt(2xln(2)) x sigma

that is used in a normal distribution (see function normal_dist).

Inputs
------

	amplitude: Array{Float64}
		peaks amplitudes
	centre: Array{Float64}
		peaks centres
	hwhm: Array{Float64}
		peaks half-width at middle heights (hwhm);
	x: Array{Float64}
		x axis values;

Options
-------

	style: ASCIIString = "None"
		see examples below.

Outputs
-------

	y_calc: Array{Float64}
		calculated y values
	y_peaks: Array{Float64}
		calculated y values of the different peaks.

----------
 Examples
----------

To have four gaussian peaks centered at 800, 900, 1000 and 1100 cm-1 with hwhm of 50 cm-1 on a Raman spectrum, you will enter:

	```julia-repl
	julia> y_calc, y_peaks = gaussiennes([1.0,1.0,1.0,1.0], [800.0,900.0,1000.0,1100.0], [50.0,50.0,50.0,50.0], x)
    ```

and y_peaks will contain in 4 columns the 4 different y values of the peaks, and y_calc their sum (the total model). Now, if you want to calculate more complex models, such as for instance contructing how the Raman peaks of water vary with pressure, you might like to parametrize the variations of the peak parameters rather than just fitting each spectrum. This will provide more robust fits of the spectra, as you will fit them together, and will also force you to find the correct underlying mathematical assumption.

The gaussiennes function allows you to do that. If you specify style = "poly", you can enter arrays for the amplitudes, centres and half-widths at half-maximum (hwhm) of the peaks, with in each column the coefficients for the polynomial variations of this parameters. The second column of x will need to contain the second variable for those polynomial functions.

Let's say for instance that we have one peak at 900 cm-1 in a pure material. It's frequency seems to linearly shift with increasing the amount of hydrogen in this material, but it's intensity is non-linearly increasing, following a quadratic variation. It's width seems constant.

How to write that with gaussiennes? Well, first you need to construct a relevant x axis: first column contains the frequency, and the second one contains the chemical variable value. In our case, we want to model the peak between 800 and 1000 cm-1, for 1 wt% H. So we have an x array build like:

	```julia-repl
    julia> frequency = collect(800:1:1000)
    julia> x = ones(length(frequency),2)
    julia> x[:,1] = frequency[:]
    julia> x[:,2] = 1.0
	```

Ok, now lets build our y peaks:

	```julia-repl
    julia> amplitudes = [1.0 0.1 0.1]
    julia> frequencies = [900.0 2.0]
    julia> hwhm = 20.0
    julia> y_calc, y_peaks = gaussiennes(amplitudes, frequencies, hwhm, x)
	```

This should provide you how the shape of the peak is as a function of both the frequency and the chemical composition there. If you want to go further, you might just want to stick gaussiennes in a loop, and play with creating various peaks with changing the chemical parameter in the x[:,2] column!
"""
function gaussiennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64};style::String = "None")
    segments = zeros(size(x)[1],size(amplitude)[1])
    if style == "None"
        for i = 1:size(amplitude)[1]
            segments[:,i] = amplitude[i] .*exp.(-log.(2.0) .* ((x[:,1].-centre[i])./hwhm[i]).^2)
        end
    elseif style == "poly"
        for i = 1:size(amplitude)[1]
            segments[:,i] = poly(vec(amplitude[i,:]),x[:,2]) .*exp.(-log.(2) .* ((x[:,1].-(poly(vec(centre[i,:]),x[:,2])))./poly(vec(hwhm[i,:]),x[:,2])).^2)
        end
    else
        error("Not implemented, see documentation")
    end
    return sum(segments, dims=2), segments
end

"""
	lorentziennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64};style::String = "None")

Inputs
------

	amplitude: Array{Float64}
		peaks amplitudes
	centre: Array{Float64}
		peaks centres
	hwhm: Array{Float64}
		peaks half-width at middle heights (hwhm)
	x: Array{Float64}
		x axis values

Options
-------

	style: ASCIIString = "None", see examples in the gaussiennes documentation.

Outputs
-------

	y_calc: Array{Float64}
		calculated y values
	y_peaks: Array{Float64}
		calculated y values of the different peaks.

"""
function lorentziennes(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},x::Array{Float64};style::String = "None")
    segments = zeros(size(x)[1],size(amplitude)[1])
    if style == "None"
        for i = 1:size(amplitude)[1]
            segments[:,i] = amplitude[i] ./ (1.0 .+ ((x[:,1].-centre[i])./hwhm[i]).^2)
        end
    elseif style == "poly"
        for i = 1:size(amplitude)[1]
            segments[:,i] = poly(vec(amplitude[i,:]),x[:,2]) ./ (1.0 .+ ((x[:,1].-(poly(vec(centre[i,:]),x[:,2])))./poly(vec(hwhm[i,:]),x[:,2])).^2)
        end
    else
        error("Not implemented, see documentation")
    end
    return sum(segments, dims=2), segments
end

"""
	pearson7(a1::Array{Float64},a2::Array{Float64},a3::Array{Float64},a4::Array{Float64},x::Array{Float64};style::String = "None")

a Pearson 7 peak with formula a1 ./ (1 + ((x-a2)./a3).^2 .* (2.0.^(1./a4) - 1.0))

Inputs
------

	a1: Array{Float64}
		parameter a1
	a2: Array{Float64}
		parameter a2
	a3: Array{Float64}
		parameter a3
	a4: Array{Float64}
		parameter a4
	x: Array{Float64}
		x axis values

Options
-------

	style: ASCIIString = "None", see examples in the gaussiennes documentation.

Outputs
-------

	y_calc: Array{Float64}
		calculated y values
	y_peaks: Array{Float64}
		y values of the different peaks.

"""
function pearson7(a1::Array{Float64},a2::Array{Float64},a3::Array{Float64},a4::Array{Float64},x::Array{Float64};style::String = "None")
    segments = zeros(size(x)[1],size(a1)[1])
    if style == "None"
        for i = 1:size(a1)[1]
            segments[:,i] = a1[i] ./ (1.0 .+ ((x[:,1].-a2[i]) ./a3[i]) .^2 .* (2.0 .^(1.0 ./a4[i]) .- 1.0))
        end
    elseif style == "poly"
        for i = 1:size(a1)[1]
            segments[:,i] = poly(vec(a1[i,:]),x[:,2]) ./ (1. .+ ((x[:,1].-(poly(vec(a2[i,:]),x[:,2]))) ./poly(vec(a3[i,:]),x[:,2])) .^2 .* (2.0 .^(1,0 ./vec(a4[i,:])) .- 1.0))
        end
    else
        error("Not implemented, see documentation")
    end
    return sum(segments, dims=2), segments
end

"""
	pseudovoigts(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},lorentzian_fraction::Array{Float64},x::Array{Float64};style::String = "None")

A mixture of gaussian and lorentzian peaks.

Inputs
------
	amplitude: Array{Float64}
		peaks amplitudes
	centre: Array{Float64}
		peaks centres
	hwhm: Array{Float64}
		peaks half-width at middle heights (hwhm)
	lorentzian_fraction: Array{Float64}
		lorentzian fraction of the pseudovoigt function. Should be comprised between 0 and 1;
	x: Array{Float64}
		x axis values
Options
-------

	style: ASCIIString = "None", see examples in the gaussiennes documentation.

Outputs
-------
	y_calc: Array{Float64}
		calculated y values
	y_peaks: Array{Float64}
		y values of the different peaks

"""
function pseudovoigts(amplitude::Array{Float64},centre::Array{Float64},hwhm::Array{Float64},lorentzian_fraction::Array{Float64},x::Array{Float64};style::String = "None")
    segments = zeros(size(x)[1],size(amplitude)[1])
	test1 = findall(lorentzian_fraction .< 0.0)
	test2 = findall(lorentzian_fraction .> 1.0)
    if size(test1)[1] != 0 || size(test2)[1] != 0
		error("lorentzian_fraction should be comprised between 0 and 1")
	end
	if style == "None"
		for i = 1:size(amplitude)[1]
			segments[:,i] =  amplitude[i].*( (1.0 .- lorentzian_fraction[i]) .*exp.(-log(2) .* ((x[:,1].-centre[i])./hwhm[i]).^2)   .+ lorentzian_fraction[i] ./ (1.0 .+ ((x[:,1].-centre[i])./hwhm[i]).^2))
        end
    elseif style == "poly"
        for i = 1:size(amplitude)[1]
			segments[:,i] =  vec(lorentzian_fraction[i,:]) .* (poly(vec(amplitude[i,:]),x[:,2]) ./ (1 .+ ((x[:,1].-(poly(vec(centre[i,:]),x[:,2])))./poly(vec(hwhm[i,:]),x[:,2])).^2))     .+      (1.0 .- vec(lorentzian_fraction[i,:])) .* poly(vec(amplitude[i,:]),x[:,2]) .*exp.(-log(2) .* ((x[:,1].-(poly(vec(centre[i,:]),x[:,2])))./poly(vec(hwhm[i,:]),x[:,2])).^2)
        end
    else
        error("Not implemented, see documentation")
	end
    return sum(segments, dims=2), segments
end

"""
	normal_dist(nd_amplitudes::Array{Float64},nd_centres::Array{Float64},nd_sigmas::Array{Float64},x::Array{Float64})

The real normal distribution / gaussian function

Inputs
------

	nd_amplitudes: Array{Float64}
		distribution probability
	nd_centres: Array{Float64}
		distribution position
	nd_sigmas: Array{Float64}
		sigma value of the distribution
	x: Array{Float64}
		x axis values

Options
-------

	style: ASCIIString = "None", see examples in the gaussiennes documentation.

Outputs
-------

	y_calc: Array{Float64}
		calculated y values
	y_peaks: Array{Float64}
		y values of the different peaks

"""
function normal_dist(nd_amplitudes::Array{Float64},nd_centres::Array{Float64},nd_sigmas::Array{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(nd_amplitudes)[1])
    for i = 1:size(nd_amplitudes)[1]
        segments[:,i] = nd_amplitudes[i]./(nd_sigmas[i].*sqrt(2.0.*pi))  .*  exp.(- (x[:,1].-nd_centres[i]).^2 ./ (2.0.*nd_sigmas[i].^2.0))
    end
    return sum(segments, dims=2), segments
end

"""
	xshift_inversion(data::Array{Float64},p::Array{Float64})

to correct for shifts in X between two spectra, used in xshift_correction.

"""
function xshift_inversion(data::Array{Float64},p::Array{Float64})
    xaxis = data[:,1]
    shifted1 = data[:,2]
    spl = Spline1D(xaxis-p[1], shifted1.*p[2] + shifted1.^2.0.*p[3])
    y = evaluate(spl,xaxis)
end

"""
	xshift_direct(original_x::Array{Float64}, original_y::Array{Float64}, p::Float64)

To correct a spectrum for a p shift in X.

Used in xshift_correction.

Inputs
------

	original_x: Array{Float64}
		x values
	original_y: Array{Float64}
		y values associated with x
	p: Array{Float64}
		value of how much x should be shifted

Outputs
-------

	original_x: Array{Float64}
		same as input
	corrected_y: Array{Float64}
		the y values corrected from the p shift in original_x
	p: Array{Float64}
		same as input.

"""
function xshift_direct(original_x::Array{Float64}, original_y::Array{Float64}, p::Float64)
    spl = Spline1D(original_x-p, original_y)
    corrected_y = evaluate(spl,original_x)
	return original_x, corrected_y, p
end

"""
	xshift_correction(full_x::Array{Float64}, full_shifted_y::Array{Float64}, ref_x::Array{Float64}, ref_y::Array{Float64},shifted_y::Array{Float64})

To correct a shift between two spectra using a reference peak.

Inputs
------

	full_x: Array{Float64}
		x values that are not good
	full_shifted_y: Array{Float64}
		y values associated with full_x
	ref_x: Array{Float64}
		x values that are good
	ref_y: Array{Float64}
		y values associated with ref_x
	shifted_y: Array{Float64}
		y values associated with a selected range of full_x that corresponds to ref_x (for instance, a specific peak that you want to use to correct the shift).

Outputs
-------

	full_x: Array{Float64}
		same as input
	corrected_y: Array{Float64}
		the full_shifted_y values corrected from the shift
	p: Array{Float64}
		same as input.

ref_x is the common X axis of two particular ref_y and shifted_y signals, that should be for instance an intense and well defined peak in your spectra. If ref_y and shifted_y do not share the same X axis, you can use first the Dierckx spline to re-sample one of them and have both sharing a common X axis. See the examples for further details.
"""
function xshift_correction(full_x::Array{Float64}, full_shifted_y::Array{Float64}, ref_x::Array{Float64}, ref_y::Array{Float64},shifted_y::Array{Float64})
	fit = curve_fit(xshift_inversion, [ref_x shifted_y], ref_y, [1.0, 1.0,1.0])
	parameters = fit.param
	return xshift_direct(full_x, full_shifted_y, parameters[1])
end

"""
	smooth(x,y;method="whittaker", window_length=5, polyorder = 2, Lambda = 10.0.^5, d=2, ese_y=1.0)

smooth the provided y signal (sampled on x)

Inputs
----------
	x: vector
		Nx1 array of x values (equally spaced).
	y: vector
		Nx1 array of y values (equally spaced).
	method: str
		Method for smoothing the signal;
		choose between savgol (Savitzky-Golay), GCVSmoothedNSpline, MSESmoothedNSpline, DOFSmoothedNSpline, whittaker, 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'.

Options
-------
	window_length: int
		The length of the filter window (i.e. the number of coefficients). window_length must be a positive odd integer.
	polyorder: int
		The order of the polynomial used to fit the samples. polyorder must be less than window_length.
	Lambda: float
		smoothing parameter of the Whittaker filter described in Eilers (2003). The higher the smoother the fit.
	d: int
		d parameter in Whittaker filter, see Eilers (2003).
	ese_y: ndarray
		errors associated with y (for the gcvspline algorithms)

Outputs
-------
	y_smo: ndarray
		smoothed signal sampled on x.

Note
====

Use of GCVSmoothedNSpline, MSESmoothedNSpline, DOFSmoothedNSpline requires installation of gcvspline. See gcvspline documentation.
See also documentation for details on GCVSmoothedNSpline, MSESmoothedNSpline, DOFSmoothedNSpline.

savgol uses the scipy.signal.savgol_filter() function.

References
==========
Eilers, P.H.C., 2003. A Perfect Smoother. Anal. Chem. 75, 3631–3636. https://doi.org/10.1021/ac034173t

Scipy Cookbook: https://scipy-cookbook.readthedocs.io/items/SignalSmooth.html?highlight=smooth

"""
function smooth(x,y;method="whittaker", window_length=5, polyorder = 2, Lambda = 10.0.^5, d=2, ese_y=1.0)
	return rampy.smooth(vec(x),vec(y),method="whittaker", window_length=window_length, polyorder = polyorder, Lambda = Lambda, d=d, ese_y=ese_y)
end

"""

	flipsp(spectra::Array{Float64})

Flip an array along the row dimension (dim = 1) if the first column values are in decreasing order.

"""
function flipsp(spectra::Array{Float64})
    if spectra[end,1] < spectra[1,1]
        spectra = spectra[end:-1:1,:]
    end
	return spectra
end

"""

	resample(x::Array{Float64},y::Array{Float64},x_new::Array{Float64})

Resample a y signal associated with x, along the x_new values.

Uses the Dierckx spline library.

"""
function resample(x::Array{Float64},y::Array{Float64},x_new::Array{Float64})
    spl = Spline1D(x,y,bc="zero")
    return evaluate(spl, x_new)
end


text:
module test_spec
using VoronoiFVM
using SparseArrays
using Test

function test_addrows(m,n,mx)
    sa=sprand(m,n,0.5)
    da=Matrix(sa)
    SparseMatrixCSC(VoronoiFVM.addzrows(da,mx))==VoronoiFVM.addzrows(sa,mx)
    sa==da
end

function test()
    @test test_addrows(0,100,1)
    @test test_addrows(1,100,1)
    @test test_addrows(1,100,2)
    @test test_addrows(1,100,4)
    @test test_addrows(10,100,11)
    @test test_addrows(10,100,1)
    true
end

end


text:
for rule in subtypes(Rule)
  definition = "==(r1::$rule,r2::$rule) = ("
  for field in fieldnames(rule)
    definition *= "r1.$field==r2.$field && "
  end
  definition *= "true)"
  eval(Meta.parse(definition))
end

## GRAMMAR
==(g1::Grammar,g2::Grammar) = g1.rules==g2.rules
function diff(g1::Grammar,g2::Grammar)
  for (key,val) in g1.rules
    if g1.rules[key] != g2.rules[key]
      println("$key:")
      println("+ $(g1.rules[key])")
      println("- $(g2.rules[key])")
    end
  end
end