#pragma once

#include "../taskflow.hpp"

#include <variant>

#include <filesystem>

namespace tf {

enum StorageLevel {

MEMORY = 0,

MEMORY_AND_DISK = 1

};

template <typename T>

class Tensor {

};

template <typename T>

Tensor<T>::Tensor(std::vector<size_t> shape) :

_shape {std::move(shape)},

_chunk_shape (_shape.size()),

_chunk_grid (_shape.size()) {

_make_chunks();

}

template <typename T>

Tensor<T>::Tensor(std::vector<size_t> shape, size_t max_chunk_size) :

_shape {std::move(shape)},

_chunk_shape (_shape.size()),

_chunk_grid (_shape.size()) {

_make_chunks(std::max(1ul, max_chunk_size));

}

template <typename T>

size_t Tensor<T>::size() const {

return std::accumulate(

_shape.begin(), _shape.end(), 1, std::multiplies<size_t>()

);

}

template <typename T>

size_t Tensor<T>::num_chunks() const {

return _chunks.size();

}

template <typename T>

size_t Tensor<T>::chunk_size() const {

return _chunks[0].data.size();

}

template <typename T>

size_t Tensor<T>::rank() const {

return _shape.size();

}

template <typename T>

const std::vector<size_t>& Tensor<T>::shape() const {

return _shape;

}

template <typename T>

const std::vector<size_t>& Tensor<T>::chunk_shape() const {

return _chunk_shape;

}

template <typename T>

template <typename... Is>

size_t Tensor<T>::flat_chunk_index(Is... rest) const {

if(sizeof...(Is) != rank()) {

TF_THROW("index rank dose not match tensor rank");

}

size_t offset;

return _flat_chunk_index(offset, rest...);

}

template <typename T>

size_t Tensor<T>::_flat_chunk_index(size_t& offset, size_t id) const {

offset = 1;

return id/_chunk_shape.back();

}

template <typename T>

template <typename... Is>

size_t Tensor<T>::_flat_chunk_index(

size_t& offset, size_t id, Is... rest

) const {

auto i = _flat_chunk_index(offset, rest...);

offset *= _chunk_grid[_chunk_shape.size() - (sizeof...(Is))];

return (id/_chunk_shape[_chunk_shape.size() - sizeof...(Is) - 1])*offset + i;

}

template <typename T>

template <typename... Is>

size_t Tensor<T>::flat_index(Is... rest) const {

if(sizeof...(Is) != rank()) {

TF_THROW("index rank dose not match tensor rank");

}

size_t offset;

return _flat_index(offset, rest...);

}

template <typename T>

size_t Tensor<T>::_flat_index(size_t& offset, size_t id) const {

offset = 1;

return id;

}

template <typename T>

template <typename... Is>

size_t Tensor<T>::_flat_index(size_t& offset, size_t id, Is... rest) const {

auto i = _flat_index(offset, rest...);

offset *= _shape[_shape.size() - (sizeof...(Is))];

return id*offset + i;

}

template <typename T>

void Tensor<T>::dump(std::ostream& os) const {

os << "Tensor<" << typeid(T).name() << "> {\n"

<< " shape=[";

for(size_t i=0; i<_shape.size(); ++i) {

if(i) os << 'x';

os << _shape[i];

}

os << "], chunk=[";

for(size_t i=0; i<_chunk_shape.size(); ++i) {

if(i) os << 'x';

os << _chunk_shape[i];

}

os << "], pgrid=[";

for(size_t i=0; i<_chunk_grid.size(); ++i) {

if(i) os << 'x';

os << _chunk_grid[i];

}

os << "]\n}\n";

}

template <typename T>

void Tensor<T>::_make_chunks(size_t M) {

size_t P = 1;

size_t N = 1;

for(int i=_shape.size()-1; i>=0; i--) {

if(M >= _shape[i]) {

_chunk_shape[i] = _shape[i];

_chunk_grid[i] = 1;

N *= _chunk_shape[i];

M /= _shape[i];

}

else {

_chunk_shape[i] = M;

_chunk_grid[i] = (_shape[i] + _chunk_shape[i] - 1) / _chunk_shape[i];

P *= _chunk_grid[i];

N *= _chunk_shape[i];

for(i--; i>=0; i--) {

_chunk_shape[i] = 1;

_chunk_grid[i] = _shape[i];

P *= _chunk_grid[i];

}

break;

}

}

_chunks.resize(P);

// we allocate the first data in memory

_chunks[0].data.resize(N);

// TODO: the rest sits in the disk

for(size_t i=1; i<_chunks.size(); ++i) {

}

// assign the storage level

_storage_level = (_chunks.size() <= 1) ? MEMORY : MEMORY_AND_DISK;

}

} // end of namespace tf -----------------------------------------------------