/**********************************************************

This software is part of J.-S. Caux's ABACUS library.

Copyright (c) J.-S. Caux.

-----------------------------------------------------------

File:  LiebLin_State_Ensemble.cc

Purpose:  State ensembles for Lieb-Liniger

***********************************************************/

#include "ABACUS.h"

using namespace std;
using namespace ABACUS;

namespace ABACUS {

  // Constructors:

  LiebLin_Diagonal_State_Ensemble::LiebLin_Diagonal_State_Ensemble ()
    : nstates(0), state(Vect<LiebLin_Bethe_State>(1)), weight(0.0, 1)
  {
  }

  LiebLin_Diagonal_State_Ensemble::LiebLin_Diagonal_State_Ensemble (const LiebLin_Bethe_State& RefState, int nstates_req)
    : nstates(nstates_req), state(Vect<LiebLin_Bethe_State>(RefState, nstates_req)), weight(1.0/nstates_req, nstates_req)
  {
  }

  // Recursively go through all arbitrary-order type 2 descendents
  void Generate_type_2_descendents (LiebLin_Bethe_State& ActualState, int* ndesc_ptr, const LiebLin_Bethe_State& OriginState)
  {
    int type_required = 3;
    Vect<string> desc_label = Descendents (ActualState, OriginState, type_required);

    if (desc_label.size() > 0) { // follow down recursively
      LiebLin_Bethe_State ActualState_here = ActualState;
      for (int idesc = 0; idesc < desc_label.size(); ++idesc) {
	ActualState_here.Set_to_Label (desc_label[idesc], OriginState.Ix2);
	// output data:
	cout << *ndesc_ptr << "\t" << desc_label[idesc] << endl;
	cout << "origin: " << OriginState.Ix2 << endl;
	cout << "shiftd: " << ActualState_here.Ix2 << endl;
	ActualState_here.Compute_All(false);
	cout << "\t\t\t\t\t\t" << ActualState.E - OriginState.E << endl;
	*ndesc_ptr += 1;
	Generate_type_2_descendents (ActualState_here, ndesc_ptr, OriginState);

	// do them mod st:
	idesc += 6;
      }
    }
    return;
  }


  LiebLin_Diagonal_State_Ensemble::LiebLin_Diagonal_State_Ensemble (DP c_int, DP L, int N, const Root_Density& rho)
  {
    // This function returns a state ensemble matching the continuous density rho.
    // The logic closely resembles the one used in Discretized_LiebLin_Bethe_State.

    //version with 4 states in ensemble

    Root_Density x(rho.Npts, rho.lambdamax);

    for(int ix=0; ix<x.Npts; ++ix) {
      x.value[ix] = x.lambda[ix];
      for( int ip=0; ip<rho.Npts; ++ip) {
	x.value[ix] += 2.0 * rho.value[ip] * rho.dlambda[ip] * atan ((x.lambda[ix] - rho.lambda[ip])/c_int);
      }
      x.value[ix] /= 2.0*PI; //normalization
    }

    // Now carry on as per Discretized_LiebLin_Bethe_State:
    // Each time N \int_{-\infty}^\lambda d\lambda' \rho(\lambda') crosses a half integer, add a particle:
    DP integral = 0.0;
    DP integral_prev = 0.0;
    int Nfound = 0;
    Vect<DP> Ix2_found(0.0, N);
    int Ix2_left, Ix2_right;
    Vect<int> Ix2(N);


    LiebLin_Bethe_State rhostate(c_int, N, L);

    nstates = 4;
    state = Vect<LiebLin_Bethe_State>(rhostate, nstates);
    weight = Vect<DP>(nstates);
    weight[0] = 1.0/nstates;
    weight[1] = 1.0/nstates;
    weight[2] = 1.0/nstates;
    weight[3] = 1.0/nstates;

    int n_moves = 0;
    bool change = true;

    for (int i = 0; i < rho.Npts; ++i) {
      integral_prev = integral;
      integral += L * rho.value[i] * rho.dlambda[i];
      if (integral > Nfound + 0.5) {
	// Subtle error: if the rho is too discontinuous, i.e. if more than one rapidity is found, must correct for this.
	if (integral > Nfound + 1.5 && integral < Nfound + 2.5) { // found two rapidities
	  ABACUSerror("The distribution of particles is too discontinous for discretisation");
	}

	else {
	  // few variables to clarify the computations, they do not need to be vectors, not used outside this loop
	  Ix2_found[Nfound] = 2.0 * L * x.Return_Value(rho.lambda[i]);

	  Ix2_left = floor(Ix2_found[Nfound]);
	  Ix2_left -= (Ix2_left + N + 1)%2 ? 1 : 0;	//adjust parity

	  Ix2_right = ceil(Ix2_found[Nfound]);
	  Ix2_right += (Ix2_right + N + 1)%2 ? 1 : 0;	//adjust parity

	  int Ix2_in = (Ix2_found[Nfound] > 0 ? Ix2_left : Ix2_right);
	  int Ix2_out = (Ix2_found[Nfound] > 0 ? Ix2_right : Ix2_left);
	  cout << rho.lambda[i] << "\t" << x.Return_Value(rho.lambda[i]) << "\t" << Ix2_found[Nfound] << endl;

	  //choose the saddle point state and remember the uncertain choices
	  if(Ix2_found[Nfound] - Ix2_left < 0.5) {
	    state[0].Ix2[Nfound] = Ix2_left;
	    state[1].Ix2[Nfound] = Ix2_left;
	    state[2].Ix2[Nfound] = Ix2_left;
	    state[3].Ix2[Nfound] = Ix2_left;
	  }
	  else if (Ix2_right - Ix2_found[Nfound] < 0.5) {
	    state[0].Ix2[Nfound] = Ix2_right;
	    state[1].Ix2[Nfound] = Ix2_right;
	    state[2].Ix2[Nfound] = Ix2_right;
	    state[3].Ix2[Nfound] = Ix2_right;
	  }
	  else {
	    //it's a kind of magic: the zeroth goes in whle the first goes out,
	    //the second goes left if the third goes right
	    state[0].Ix2[Nfound] = Ix2_in;
	    state[1].Ix2[Nfound] = Ix2_out;
	    state[2+change].Ix2[Nfound] = Ix2_left;
	    state[2+!change].Ix2[Nfound] = Ix2_right;
	    change = !change;
	    ++n_moves;
	  }
	  Nfound++;
	}
      }
    }

    //fix state[3] and state[4]
    for(int i=0; i<N-1; ++i) {
      if(state[2].Ix2[i] == state[2].Ix2[i+1]) {
	if(state[2].Ix2[i] != state[2].Ix2[i-1] + 2) state[2].Ix2[i] -= 2;
	else state[2].Ix2[i+1] += 2;
      }
      if(state[3].Ix2[i] == state[3].Ix2[i+1]) {
	if(state[3].Ix2[i] != state[3].Ix2[i-1] + 2) state[3].Ix2[i] -= 2;
	else state[3].Ix2[i+1] += 2;
      }
    }

    bool all_Diff = true;
    //check that all 'ns' states are different, we assume that Ix2 are ordered
    for(int i=0; i<nstates; ++i)
      for(int j=i+1; j<nstates; ++j) {
	if(state[i].Ix2 == state[j].Ix2)
	  all_Diff = false;
      }

    if(!all_Diff) {
      //check if the first two states are different
      if(state[0].Ix2 == state[1].Ix2)
	ABACUSerror("Cannot create 4 (nor 2) different states in LiebLin_Diagonal_State_Ensemble. "
		    "Consider increasing system size");
      else {
	nstates = 2;
	weight[0] = 0.5;
	weight[1] = 0.5;
      }
    }

    for(int i=0; i<nstates; ++i) {
      state[i].Set_Label_from_Ix2(state[0].Ix2);
      state[i].Compute_All(true);
    }


    cout << weight[0] << "\t" << state[0].Ix2 << endl << weight[0] << "\t" << state[1].Ix2 << endl;

    return;
  }

  LiebLin_Diagonal_State_Ensemble& LiebLin_Diagonal_State_Ensemble::operator= (const LiebLin_Diagonal_State_Ensemble& rhs)
  {
    if (this != &rhs) {
      nstates = rhs.nstates;
      state = rhs.state;
      weight = rhs.weight;
    }

    return(*this);
  }

  void LiebLin_Diagonal_State_Ensemble::Load (DP c_int, DP L, int N, const char* ensfile_Cstr)
  {
    ifstream infile(ensfile_Cstr);

    // Count nr of lines in file:
    string dummy;
    int nrlines = 0;
    while (getline(infile, dummy))  ++nrlines;
    infile.close();

    nstates = nrlines;
    LiebLin_Bethe_State examplestate (c_int, L, N);
    state = Vect<LiebLin_Bethe_State> (examplestate, nstates);
    weight = Vect<DP> (0.0, nstates);

    infile.open(ensfile_Cstr);

    string dummylabel;
    for (int ns = 0; ns < nstates; ++ns) {
      infile >> weight[ns] >> dummylabel;
      for (int i = 0; i < state[ns].N; ++i) infile >> state[ns].Ix2[i];
    }
    infile.close();

    for (int ns = 0; ns < nstates; ++ns) {
      state[ns].Set_Label_from_Ix2 (state[0].Ix2); // labels are ref to the saddle-point state
      state[ns].Compute_All(true);
    }
  }

  void LiebLin_Diagonal_State_Ensemble::Save (const char* ensfile_Cstr)
  {
    ofstream outfile;

    outfile.open(ensfile_Cstr);

    for (int ns = 0; ns < nstates; ++ns) {
      if (ns > 0) outfile << endl;
      outfile << setprecision(16) << weight[ns] << "\t" << state[ns].label << "\t";
      for (int i = 0; i < state[ns].N; ++i) outfile << " " << state[ns].Ix2[i];
    }
  }


  LiebLin_Diagonal_State_Ensemble LiebLin_Thermal_Saddle_Point_Ensemble (DP c_int, DP L, int N, DP kBT)
  {
    // This function constructs a manifold of states around the thermal saddle-point.

    // Start as per Canonical_Saddle_Point_State:

    // This function returns the discretized state minimizing the canonical free energy
    // F = E - T S.

    // This is obtained by rediscretizing the solution coming from TBA.


    // Otherwise, return the discretized TBA saddle-point state:
    LiebLin_TBA_Solution TBAsol = LiebLin_TBA_Solution_fixed_nbar (c_int, N/L, kBT, 1.0e-4, 100);

    LiebLin_Diagonal_State_Ensemble ensemble(c_int, L, N, TBAsol.rho);

    cout << "nbar: " << TBAsol.nbar << endl;
    cout << "ebar: " << TBAsol.ebar << endl;

    cout << ensemble.state[0].E << endl;
    return(ensemble);
  }

} // namespace ABACUS