datautils.jl

module DataUtils

using ProgressMeter
using JuLIP
using ACEfriction.mUtils: reinterpret
using StaticArrays, SparseArrays
export FrictionData, BlockDenseArray 
export save_h5fdata, load_h5fdata
using HDF5


function _array2svector(x::Array{T,2}) where {T}
    return [ SVector{3}(x[i,:]) for i in 1:size(x)[1] ]
end
function _svector2array(c_vec::Vector{SVector{N_rep, T}}) where {N_rep,T<:Number}
    """
    input: c_vec each 
    N_basis = length(c_vec)
    N_rep = numbero of channels
    """
    c_matrix = Array{T}(undef,length(c_vec),N_rep)
    for j in eachindex(c_vec)
        c_matrix[j,:] = c_vec[j]
    end
    return c_matrix
end

"""
    save_h5fdata(rdata, filename )

Saves a friction tensor dataset with the HDF5 file structure expected by `load_h5fdata`, for use in training scripts. 

# Arguments

## `rdata`:
A `Vector{NamedTuple}` containing the training data in the following format.

```
(
    at = ::JuLIP.Atoms,
    friction_tensor = sparse(::Matrix),
    friction_indices = ::Vector{Int}
)
```

## `filename`

Name of the file to save to (including h5 extension). 

"""
function save_h5fdata(rdata, filename )
    fid = h5open(filename, "w")
    try
        # iterate over each data entry
        write_attribute(fid, "N_data", length(rdata))
        @showprogress for (i,d) in enumerate(rdata)
            g = create_group(fid, "$i")
            # write atoms data
            ag = create_group(g, "atoms")
            dset_pos = create_dataset(ag, "positions", Float64, (length(d.at.X), 3))
            for (k,x) in enumerate(d.at.X)
                dset_pos[k,:] = x
            end
            write(ag, "atypes", Int.(d.at.Z))
            write(ag, "cell", Matrix(d.at.cell))
            write(ag, "pbc", Array(d.at.pbc))
            # write friction data
            fg = create_group(g, "friction_tensor")
            (I,J,V) = findnz(d.friction_tensor)
            write(fg, "ft_I", I)
            write(fg, "ft_J", J)
            dset_ft = create_dataset(fg, "ft_val", Float64, (length(V), 3, 3))
            for (k,v) in enumerate(V)
                dset_ft[k,:,:] = v
            end
            write(fg, "ft_mask", d.friction_indices)
        end
    catch 
        close(fid)
        rethrow(e)
    end
    HDF5.close(fid)
end

_hdf52Atoms( ag::HDF5.Group ) = JuLIP.Atoms(;
                X=[SVector{3}(d) for d in eachslice(read(ag["positions"]); dims=1)],
                Z=read(ag["atypes"]), 
                cell= read(ag["cell"]),
                pbc=read(ag["pbc"])
            )

        
function _hdf52ft( ftg::HDF5.Group ) 
    spft = sparse( read(ftg["ft_I"]),read(ftg["ft_J"]), [SMatrix{3,3}(d) for d in eachslice(read(ftg["ft_val"]); dims=1)] )
    ft_mask = read(ftg["ft_mask"])
    return (friction_tensor = spft, mask = ft_mask)
end

function load_h5fdata(filename)
    fid = h5open(filename, "r")
    N_data = read_attribute(fid, "N_data")
    rdata = @showprogress [begin
                at = _hdf52Atoms( fid["$i/atoms/"]) 
                spft, ft_mask = _hdf52ft( fid["$i/friction_tensor/"])
                (at=at, friction_tensor=spft, friction_indices=ft_mask)
            end
            for i=1:N_data]
    HDF5.close(fid)
    return rdata
end

"""
Semi-sparse matrix representation of a square matrix M of unspecified dimension. Entries of M are 
specified by the fields `values::Matrix{T}` and `indexmap` assuming

    1. `M[k,l] = 0` if either of the indices `k`, `l` is not contained in `indexmap`.
    2. values[i,j] = M[indexmap[i],indexmap[j]]

"""
struct BlockDenseMatrix{T} <: AbstractArray{Float64,2}
    tensor::Matrix{T}
    indices
end

function BlockDenseArray(full_tensor::Matrix; indices=1:size(full_tensor,1)) 
    @assert size(full_tensor,1) == size(full_tensor,2)
    return BlockDenseMatrix(full_tensor[indices,indices], indices)
end

struct FrictionData{A} 
    atoms::Atoms
    friction_tensor::A
    friction_indices
    friction_tensor_ref
end

function FrictionData(d::NamedTuple{(:at, :friction_tensor, :friction_indices)})
    return FrictionData(d.at, d.friction_tensor, d.friction_indices, nothing)
end
function FrictionData(d::NamedTuple{(:at, :friction_tensor, :friction_indices,:friction_indices_ref)})
    return FrictionData(d.at, d.friction_tensor, d.friction_indices, d.friction_tensor_ref)
end
# function FrictionData(atoms::Atoms, friction_tensor, friction_indices;  
#                                                                 friction_tensor_ref=nothing)
#     return FrictionData(atoms, friction_tensor, friction_indices, weights, friction_tensor_ref)
# end


end