/*! Type II quasisolitons in Lieb-Liniger: rhoxt

  Purpose:
  Produce the .x, .t and .rhoxt files
  respectively containing the coordinate and time vectors, and the rho(x,t) data matrix
  for (multi-)type II hole wavepackets in Lieb-Liniger.

  It automatically makes use of available hardware concurrency.

  This executable requires Abacus version 2.

  See README for compilation instructions.

  Copyright © Jean-Sébastien Caux, Anahita Sarvi and Cesare Vianello.

  This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.

  This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
*/


import std;

import abacus;


class Orchestrator
{
 public:
  Orchestrator(Real c, Real L, int N,
	       int nholes, int width, int offset, std::string protocol,
	       int Nx, int Nt, Real tmax)
    : c_(c), L_(L), N_(N)
    , nholes_(nholes), width_(width), offset_(offset), protocol_(protocol)
    , Nx_(Nx), Nt_(Nt), tmax_(tmax)
    , rho_expval_(std::vector<std::vector<Real>>(Nx+1, std::vector<Real>(Nt+1)))
    , x_(std::vector<Real>(Nx+1))
    , t_(std::vector<Real>(Nt+1))
  {
    set_up();
  }

public:
  Real c_;
  Real L_;
  int N_;
  int nholes_;
  int width_;
  int offset_;
  std::string protocol_;
  int Nx_;
  int Nt_;
  Real tmax_;

  int nr_states_;
  std::vector<std::string> label_;
  std::vector<int> iK_;
  std::vector<Real> E_;

  std::vector<std::complex<Real>> amplitude;

  std::map<std::string, std::map<std::string, std::complex<Real>>> rho_ME_;

  std::vector<std::vector<Real>> rho_expval_;

  std::vector<Real> x_;
  std::vector<Real> t_;

  std::queue<int> ix_queue_;
  std::mutex ix_mutex_;
  std::mutex rho_mutex_;

  void set_up();
  void worker();
  void save();
};

void Orchestrator::set_up()
{
  for (int i { 0 }; i <= Nx_; ++i) ix_queue_.push(i);

  // Define the input files

  std::stringstream filename_base;
  filename_base << "c_" << c_ << "_N_" << N_ << "_L_" << L_
		<< "_nholes_" << nholes_ << "_width_" << width_ << "_offset_" << offset_;

  std::stringstream states_filename;
  states_filename << filename_base.str() << ".states";
  std::ifstream states_file;
  states_file.open(states_filename.str());
  states_file >> std::setprecision(std::numeric_limits<Real>::digits10 + 1);

  std::stringstream amplitudes_filename;
  amplitudes_filename << filename_base.str() << "_" << protocol_ << ".amplitudes";
  std::ifstream amplitudes_file;
  amplitudes_file.open(amplitudes_filename.str());
  amplitudes_file >> std::setprecision(std::numeric_limits<Real>::digits10 + 1);

  std::stringstream rho_ME_filename;
  rho_ME_filename << filename_base.str() << ".rho";
  std::ifstream rho_ME_file;
  rho_ME_file.open(rho_ME_filename.str());
  rho_ME_file >> std::setprecision(std::numeric_limits<Real>::digits10 + 1);


  // Input the states info
  std::string tmp_label;
  int tmp_iK;
  Real tmp_E;

  states_file >> tmp_label;
  do {
    states_file >> tmp_iK >> tmp_E;
    label_.push_back(tmp_label);
    iK_.push_back(tmp_iK);
    E_.push_back(tmp_E);
  } while (states_file >> tmp_label);
  states_file.close();

  nr_states_ = int(label_.size());


  // Input the amplitudes

  std::complex<Real> tmp_amplitude;
  for (int is { 0 }; is < nr_states_; ++is) {
    amplitudes_file >> tmp_amplitude;
    amplitude.push_back(tmp_amplitude);
  }
  amplitudes_file.close();


  // Input the density operator matrix elements

  for (int ibra { 0 }; ibra < nr_states_; ++ibra) {
    for (int iket { 0 }; iket <= ibra; ++iket) {
      rho_ME_file >> rho_ME_[label_[ibra]][label_[iket]];
    }
  }
  rho_ME_file.close();


  // Define the output files

  std::stringstream x_filename;
  x_filename << filename_base.str() << "_" << protocol_
	     << "_Nx_" << Nx_ << "_Nt_" << Nt_ << "_tmax_" << tmax_ << ".x";
  std::ofstream x_file;
  x_file.open(x_filename.str(), std::ios::out | std::ios::trunc);
  x_file << std::setprecision(std::numeric_limits<Real>::digits10 + 1);

  std::stringstream t_filename;
  t_filename << filename_base.str() << "_" << protocol_
	     << "_Nx_" << Nx_ << "_Nt_" << Nt_ << "_tmax_" << tmax_ << ".t";
  std::ofstream t_file;
  t_file.open(t_filename.str(), std::ios::out | std::ios::trunc);
  t_file << std::setprecision(std::numeric_limits<Real>::digits10 + 1);


  // Density as function of space and time

  for (int ix { 0 }; ix <= Nx_; ++ix) {
    x_[ix] = -L_/2 + L_ * ix/Nx_;
    x_file << std::endl << -L_/2 + L_ * ix/Nx_;
  }
  x_file.close();

  for (int it { 0 }; it <= Nt_; ++it) {
    t_[it] = tmax_ * it/Nt_;
    t_file << std::endl << tmax_ * it/Nt_;
  }
  t_file.close();

}

void Orchestrator::save()
{
  std::stringstream filename_base;
  filename_base << "c_" << c_ << "_N_" << N_ << "_L_" << L_ << "_nholes_" << nholes_;
  filename_base << "_width_" << width_;
  filename_base << "_offset_" << offset_;

  std::stringstream rhoxt_filename;
  rhoxt_filename << filename_base.str() << "_" << protocol_
		 << "_Nx_" << Nx_ << "_Nt_" << Nt_ << "_tmax_" << tmax_ << ".rhoxt";
  std::ofstream rhoxt_file;
  rhoxt_file.open(rhoxt_filename.str(), std::ios::out | std::ios::trunc);
  rhoxt_file << std::setprecision(std::numeric_limits<Real>::digits10 + 1);

  for (int ix { 0 }; ix <= Nx_; ++ix) {
    for (IndexU it { 0 }; it <= IndexU(Nt_); ++it) {
      rhoxt_file << rho_expval_[ix][it] << "\t";
    }
    rhoxt_file << std::endl;
  }
  rhoxt_file.close();
}

void Orchestrator::worker ()
{
  while (!ix_queue_.empty()) {

    ix_mutex_.lock();
    int ix_here { ix_queue_.front() };
    ix_queue_.pop();
    ix_mutex_.unlock();

    std::vector<Real> computed_rho_expval;

    for (int it { 0 }; it < t_.size(); ++it) {
      Real rho_value { 0 };
      for (int ibra { 0 }; ibra < nr_states_; ++ibra) {
	for (int iket { 0 }; iket < ibra; ++iket) {
	  // first add off-diagonal values
	  rho_value +=
	    2*std::real(std::conj(amplitude[ibra]) * amplitude[iket]
			* rho_ME_[label_[ibra]][label_[iket]]
			* std::exp(1_ir * ((E_[ibra] - E_[iket]) * t_[it] -
					   (twopi_r * (iK_[ibra] - iK_[iket]) * x_[ix_here]/L_))));
	}
	// then diagonal one (no factor of 2 here)
	rho_value += std::norm(amplitude[ibra]) * std::real(rho_ME_[label_[ibra]][label_[ibra]]);
      }
      computed_rho_expval.push_back(rho_value);
    }
    rho_mutex_.lock();
    rho_expval_[ix_here].swap(computed_rho_expval);
    rho_mutex_.unlock();
  }
}


int main(int argc, char* argv[])
{
  using namespace std::complex_literals;

  if (argc != 11) {
    std::cout << "Executable rhoxt\n"
	      << " part of the Type II quasisolitons in Lieb-Liniger software suite\n"
	      << " copyright © Jean-Sébastien Caux, Anahita Sarvi and Cesare Vianello.\n";
    std::cout << "\nPurpose:\n"
	      << " Produce the .x, .t and .rhoxt files\n"
	      << " respectively containing the coordinate and time vectors, and the rho(x,t) data matrix\n"
	      << " for (multi-)type II hole wavepackets in Lieb-Liniger.\n";

    std::cout << "\nPrerequirements:\n----------------\n"
	      << " - .states and .rho files produced by executable rho\n"
	      << " - .amplitudes file for the required protocol, produced by executable amplitudes_[protocol type]\n";

    std::cout << "\nUsage:\n------\n";
    std::cout << "rhoxt <c> <L> <N> <nr holes> <width> <offset> <protocol> <Nx> <Nt> <tmax>\n\n";

    int warg { 16 }, wtype { 10 }, wcons { 32 };
    std::cout << std::left << std::setw(warg) << "Argument" << std::setw(wtype) << "Type"
	      << std::setw(wcons) << "Constraints" << "Description\n";
    std::cout << std::left << std::setw(warg) << "--------" << std::setw(wtype) << "----"
	      << std::setw(wcons) << "-----------" << "-----------\n";
    std::cout << std::left << std::setw(warg) << "c" << std::setw(wtype) << "Real"
	      << std::setw(wcons) << "> 0"
	      << "Value of the interaction parameter\n";
    std::cout << std::left << std::setw(warg) << "L" << std::setw(wtype) << "Real"
	      << std::setw(wcons) << "> 0"
	      << "System size\n";
    std::cout << std::left << std::setw(warg) << "N" << std::setw(wtype) << "int"
	      << std::setw(wcons) << "> 0"
	      << "Number of particles\n";
    std::cout << std::left << std::setw(warg) << "nr holes" << std::setw(wtype) << "int"
	      << std::setw(wcons) << "1 <= nr holes <= N"
	      << "Number of holes (Type II modes)\n";
    std::cout << std::left << std::setw(warg) << "width" << std::setw(wtype) << "int"
	      << std::setw(wcons) << "nr holes < width <= N"
	      << "Width of the hole window\n";
    std::cout << std::left << std::setw(warg) << "offset" << std::setw(wtype) << "int"
	      << std::setw(wcons) << "0 <= offset <= N-width"
	      << "Offset of the hole window w/r to the right Fermi edge\n";
    std::cout << std::left << std::setw(warg) << "protocol" << std::setw(wtype) << "string"
	      << std::setw(wcons) << "(see amplitudes executables)"
	      << "Protocol used for defining the amplitudes in the wavepacket\n";
    std::cout << std::left << std::setw(warg) << "Nx" << std::setw(wtype) << "int"
	      << std::setw(wcons) << "> 0"
	      << "Number of spatial sampling points\n";
    std::cout << std::left << std::setw(warg) << "Nt" << std::setw(wtype) << "int"
	      << std::setw(wcons) << "> 0"
	      << "Number of time slices\n";
    std::cout << std::left << std::setw(warg) << "tmax" << std::setw(wtype) << "Real"
	      << std::setw(wcons) << "> 0"
	      << "Maximal time to reach\n";

    return 0;
  }

  std::cout << std::setprecision(std::numeric_limits<Real>::digits10 + 1);

  Real c { std::stold(argv[1]) };
  Real L { std::stold(argv[2]) };
  int N { std::stoi(argv[3]) };
  int nholes { std::stoi(argv[4]) };
  int width { std::stoi(argv[5]) };
  int offset { std::stoi(argv[6]) };
  std::string protocol { argv[7] };   // protocol used to define the amplitudes
  int Nx { std::stoi(argv[8]) };      // number of points for x axis
  int Nt { std::stoi(argv[9]) };      // number of points for t slices
  Real tmax { std::stold(argv[10]) };  // time up to which movie is produced

  Orchestrator orchestrator(c, L, N, nholes, width, offset, protocol, Nx, Nt, tmax);

  int nr_threads_used { ι(std::thread::hardware_concurrency()*2/3) + 1 };

  std::vector<std::thread> threads;

  for (int i { 0 }; i < nr_threads_used; ++i)
    threads.emplace_back(&Orchestrator::worker, &orchestrator);

  for (auto& t : threads) t.join();

  orchestrator.save();

  return 0;
}