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

  Purpose:
  Produce the .states and .rho files
  respectively containing the states and density matrix elements
  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;


std::mutex bra_queue_mutex_;
std::mutex rho_mutex_;


class Orchestrator
{
public:
  Orchestrator(std::vector<LiebLinigerBetheState>& subbasis_states,
	       std::vector<std::vector<std::complex<Real>>>& rho_ME)
    : subbasis_states_(subbasis_states)
    , rho_ME_(rho_ME)
  {
    // Set up the queue
    for (IndexU ibra { 0 }; ibra < subbasis_states_.get().size(); ++ibra)
      bra_queue_.push_back(ibra);
  }

public:
  std::reference_wrapper<std::vector<LiebLinigerBetheState>> subbasis_states_;
  std::reference_wrapper<std::vector<std::vector<std::complex<Real>>>> rho_ME_;
  std::vector<int> bra_queue_;

public:
  void worker();
};

void Orchestrator::worker()
{
  bool working { true };
  int ibra;

  while (working) {
    bra_queue_mutex_.lock();
    working = !bra_queue_.empty();
    if (working) {
      ibra = bra_queue_.back();  // proceed from last (longest) to first
      bra_queue_.pop_back();
    }
    bra_queue_mutex_.unlock();

    if (!working) break;

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

    for (int iket { 0 }; iket <= ibra; ++iket) {
      computed_rho_ME.push_back(matrix_element_ρ (subbasis_states_.get()[ibra],
						  subbasis_states_.get()[iket]));
    }

    rho_mutex_.lock();
    rho_ME_.get()[ibra].swap(computed_rho_ME);
    rho_mutex_.unlock();
  }
}


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

  if (argc != 7) {
    std::cout << "Executable rho\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 .states and .rho files\n"
	      << " respectively containing the states and density matrix elements\n"
	      << " for (multi-)type II hole wavepackets in Lieb-Liniger.\n";

    std::cout << "\nUsage:\n------\n";
    std::cout << "rho <c> <L> <N> <nr holes> <width> <offset>\n\n";

    int warg { 16 }, wtype { 10 }, wcons { 26 };
    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";

    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]) };

  if (nholes > width) throw std::invalid_argument("nholes > width, bailing out");
  if (nholes > N) throw std::invalid_argument("nholes > N, bailing out");
  if (width + offset > N) throw std::invalid_argument("width+offset > N, bailing out");


  // Build the states defining the sub-basis of the Hilbert space we work in.
  // We define a Young tableau to specify the displacements.
  // The empty tableau means all holes are to the right.
  // The row length defines the particle right-displacement as usual.
  Tableau tableau(width-nholes, nholes);

  BosonicContinuum bc1(L);
  LiebLinigerModel LL1(bc1, c);
  LiebLinigerBetheState llbs(LL1, N);

  std::vector<int> Ix2_gs{llbs.Ix2_};
  std::vector<int> Ix2;


  // Build the sub-basis

  std::vector<LiebLinigerBetheState> subbasis_states;

  for (std::size_t id { 0 }; id <= tableau.maxid(); ++id) {
    tableau.set_to_id(id);
    Ix2 = Ix2_gs;
    llbs.set_ground_state_Ix2();
    // The rightmost nholes + offset q# are each moved nholes units to the right
    for (int i { N - nholes - offset }; i < N; ++i) Ix2[i] += 2* nholes;
    // The next nholes q# are moved right by the length of the tableau row
    for (int i { 0 }; i < tableau.nr_rows(); ++i) Ix2[N-1 - nholes - offset - i] += 2* tableau.row_l(i);
    llbs.set_g_Ix2(Ix2);
    llbs.initialize();
    llbs.compute_Momentum();
    llbs.polish();
    if (!llbs.converged()) {
      std::cout << "state " << llbs.label() << " did not converge,\n"
		<< llbs.iter_details() << std::endl;
      throw AbacusException("Bethe equations solution failed to converge");
    }
    subbasis_states.push_back(llbs);
  }


  // Define the output files

  std::stringstream filename_base;
  filename_base << "c_" << c << "_N_" << N << "_L_" << L << "_nholes_" << nholes;
  filename_base << "_width_" << width;
  filename_base << "_offset_" << offset;

  std::stringstream states_filename;
  states_filename << filename_base.str() << ".states";
  std::ofstream states_file;
  states_file.open(states_filename.str(), std::ios::out | std::ios::trunc);
  states_file << std::setprecision(std::numeric_limits<Real>::digits10 + 1);

  std::stringstream rho_ME_filename;
  rho_ME_filename << filename_base.str() << ".rho";
  std::ofstream rho_ME_file;
  rho_ME_file.open(rho_ME_filename.str(), std::ios::out | std::ios::trunc);
  rho_ME_file << std::setprecision(std::numeric_limits<Real>::digits10 + 1);


  // Output the states

  for (auto state : subbasis_states) {
    states_file << std::endl;
    states_file << state.label() << "\t" << state.iK() << "\t" << state.E();
  }
  states_file.close();


  // Compute the matrix elements

  // Allocate rho_ME_ triangular matrix
  std::vector<std::vector<std::complex<Real>>> rho_ME;
  rho_ME = std::vector<std::vector<std::complex<Real>>>(subbasis_states.size());
  for (IndexU i { 0 }; i < subbasis_states.size(); ++ i)
    rho_ME[i] = std::vector<std::complex<Real>>(i+1);

  Orchestrator orchestrator(subbasis_states, rho_ME);

  // Use nr processors-2 (and at least one) threads
  int nr_threads_used { std::max(1, ι(std::thread::hardware_concurrency() - 2)) };

  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();

  // Output to ME
  for (IndexU ibra { 0 }; ibra < subbasis_states.size(); ++ibra) {
    for (IndexU iket { 0 }; iket <= ibra; ++iket) {
      rho_ME_file << rho_ME[ibra][iket] << "\t";
    }
    rho_ME_file << std::endl;
  }
  rho_ME_file.close();

  return 0;
}