using DFWannier
T= Float32
x = WannierModel{T}("/home/ponet/Documents/PhD/GeTe/NSOC/paperxsf/","/home/ponet/Documents/PhD/GeTe/SOC/GeTe_bands.out",[[Atom(T[0.0,0.0,-0.0239129,-2*0.155854]...) for i=1:4]...,[Atom(T[0.0,0.0,5.5540692,2*0.318205]...) for i=1:4]...],true);
x2 = WannierModel{T}("/home/ponet/Documents/PhD/GeTe/NSOC/paperxsf/","/home/ponet/Documents/PhD/GeTe/SOC/GeTe_bands.out",[[Atom(T[0.0,0.0,-0.0239129,-0.155854]...) for i=1:4]...,[Atom(T[0.0,0.0,5.5540692,0.318205]...) for i=1:4]...]);
using Plots
@time test = calculate_eig_angmom_soc_bloch_gpu(x,90:0.1:110.);
@time test1 = calculate_eig_angmom_soc_bloch_gpu(x,90:0.1:110.);
@time test1 = calculate_eig_angmom_soc_bloch(x2,90:0.1:110.);
plot(plot(test1[8],:angmom2_x),plot(test[8],:angmom2_x))

@time benchmark = construct_bloch_sum(x.wfcs[1],x.kpoints[1]);

begin
test1 = construct_bloch_sum(x.wfcs[1],x.kpoints[1])
assert(Array(test1.values)==Array(benchmark.values))
end

test2 = calculate_eig_angmom_soc_bloch(x2,90:110.);
using Plots
test = construct_bloch_sum(x.wfcs[1],x.kpoints[1])
Array(test.values)
test = construct_bloch_sum(x2.wfcs[1],x2.kpoints[1])
T=Float64
test_wfc1 =  DFWannier.host2gpu(read_xsf_file("/home/ponet/Documents/PhD/GeTe/NSOC/paperxsf/wan_00003.xsf",Atom{T}(0.0,0.0,0.1,0.1),T))
test_wfc2=  DFWannier.host2gpu(read_xsf_file("/home/ponet/Documents/PhD/GeTe/NSOC/paperxsf/wan_00004.xsf",Atom{T}(0.0,0.0,0.1,0.1),T))
dims = size(test_wfc1.values)
cu_dims = CuArray(UInt32[dims...])
Lx = CuArray(zeros(Complex{T},dims))
Ly = CuArray(zeros(Complex{T},dims))
Lz = CuArray(zeros(Complex{T},dims))
n1 = CuArray(zeros(Complex{T},dims))
n2 = CuArray(zeros(Complex{T},dims))
grid = Array(test_wfc1.grid)
origin = grid[1,1,1]
a = grid[2,1,1] .- origin
b = grid[1,2,1] .- origin
c = grid[1,1,2] .- origin
V = CuArray(inv([[a...] [b...] [c...]]))
begin
  # indices = Array{Tuple{Tuple{Int32,Int32,Int32},Tuple{Int32,Int32,Int32}},1}()
  # coeff = Array{Complex{T},1}()
  indices = Array{Tuple{CuArray{Int32,1},CuArray{Int32,1}},1}()
  coeff = Array{CuArray{Complex{T},1},1}()
  k=T[0.2,0.2,0.3]
  for R1=-1:0,R2=-1:0,R3=-1:0
    if R1+R2+R3 == 0
      continue
    end
    R= R1*test_wfc1.cell[1]+R2*test_wfc1.cell[2]+R3*test_wfc1.cell[3]
    ind1,ind2 = DFWannier.find_start(test_wfc1,R,27)
    # push!(indices,((Int32(ind1[1]),Int32(ind1[2]),Int32(ind1[3])),(Int32(ind1[1]-ind2[1]),Int32(ind1[2]-ind2[2]),Int32(ind1[3]-ind2[3]))))
    push!(indices,(CuArray{Int32}([ind1...]),CuArray{Int32}([ind2...])))
    # push!(coeff,Complex{T}(exp(dot(-2*pi*k,[R1,R2,R3])*1im)))
    push!(coeff,CuArray(Complex{T}(exp(dot(-2*pi*k,[R1,R2,R3])*1im))))
  end
  dfprintln(length(indices),length(coeff))
  k_wfcs = Array{Wfc3D_gpu{T},1}()
  #optimize so we dont redo the zero index ...
  for wfc in [test_wfc1,test_wfc2]
    push!(k_wfcs,Wfc3D_gpu(wfc.grid,CuArray(zeros(Complex{T},size(wfc.values))),wfc.cell,wfc.atom))
  end
end
@time  for i =1:2000
  construct_bloch_sums([test_wfc1,test_wfc2],k_wfcs,k,CuArray(indices),CuArray(coeff))
end
calculate_angmom(test_wfc1,test_wfc2,V,CuArray([test_wfc1.atom.center.x,test_wfc1.atom.center.y,test_wfc1.atom.center.z]),dims,Lx,Ly,Lz,n2,n2)

using CUDAnative, CUDAdrv

function haversine_cpu(lat1::Float32, lon1::Float32, lat2::Float32, lon2::Float32, radius::Float32)
  c1 = cospi(lat1 / 180.0f0)
  c2 = cospi(lat2 / 180.0f0)
  dlat = lat2 - lat1
  dlon = lon2 - lon1
  d1 = sinpi(dlat / 360.0f0)
  d2 = sinpi(dlon / 360.0f0)
  t = d2 * d2 * c1 * c2
  a = d1 * d1 + t
  c = 2.0f0 * asin(min(1.0f0, sqrt(a)))
  return radius * c
end

function pairwise_dist_cpu(lat::Vector{Float32}, lon::Vector{Float32})
  # allocate
  n = length(lat)
  rowresult = Array{Float32}(n, n)

  # brute force fill in each cell
  for i in 1:n, j in 1:n
    @inbounds rowresult[i, j] = haversine_cpu(lat[i], lon[i], lat[j], lon[j] , 6372.8f0)
  end

  return rowresult
end

# from https://devblogs.nvidia.com/parallelforall/fast-great-circle-distance-calculation-cuda-c/
function haversine_gpu(lat1::Float32, lon1::Float32, lat2::Float32, lon2::Float32, radius::Float32)
  # XXX: need to prefix math intrinsics with CUDAnative
  c1 = CUDAnative.cospi(lat1 / 180.0f0)
  c2 = CUDAnative.cospi(lat2 / 180.0f0)
  dlat = lat2 - lat1
  dlon = lon2 - lon1
  d1 = CUDAnative.sinpi(dlat / 360.0f0)
  d2 = CUDAnative.sinpi(dlon / 360.0f0)
  t = d2 * d2 * c1 * c2
  a = d1 * d1 + t
  c = 2.0f0 * CUDAnative.asin(CUDAnative.min(1.0f0, CUDAnative.sqrt(a)))
  return radius * c
end

# pairwise distance calculation kernel
function pairwise_dist_kernel(lat::CuDeviceVector{Float32}, lon::CuDeviceVector{Float32},
  rowresult::CuDeviceMatrix{Float32}, n)
  i = (blockIdx().x-1) * blockDim().x + threadIdx().x
  j = (blockIdx().y-1) * blockDim().y + threadIdx().y

  if i <= n && j <= n
    # store to shared memory
    shmem = @cuDynamicSharedMem(Float32, 2*blockDim().x + 2*blockDim().y)
    if threadIdx().y == 1
      shmem[threadIdx().x] = lat[i]
      shmem[blockDim().x + threadIdx().x] = lon[i]
    end
    if threadIdx().x == 1
      shmem[2*blockDim().x + threadIdx().y] = lat[j]
      shmem[2*blockDim().x + blockDim().y + threadIdx().y] = lon[j]
    end
    sync_threads()

    # load from shared memory
    lat_i = shmem[threadIdx().x]
    lon_i = shmem[blockDim().x + threadIdx().x]
    lat_j = shmem[2*blockDim().x + threadIdx().y]
    lon_j = shmem[2*blockDim().x + blockDim().y + threadIdx().y]

    @inbounds rowresult[i, j] = haversine_gpu(lat_i, lon_i, lat_j, lon_j, 6372.8f0)
  end
end

function pairwise_dist_gpu(lat::Vector{Float32}, lon::Vector{Float32})
  # upload
  lat_gpu = CuArray(lat)
  lon_gpu = CuArray(lon)

  # allocate
  n = length(lat)
  rowresult_gpu = CuArray{Float32}(n, n)

  # calculate launch configuration
  # NOTE: we want our launch configuration to be as square as possible,
  #       because that minimizes shared memory usage
  ctx = CuCurrentContext()
  dev = device(ctx)
  total_threads = min(n, attribute(dev, CUDAdrv.MAX_THREADS_PER_BLOCK))
  threads_x = floor(Int, sqrt(total_threads))
  threads_y = total_threads ÷ threads_x
  threads = (threads_x, threads_y)
  blocks = ceil.(Int, n ./ threads)

  # calculate size of dynamic shared memory
  shmem = 2 * sum(threads) * sizeof(Float32)
  dfprintln(shmem)
  @cuda (blocks, threads, shmem) pairwise_dist_kernel(lat_gpu, lon_gpu, rowresult_gpu, n)

  return Array(rowresult_gpu)
end

# generate reasonable data
const n = 10000
const lat = rand(Float32, n) .* 45
const lon = rand(Float32, n) .* -120

@test pairwise_dist_cpu(lat, lon) ≈ pairwise_dist_gpu(lat, lon)


using Test
import Omega.Lang: AppExpr, LetExpr
import Omega.Lang: unparse

function test1()
  op = """
  μ : uniform 0 1
  x : normal μ 1
  rand x + μ
  """
  oe = LetExpr(:μ, AppExpr(:uniform, [0, 1]),
              LetExpr(:x, AppExpr(:uniform, [0, 1]),
              AppExpr(:rand, )
              ))

  @test unparse(oe) == op
end

function test2()
  expr = quote
  function f(w)
    x = sqrt(w) + 1
    y = sqrt(x) + 3
    if

  end
end


export pooling

const MAXPOOLING2D_F32 = Libdl.dlsym(libmerlin, :maxpooling2d_f32)
const AVGPOOLING2D_F32 = Libdl.dlsym(libmerlin, :avgpooling2d_f32)
const ∇MAXPOOLING2D_F32 = Libdl.dlsym(libmerlin, :maxpooling2d_f32_grad)
const ∇AVGPOOLING2D_F32 = Libdl.dlsym(libmerlin, :avgpooling2d_f32_grad)

maxpooling2d_handle(::Type{Float32}) = MAXPOOLING2D_F32
avgpooling2d_handle(::Type{Float32}) = AVGPOOLING2D_F32
∇maxpooling2d_handle(::Type{Float32}) = ∇MAXPOOLING2D_F32
∇avgpooling2d_handle(::Type{Float32}) = ∇AVGPOOLING2D_F32

"""
    pooling(mode, x::Var, dims::Tuple, [pads], [stride])

N-dimensional pooling function.

* mode: `:max` or `:average`
* dims: kernel size.
* [pads=(0,0,...)]: spatial padding size.
* [strides=(1,1,...)]: kernel strides.

```julia
x = Var(rand(Float32,5,4,3,2))
y = pooling(:max, x, (2,2), pads=(0,0), strides=(1,1))
```
"""
function pooling{N}(mode::Symbol, x::Var, dims::NTuple{N,Int}; pads=nothing, strides=nothing)
    pads == nothing && (pads = ntuple(_ -> 0, N))
    strides == nothing && (strides = ntuple(_ -> 1, N))
    if mode == :max
        f = maxpooling
    elseif mode == :average
        f = avgpooling
    else
        throw("Invalid mode is specified: $(mode).")
    end
    isvoid(x.data) && return Var(nothing, f, (x,dims,pads,strides))
    iscuda(x.data) && return CUDA.f(x, dims, pads, strides)
    f(x, dims, pads, strides)
end

function maxpooling(x::Var, dims::NTuple{2,Int}, pads, strides)
    T = eltype(x.data)
    h = maxpooling2d_handle(T)
    spdims = ntuple(i -> (size(x.data,i) + 2pads[i] - dims[i]) ÷ strides[i] + 1, 2)
    y = Array(T, spdims..., size(x,3), size(x,4))
    inds = Array(Cint, length(y))
    function f{T}(x::Array{T})
        ccall(h, Void, (Ptr{T},Ptr{T},Ptr{Cint},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
            x, y, inds, size(x,1), size(x,2), size(x,3)*size(x,4),
            dims[1], dims[2], pads[1], pads[2], strides[1], strides[2])
    end
    f(x.data)

    function df{T}(gy::Array{T})
        isvoid(x.grad) && return
        ccall(∇maxpooling2d_handle(T), Void, (Ptr{T},Ptr{T},Ptr{Cint},Cint),
            gy, x.grad, inds, length(gy))
    end
    Var(y, df, (x,))
end

function avgpooling(x::Var, dims::NTuple{2,Int}, pads, strides)
    T = eltype(x.data)
    h = avgpooling2d_handle(T)
    spdims = ntuple(i -> (size(x.data,i) + 2pads[i] - dims[i]) ÷ strides[i] + 1, 2)
    y = Array(T, spdims..., size(x,3), size(x,4))
    function f{T}(x::Array{T})
        ccall(h, Void, (Ptr{T},Ptr{T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
            x, y, size(x,1), size(x,2), size(x,3)*size(x,4),
            dims[1], dims[2], pads[1], pads[2], strides[1], strides[2])
    end
    f(x.data)

    function df{T}(gy::Array{T})
        isvoid(x.grad) && return
        gx = x.grad
        h = ∇avgpooling2d_handle(T)
        ccall(h, Void, (Ptr{T},Ptr{T},Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint,Cint),
            gy, gx, size(gx,1), size(gx,2), size(gx,3)*size(gx,4),
            dims[1], dims[2], pads[1], pads[2], strides[1], strides[2])
    end
    Var(y, df, (x,))
end


@testset "Counting" begin

    @testset "GC content" begin
        @test gc_content(dna"") === 0.0
        @test gc_content(dna"AATA") === 0.0
        @test gc_content(dna"ACGT") === 0.5
        @test gc_content(dna"CGGC") === 1.0
        @test gc_content(dna"ACATTGTGTATAACAAAAGG") === 6 / 20
        @test gc_content(dna"GAGGCGTTTATCATC"[2:end]) === 6 / 14

        @test gc_content(mer"AATA"dna) === 0.0
        @test gc_content(mer"ACGT"dna) === 0.5
        @test gc_content(mer"CGGC"dna) === 1.0
        @test gc_content(mer"ACATTGTGTATAACAAAAGG"dna) === 6 / 20

        @test gc_content(rna"") === 0.0
        @test gc_content(rna"AAUA") === 0.0
        @test gc_content(rna"ACGU") === 0.5
        @test gc_content(rna"CGGC") === 1.0
        @test gc_content(rna"ACAUUGUGUAUAACAAAAGG") === 6 / 20
        @test gc_content(rna"GAGGCGUUUAUCAUC"[2:end]) === 6 / 14

        @test gc_content(mer"AAUA"rna) === 0.0
        @test gc_content(mer"ACGU"rna) === 0.5
        @test gc_content(mer"CGGC"rna) === 1.0
        @test gc_content(mer"ACAUUGUGUAUAACAAAAGG"rna) === 6 / 20

        @test_throws Exception gc_content(aa"ARN")

        Random.seed!(1234)
        for _ in 1:200
            s = randdnaseq(rand(1:100))
            @test gc_content(s) === count(isGC, s) / length(s)
            @test gc_content(LongSequence{DNAAlphabet{2}}(s)) === count(isGC, s) / length(s)
            i = rand(1:lastindex(s))
            j = rand(i-1:lastindex(s))
            @test gc_content(s[i:j]) === (j < i ? 0.0 : count(isGC, s[i:j]) / (j - i + 1))
        end
    end

    function testcounter(pred::Function, alias::Function, seqa::BioSequence, seqb::BioSequence)
        # Test that order does not matter.
        @test count(pred, seqa, seqb) == count(pred, seqb, seqa)
        @test BioSequences.count_naive(pred, seqa, seqb) == BioSequences.count_naive(pred, seqb, seqa)
        @test alias(seqa, seqb) == alias(seqb, seqa)
        # Test that result is the same as counting naively.
        @test count(pred, seqa, seqb) == BioSequences.count_naive(pred, seqa, seqb)
        @test count(pred, seqb, seqa) == BioSequences.count_naive(pred, seqb, seqa)
        # Test that the alias function works.
        @test count(pred, seqa, seqb) == alias(seqa, seqb)
        @test count(pred, seqb, seqa) == alias(seqb, seqa)
    end

    function counter_random_tests(pred::Function, alias::Function, alphx::Type{<:Alphabet}, alphy::Type{<:Alphabet}, subset::Bool)
        for _ in 1:50
            seqA = random_seq(alphx, rand(10:100))
            seqB = random_seq(alphy, rand(10:100))
            sa = seqA
            sb = seqB
            if subset
                intA = random_interval(1, length(seqA))
                intB = random_interval(1, length(seqB))
                subA = seqA[intA]
                subB = seqB[intB]
                sa = subA
                sb = subB
            end
            testcounter(pred, alias, sa, sb)
        end
    end

    @testset "Mismatches" begin
        for a in (DNAAlphabet, RNAAlphabet)
            for sub in (true, false)
                for n in (4, 2)
                    counter_random_tests(!=, mismatches, a{n}, a{n}, sub)
                end
                counter_random_tests(!=, mismatches, a{4}, a{2}, sub)
                counter_random_tests(!=, mismatches, a{2}, a{4}, sub)
            end
        end
    end

    @testset "Matches" begin
        for a in (DNAAlphabet, RNAAlphabet)
            for sub in (true, false)
                for n in (4, 2)
                    counter_random_tests(==, matches, a{n}, a{n}, sub)
                end
                counter_random_tests(==, matches, a{4}, a{2}, sub)
                counter_random_tests(==, matches, a{2}, a{4}, sub)
            end
        end
    end

    @testset "Ambiguous" begin
        for a in (DNAAlphabet, RNAAlphabet)
            for sub in (true, false)
                for n in (4, 2)
                    counter_random_tests(isambiguous, n_ambiguous, a{n}, a{n}, sub)
                end
                counter_random_tests(isambiguous, n_ambiguous, a{4}, a{2}, sub)
                counter_random_tests(isambiguous, n_ambiguous, a{2}, a{4}, sub)
            end
        end
    end

    @testset "Certain" begin
        for a in (DNAAlphabet, RNAAlphabet)
            for sub in (true, false)
                for n in (4, 2)
                    counter_random_tests(iscertain, n_certain, a{n}, a{n}, sub)
                end
                counter_random_tests(iscertain, n_certain, a{4}, a{2}, sub)
                counter_random_tests(iscertain, n_certain, a{2}, a{4}, sub)
            end
        end
    end

    @testset "Gap" begin
        for a in (DNAAlphabet, RNAAlphabet)
            for sub in (true, false)
                for n in (4, 2)
                    counter_random_tests(isgap, n_gaps, a{n}, a{n}, sub)
                end
                counter_random_tests(isgap, n_gaps, a{4}, a{2}, sub)
                counter_random_tests(isgap, n_gaps, a{2}, a{4}, sub)
            end
        end
    end

end


# GaussDCA
# ========

"""
Wrapper function to `GaussDCA.gDCA`.
You need to install GaussDCA:

```julia
using Pkg
Pkg.add(PackageSpec(url="https://github.com/carlobaldassi/GaussDCA.jl"))
```
Look into [GaussDCA.jl README](https://github.com/carlobaldassi/GaussDCA.jl) for further information.
If you use this wrapper, **please cite the GaussDCA publication and the package's doi**.

It's possible to indicate the path to the julia binary where GaussDCA is installed.
However, it's recommended to use the same version where MIToS is installed. That is
because this function use `serialize`/`deserialize` to transfer data between the processes.

**GaussDCA Publication:**
Baldassi, Carlo, Marco Zamparo, Christoph Feinauer, Andrea Procaccini, Riccardo Zecchina, Martin Weigt, and Andrea Pagnani.
"Fast and accurate multivariate Gaussian modeling of protein families: predicting residue contacts and protein-interaction partners."
PloS one 9, no. 3 (2014): e92721.
"""
function gaussdca(msa; juliapath::String=joinpath(Sys.BINDIR,Base.julia_exename()), kargs...)
    base_name = tempname()
    if Sys.iswindows()
        base_name = escape_string(base_name)
        juliapath = escape_string(juliapath)
        if !endswith(juliapath, ".exe")
            juliapath = juliapath * ".exe"
        end
    end
    script_file = base_name * ".jl"
    msa_file = base_name * ".fasta"
    jdl_file = base_name * ".jls"
    plm = fill!(columnpairsmatrix(msa), NaN)
    write(msa_file, msa, FASTA)
    try
        _create_script(script_file, msa_file, jdl_file; kargs...)
        run(`$juliapath $script_file`)
        pairedvalues = open(deserialize, jdl_file, "r")
        for (i,j,value) in pairedvalues
           plm[i,j] = value
        end
    finally
        isfile(script_file) && rm(script_file)
        isfile(msa_file) && rm(msa_file)
        isfile(jdl_file) && rm(jdl_file)
    end
    plm
end

function _create_script(script_file::String, msa_file::String, jdl_file::String; kargs...)
    if length(kargs) > 0
        for (k,v) in kargs
            @assert isa(v,Number) || isa(v,Symbol) "Argument values must be numbers or symbols"
        end
        str_kargs = "," * join([ isa(v,Symbol) ? "$k=:$v" : "$k=$v" for (k,v) in kargs],',')
    else
        str_kargs = ""
    end
    open(script_file, "w") do fh
        println(fh, "using Serialization;")
        println(fh, "using GaussDCA;")
        println(fh, "values = GaussDCA.gDCA(\"$msa_file\"$str_kargs);")
        println(fh, "open(fh -> serialize(fh, values), \"$(jdl_file)\", \"w\")")
    end
end


############################################################################################
# Define the basic type system
############################################################################################
"Most abstract supertype for all SEMs"
abstract type AbstractSem end

"Supertype for all single SEMs, e.g. SEMs that have at least the fields `observed`, `imply`, `loss` and `optimizer`"
abstract type AbstractSemSingle{O, I, L, D} <: AbstractSem end

"Supertype for all collections of multiple SEMs"
abstract type AbstractSemCollection <: AbstractSem end

"Supertype for all loss functions of SEMs. If you want to implement a custom loss function, it should be a subtype of `SemLossFunction`."
abstract type SemLossFunction end

"""
    SemLoss(args...; loss_weights = nothing, ...)

Constructs the loss field of a SEM. Can contain multiple `SemLossFunction`s, the model is optimized over their sum.
See also [`SemLossFunction`](@ref).

# Arguments
- `args...`: Multiple `SemLossFunction`s.
- `loss_weights::Vector`: Weights for each loss function. Defaults to unweighted optimization.

# Examples
```julia
my_ml_loss = SemML(...)
my_ridge_loss = SemRidge(...)
my_loss = SemLoss(SemML, SemRidge; loss_weights = [1.0, 2.0])
```
"""
mutable struct SemLoss{F <: Tuple, T}
    functions::F
    weights::T
end

function SemLoss(functions...; loss_weights = nothing, kwargs...)

    if !isnothing(loss_weights)
        loss_weights = SemWeight.(loss_weights)
    else
        loss_weights = Tuple(SemWeight(nothing) for _ in 1:length(functions))
    end

    return SemLoss(
        functions,
        loss_weights
        )
end

# weights for loss functions or models. If the weight is nothing, multiplication returs second argument
struct SemWeight{T}
    w::T
end

Base.:*(x::SemWeight{Nothing}, y) = y
Base.:*(x::SemWeight, y) = x.w*y

"""
Supertype of all objects that can serve as the `optimizer` field of a SEM.
Connects the SEM to its optimization backend and controls options like the optimization algorithm.
If you want to connect the SEM package to a new optimization backend, you should implement a subtype of SemOptimizer.
"""
abstract type SemOptimizer end

"""
Supertype of all objects that can serve as the observed field of a SEM.
Pre-processes data and computes sufficient statistics for example.
If you have a special kind of data, e.g. ordinal data, you should implement a subtype of SemObserved.
"""
abstract type SemObserved end

"""
Supertype of all objects that can serve as the imply field of a SEM.
Computed model-implied values that should be compared with the observed data to find parameter estimates,
e. g. the model implied covariance or mean.
If you would like to implement a different notation, e.g. LISREL, you should implement a subtype of SemImply.
"""
abstract type SemImply end

"Subtype of SemImply for all objects that can serve as the imply field of a SEM and use some form of symbolic precomputation."
abstract type SemImplySymbolic <: SemImply end

"""
    Sem(;observed = SemObservedData, imply = RAM, loss = SemML, optimizer = SemOptimizerOptim, kwargs...)

Constructor for the basic `Sem` type.
All additional kwargs are passed down to the constructors for the observed, imply, loss and optimizer fields.

# Arguments
- `observed`: object of subtype `SemObserved` or a constructor.
- `imply`: object of subtype `SemImply` or a constructor.
- `loss`: object of subtype `SemLossFunction`s or constructor; or a tuple of such.
- `optimizer`: object of subtype `SemOptimizer` or a constructor.

Returns a Sem with fields
- `observed::SemObserved`: Stores observed data, sample statistics, etc. See also [`SemObserved`](@ref).
- `imply::SemImply`: Computes model implied statistics, like Σ, μ, etc. See also [`SemImply`](@ref).
- `loss::SemLoss`: Computes the objective and gradient of a sum of loss functions. See also [`SemLoss`](@ref).
- `optimizer::SemOptimizer`: Connects the model to the optimizer. See also [`SemOptimizer`](@ref).
"""
mutable struct Sem{O <: SemObserved, I <: SemImply, L <: SemLoss, D <: SemOptimizer} <: AbstractSemSingle{O, I, L, D}
    observed::O
    imply::I
    loss::L
    optimizer::D
end

############################################################################################
# automatic differentiation
############################################################################################
"""
    SemFiniteDiff(;observed = SemObservedData, imply = RAM, loss = SemML, optimizer = SemOptimizerOptim, kwargs...)

Constructor for `SemFiniteDiff`.
All additional kwargs are passed down to the constructors for the observed, imply, loss and optimizer fields.

# Arguments
- `observed`: object of subtype `SemObserved` or a constructor.
- `imply`: object of subtype `SemImply` or a constructor.
- `loss`: object of subtype `SemLossFunction`s or constructor; or a tuple of such.
- `optimizer`: object of subtype `SemOptimizer` or a constructor.

Returns a Sem with fields
- `observed::SemObserved`: Stores observed data, sample statistics, etc. See also [`SemObserved`](@ref).
- `imply::SemImply`: Computes model implied statistics, like Σ, μ, etc. See also [`SemImply`](@ref).
- `loss::SemLoss`: Computes the objective and gradient of a sum of loss functions. See also [`SemLoss`](@ref).
- `optimizer::SemOptimizer`: Connects the model to the optimizer. See also [`SemOptimizer`](@ref).
"""
struct SemFiniteDiff{O <: SemObserved, I <: SemImply, L <: SemLoss, D <: SemOptimizer} <: AbstractSemSingle{O, I, L, D}
    observed::O
    imply::I
    loss::L
    optimizer::D
end

############################################################################################
# ensemble models
############################################################################################
"""
    SemEnsemble(models..., optimizer = SemOptimizerOptim, weights = nothing, kwargs...)

Constructor for ensemble models.

# Arguments
- `models...`: `AbstractSem`s.
- `optimizer`: object of subtype `SemOptimizer` or a constructor.
- `weights::Vector`:  Weights for each model. Defaults to the number of observed data points.

All additional kwargs are passed down to the constructor for the optimizer field.

Returns a SemEnsemble with fields
- `n::Int`: Number of models.
- `sems::Tuple`: `AbstractSem`s.
- `weights::Vector`: Weights for each model.
- `optimizer::SemOptimizer`: Connects the model to the optimizer. See also [`SemOptimizer`](@ref).
- `identifier::Dict`: Stores parameter labels and their position.
"""
struct SemEnsemble{N, T <: Tuple, V <: AbstractVector, D, I} <: AbstractSemCollection
    n::N
    sems::T
    weights::V
    optimizer::D
    identifier::I
end

function SemEnsemble(models...; optimizer = SemOptimizerOptim, weights = nothing, kwargs...)
    n = length(models)
    npar = n_par(models[1])

    # default weights

    if isnothing(weights)
        nobs_total = sum(n_obs.(models))
        weights = [n_obs(model)/nobs_total for model in models]
    end

    # check identifier equality
    id = identifier(models[1])
    for model in models
        if id != identifier(model)
            throw(ErrorException("The identifier of your models do not match. \n
            Maybe you tried to specify models of an ensemble via ParameterTables. \n
            In that case, you may use RAMMatrices instead."))
        end
    end

    # optimizer
    if !isa(optimizer, SemOptimizer)
        optimizer = optimizer(;kwargs...)
    end

    return SemEnsemble(
        n,
        models,
        weights,
        optimizer,
        id
        )
end

"""
    n_models(ensemble::SemEnsemble) -> Integer

Returns the number of models in an ensemble model.
"""
n_models(ensemble::SemEnsemble) = ensemble.n
"""
    models(ensemble::SemEnsemble) -> Tuple{AbstractSem}

Returns the models in an ensemble model.
"""
models(ensemble::SemEnsemble) = ensemble.sems
"""
    weights(ensemble::SemEnsemble) -> Vector

Returns the weights of an ensemble model.
"""
weights(ensemble::SemEnsemble) = ensemble.weights
"""
    optimizer(ensemble::SemEnsemble) -> SemOptimizer

Returns the optimizer part of an ensemble model.
"""
optimizer(ensemble::SemEnsemble) = ensemble.optimizer

############################################################################################
# additional methods
############################################################################################
"""
    observed(model::AbstractSemSingle) -> SemObserved

Returns the observed part of a model.
"""
observed(model::AbstractSemSingle) = model.observed

"""
    imply(model::AbstractSemSingle) -> SemImply

Returns the imply part of a model.
"""
imply(model::AbstractSemSingle) = model.imply

"""
    loss(model::AbstractSemSingle) -> SemLoss

Returns the loss part of a model.
"""
loss(model::AbstractSemSingle) = model.loss

"""
    optimizer(model::AbstractSemSingle) -> SemOptimizer

Returns the optimizer part of a model.
"""
optimizer(model::AbstractSemSingle) = model.optimizer


"""
    Reduced_lattice

Struct containing symmetry irreducible sites of a lattice graph.
* `name     :: String`                  : name of the original lattice
* `size     :: Int64`                   : bond truncation of the original lattice
* `model    :: String`                  : name of the initialized model
* `J        :: Vector{Vector{Float64}}` : coupling vector of the initialized model
* `sites    :: Vector{Site}`            : list of symmetry irreducible sites
* `overlap  :: Vector{Matrix{Int64}}`   : pairs of irreducible sites with their respective multiplicity in range of origin and another irreducible site
* `mult     :: Vector{Int64}`           : multiplicities of irreducible sites
* `exchange :: Vector{Int64}`           : images of the pair (origin, irreducible site) under site exchange
* `project  :: Matrix{Int64}`           : projections of pairs (site1, site2) of the original lattice to pair (origin, irreducible site)
"""
struct Reduced_lattice
    name     :: String
    size     :: Int64
    model    :: String
    J        :: Vector{Vector{Float64}}
    sites    :: Vector{Site}
    overlap  :: Vector{Matrix{Int64}}
    mult     :: Vector{Int64}
    exchange :: Vector{Int64}
    project  :: Matrix{Int64}
end

# check if matrix in list of matrices within numerical tolerance
function is_in(
    e    :: SMatrix{3, 3, Float64},
    list :: Vector{SMatrix{3, 3, Float64}}
    )    :: Bool

    in = false

    for item in list
        if maximum(abs.(item .- e)) <= 1e-8
            in = true
            break
        end
    end

    return in
end

# rotate vector onto a reference vector using Rodrigues formula
function get_rotation(
    vec :: SVector{3, Float64},
    ref :: SVector{3, Float64}
    )   :: SMatrix{3, 3, Float64}

    # buffer geometric information
    a = vec ./ norm(vec)
    b = ref ./ norm(ref)
    n = cross(a, b)
    s = norm(n)
    c = dot(a, b)

    # check if vectors are collinear
    if s < 1e-8
        # if vector are antiparallel do inversion
        if c < 0.0
            return SMatrix{3, 3, Float64}(-1.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, -1.0)
        # if vectors are parallel return unity
        else
            return SMatrix{3, 3, Float64}(+1.0, 0.0, 0.0, 0.0, +1.0, 0.0, 0.0, 0.0, +1.0)
        end
    # if vectors are non-collinear use Rodrigues formula
    else
        n    = n ./ s
        temp = SMatrix{3, 3, Float64}(0.0, n[3], -n[2], -n[3], 0.0, n[1], n[2], -n[1], 0.0)
        mat  = SMatrix{3, 3, Float64}(+1.0, 0.0, 0.0, 0.0, +1.0, 0.0, 0.0, 0.0, +1.0)
        mat  = mat .+ s .* temp .+ (1.0 - c) .* temp * temp

        return mat
    end
end

# try to rotate vector onto a reference vector around a given axis (return matrix of zeros in case of failure)
function get_rotation_for_axis(
    vec  :: SVector{3, Float64},
    ref  :: SVector{3, Float64},
    axis :: SVector{3, Float64}
    )    :: SMatrix{3, 3, Float64}

    # normalize vectors and axis
    a  = vec ./ norm(vec)
    b  = ref ./ norm(ref)
    ax = axis ./ norm(axis)

    # determine othogonal projections onto axis
    ovec = vec .- dot(vec, ax) .* ax
    ovec = ovec ./ norm(ovec)
    oref = ref .- dot(ref, ax) .* ax
    oref = oref ./ norm(oref)

    # buffer geometric information
    s = norm(cross(ovec, oref))
    c = dot(ovec, oref)

    # allocate rotation matrix
    mat = SMatrix{3, 3, Float64}(ax[1]^2 + (1.0 - ax[1]^2) * c,
                                 ax[1] * ax[2] * (1.0 - c) + ax[3] * s,
                                 ax[1] * ax[3] * (1.0 - c) - ax[2] * s,
                                 ax[1] * ax[2] * (1.0 - c) - ax[3] * s,
                                 ax[2]^2 + (1.0 - ax[2]^2) * c,
                                 ax[2] * ax[3] * (1.0 - c) + ax[1] * s,
                                 ax[1] * ax[3] * (1.0 - c) + ax[2] * s,
                                 ax[2] * ax[3] * (1.0 - c) - ax[1] * s,
                                 ax[3]^2 + (1.0 - ax[3]^2) * c)

    # check if rotation works as expected
    if norm(mat * a .- b) > 1e-8
        # check if sense of rotation has to be inverted
        if norm(transpose(mat) * a .- b) < 1e-8
            return transpose(mat)
        # if algorithm fails return matrix of zeros
        else
            return SMatrix{3, 3, Float64}(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
        end
    # if sanity check is passed return rotation matrix
    else
        return mat
    end
end

# try to obtain a rotation of (vec1, vec2) onto (ref1, ref2) (return matrix of zeros in case of failure)
function rotate(
    vec1 :: SVector{3, Float64},
    vec2 :: SVector{3, Float64},
    ref1 :: SVector{3, Float64},
    ref2 :: SVector{3, Float64}
    )    :: SMatrix{3, 3, Float64}

    # rotate vec1 onto ref1
    mat = get_rotation(vec1, ref1)

    # try to rotate vec2 around ref1 onto ref2
    mat = get_rotation_for_axis(mat * vec2, ref2, ref1) * mat

    return mat
end

"""
    get_trafos_orig(
        l :: Lattice
        ) :: Vector{SMatrix{3, 3, Float64}}

Compute transformations which leave the origin of the lattice invariant (point group symmetries).
"""
function get_trafos_orig(
    l :: Lattice
    ) :: Vector{SMatrix{3, 3, Float64}}

    # allocate list for transformations, set reference site and its bonds
    trafos = SMatrix{3, 3, Float64}[]
    ref    = l.sites[1]
    con    = l.uc.bonds[1]

    # get a pair of non-collinear neighbors of origin
    for i1 in eachindex(con)
        for i2 in eachindex(con)
            # check that sites are unequal
            if i1 == i2
                continue
            end

            # get connected sites in Bravais representation
            int_i1 = ref.int .+ con[i1]
            int_i2 = ref.int .+ con[i2]

            # get connected sites in real space representation
            vec_i1 = get_vec(int_i1, l.uc)
            vec_i2 = get_vec(int_i2, l.uc)

            # check that sites are non-collinear
            if norm(cross(vec_i1, vec_i2)) <= 1e-8
                continue
            end

            # get site structs
            site_i1 = Site(int_i1, vec_i1)
            site_i2 = Site(int_i2, vec_i2)

            # get bonds
            bond_i1 = get_bond(ref, site_i1, l)
            bond_i2 = get_bond(ref, site_i2, l)

            # get another pair of non-collinear neighbors of origin
            for j1 in eachindex(con)
                for j2 in eachindex(con)
                    # check that sites are unequal
                    if j1 == j2
                        continue
                    end

                    # get connected sites in Bravais representation
                    int_j1 = ref.int .+ con[j1]
                    int_j2 = ref.int .+ con[j2]

                    # get connected sites in real space representation
                    vec_j1 = get_vec(int_j1, l.uc)
                    vec_j2 = get_vec(int_j2, l.uc)

                    # check that sites are non-collinear
                    if norm(cross(vec_j1, vec_j2)) <= 1e-8
                        continue
                    end

                    # get site structs
                    site_j1 = Site(int_j1, vec_j1)
                    site_j2 = Site(int_j2, vec_j2)

                    # get bonds
                    bond_j1 = get_bond(ref, site_j1, l)
                    bond_j2 = get_bond(ref, site_j2, l)

                    # try to obtain rotation
                    mat = rotate(site_i1.vec, site_i2.vec, site_j1.vec, site_j2.vec)

                    # if successful, verify rotation on test set before saving it
                    if maximum(abs.(mat)) > 1e-8
                        # check if rotation is already known
                        if is_in(mat, trafos) == false
                            # verify rotation
                            valid = true

                            for n in eachindex(l.test_sites)
                                # apply transformation to lattice site
                                mapped_vec = mat * l.test_sites[n].vec
                                orig_bond  = get_bond(ref, l.test_sites[n], l)
                                mapped_ind = get_site(mapped_vec, l)

                                # check if resulting site is in lattice and if the bonds match
                                if mapped_ind == 0
                                    valid = false
                                else
                                    valid = are_equal(orig_bond, get_bond(ref, l.sites[mapped_ind], l))
                                end

                                # break if test fails for one test site
                                if valid == false
                                    break
                                end
                            end

                            # save transformation
                            if valid
                                push!(trafos, mat)
                            end
                        end

                        # check if rotation combined with inversion is already known
                        if is_in(-mat, trafos) == false
                            # verify rotation
                            valid = true

                            for n in eachindex(l.test_sites)
                                # apply transformation to lattice site
                                mapped_vec = -mat * l.test_sites[n].vec
                                orig_bond  = get_bond(ref, l.test_sites[n], l)
                                mapped_ind = get_site(mapped_vec, l)

                                # check if resulting site is in lattice and if the bonds match
                                if mapped_ind == 0
                                    valid = false
                                else
                                    valid = are_equal(orig_bond, get_bond(ref, l.sites[mapped_ind], l))
                                end

                                # break if test fails for one test site
                                if valid == false
                                    break
                                end
                            end

                            # save transformation
                            if valid
                                push!(trafos, -mat)
                            end
                        end
                    end
                end
            end
        end
    end

    return trafos
end

# compute reduced representation of the lattice
function get_reduced(
    l :: Lattice
    ) :: Vector{Int64}

    # allocate a list of indices
    reduced = Int64[i for i in eachindex(l.sites)]

    # get transformations
    trafos = get_trafos_orig(l)

    # iterate over sites and try to find sites which are symmetry equivalent
    for i in 2 : length(reduced)
        # continue if index has already been replaced
        if reduced[i] < i
            continue
        end

        # apply all available transformations to current site and see where it is mapped
        for j in eachindex(trafos)
            mapped_vec = trafos[j] * l.sites[i].vec

            # check that site is not mapped to itself
            if norm(mapped_vec .- l.sites[i].vec) > 1e-8
                # determine image of site
                index = get_site(mapped_vec, l)

                # assert that the image is valid, otherwise this is not a valid transformation and our algorithm failed
                @assert index > 0 "Validity on test set could not be generalized."

                # replace site in list if bonds match
                if are_equal(l.bonds[1, i], get_bond(l.sites[1], l.sites[index], l))
                    reduced[index] = i
                end
            end
        end
    end

    return reduced
end

"""
    get_trafos_uc(
        l :: Lattice
        ) :: Vector{Tuple{SMatrix{3, 3, Float64}, Bool}}

Compute mappings of a lattice's basis sites to the origin.
The mappings consist of a transformation matrix and a boolean indicating if an inversion was used or not.
"""
function get_trafos_uc(
    l :: Lattice
    ) :: Vector{Tuple{SMatrix{3, 3, Float64}, Bool}}

    # allocate list for transformations, set reference site and its bonds
    trafos = Vector{Tuple{SMatrix{3, 3, Float64}, Bool}}(undef, length(l.uc.basis) - 1)
    ref    = l.sites[1]
    con    = l.uc.bonds[1]

    # iterate over basis sites and find a symmetry for each of them
    for b in 2 : length(l.uc.basis)
        # set basis site and connections
        int       = SVector{4, Int64}(0, 0, 0, b)
        basis     = Site(int, get_vec(int, l.uc))
        con_basis = l.uc.bonds[b]

        # get a pair of non-collinear neighbors of basis
        for b1 in eachindex(con_basis)
            for b2 in eachindex(con_basis)
                # check that sites are unequal
                if b1 == b2
                    continue
                end

                # get connected sites in Bravais representation
                int_b1 = basis.int .+ con_basis[b1]
                int_b2 = basis.int .+ con_basis[b2]

                # get connected sites in real space representation
                vec_b1 = get_vec(int_b1, l.uc)
                vec_b2 = get_vec(int_b2, l.uc)

                # check that sites are non-collinear
                if norm(cross(vec_b1 .- basis.vec, vec_b2 .- basis.vec)) <= 1e-8
                    continue
                end

                # get site structs
                site_b1 = Site(int_b1, vec_b1)
                site_b2 = Site(int_b2, vec_b2)

                # get bonds
                bond_b1 = get_bond(basis, site_b1, l)
                bond_b2 = get_bond(basis, site_b2, l)

                # get a pair of non-collinear neigbors of origin
                for ref1 in eachindex(con)
                    for ref2 in eachindex(con)
                        # check that sites are unequal
                        if ref1 == ref2
                            continue
                        end

                        # get connected sites in Bravais representation
                        int_ref1 = ref.int .+ con[ref1]
                        int_ref2 = ref.int .+ con[ref2]

                        # get connected sites in real space representation
                        vec_ref1 = get_vec(int_ref1, l.uc)
                        vec_ref2 = get_vec(int_ref2, l.uc)

                        # check that sites are non-collinear
                        if norm(cross(vec_ref1, vec_ref2)) <= 1e-8
                            continue
                        end

                        # get site structs
                        site_ref1 = Site(int_ref1, vec_ref1)
                        site_ref2 = Site(int_ref2, vec_ref2)

                        # get bonds
                        bond_ref1 = get_bond(ref, site_ref1, l)
                        bond_ref2 = get_bond(ref, site_ref2, l)

                        # check that the bonds match
                        if are_equal(bond_b1, bond_ref1) == false || are_equal(bond_b2, bond_ref2) == false
                            continue
                        end

                        # try shift -> rotation
                        mat = rotate(site_b1.vec .- basis.vec, site_b2.vec .- basis.vec, site_ref1.vec, site_ref2.vec)

                        # if successful, verify tranformation on test set before saving it
                        if maximum(abs.(mat)) > 1e-8
                            valid = true

                            for n in eachindex(l.test_sites)
                                # test only those sites which are in range of basis
                                if get_metric(l.test_sites[n], basis, l.uc) <= l.size
                                    # apply transformation to lattice site
                                    mapped_vec = mat * (l.test_sites[n].vec .- basis.vec)
                                    orig_bond  = get_bond(basis, l.test_sites[n], l)
                                    mapped_ind = get_site(mapped_vec, l)

                                    # check if resulting site is in lattice and if bonds match
                                    if mapped_ind == 0
                                        valid = false
                                    else
                                        valid = are_equal(orig_bond, get_bond(ref, l.sites[mapped_ind], l))
                                    end

                                    # break if test fails for one test site
                                    if valid == false
                                        break
                                    end
                                end
                            end

                            # save transformation
                            if valid
                                trafos[b - 1] = (mat, false)
                                break
                                break
                                break
                                break
                            end
                        end

                        # try inversion -> shift -> rotation
                        mat = rotate(-(site_b1.vec .- basis.vec), -(site_b2.vec .- basis.vec), site_ref1.vec, site_ref2.vec)

                        # if successful, verify tranformation on test set before saving it
                        if maximum(abs.(mat)) > 1e-8
                            valid = true

                            for n in eachindex(l.test_sites)
                                # test only those sites which are in range of basis
                                if get_metric(l.test_sites[n], basis, l.uc) <= l.size
                                    # apply transformation to lattice site
                                    mapped_vec = mat * (-l.test_sites[n].vec .+ basis.vec)
                                    orig_bond  = get_bond(basis, l.test_sites[n], l)
                                    mapped_ind = get_site(mapped_vec, l)

                                    # check if resulting site is in lattice and if bonds match
                                    if mapped_ind == 0
                                        valid = false
                                    else
                                        valid = are_equal(orig_bond, get_bond(ref, l.sites[mapped_ind], l))
                                    end

                                    # break if test fails for one test site
                                    if valid == false
                                        break
                                    end
                                end
                            end

                            # save transformation
                            if valid
                                trafos[b - 1] = (mat, true)
                                break
                                break
                                break
                                break
                            end
                        end
                    end
                end
            end
        end
    end

    return trafos
end

# apply transformation (i, j) -> (i0, j*), where i0 is the origin, to site j
function apply_trafo(
    s      :: Site,
    b      :: Int64,
    shift  :: SVector{3, Float64},
    trafos :: Vector{Tuple{SMatrix{3, 3, Float64}, Bool}},
    l      :: Lattice
    )      :: SVector{3, Float64}

    # check if equivalent to origin
    if b != 1
        # if inequivalent to origin use transformation inside unitcell
        if trafos[b - 1][2]
            return trafos[b - 1][1] * (shift .- s.vec .+ l.uc.basis[b])
        else
            return trafos[b - 1][1] * (s.vec .- shift .- l.uc.basis[b])
        end
    # if equivalent to origin, only shift along translation vectors needs to be performed
    else
        return s.vec .- shift
    end
end

# compute mappings onto reduced lattice
function get_mappings(
    l       :: Lattice,
    reduced :: Vector{Int64}
    )       :: Matrix{Int64}

    # allocate matrix
    num = length(l.sites)
    mat = Matrix{Int64}(I, num, num)

    # get transformations inside unitcell
    trafos = get_trafos_uc(l)

    # get distances to origin
    dists = Float64[norm(l.sites[i].vec) for i in eachindex(l.sites)]
    d_max = maximum(dists)

    # group sites in shells
    shell_dists = unique(trunc.(dists, digits = 8))
    shells      = Vector{Vector{Int64}}(undef, length(shell_dists))

    for i in eachindex(shells)
        shell = Int64[]

        for j in eachindex(dists)
            if abs(dists[j] - shell_dists[i]) < 1e-8
                push!(shell, j)
            end
        end

        shells[i] = shell
    end

    # compute entries of matrix
    Threads.@threads for i in eachindex(l.sites)
        # determine shift for second site
        si    = l.sites[i]
        b     = si.int[4]
        shift = get_vec(SVector{4, Int64}(si.int[1], si.int[2], si.int[3], 1), l.uc)

        for j in eachindex(l.sites)
            # compute only off-diagonal entries
            if i != j
                # apply symmetry transformation
                mapped_vec = apply_trafo(l.sites[j], b, shift, trafos, l)

                # locate transformed site in lattice
                index = 0
                d_map = norm(mapped_vec)

                # check if transformed site can be in lattice
                if d_max - d_map > -1e-8
                    # find respective shell
                    shell = shells[argmin(abs.(shell_dists .- d_map))]

                    # find matching site in shell
                    for k in shell
                        if norm(mapped_vec .- l.sites[k].vec) < 1e-8
                            index = k
                            break
                        end
                    end

                    if index != 0
                        mat[i, j] = reduced[index]
                    end
                end
            end
        end
    end

    return mat
end

# compute irreducible sites in overlap of two sites
function get_overlap(
    l           :: Lattice,
    reduced     :: Vector{Int64},
    irreducible :: Vector{Int64},
    mappings    :: Matrix{Int64}
    )           :: Vector{Matrix{Int64}}

    # allocate overlap
    overlap = Vector{Matrix{Int64}}(undef, length(irreducible))

    for i in eachindex(irreducible)
        # collect all sites in range of irreducible and origin
        temp = NTuple{2, Int64}[]

        for j in eachindex(l.sites)
            if mappings[irreducible[i], j] != 0
                push!(temp, (reduced[j], mappings[j, irreducible[i]]))
            end
        end

        # determine how often a certain pair occurs
        pairs = unique(temp)
        table = zeros(Int64, length(pairs), 3)

        for j in eachindex(pairs)
            # convert from original lattice index to new "irreducible" index
            table[j, 1] = findall(index -> index == pairs[j][1], irreducible)[1]
            table[j, 2] = findall(index -> index == pairs[j][2], irreducible)[1]

            # count multiplicity
            for k in eachindex(temp)
                if pairs[j] == temp[k]
                    table[j, 3] += 1
                end
            end
        end

        # save into overlap
        overlap[i] = table
    end

    return overlap
end

# compute multiplicity of irreducible sites
function get_mult(
    reduced     :: Vector{Int64},
    irreducible :: Vector{Int64}
    )           :: Vector{Int64}

    # allocate vector to store multiplicities
    mult = zeros(Float64, length(irreducible))

    for i in eachindex(mult)
        for j in eachindex(reduced)
            if reduced[j] == irreducible[i]
                mult[i] += 1
            end
        end
    end

    return mult
end

# compute mappings under site exchange
function get_exchange(
    irreducible :: Vector{Int64},
    mappings    :: Matrix{Int64}
    )           :: Vector{Int64}

    # allocate exchange list
    exchange = zeros(Int64, length(irreducible))

    # determine mappings under site exchange
    for i in eachindex(exchange)
        exchange[i] = findall(index -> index == mappings[irreducible[i], 1], irreducible)[1]
    end

    return exchange
end

# convert mapping table entries to irreducible site indices
function get_project(
    l           :: Lattice,
    irreducible :: Vector{Int64},
    mappings    :: Matrix{Int64}
    )           :: Matrix{Int64}

    # allocate projections
    project = zeros(Int64, length(l.sites), length(l.sites))

    for i in eachindex(l.sites)
        for j in eachindex(l.sites)
            if mappings[i, j] != 0
                project[i, j] = findall(index -> index == mappings[i, j], irreducible)[1]
            end
        end
    end

    return project
end

"""
    get_reduced_lattice(
        model   :: String,
        J       :: Vector{Vector{Float64}},
        l       :: Lattice
        ;
        verbose :: Bool = true
        )       :: Reduced_lattice

Compute symmetry reduced representation of a given lattice graph with spin interactions between sites.
The interactions are defined by passing a model's name and coupling vector.
"""
function get_reduced_lattice(
    model   :: String,
    J       :: Vector{Vector{Float64}},
    l       :: Lattice
    ;
    verbose :: Bool = true
    )       :: Reduced_lattice

    if verbose
        println("Performing symmetry reduction ...")
    end

    # initialize model by modifying bond matrix of lattice
    init_model!(model, J, l)

    # get reduced representation of lattice
    reduced     = get_reduced(l)
    irreducible = unique(reduced)
    sites       = Site[Site(l.sites[i].int, get_vec(l.sites[i].int, l.uc)) for i in irreducible]

    # get mapping table
    mappings = get_mappings(l, reduced)

    # get overlap
    overlap = get_overlap(l, reduced, irreducible, mappings)

    # get multiplicities
    mult = get_mult(reduced, irreducible)

    # get exchanges
    exchange = get_exchange(irreducible, mappings)

    # get projections
    project = get_project(l, irreducible, mappings)

    # build reduced lattice
    r = Reduced_lattice(l.name, l.size, model, J, sites, overlap, mult, exchange, project)

    if verbose
        println("Done. Reduced lattice has $(length(r.sites)) sites.")
    end

    return r
end


module TestDepthFirst

using TapirSchedulers
using TapirSchedulers:
    DEPTH_FIRST_SCHEDULER, Priority, rollback_priority, get_current_priority_range
using Test

current_priority() = first(get_current_priority_range())

function record_priority(n, sch = DEPTH_FIRST_SCHEDULER[])
    priorities = fill!(Vector{Priority}(undef, n), typemax(Priority))
    workers = zeros(Int, n)
    function recur(indices)
        if length(indices) == 0
        elseif length(indices) == 1
            @inbounds priorities[indices[1]] = current_priority()
            @inbounds workers[indices[1]] = Threads.threadid()
        else
            f = first(indices)
            l = last(indices)
            m = (l - f + 1) ÷ 2 + f - 1
            work = TapirSchedulers.spawn!(sch) do
                recur(m+1:l)
            end
            recur(f:m)
            fetch(work)::Nothing
        end
        return nothing
    end
    rollback_priority() do
        recur(eachindex(priorities))
    end
    return (; priorities, workers)
end

function test_priority()
    e = 5
    @assert e < sizeof(Priority) * 8
    (; priorities) = record_priority(2^e)
    @test all(<(typemax(Priority)), priorities)
    @test issorted(priorities)
end

end  # module


#===============================================================================

    RawData.jl

    struct containing 3D array of raw data
    parameterized on unit type

    Dimensions:
    1. Unit: `cycle` for amplification data (type Int_T)
       `temperature` for melting curve data (type Float_T).
    2. Well: enumerated type, maximum 16.
    3. Channel: 1 or 2.

    Author: Tom Price
    Date:   July 2019

===============================================================================#

abstract type AbstractRaw end

struct RawData{F <: Union{Int_T,Float_T}} <: AbstractRaw
    data ::Array{F,3}
end


module Foo
"""
```jldoctest
julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8

julia> reshape(Int64[1, 2, 3, 4], (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 1
 2
 3
 4
```
```jldoctest
julia> reshape(Int64[1, 2, 3, 4], (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 1
 2
 3
 4

julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest
julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8

julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest
julia> begin
           reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
       end
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest
reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))

# output

4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest; filter = r"foo"
julia> println("  foobar")
  foobar
```
```jldoctest
julia> 1 + 2
3

julia> 3 + 4
7
```
"""
foo() = 1

    """
    ```jldoctest
    julia> begin
               reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
           end
    4×1×1 Array{Int64,3}:
    [:, :, 1] =
     2
     4
     6
     8
    ```
    ```jldoctest
    julia> println(); println("foo")

    foo
    ```
    """
    foo(x) = 1

end # module


#####
##### statistics and diagnostics
#####

module Diagnostics

export EBFMI, summarize_tree_statistics, explore_log_acceptance_ratios, leapfrog_trajectory,
    InvalidTree, REACHED_MAX_DEPTH, is_divergent

using DynamicHMC: GaussianKineticEnergy, Hamiltonian, evaluate_ℓ, InvalidTree,
    REACHED_MAX_DEPTH, is_divergent, log_acceptance_ratio, PhasePoint, rand_p, leapfrog,
    logdensity, MAX_DIRECTIONS_DEPTH

using ArgCheck: @argcheck
using DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF
using LogDensityProblems: dimension
using Parameters: @unpack
import Random
using Statistics: mean, quantile, var

"""
$(SIGNATURES)

Energy Bayesian fraction of missing information. Useful for diagnosing poorly
chosen kinetic energies.

Low values (`≤ 0.3`) are considered problematic. See Betancourt (2016).
"""
function EBFMI(tree_statistics)
    πs = map(x -> x.π, tree_statistics)
    mean(abs2, diff(πs)) / var(πs)
end

"Acceptance quantiles for [`TreeStatisticsSummary`](@ref) diagnostic summary."
const ACCEPTANCE_QUANTILES = [0.05, 0.25, 0.5, 0.75, 0.95]

"""
$(TYPEDEF)

Storing the output of [`NUTS_statistics`](@ref) in a structured way, for pretty
printing. Currently for internal use.

# Fields

$(FIELDS)
"""
struct TreeStatisticsSummary{T <: Real, C <: NamedTuple}
    "Sample length."
    N::Int
    "average_acceptance"
    a_mean::T
    "acceptance quantiles"
    a_quantiles::Vector{T}
    "termination counts"
    termination_counts::C
    "depth counts (first element is for `0`)"
    depth_counts::Vector{Int}
end

"""
$(SIGNATURES)

Count termination reasons in `tree_statistics`.
"""
function count_terminations(tree_statistics)
    max_depth = 0
    divergence = 0
    turning = 0
    for tree_statistic in tree_statistics
        it = tree_statistic.termination
        if it == REACHED_MAX_DEPTH
            max_depth += 1
        elseif is_divergent(it)
            divergence += 1
        else
            turning += 1
        end
    end
    (max_depth = max_depth, divergence = divergence, turning = turning)
end

"""
$(SIGNATURES)

Count depths in tree statistics.
"""
function count_depths(tree_statistics)
    c = zeros(Int, MAX_DIRECTIONS_DEPTH + 1)
    for tree_statistic in tree_statistics
        c[tree_statistic.depth + 1] += 1
    end
    c[1:something(findlast(!iszero, c), 0)]
end

"""
$(SIGNATURES)

Summarize tree statistics. Mostly useful for NUTS diagnostics.
"""
function summarize_tree_statistics(tree_statistics)
    As = map(x -> x.acceptance_rate, tree_statistics)
    TreeStatisticsSummary(length(tree_statistics),
                          mean(As), quantile(As, ACCEPTANCE_QUANTILES),
                          count_terminations(tree_statistics),
                          count_depths(tree_statistics))
end

function Base.show(io::IO, stats::TreeStatisticsSummary)
    @unpack N, a_mean, a_quantiles, termination_counts, depth_counts = stats
    println(io, "Hamiltonian Monte Carlo sample of length $(N)")
    print(io, "  acceptance rate mean: $(round(a_mean; digits = 2)), 5/25/50/75/95%:")
    for aq in a_quantiles
        print(io, " ", round(aq; digits = 2))
    end
    println(io)
    function print_percentages(pairs)
        is_first = true
        for (key, value) in sort(collect(pairs), by = first)
            if is_first
                is_first = false
            else
                print(io, ",")
            end
            print(io, " $(key) => $(round(Int, 100*value/N))%")
        end
    end
    print(io, "  termination:")
    print_percentages(pairs(termination_counts))
    println(io)
    print(io, "  depth:")
    print_percentages(zip(axes(depth_counts, 1) .- 1, depth_counts))
end

####
#### Acceptance ratio diagnostics
####

"""
$(SIGNATURES)

From an initial position, calculate the uncapped log acceptance ratio for the given log2
stepsizes and momentums `ps`, `N` of which are generated randomly by default.
"""
function explore_log_acceptance_ratios(ℓ, q, log2ϵs;
                                       rng = Random.GLOBAL_RNG,
                                       κ = GaussianKineticEnergy(dimension(ℓ)),
                                       N = 20, ps = [rand_p(rng, κ) for _ in 1:N])
    H = Hamiltonian(κ, ℓ)
    Q = evaluate_ℓ(ℓ, q)
    [log_acceptance_ratio(H, PhasePoint(Q, p), 2.0^log2ϵ) for log2ϵ in log2ϵs, p in ps]
end

"""
$(TYPEDEF)

Implements an iterator on a leapfrog trajectory until the first non-finite log density.

# Fields

$(FIELDS)
"""
struct LeapfrogTrajectory{TH,TZ,TF,Tϵ}
    "Hamiltonian"
    H::TH
    "Initial position"
    z₀::TZ
    "Negative energy at initial position."
    π₀::TF
    "Stepsize (negative: move backward)."
    ϵ::Tϵ
end

Base.IteratorSize(::Type{<:LeapfrogTrajectory}) = Base.SizeUnknown()

function Base.iterate(lft::LeapfrogTrajectory, zi = (lft.z₀, 0))
    @unpack H, ϵ, π₀ = lft
    z, i = zi
    if isfinite(z.Q.ℓq)
        z′ = leapfrog(H, z, ϵ)
        i′ = i + sign(ϵ)
        _position_information(lft, z′, i′), (z′, i′)
    else
        nothing
    end
end

"""
$(SIGNATURES)

Position information returned by [`leapfrog_trajectory`](@ref), see documentation there.
Internal function.
"""
function _position_information(lft::LeapfrogTrajectory, z, i)
    @unpack H, π₀ = lft
    (z = z, position = i, Δ = logdensity(H, z) - π₀)
end

"""
$(SIGNATURES)

Calculate a leapfrog trajectory visiting `positions` (specified as a `UnitRange`, eg `-5:5`)
relative to the starting point `q`, with stepsize `ϵ`. `positions` has to contain `0`, and
the trajectories are only tracked up to the first non-finite log density in each direction.

Returns a vector of `NamedTuple`s, each containin

- `z`, a [`PhasePoint`](@ref) object,

- `position`, the corresponding position,

- `Δ`, the log density + the kinetic energy relative to position `0`.
"""
function leapfrog_trajectory(ℓ, q, ϵ, positions::UnitRange{<:Integer};
                             rng = Random.GLOBAL_RNG,
                             κ = GaussianKineticEnergy(dimension(ℓ)), p = rand_p(rng, κ))
    A, B = first(positions), last(positions)
    @argcheck A ≤ 0 ≤ B "Positions has to contain `0`."
    Q = evaluate_ℓ(ℓ, q)
    H = Hamiltonian(κ, ℓ)
    z₀ = PhasePoint(Q, p)
    π₀ = logdensity(H, z₀)
    lft_fwd = LeapfrogTrajectory(H, z₀, π₀, ϵ)
    fwd_part = collect(Iterators.take(lft_fwd, B))
    bwd_part = collect(Iterators.take(LeapfrogTrajectory(H, z₀, π₀, -ϵ), -A))
    vcat(reverse!(bwd_part), _position_information(lft_fwd, z₀, 0), fwd_part)
end

end


"""
    ReversedEasing

Will be used to reverse an easing inside [`easing_to_animation`](@ref).
Can be constructed from an easing function using [`rev`](@ref).
"""
struct ReversedEasing
    easing::Easing
end

"""
    rev(e::Easing)

Reverse an easing function such that `easing_to_animation` maps it to `[1.0, 0.0]` instead of `[0.0, 1.0]`.
An example can be seen in [`rotate`](@ref)
"""
rev(e::Easing) = ReversedEasing(e)

"""
    easing_to_animation(easing)

Converts an easing to an Animation with time goes from `0.0` to `1.0` and value from `0` to `1`.
"""
easing_to_animation(easing) = Animation(0.0, 0.0, easing, 1.0, 1.0)

"""
    easing_to_animation(rev_easing::ReversedEasing)

Converts an easing to an Animation with time goes from `0.0` to `1.0` and value from `1` to `0`.
"""
easing_to_animation(rev_easing::ReversedEasing) =
    Animation(0.0, 1.0, rev_easing.easing, 1.0, 0.0)


## BOUNDEDNESS TRAIT

# For random search and perhaps elsewhere, we need a variation on the
# built-in boundedness notions:
abstract type PositiveUnbounded <: Unbounded end
abstract type Other <: Unbounded end

boundedness(::NumericRange{<:Any,<:Bounded}) = Bounded
boundedness(::NumericRange{<:Any,<:LeftUnbounded}) = Other
boundedness(::NumericRange{<:Any,<:DoublyUnbounded}) = Other
function boundedness(r::NumericRange{<:Any,<:RightUnbounded})
    if r.lower >= 0
        return PositiveUnbounded
    end
    return Other
end

"""
    MLJTuning.grid([rng, ] prototype, fields, iterators)

Given a vector of (nested) field names `fields`, and a vector of
`iterators` (each element being a vector of values for the
corresponding field) return a vector of models generated by cloning
and mutating the fields of `prototype` through all possible
combinations of values given by the iterator values.

Specification of an `AbstractRNG` object `rng` implies shuffling of
the results. Otherwise models are ordered, with the first
hyperparameter referenced cycling fastest.

"""
grid(rng::AbstractRNG, prototype::Model, fields, iterators) =
    shuffle(rng, grid(prototype, fields, iterators))

function grid(prototype::Model, fields, iterators)

    A = MLJBase.unwind(iterators...)

    N = size(A, 1)
    map(1:N) do i
        clone = deepcopy(prototype)
        for k in eachindex(fields)
            recursive_setproperty!(clone, fields[k], A[i,k])
        end
        clone
    end
end

## PRE-PROCESSING OF USER-SPECIFIED CARTESIAN RANGE OBJECTS

"""
    process_grid_range(user_specified_range, resolution, verbosity)

Convert a user-specified range (see [`Grid`](@ref)) into a tuple  of
tuples `(ranges, resolutions)`.

For example, if `r1`, `r2` are `NumericRange`s and `s` is a
NominalRange` with 5 values, then we have:

    julia> MLJTuning.process_grid_range([(r1, 3), r2, s], 42, 1) ==
                            ((r1, r2, s), (3, 42, 5))
    true

If `verbosity` > 0, then a warning is issued if a `Nominal` range is
paired with a resolution.

"""
process_grid_range(user_specified_range, args...) =
    throw(ArgumentError("Unsupported range. "))

process_grid_range(usr::Union{ParamRange,Tuple{ParamRange,Int}}, args...) =
    process_grid_range([usr, ], args...)

function process_grid_range(user_specified_range::AbstractVector,
                    resolution, verbosity)
    # r unpaired:
    stand(r) = throw(ArgumentError("Unsupported range. "))
    stand(r::NumericRange) = (r, resolution)
    stand(r::NominalRange) = (r, length(r.values))

    # (r, res):
    stand(t::Tuple{NumericRange,Integer}) = t
    function stand(t::Tuple{NominalRange,Integer})
        verbosity < 0 ||
            @warn  "Ignoring a resolution specified for a `NominalRange`. "
        return (first(t), length(first(t).values))
    end

    ret = zip(stand.(user_specified_range)...) |> collect
    return first(ret), last(ret)
end

"""
    process_random_range(user_specified_range,
                         bounded,
                         positive_unbounded,
                         other)

Utility to convert a user-specified range (see [`RandomSearch`](@ref))
into an n-tuple of `(field, sampler)` pairs.

"""
process_random_range(user_specified_range, args...) =
    throw(ArgumentError("Unsupported range #1. "))

const DIST = Distributions.Distribution

process_random_range(user_specified_range::Union{ParamRange, Tuple{Any,Any}},
                     args...) =
    process_random_range([user_specified_range, ], args...)

function process_random_range(user_specified_range::AbstractVector,
                              bounded,
                              positive_unbounded,
                              other)
    # r not paired:
    stand(r) = throw(ArgumentError("Unsupported range #2. "))
    stand(r::NumericRange) = stand(r, boundedness(r))
    stand(r::NumericRange, ::Type{<:Bounded}) = (r.field, sampler(r, bounded))
    stand(r::NumericRange, ::Type{<:Other}) = (r.field, sampler(r, other))
    stand(r::NumericRange, ::Type{<:PositiveUnbounded}) =
        (r.field, sampler(r, positive_unbounded))
    stand(r::NominalRange) = (n = length(r.values);
                              (r.field, sampler(r, fill(1/n, n))))
    # (r, d):
    stand(t::Tuple{ParamRange,Any}) = stand(t...)
    stand(r, d) = throw(ArgumentError("Unsupported range #3. "))
    stand(r::NominalRange, d::AbstractVector{Float64}) =  _stand(r, d)
    stand(r::NumericRange, d:: Union{DIST, Type{<:DIST}}) = _stand(r, d)
    _stand(r, d) = (r.field, sampler(r, d))

    # (field, s):
    stand(t::Tuple{Union{Symbol,Expr},Any}) = t

    return  Tuple(stand.(user_specified_range))

    # ret = zip(stand.(user_specified_range)...) |> collect
    # return first(ret), last(ret)
end


module CurrencyFormatting

using ..Baskets
using ..CurrenciesBase
using ..CurrencyData
using ..DeclarativeFormatting
import ..DeclarativeFormatting: conflict, reconcile

export format

include("rules.jl")
include("decimals.jl")
include("templates.jl")
include("render.jl")

# Single-line output in two-argument `show`
Base.show(io::IO, m::Monetary) =
        show(IOContext(io, :compact => true), "text/plain", m)

function Base.show(io::IO, b::Basket)
    write(io, "$(typeof(b))([")
    write(io, join([sprint(show, c; context=:compact => true) for c in b], ","))
    print(io, "])")
end

# Multi-line output with MIME types
function Base.show(io::IO, ::MIME"text/plain", m::Monetary)
    if get(io, :compact, false)
        print(io, m / majorunit(m), currency(m))
    else
        print(io, format(m; styles=[:plain]))
    end
end

function Base.show(io::IO, ::MIME"text/latex", m::Monetary)
    print(io, string('$', format(m; styles=[:latex]), '$'))
end

function Base.show(io::IO, ::MIME"text/plain", b::Basket)
    len = length(b)
    write(io, "$len-currency $(typeof(b)):")
    for val in b
        write(io, "\n ")
        show(io, "text/plain", val)
    end
end

function Base.show(io::IO, ::MIME"text/markdown", b::Basket)
    len = length(b)
    println(io, "\$$len\$-currency `$(typeof(b))`:")
    for val in b
        write(io, "\n - ")
        show(io, "text/latex", val)
    end
end

end  # module CurrencyFormatting


const required_fields = [
    "byr",  # Birth Year
    "ecl",  # Eye Color
    "eyr",  # Expiration Year
    "hcl",  # Hair Color
    "hgt",  # Height
    "iyr",  # Issue Year
    "pid"   # Passport ID
]

function new_passport(fields)

    p = Dict{String,Any}()
    for field in vcat(split.(fields)...)
       key, value = split(field,":")
       p[key] = value
    end
    return p

end

function read_passports(filename)

   passports = []
   open(filename) do f

       fields  = []
       for line in eachline(f)

           if isempty(line)
               push!(passports, new_passport(fields))
               fields = []
           else
               push!(fields, line)
           end
       end

       # end of the file, last passport
       push!(passports, new_passport(fields))

   end

   return passports

end

function valid1( p )

    valid_fields = collect(filter(k->(k!="cid"), keys(p)))

    if length(valid_fields) == 7
        return all( sort(valid_fields) .== required_fields)
    else
        return false
    end

end

const eye_colors = ["amb","blu","brn","gry","grn","hzl","oth"]

function valid2( p ) # not finish...

    if valid1(p)

        1920 <= parse(Int, p["byr"]) <= 2002 || (return false)
        2010 <= parse(Int, p["iyr"]) <= 2020 || (return false)
        2020 <= parse(Int, p["eyr"]) <= 2030 || (return false)
        hgt = p["hgt"]
        if endswith(hgt, "cm")
            150 <= parse(Int,hgt[1:end-2]) <= 193 || (return false)
        elseif endswith(hgt, "in")
            59 <= parse(Int,hgt[1:end-2]) <= 76 || (return false)
        else
            return false
        end
        startswith(p["hcl"], "#") || (return false)
        tryparse(Int, p["hcl"][2:end], base=16) == nothing && (return false)
        p["ecl"] in eye_colors || (return false)
        length(p["pid"]) == 9 || (return false)
        tryparse(Int, p["pid"]) == nothing && (return false)

        return true

    else
        return false
    end

end

passports = read_passports( "input04.txt")
println(count(valid1.(passports)))
println(count(valid2.(passports)))


"""
    AbstractSize{N}

Supertype for all array sizes.
"""
abstract type AbstractSize{N} end

const SizeType{N} = Union{AbstractSize{N},Dims{N}}

"""
    Size{N} <: AbstractSize{N}

Size type for resizable arrays, which is a mutable wrapper of `Dims{N}` to
represent the dimension of an resizable array. Mutate `d`th dimension to `n` by
`sz[d] = n` mutate the whole dimensions to `nsz` by `set!(sz, nsz)`.
"""
mutable struct Size{N} <: AbstractSize{N}
    sz::Dims{N}
end
@inline Size(I::Int...) = Size(I)
@inline Size(A::AbstractArray) = Size(size(A))

# use SizeOrTuple instead of SizeType{N} to avoid override the methods for tuple
const SizeOrTuple = Union{Size,Tuple}
_totuple(sz::Size) = sz.sz
_totuple(tp::Tuple) = tp

@inline Base.length(::Size{N}) where {N} = N
@inline Base.convert(::Type{T}, sz::Size) where {T<:Tuple} = convert(T, sz.sz)
@inline Base.:(==)(t1::SizeOrTuple, t2::SizeOrTuple) = _totuple(t1) == _totuple(t2)

"""
    set!(sz::Size{N}, nsz::NTuple{N,Integer})

Set `sz` to `nsz`.

# Example

```jldoctest
julia> sz = Size(1, 2, 3)
Size{3}((1, 2, 3))

julia> set!(sz, (3, 2, 1))
Size{3}((3, 2, 1))
```
"""
set!(sz::Size{N}, nsz::NTuple{N,Integer}) where {N} = set!(sz, convert(Dims{N}, nsz))
set!(sz::Size{1}, nsz::Dims{1}) = (sz[1] = nsz[1]; sz)
set!(sz::Size{2}, nsz::Dims{2}) = (sz[1] = nsz[1]; sz[2] = nsz[2]; sz)
set!(sz::Size{3}, nsz::Dims{3}) = (sz[1] = nsz[1]; sz[2] = nsz[2]; sz[3] = nsz[3]; sz)
set!(sz::Size{N}, nsz::Dims{N}) where {N} = sz.sz = nsz

# The below two methods is a modification of `MArray` in `StaticArrays.jl`
# https://github.com/JuliaArrays/StaticArrays.jl/blob/master/src/MArray.jl#L80
function Base.getindex(sz::Size{N}, i::Int) where {N}
    @boundscheck 1 <= i <= N || throw(BoundsError(sz, i))
    return GC.@preserve sz unsafe_load(
        Base.unsafe_convert(Ptr{Int}, pointer_from_objref(sz)),
        i,
    )
end

function Base.setindex!(sz::Size{N}, v, i::Int) where {N}
    @boundscheck 1 <= i <= N || throw(BoundsError(sz, i))
    GC.@preserve sz unsafe_store!(
        Base.unsafe_convert(Ptr{Int}, pointer_from_objref(sz)),
        convert(Int, v),
        i,
    )
    return sz
end

# getsize
"""
    getsize(A::AbstractArray, [dim])

Returns the dimensions of `A` unlike `size` which may not return a
`Dims{N}`.

!!! warning
    You should never call this function directly. Use `size` instead.
"""
getsize

Base.@propagate_inbounds getsize(A::AbstractArray, d::Integer) = getsize(A, Int(d))
Base.@propagate_inbounds getsize(A::AbstractArray, d::Int) = getsize(A)[d]

"""
    size_type(A::AbstractArray)

Get the size type of `A`, determine the methods of `setsize!`. The default
`size_type` is `NoneSize`, which means `setsize!` will "do nothing".
"""
size_type(A::AbstractArray) = size_type(typeof(A))
size_type(::Type{T}) where {T<:AbstractArray} = Dims{ndims(T)}

# setsize!(A, sz)
"""
    setsize!(A::AbstractArray{T,N}, sz) where {T,N} -> AbstractArray

Set the size of `A` to `sz`.
"""
setsize!(A::AbstractArray{T,N}, sz::NTuple{N,Any}) where {T,N} =
    setsize!(size_type(A), A, to_dims(sz))

@inline setsize!(::Type{S}, A::AbstractArray, ::Dims{N}) where {N,S<:Dims{N}} = A
@inline setsize!(::Type{S}, A::AbstractArray{T,N}, sz::Dims{N}) where {T,N,S<:Size{N}} =
    (set!(getsize(A)::S, sz); A)

# setsize!(A, d, i)
"""
    setsize!(A::AbstractArray, d::Integer, n) -> AbstractArray

Set the `d`th dimension to `n`.
"""
Base.@propagate_inbounds setsize!(A::AbstractArray, d::Integer, n) =
    setsize!(size_type(A), A, Int(d), to_dims(n))

@inline setsize!(::Type{S}, A::AbstractArray, ::Int, ::Int) where {S<:Dims} = A
setsize!(::Type{S}, A::AbstractArray{T,N}, d::Int, n::Int) where {T,N,S<:Size{N}} =
    (Base.@_propagate_inbounds_meta; getsize(A)[d] = n; A) # meta to avoid too long line

"""
    to_dims(inds::Tuple) -> Dims

Convert the given indices to `Dims`.

!!! note

    The given indices should be a return value of `to_indices`.
    If `inds[i]` is an `Integer`, this function would converted it to `Int`;
    If `inds[i]` is an `AbstractVector`, this function would return its length.
"""
to_dims(inds::Tuple) = map(to_dims, inds)::Dims
to_dims(inds::Dims) = inds
to_dims(ind::AbstractVector) = length(ind)::Int
to_dims(ind::Integer) = Int(ind)