#include "psi.h"

#include "source_base/global_variable.h"
#include "source_base/module_device/device.h"
#include "source_base/tool_quit.h"
#include "source_io/module_parameter/parameter.h"

#include <cassert>
#include <complex>
#include <type_traits>

namespace psi
{

Range::Range(const size_t range_in)
{
 k_first = true;
 index_1 = 0;
 range_1 = range_in;
 range_2 = range_in;
}

Range::Range(const bool k_first_in, const size_t index_1_in, const size_t range_1_in, const size_t range_2_in)
{
 k_first = k_first_in;
 index_1 = index_1_in;
 range_1 = range_1_in;
 range_2 = range_2_in;
}

// Constructor 0: basic
template <typename T, typename Device>
Psi<T, Device>::Psi()
{
}

template <typename T, typename Device>
Psi<T, Device>::~Psi()
{
 if (this->allocate_inside)
 {
 delete_memory_op()(this->psi);
 }
}

// Constructor 1:
template <typename T, typename Device>
Psi<T, Device>::Psi(const int nk_in,
 const int nbd_in,
 const int nbs_in,
 const std::vector<int>& ngk_in,
 const bool k_first_in)
{
 assert(nk_in > 0);
 assert(nbd_in >= 0);
 assert(nbs_in > 0);

 this->k_first = k_first_in;
 this->allocate_inside = true;

 this->ngk = ngk_in.data(); // modify later
 // This function will delete the psi array first(if psi exist), then malloc a new memory for it.
 resize_memory_op()(this->psi, nk_in * static_cast<std::size_t>(nbd_in) * nbs_in, "no_record");

 this->nk = nk_in;
 this->nbands = nbd_in;
 this->nbasis = nbs_in;

 this->current_b = 0;
 this->current_k = 0;
 this->current_nbasis = nbs_in;
 this->psi_current = this->psi;
 this->psi_bias = 0;

 // Currently only GPU's implementation is supported for device recording!
 base_device::information::print_device_info<Device>(this->ctx, GlobalV::ofs_device);
 base_device::information::record_device_memory<Device>(this->ctx,
 GlobalV::ofs_device,
 "Psi->resize()",
 sizeof(T) * nk_in * nbd_in * nbs_in);
}

// Constructor 3-1: 2D Psi version
template <typename T, typename Device>
Psi<T, Device>::Psi(T* psi_pointer,
 const int nk_in,
 const int nbd_in,
 const int nbs_in,
 const int current_nbasis_in,
 const bool k_first_in)
{
 // Currently this function only supports nk_in == 1 when called within diagH_subspace_init.
 // assert(nk_in == 1); // NOTE because lr/utils/lr_uril.hpp func & get_psi_spin func

 this->k_first = k_first_in;
 this->allocate_inside = false;

 this->ngk = nullptr;
 this->psi = psi_pointer;

 this->nk = nk_in;
 this->nbands = nbd_in;
 this->nbasis = nbs_in;

 this->current_k = 0;
 this->current_b = 0;
 this->current_nbasis = current_nbasis_in;
 this->psi_current = psi_pointer;
 this->psi_bias = 0;

 // Currently only GPU's implementation is supported for device recording!
 base_device::information::print_device_info<Device>(this->ctx, GlobalV::ofs_device);
}

// Constructor 3-2: 2D Psi version
template <typename T, typename Device>
Psi<T, Device>::Psi(const int nk_in,
 const int nbd_in,
 const int nbs_in,
 const int current_nbasis_in,
 const bool k_first_in)
{
 // Currently this function only supports nk_in == 1 when called within diagH_subspace_init.
 // assert(nk_in == 1);

 this->k_first = k_first_in;
 this->allocate_inside = true;

 this->ngk = nullptr;
 assert(nk_in > 0 && nbd_in >= 0 && nbs_in > 0);
 resize_memory_op()(this->psi, nk_in * static_cast<std::size_t>(nbd_in) * nbs_in, "no_record");

 this->nk = nk_in;
 this->nbands = nbd_in;
 this->nbasis = nbs_in;

 this->current_k = 0;
 this->current_b = 0;
 this->current_nbasis = current_nbasis_in;
 this->psi_current = this->psi;
 this->psi_bias = 0;

 // Currently only GPU's implementation is supported for device recording!
 base_device::information::print_device_info<Device>(this->ctx, GlobalV::ofs_device);
 base_device::information::record_device_memory<Device>(this->ctx,
 GlobalV::ofs_device,
 "Psi->resize()",
 sizeof(T) * nk_in * nbd_in * nbs_in);
}

// Constructor 2-1:
template <typename T, typename Device>
Psi<T, Device>::Psi(const Psi& psi_in)
{

 this->ngk = psi_in.ngk;
 this->nk = psi_in.get_nk();
 this->nbands = psi_in.get_nbands();
 this->nbasis = psi_in.get_nbasis();
 this->current_k = psi_in.get_current_k();
 this->current_b = psi_in.get_current_b();
 this->k_first = psi_in.get_k_first();
 // this function will copy psi_in.psi to this->psi no matter the device types of each other.

 this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis());
 base_device::memory::synchronize_memory_op<T, Device, Device>()(this->psi,
 psi_in.get_pointer() - psi_in.get_psi_bias(),
 psi_in.size());
 this->psi_bias = psi_in.get_psi_bias();
 this->current_nbasis = psi_in.get_current_nbas();
 this->psi_current = this->psi + psi_in.get_psi_bias();
}

// Constructor 2-2:
template <typename T, typename Device>
template <typename T_in, typename Device_in>
Psi<T, Device>::Psi(const Psi<T_in, Device_in>& psi_in)
{

 this->ngk = psi_in.get_ngk_pointer();
 this->nk = psi_in.get_nk();
 this->nbands = psi_in.get_nbands();
 this->nbasis = psi_in.get_nbasis();
 this->current_k = psi_in.get_current_k();
 this->current_b = psi_in.get_current_b();
 this->k_first = psi_in.get_k_first();
 // this function will copy psi_in.psi to this->psi no matter the device types of each other.

 this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis());

 // Specifically, if the Device_in type is CPU and the Device type is GPU.
 // Which means we need to initialize a GPU psi from a given CPU psi.
 // We first malloc a memory in CPU, then cast the memory from T_in to T in CPU.
 // Finally, synchronize the memory from CPU to GPU.
 // This could help to reduce the peak memory usage of device.
 if (std::is_same<Device, base_device::DEVICE_GPU>::value && std::is_same<Device_in, base_device::DEVICE_CPU>::value)
 {
 auto* arr = (T*)malloc(sizeof(T) * psi_in.size());
 // cast the memory from T_in to T in CPU
 base_device::memory::cast_memory_op<T, T_in, Device_in, Device_in>()(arr,
 psi_in.get_pointer()
 - psi_in.get_psi_bias(),
 psi_in.size());
 // synchronize the memory from CPU to GPU
 base_device::memory::synchronize_memory_op<T, Device, Device_in>()(this->psi,
 arr,
 psi_in.size());
 free(arr);
 }
 else
 {
 base_device::memory::cast_memory_op<T, T_in, Device, Device_in>()(this->psi,
 psi_in.get_pointer() - psi_in.get_psi_bias(),
 psi_in.size());
 }
 this->psi_bias = psi_in.get_psi_bias();
 this->current_nbasis = psi_in.get_current_nbas();
 this->psi_current = this->psi + psi_in.get_psi_bias();
}

template <typename T, typename Device>
void Psi<T, Device>::set_all_psi(const T* another_pointer, const std::size_t size_in)
{
 assert(size_in == this->size());
 synchronize_memory_op()(this->psi, another_pointer, this->size());
}

template <typename T, typename Device>
Psi<T, Device>& Psi<T, Device>::operator=(const Psi<T, Device>& psi_in)
{
// printf("%d\n", &psi_in);
 this->ngk = psi_in.ngk;
 this->nk = psi_in.get_nk();
 this->nbands = psi_in.get_nbands();
 this->nbasis = psi_in.get_nbasis();
 this->current_k = psi_in.get_current_k();
 this->current_b = psi_in.get_current_b();
 this->k_first = psi_in.get_k_first();
 // this function will copy psi_in.psi to this->psi no matter the device types of each other.

 this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis());
 base_device::memory::synchronize_memory_op<T, Device, Device>()(this->psi,
 psi_in.psi,
 psi_in.size());
 this->psi_bias = psi_in.get_psi_bias();
 this->current_nbasis = psi_in.get_current_nbas();
 this->psi_current = this->psi + psi_in.get_psi_bias();

 return *this;
}

template <typename T, typename Device>
void Psi<T, Device>::resize(const int nks_in, const int nbands_in, const int nbasis_in)
{
 assert(nks_in > 0 && nbands_in >= 0 && nbasis_in > 0);

 // This function will delete the psi array first(if psi exist), then malloc a new memory for it.
 resize_memory_op()(this->psi, nks_in * static_cast<std::size_t>(nbands_in) * nbasis_in, "no_record");

 // this->zero_out();

 this->nk = nks_in;
 this->nbands = nbands_in;
 this->nbasis = nbasis_in;
 this->current_nbasis = nbasis_in;
 this->psi_current = this->psi;
 // GlobalV::ofs_device << "allocated xxx MB memory for psi" << std::endl;
}

template <typename T, typename Device>
T* Psi<T, Device>::get_pointer() const
{
 return this->psi_current;
}

template <typename T, typename Device>
T* Psi<T, Device>::get_pointer(const int& ikb) const
{
 assert(ikb >= 0);
 assert(this->k_first ? ikb < this->nbands : ikb < this->nk);
 return this->psi_current + ikb * this->nbasis;
}

template <typename T, typename Device>
const bool& Psi<T, Device>::get_k_first() const
{
 return this->k_first;
}

template <typename T, typename Device>
const Device* Psi<T, Device>::get_device() const
{
 return this->ctx;
}

template <typename T, typename Device>
const int* Psi<T, Device>::get_ngk_pointer() const
{
 return this->ngk;
}

template <typename T, typename Device>
const size_t& Psi<T, Device>::get_psi_bias() const
{
 return this->psi_bias;
}

template <typename T, typename Device>
const int& Psi<T, Device>::get_current_ngk() const
{
 if (this->get_npol() == 1)
 {
 return this->current_nbasis;
 }
 else
 {
 return this->nbasis;
 }
}

template <typename T, typename Device>
int Psi<T, Device>::get_npol() const
{
 if (PARAM.inp.nspin == 4)
 {
 return 2;
 }
 else
 {
 return 1;
 }
}

template <typename T, typename Device>
const int& Psi<T, Device>::get_nk() const
{
 return this->nk;
}

template <typename T, typename Device>
const int& Psi<T, Device>::get_nbands() const
{
 return this->nbands;
}

template <typename T, typename Device>
const int& Psi<T, Device>::get_nbasis() const
{
 return this->nbasis;
}

template <typename T, typename Device>
std::size_t Psi<T, Device>::size() const
{
 if (this->psi == nullptr)
 {
 return 0;
 }
 return this->nk * static_cast<std::size_t>(this->nbands) * this->nbasis;
}

template <typename T, typename Device>
void Psi<T, Device>::fix_k(const int ik) const
{
 assert(ik >= 0);
 this->current_k = ik;
 if (this->ngk != nullptr)
 {
 this->current_nbasis = this->ngk[ik];
 }
 else
 {
 this->current_nbasis = this->nbasis;
 }

 if (this->k_first)
 {
 this->current_b = 0;
 }
 int base = this->current_b * this->nk * this->nbasis;
 if (ik >= this->nk)
 {
 // mem_saver: fix to base
 this->psi_bias = base;
 this->psi_current = const_cast<T*>(&(this->psi[base]));
 }
 else
 {
 this->psi_bias = k_first ? ik * this->nbands * this->nbasis : base + ik * this->nbasis;
 this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
 }
}
template <typename T, typename Device>
void Psi<T, Device>::fix_b(const int ib) const
{
 assert(ib >= 0);
 this->current_b = ib;

 if (!this->k_first)
 {
 this->current_k = 0;
 }
 int base = this->current_k * this->nbands * this->nbasis;
 if (ib >= this->nbands)
 {
 // mem_saver: fix to base
 this->psi_bias = base;
 this->psi_current = const_cast<T*>(&(this->psi[base]));
 }
 else
 {
 this->psi_bias = k_first ? base + ib * this->nbasis : ib * this->nk * this->nbasis;
 this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
 }
}

template <typename T, typename Device>
void Psi<T, Device>::fix_kb(const int ik, const int ib) const
{
 assert(ik >= 0 && ib >= 0);
 this->current_k = ik;
 this->current_b = ib;
 if (ik >= this->nk || ib >= this->nbands)
 { // fix to 0
 this->psi_bias = 0;
 this->psi_current = const_cast<T*>(&(this->psi[0]));
 }
 else
 {
 this->psi_bias = k_first ? (ik * this->nbands + ib) * this->nbasis : (ib * this->nk + ik) * this->nbasis;
 this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
 }
}

template <typename T, typename Device>
T& Psi<T, Device>::operator()(const int ikb1, const int ikb2, const int ibasis) const
{
 assert(ikb1 >= 0 && ikb2 >= 0 && ibasis >= 0);
 assert(this->k_first ? ikb1 < this->nk && ikb2 < this->nbands : ikb1 < this->nbands && ikb2 < this->nk);
 return this->k_first ? this->psi[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis]
 : this->psi[(ikb1 * this->nk + ikb2) * this->nbasis + ibasis];
}

template <typename T, typename Device>
T& Psi<T, Device>::operator()(const int ikb2, const int ibasis) const
{
 assert(this->k_first ? this->current_b == 0 : this->current_k == 0);
 assert(this->k_first ? ikb2 >= 0 && ikb2 < this->nbands : ikb2 >= 0 && ikb2 < this->nk);
 assert(ibasis >= 0 && ibasis < this->nbasis);
 return this->psi_current[ikb2 * this->nbasis + ibasis];
}

template <typename T, typename Device>
T& Psi<T, Device>::operator()(const int ibasis) const
{
 assert(ibasis >= 0 && ibasis < this->nbasis);
 return this->psi_current[ibasis];
}

template <typename T, typename Device>
int Psi<T, Device>::get_current_k() const
{
 return this->current_k;
}

template <typename T, typename Device>
int Psi<T, Device>::get_current_b() const
{
 return this->current_b;
}

template <typename T, typename Device>
int Psi<T, Device>::get_current_nbas() const
{
 return this->current_nbasis;
}

template <typename T, typename Device>
const int& Psi<T, Device>::get_ngk(const int ik_in) const
{
 assert(this->ngk != nullptr);
 return this->ngk[ik_in];
}

template <typename T, typename Device>
void Psi<T, Device>::zero_out()
{
 // this->psi.assign(this->psi.size(), T(0));
 set_memory_op()(this->psi, 0, this->size());
}

template <typename T, typename Device>
std::tuple<const T*, int> Psi<T, Device>::to_range(const Range& range) const
{
 const int& i1 = range.index_1;
 const int& r1 = range.range_1;
 const int& r2 = range.range_2;

 if (range.k_first != this->k_first || r1 < 0
 || r2 < r1
 // || (range.k_first && (r2 >= this->nbands || i1 >= this->nk))
 // || (!range.k_first && (r2 >= this->nk || i1 >= this->nbands)))
 || (range.k_first ? (i1 >= this->nk) : (i1 >= this->nbands)) // illegal index 1
 || (range.k_first ? (i1 > 0 && r2 >= this->nbands) : (i1 > 0 && r2 >= this->nk)) // illegal range of index 2
 || (range.k_first ? (i1 < 0 && r2 >= this->nk) : (i1 < 0 && r2 >= this->nbands))) // illegal range of index 1
 {
 return std::tuple<const T*, int>(nullptr, 0);
 }
 else if (i1 < 0) // [r1, r2] is the range of index1 with length m
 {
 const T* p = &this->psi[r1 * (k_first ? this->nbands : this->nk) * this->nbasis];
 int m = (r2 - r1 + 1) * this->get_npol();
 return std::tuple<const T*, int>(p, m);
 }
 else // [r1, r2] is the range of index2 with length m
 {
 const T* p = &this->psi[(i1 * (k_first ? this->nbands : this->nk) + r1) * this->nbasis];
 int m = (r2 - r1 + 1) * this->get_npol();
 return std::tuple<const T*, int>(p, m);
 }
}

template class Psi<float, base_device::DEVICE_CPU>;
template class Psi<std::complex<float>, base_device::DEVICE_CPU>;
template class Psi<double, base_device::DEVICE_CPU>;
template class Psi<std::complex<double>, base_device::DEVICE_CPU>;
template Psi<std::complex<float>, base_device::DEVICE_CPU>::Psi(
 const Psi<std::complex<double>, base_device::DEVICE_CPU>&);
template Psi<std::complex<double>, base_device::DEVICE_CPU>::Psi(
 const Psi<std::complex<float>, base_device::DEVICE_CPU>&);
#if ((defined __CUDA) || (defined __ROCM))
template class Psi<float, base_device::DEVICE_GPU>;
template class Psi<std::complex<float>, base_device::DEVICE_GPU>;
template Psi<float, base_device::DEVICE_CPU>::Psi(const Psi<float, base_device::DEVICE_GPU>&);
template Psi<float, base_device::DEVICE_GPU>::Psi(const Psi<float, base_device::DEVICE_CPU>&);
template Psi<std::complex<float>, base_device::DEVICE_CPU>::Psi(
 const Psi<std::complex<float>, base_device::DEVICE_GPU>&);
template Psi<std::complex<float>, base_device::DEVICE_GPU>::Psi(
 const Psi<std::complex<float>, base_device::DEVICE_CPU>&);

template class Psi<double, base_device::DEVICE_GPU>;
template class Psi<std::complex<double>, base_device::DEVICE_GPU>;
template Psi<double, base_device::DEVICE_CPU>::Psi(const Psi<double, base_device::DEVICE_GPU>&);
template Psi<double, base_device::DEVICE_GPU>::Psi(const Psi<double, base_device::DEVICE_CPU>&);
template Psi<std::complex<double>, base_device::DEVICE_CPU>::Psi(
 const Psi<std::complex<double>, base_device::DEVICE_GPU>&);
template Psi<std::complex<double>, base_device::DEVICE_CPU>::Psi(
 const Psi<std::complex<float>, base_device::DEVICE_GPU>&);
template Psi<std::complex<double>, base_device::DEVICE_GPU>::Psi(
 const Psi<std::complex<double>, base_device::DEVICE_CPU>&);
template Psi<std::complex<float>, base_device::DEVICE_GPU>::Psi(
 const Psi<std::complex<double>, base_device::DEVICE_CPU>&);
template Psi<std::complex<double>, base_device::DEVICE_GPU>::Psi(
 const Psi<std::complex<float>, base_device::DEVICE_GPU>&);
#endif
} // namespace psi