Untitled

//#include "/home/lipegal/Documents/OIST/code/Modular/kp_lib.h"
#include "/flash/BuschU/andres/Transfer_protocol/kp_lib.h"
#include <vector>
#include <functional>
#include <iostream>
#include <fstream>
#include <cmath>
#include <boost/numeric/odeint.hpp>
//signature
using namespace boost::numeric::odeint;
using namespace std;

int in_trunc_dim = 100, num_comp = 300;
int evo_trunc_d = 20;
double tol = 1e-8;
double dt = 1e-4;

// Blackman window
double a0 = 0.42, a1 = 0.5, a2 = 0.08, a3 = 0.012604;

double dip_(int n, int m, double L){
  if(n == m) return ((L * L)/ 12.) - (L * L) / (2 * PI * PI * n * n);
  if((n + m) % 2 == 0) return (8 * L * L * m * n) / (PI * PI * pow(n * n - m * m, 2));
  return 0.0;
}

double dip_x(int n, int m, double L)
{
  if (n == m) return 0.0;
  if ((n + m) % 2 == 0) return 0.0;

  double sign = pow(-1.0, (n + m + 1) / 2.0);
  return (-4.0 * L / (M_PI * M_PI)) * sign *
    (1.0 / pow(n - m, 2) + 1.0 / pow(n + m, 2));
}

// Contrasted with S. Iqbal - F.saif paper
double dip_x2(int n, int m, double a)
{
  // ---- Diagonal term ----
  if (n == m) {
    return (a * a) / 3.0
      + (2.0 * a * a * pow(-1.0, n))
      / (n * n * M_PI * M_PI);
  }

  // ---- Selection rule: same parity only ----
  if ( (n + m) % 2 != 0 )
    return 0.0;

  // ---- Off-diagonal term ----
  double sign = pow(-1.0, (m - n) / 2.0);

  return (8.0 * a * a / (M_PI * M_PI)) * sign *
    ( 1.0 / pow(m - n, 2)
      - 1.0 / pow(m + n, 2) );
}

// Nuttall envelope
double nuttall(double t, double t_pi) {
  double n = t / t_pi;
  return a0 - a1 * cos(2 * PI * n) + a2 * cos(4 * PI * n) - a3 * cos(6 * PI * n);
}

// Blackman envelope
double blackman(double t, double t_pi) {
  double n = t / t_pi;
  return a0 - a1 * cos(2. * PI * n) + a2 * cos(4. * PI * n);
}

// Gaussian envelope
double gaussian(double t, double t_pi, double sigma) {
  return exp(- 0.5 * pow( (t - (0.5 * t_pi)) / sigma, 2));
}

struct EvoResult {
  ComplexVector final_state;
  double max_population;
  int steps;
};

EvoResult evolve_state_fixed_dt(ComplexVector psi_init, double t_tot, double omega, int evo_trunc_d,
                                auto& schrodinger_rhs, const vector<ComplexVector>& eigs_in, int target_level,
                                ofstream& fidelFile, int L, ofstream* wfFile=nullptr) {
  ComplexVector psi_t = psi_init;
  double t = 0.0;
  int step_count = 0;

  // ---- FIXED DT FROM OMEGA ----
  const int Npp = 5000; // points per period
  const double dt = 2.0 * M_PI / (omega * Npp);

  runge_kutta4<ComplexVector> stepper;

  double output_interval = 0.1;
  double next_output_time = 0.0;

  cout << "Fixed-step RK4 evolution\n";
  cout << "omega = " << omega << "\n";
  cout << "dt    = " << dt << "\n";
  cout << "steps = " << int(t_tot / dt) << endl;

  while (t < t_tot) {
    stepper.do_step(schrodinger_rhs, psi_t, t, dt);
    t += dt;
    step_count++;

    normalize(psi_t);
    // ---- OUTPUT ----
    while (t >= next_output_time) {
      //wf_data wf_t;
      //build_wf(psi_t, wf_t, L);
      //save_wf_par(wfFile, wf_t, t);

      fidelFile << t;

      // Eigenstates overlaps
      for (int n = 0; n <= target_level + 2; ++n) {
        if (n < 0 || n >= evo_trunc_d) {
          fidelFile << " 0.0";
          continue;
        }
        Complex ovlp = 0.0;
        for (int i = 0; i < evo_trunc_d; ++i)
          ovlp += conj(eigs_in[n](i)) * psi_t(i);
        fidelFile << " " << norm(ovlp);
      }

      fidelFile << "\n";

      next_output_time += output_interval;
    }

    if (step_count % 1000 == 0) {
      cout << "Progress: " << (t / t_tot) * 100.0 << "%\r" << flush;
    }
  }

  cout << "\nEvolution completed: " << step_count << " steps\n";

  return {psi_t, 0.0, step_count};
}


RealMatrix build_H_with(double pref, double L, double trunc_dim) {
  HamKernel kernel = [&](int n, int m) {
    return kin(n , m , L) + pref * dip_x2(n, m , L/2);
  };
  return build_H(trunc_dim, kernel);
}

// ------------------------------------------------------------------
// Main
// ------------------------------------------------------------------
int main(int argc, char* argv[]) {
  if (argc < 2) {
    cerr << "Usage: ./run [target_level] [alpha]" << endl;
    return 1;
  }

  int target_level = atoi(argv[1]);
  cout << "Target level: " << target_level << endl;

  // -------------------- System setup ------------------------------
  double doms_num = 1;
  std::vector<Domain> domains;
  double L_[] = {1};
  int    M_[] = {0};
  double B_[] = {0.0};
  double W_[] = {0.0}; // Empty well
  double D_[] = {0.0};
  double h_[] = {0.0};

  dom_load(domains, L_, M_, B_, W_, D_, h_, doms_num);
  print_syst_info(domains);

  double L = syst_L(domains);
  RealVector offsets = offsets_(domains);

  //--------------------------- EMPTY WELL HAMILTONIAN ----------------------------
  dom_load (domains, L_, M_, B_, W_, D_, h_, doms_num);

  // Generate Hamiltonian function
  auto H_in = [&](int n, int m) {
    return Ham_uni(n, m, domains, offsets, L);
  };

  Ham_data in_data = H_diag (H_in, in_trunc_dim, tol, num_comp);

  cout << "--Initial state succesfully extracted--" << endl;
  cout << "in_dim = " << in_data.dim << endl;

  // Ground state of the empty well
  ComplexVector psi_in = eigen_extract(in_data, 0);
  normalize(psi_in);

  //----------------------------------------------------------------------

  // Initial and final eigenstates
  vector<ComplexVector> eigs_in;
  for (int k = 0; k < evo_trunc_d; ++k) {
    eigs_in.push_back(eigen_extract(in_data, k));
  }
  in_data = truncate_Ham_data(in_data, evo_trunc_d);

  double omega = in_data.eigvals[target_level - 1] - in_data.eigvals[0];
  double dipole_val = dip_x2(target_level, 1, L/2);

  // Adaptive parameter
  double alpha =  atof(argv[2]);

  // Gaussian parameters
  double lambda = alpha * (PI * PI / (L * L * dipole_val));
  double sigma = (L * L) / (sqrt(2 * PI) * alpha * PI);
  double t_pi = 10 * sigma;
  double t_tot = t_pi + 0.5 * t_pi;

  if (fabs(dipole_val) < 1e-6) {
    cerr << "Dipole too small for target level " << target_level << endl;
    return 1;
  }

  cout << "t_pi = " << t_pi << " , sigma = " << sigma << " ,lambda = " << lambda << " ,omega = "<< omega << endl;

  //for (int n = 1; n < 50; ++n) cout << "<1|x|" << n + 1 << "> = " << dip(n + 1, 1,L) << endl;

  // To print the gaussian
  /*
    ofstream gaussFile("gaussian.dat", ios::trunc);
    double t_curr = 0;
    double dtg= 0.01;
    while (t_curr < t_pi) {
    gaussFile << t_curr << " " << gaussian(t_curr, t_pi,sigma) << endl;
    t_curr += dtg;
    }

    // Blackman - Nutall parameters
    double t_pi = (L * L) / (a0 * alpha * PI);
    double lambda = PI / (dipole_val * a0 * t_pi);

    cout << "  omega = " << omega << ", dipole = " << dipole_val << endl;
    cout << "  alpha = " << alpha << ", t_pi = " << t_pi << ", lambda = " << lambda << endl;
  */

  // ----------------- Files definition ---------------------
  std::ostringstream fidel_th_ss;
  //std::ostringstream wf_th_ss;

  fidel_th_ss << "data/fidelity_" << std::fixed << std::setprecision(3) << alpha <<"_" << target_level << ".dat";
  //wf_th_ss << "data/wf_" << std::fixed << std::setprecision(3) << alpha << "_" << target_level << ".dat";

  //fidel_th_ss << "data/fidelity.dat";
  //wf_th_ss << "data/wf_prep.dat";

  ofstream fidel_thFile(fidel_th_ss.str(), ios::trunc);
  //ofstream wf_thFile(wf_th_ss.str(), ios::trunc);

  ofstream summaryFile("data/summary_fidelity.dat", ios::app);

  // --------------------------TIME EVOLUTION------------------------------
  ComplexVector psi_t = psi_in;

  RealMatrix H_kin = build_H_with(0.0, L, evo_trunc_d);
  RealMatrix H_drive = build_H_with(1.0, L, evo_trunc_d);

  auto schrodinger_rhs = [&](const ComplexVector &psi, ComplexVector &dpsi, double t) {
    double pref = 0.0;
    if (t <= t_pi) {
      pref = 1 * gaussian(t, t_pi, sigma) * lambda * cos(omega * t);
    }
    for (int i = 0; i < evo_trunc_d; ++i) {
      Complex sum = 0.;
      for (int j = 0; j < evo_trunc_d; ++j) {
        sum += (H_kin(i,j) + pref * H_drive(i,j) ) * psi(j);
      }
      dpsi(i) = Complex(0.0, -1.0) * sum;
    }
  };

  auto result_theo = evolve_state_fixed_dt(psi_in, t_tot, omega, evo_trunc_d, schrodinger_rhs,
                                           eigs_in, target_level, fidel_thFile, L);
  cout << "\n=== PROTOCOL COMPLETE ===\n";

  // --- Compute final fidelity with target eigenstate
  ComplexVector psi_final = result_theo.final_state;

  Complex overlap = 0.0;
  for (int i = 0; i < evo_trunc_d; ++i) {
    overlap += conj(eigs_in[target_level - 1](i)) * psi_final(i);
  }

  double final_fidelity = norm(overlap);
  summaryFile << alpha << " " << target_level << " " << t_pi << " " << final_fidelity << endl;

  return 0;
}