#include "itensor/all.h"
#include <iostream>
#include <iomanip>

using namespace itensor;

std::vector<double> test1(int argc, char* argv[], Real lambda){

  int N = 50;
  auto sweeps = Sweeps(5);
  sweeps.maxm() = 50,80,100,150,150;
  sweeps.minm() = 10,10,30,50,90;
  sweeps.cutoff() = 1E-6,1E-8,1E-10,1E-12,1E-12;
  sweeps.niter() = 4,3,2,2,2;
  sweeps.noise() = 1E-7,1E-8,1E-10,0,0;
  /*auto sweeps = Sweeps(12);
  sweeps.maxm() = 10,10,10,20,30,50,80,100,150,150,200,250;
  sweeps.minm() = 10,10,10,10,10,10,10,30,50,90,150,200;
  sweeps.cutoff() = 1E-6,1E-6,1E-6,1E-6,1E-6,1E-6,1E-6,1E-8,1E-10,1E-12,1E-12,1E-14,1E-15;
  sweeps.niter() = 4,4,4,4,4,4,3,2,2,2,2,2;
  sweeps.noise() = 1E-7,1E-7,1E-7,1E-7,1E-7,1E-7,1E-8,1E-10,0,0,0,0;*/

  bool use_Z3 = true;
  auto sites = SpinHalf(N);
  auto ampo = AutoMPO(sites);
  for(int j = 1; j < N; ++j){
    ampo += -4.0*lambda,"Sz",j,"Sz",j+1;
  }
  for(int j = 1; j <= N; ++j){
    ampo += -2*(1.0-lambda),"Sx",j;
  }

  // dmrg
  auto H = MPO(ampo);
  auto psi = MPS(sites);
  auto energy = dmrg(psi,H,sweeps);

  // calculate the entanglement entropy
  Real s = 0;
  ITensor S, V;
  int b = N/2;
  psi.position(b);
  ITensor wf = psi.A(b) * psi.A(b+1);
  auto U = psi.A(b);
  auto spectrum = svd(wf,U,S,V);
  for (auto p : spectrum.eigs()){
    if(p > 1E-12) s += -p*log(p);
  }

  auto psi1 = MPS(sites);
  auto wfs = std::vector<MPS>(1);
  wfs.at(0) = psi;
  auto energy1 = dmrg(psi1, H, wfs, sweeps, {"Quiet=", true, "Weight=", 20.0});

  return std::vector<double>{s, energy, energy1};
}

int main(int argc, char* argv[]){
  int N_parameter_sets = 15;
  double lambda_min = 0.01;
  double lambda_max = 0.99;
  double lambda_grad = (lambda_max - lambda_min) / (N_parameter_sets - 1);

  std::vector<double> lambdas;
  for (int i = 0; i < N_parameter_sets; i++){
    double lambda = lambda_min + lambda_grad * i;
    lambdas.push_back(lambda);
  }

  std::vector<std::vector<double> > results;
  for (auto lambda: lambdas){
    auto result = std::vector<double>();
    result.push_back(lambda);
    for (auto x : test1(argc, argv, lambda)){
      result.push_back(x);
    }
    results.push_back(result);
  }

  for (auto r : results){
    std::stringstream ss;
    for (auto y : r){
      //ss << std::fixed << std::setprecision(12) << y << ", ";
      ss << std::setw(12) << std::fixed << std::setprecision(12) << y << ", ";
    }
    /* Output lambda, entropy, E_ground, E_excited */
    std::cout << ss.str() << std::endl;
  }
  return 0;
}