#include <iostream>
#include <iomanip>      // std::setw

#include "softwareAstra.h"
#include "generatorParticle.h"
#include "mixedTools.h"
#include "PhysicalConstants.h"
#include "INIReader.h"

softwareAstra::softwareAstra() : abstractSoftware()
{
  INIReader reader("../../pspa_config.ini");
  debug = reader.GetInteger("Debugging", "debug_mode", -1);
  
  nameOfSoftware_ = new nomDeLogiciel("astra");
  userDir_ = string();
}

softwareAstra::softwareAstra(computingBlock *cmpb,dataManager *dt) : abstractSoftware(cmpb,dt)
{
  INIReader reader("../../pspa_config.ini");
  debug = reader.GetInteger("Debugging", "debug_mode", -1);
  
  nameOfSoftware_ = new nomDeLogiciel("astra");
  userDir_ = dataManager_->getUserDir();
}

bool softwareAstra::createInputFile(particleBeam *beamBefore)
{
  if(debug > 0) cout << "softwareAstra::createInputFile()\n";

  const vector< pair< string,vector<smap> > >& commands= getComputingBlock()->actions();
  
  string action= commands.at( 0 ).first;
  if(action != "track") {
    dataManager_->consoleMessage("softwareAstra::createInputFile: program astra used only for tracking particles");
    return false;
  }
  
  // input file for astra has to have the extension .in
  string name= userDir_ + simulationId_ + ".in";
  ofstream outfile;
  outfile.open(name.c_str(),ios::out);
  if (!outfile) {
    dataManager_->consoleMessage("softwareAstra::createInputFile: error opening output stream " + name +"\n");
    return false;
  }

  const vector<smap>& v= commands.at( 0 ).second;
  string fname= v[0].find("fileName")->second.at(0);
  if (fname.empty()) {
    string wDialog= "astra requires an input particle distribution";
    dataManager_->messageEcran("softwareAstra::createInputFile: WARNING ",wDialog);
    return false;
  }
 
  // The namelist NEWRUN contains general instructions for the tracking
  outfile << newrun( v[0] ) << endl;
  
  // namelist OUTPUT
  outfile << output( v[1] ) << endl;

  // settings for the space charge calculation
  outfile << charge( v[2] ) << endl;

  // include apertures
  outfile << aperture( v[3] ) << endl;

  // cavity & solenoid
  outfile << cavity() << endl;
  outfile << solenoid() << endl;
  
  return true;
}

// namelist[0] NEWRUN
string softwareAstra::newrun(const smap& v)
{
  if(debug > 0) cout << "softwareAstra::newrun()\n";
  
  ostringstream os;
  os << "&NEWRUN\n";
  os << "\tHead=" << simulationId_ << "\n";
  os << "\tRUN=1\n";

  os << "\tDistribution=" << "'" + v.find("fileName")->second.at(0) + "'\n";
  os << "\tQbunch=" << v.find("qBunch")->second.at(0) << "\n";
  os << "\tXoff=" << v.find("xOff")->second.at(0) << "\n";
  os << "\tYoff=" << v.find("yOff")->second.at(0) << "\n";
  os << "\tTRACK_ALL=" << v.find("trackAll")->second.at(0) << "\n";
  os << "\tcheck_ref_part=.T\n"; //if true, the run will be interrupted if the reference particle is lost during the onand off-axis reference particle tracking.
  os << "\tPHASE_SCAN=.F\n";
  os << "\tAUTO_PHASE=" << v.find("autoPhase")->second.at(0) << "\n";
  os << "\tH_max=0.001\n";

  os << "/" << endl;
  return os.str();
}

// namelist[1] OUTPUT
string softwareAstra::output(const smap& v)
{
  if(debug > 0) cout << "softwareAstra::output()\n";

  ostringstream os;
  os << "&OUTPUT\n";
  string txt;

  txt= v.find("zStart")->second.at(0);
  double zstart= atof( txt.c_str() );
  os << "\tZSTART=" << zstart << "\n";
  txt= v.find("zStop")->second.at(0);
  double zstop= atof( txt.c_str() );
  os << "\tZSTOP=" << zstop << "\n";
  txt= v.find("zEmit")->second.at(0);
  int zemit= atoi( txt.c_str() );
  os << "\tZemit=" << zemit << "\n";
  txt= v.find("zPhase")->second.at(0);
  int zphase= atoi( txt.c_str() );
  os << "\tZphase=" << zphase << "\n";

  // tracking will stop when the bunch center passes ZSTOP
  ZStop_= zstop;
  
  /* ************************************************************* 
  // additional position for the generation of output
  double ln= 0.0;
  for(int k = 0; k < getComputingBlock()->getNumberOfElements(); k++)
  {
  abstractElement* ptr = getComputingBlock()->getElement(k);
  ln+= ptr->getLengthOfElement();
  os << "\tScreen(" << k+1 << ")="  << ln << ",\n";
  }
  ************************************************************* */
  os << "\tLmagnetized=.F\n";
  os << "\tRefS=.T\n";
  os << "\tEmitS=.T\n";
  os << "\tPhaseS=.T\n";

  os << "/" << endl;
  return os.str();
}

// namelist[2] CHARGE
string softwareAstra::charge(const smap& v)
{
  if(debug > 0) cout << "softwareAstra::charge()\n";

  ostringstream os;
  os << "&CHARGE\n";

  os << "\tLSPCH=" << v.find("LSPCH")->second.at(0) << "\n";
  if( v.size() > 1) {
    string txt;
    txt= v.find("nRad")->second.at(0);
    os << "\tNrad=" << atoi( txt.c_str() ) << "\n";
    txt= v.find("cellVar")->second.at(0);
    os << "\tCell_var=" << atof( txt.c_str() ) << "\n";
    txt= v.find("nlongIn")->second.at(0);
    os << "\tNlong_in=" << atoi( txt.c_str() ) << "\n";
    txt= v.find("nMin")->second.at(0);
    os << "\tN_min=" << atof( txt.c_str() ) << "\n";
    txt= v.find("minGrid")->second.at(0);
    os << "\tmin_grid=" << atof( txt.c_str() ) << "\n";
    txt= v.find("maxScale")->second.at(0);
    os << "\tMax_scale=" << atof( txt.c_str() ) << "\n";
    txt= v.find("lMirror")->second.at(0);
    os << "\tLmirror=" << txt << "\n";
  }
  
  os << "/" << endl;
  return os.str();
}

// namelist[3] APERTURE
string softwareAstra::aperture(const smap& v)
{
  if(debug > 0) cout << "softwareAstra::aperture()\n";

  ostringstream os;
  os << "&APERTURE\n";

  os << "\tLApert=" << v.find("LApert")->second.at(0) << "\n";
  for (int k = 1; k < v.size(); ++k)
    {
      smap::const_iterator it= v.find( mixedTools::intToString(k) );
      if( it == v.end() ) continue;
      os << "\tFile_Aperture(" << k << ")= " << it->second.at(0) << "\n";
      os << "\tAp_Z1(" << k << ")= " << it->second.at(1) << "\n"; // m
      os << "\tAp_Z2(" << k << ")= " << it->second.at(2) << "\n"; // m
      os << "\tAp_R("  << k << ")= " << it->second.at(3) << "\n"; // mm
    }
  
  os << "/" << endl;
  return os.str();
}

// namelist[4] ERROR
string softwareAstra::error(const smap& v)
{
  if(debug > 0) cout << "softwareAstra::error()\n";

  ostringstream os;
  os << "&ERROR\n";
  if( !v.empty() ) {
    os << "\tLError=.T\n";
  }
  os << "/" << endl;
  return os.str();
}

// namelist[5] SCAN
string softwareAstra::scan(const smap& v)
{
  if(debug > 0) cout << "softwareAstra::scan()\n";

  ostringstream os;
  os << "&SCAN\n";
  if( !v.empty() ) {
    os << "\tLSCan=.T\n";
  }
  os << "/" << endl;
  return os.str();
}

// namelist[6] MODULES

// namelist FEM
string softwareAstra::fem()
{
  if(debug > 0) cout << "softwareAstra::fem()\n";

  ostringstream os;
  os << "&FEM\n";
  os << "/" << endl;
  return os.str();
}

// namelist CAVITY
string softwareAstra::cavity()
{
  vector<abstractElement*> LFields;
  unsigned nElts= getComputingBlock()->getNumberOfElements();
  for(unsigned k = 0; k < nElts; k++)
    {
      abstractElement *ptr= getComputingBlock()->getElement( k );
      string eType= ptr->getGenericName();
      if(eType == "linac_cavity") LFields.push_back( ptr );
    }

  ostringstream os;
  os << "&CAVITY\n";

  ostringstream osx;
  for(unsigned k = 0; k < LFields.size(); k++)
    {
      abstractElement *ptr= LFields.at( k );
      vector<psvs> v= ptr->parametersToSoftware();

      // filename of the rf field
      string efield= "'" + v.at(2).second.at(0) + "'";
      osx << "\tFile_Efield(" << k+1 << ")=" << efield << "\n";
      // frequency of the rf field [GHz]
      osx << "\tNue("<< k+1 << ")=" << v.at(2).second.at(2) << "\n";
      // maximum field amplitude [MV/m]
      osx << "\tMaxE(" << k+1 << ")=" << v.at(2).second.at(3) << "\n";
      // phase of the rf field
      osx << "\tPhi(" << k+1 << ")=" << v.at(2).second.at(4) << "\n";
      // shifts the longitudinal cavity position
      osx << "\tC_pos(" << k+1 << ")=" << v.at(2).second.at(5) << "\n";
      // number of cells for TWS
      osx << "\tC_numb(" << k+1 << ")=" << v.at(2).second.at(1) << "\n";
    }
  
  if (LFields.empty()) {
    os << "\tLEField=.F\n";
  } else {
    os << "\tLEField=.T\n";
    os << osx.str();
  }
  
  os << "/" << endl;
  return os.str();
}

// namelist SOLENOID
string softwareAstra::solenoid()
{
  vector<abstractElement*> LFields;
  unsigned nElts= getComputingBlock()->getNumberOfElements();
  for(unsigned k = 0; k < nElts; k++)
    {
      abstractElement *ptr= getComputingBlock()->getElement( k );
      string eType= ptr->getGenericName();
      if(eType == "linac_cavity") LFields.push_back( ptr );
    }

  ostringstream os;
  os << "&SOLENOID\n";
 
  // if false, all solenoid fields are turned off
  bool LBfield= false;
  
  ostringstream osx;
  for(unsigned k = 0; k < LFields.size(); k++)
    {
      abstractElement *ptr= LFields.at( k );
      vector<psvs> v= ptr->parametersToSoftware();
     
      string bfield= v.at(3).second.at(0);
      if( bfield.empty() ) continue;

      LBfield= true;
      // filename of the solenoid field
      osx << "\tFile_Bfield(" << k+1 << ")=" << "'"+bfield+"'" << "\n";
      // maximum field value [T]
      osx << "\tMaxB(" << k+1 << ")=" << v.at(3).second.at(1) << "\n";
      // shifts the longitudinal solenoid position [m]
      osx << "\tS_pos(" << k+1 << ")=" << v.at(3).second.at(2) << "\n";
      // horizontal offset of the solenoid [m]
      osx << "\tS_xoff(" << k+1 << ")=" << v.at(3).second.at(3) << "\n";
      // vertical offset of the solenoid [m]
      osx << "\tS_yoff(" << k+1 << ")=" << v.at(3).second.at(4) << "\n";
      // field expansion extends only to 1st order
      osx << "\tS_higher_order(" << k+1 << ")=.F\n";
    }

  if (LBfield) {
    os << "\tLBField=.T\n";
    os << osx.str();
  } else {
    os << "\tLBField=.F\n";
  }
   
  os << "/" << endl;
  return os.str();
}

bool softwareAstra::execute()
{
  cout << "softwareAstra::execute()\n";
  
  ostringstream sortie;
  bool status= true;

  string wrkDir = dataManager_->getSoftwareDir();
  string mjob = wrkDir.append("astra");
  if (!mixedTools::exists( mjob.c_str() )) {
    
    sortie << "Error in softwareAstra::execute:: " << mjob << " does not exist\n";
    status= false;
    
  } else {

    string nameInput= simulationId_ + ".in";
    //sortie << " run astra " << nameInput << endl;
    dataManager_->consoleMessage("run astra => "+nameInput+"\n");
    
    mjob += string(" ") + nameInput;
    string resultOfRun;
    bool success= launchJob(mjob,resultOfRun);

    if ( !success ) {
      status= false;
    } else {
      sortie << "astra has finished normally\n";
      dataManager_->consoleOutput(resultOfRun);

      // creates the output file => xxx-output.txt;
      string nameOut= userDir_ + getFileName("output",false);
      ofstream outfile;
      outfile.open(nameOut.c_str(),ios::out);
      if ( !outfile ) {
	sortie << "error opening astra output " << nameOut << endl;
	status= false;
      } else {
	// copy output in 'xxx-output.txt"
	outfile << resultOfRun << endl;
	outfile.close();
      }
    }
  }
  
  dataManager_->consoleMessage(sortie.str());
  return status;
}

bool softwareAstra::buildBeamAfterElements()
{
  cout << "softwareAstra::buildBeamAfterElements()\n";
  
  // The approximate longitudinal position of a saved particle distribution is indicated in the file name as a four digit number, which corresponds in general to the rounded beam position in cm.
  ostringstream os;
  os.fill('0');
  os << std::setw(4) << 100*ZStop_;
  string fname= userDir_ + simulationId_;
  string output= fname + "." + os.str() + ".001";
  
  cout << "ASTRA:: output= " << output << endl;

  bareParticle refPart;
  vector<bareParticle> particles;
  bool ok;
  ok= beamFromDistribution(output,refPart,particles);
  
  if (ok) {   
    particleBeam *newDiag= dataManager_->updateCurrentDiagnostic(true);
    // pour l'instant on choisit un centroid nul;
    vector<double> centroid= vector<double>(6,0.0);
    newDiag->setWithParticles(centroid,refPart,particles); 
  } else {
    // si echec
    dataManager_->consoleMessage("softwareAstra::buildBeamAfterElements: error in reading particle distribution file" + output);
    return false;
  }

  return true;
}

bool softwareAstra::beamFromDistribution(string fname,bareParticle& refPart,vector<bareParticle>& particles)
{
  cout << "softwareAstra::beamFromDistribution()\n";
  
  FILE *filefais= fopen(fname.c_str(),"r");  if ( filefais == (FILE*)0 ) {
    dataManager_->consoleMessage("softwareAstra::beamFromDistribution: error opening file => " + fname);
    return false;
  }
  cout << "softwareAstra::beamFromDistribution: open " << fname << endl;

  generatorParticle partic;
  std::vector<generatorParticle> faisceau;
  
  // the first line of the file defines the reference particle
  double timeRef= 0.0;
  double bgZRef = 0.0;
  if (partic.readFromGeneratorFile(filefais) > 0) {
    timeRef= partic.clock;
    // position in [m]
    TRIDVECTOR posRef(partic.xx,partic.yy,partic.zz);
    bgZRef= partic.pz/EREST_eV;
    // impulsion [eV/c]
    TRIDVECTOR betagammaRef(partic.px/EREST_eV,partic.py/EREST_eV,bgZRef);
    refPart= bareParticle(posRef,betagammaRef);
    
    // the coordinates of the reference particle are in absolute coordinates. Longitudinal particle coordinates, i.e. z, pz and t are given relative to the reference particle.
    partic.pz= 0.0;
    faisceau.push_back( partic );
  }

  int nbProbPart= 0;
  while (partic.readFromGeneratorFile(filefais) > 0) 
    {
      faisceau.push_back( partic );
      if ( partic.flag != -1 ) nbProbPart++;
    }

  if (faisceau.size() == 0) {
    dataManager_->consoleMessage("%%ERROR in softwareAstra::beamFromDistribution => no particle found");
      return false;
    }

  // particlesPassives : on ne fait rien de ces particules pour l'instant
  vector<bareParticle> 	passiveParticles;
  
  particles.clear();
  for (unsigned k = 0; k < faisceau.size(); k++) 
    {
      double bgx = faisceau.at(k).px / EREST_eV;
      double bgy = faisceau.at(k).py / EREST_eV;
      double bgz = faisceau.at(k).pz / EREST_eV;
      bgz += bgZRef;
      TRIDVECTOR betagam(bgx,bgy,bgz);
      // ? double gamma = sqrt(1.0 + betagam.norm2());

      // decalage temporel par rapport a la reference  ?
      // on prend le faisceau sous la forme "a un instant donne"
      double cdt= CLIGHT_m_per_ns*(faisceau.at(k).clock-timeRef); // [m]
      double x = faisceau.at(k).xx;
      double y = faisceau.at(k).yy;
      TRIDVECTOR pos(x,y,cdt); // [m]
      particles.push_back( bareParticle(pos,betagam) );
    } //k

  return true;
}

json softwareAstra::readOpticalParameters()
{
  cout << "softwareAstra::readOpticalParameters()\n";
 
  ifstream infile;
  // get the file xxx.Xemit.001
  string f1= userDir_ + simulationId_ + ".Xemit.001";
  infile.open(f1.c_str(),ios::in);
  const smap& xemit= getTableContents( infile );
  
  infile.close();
  
  // get the file xxx.Yemit.001
  string f2= userDir_ + simulationId_ + ".Yemit.001";
  infile.open(f2.c_str(),ios::in);
  const smap& yemit= getTableContents( infile );
  
  infile.close();
  
  // get the file xxx.Zemit.001
  string f3= userDir_ + simulationId_ + ".Zemit.001";
  infile.open(f3.c_str(),ios::in);
  const smap& zemit= getTableContents( infile );

  infile.close();

  json j;
  j["software"]= "astra";

  const vector< pair< string,vector<smap> > >& commands= getComputingBlock()->actions();
  j["action"]= commands.at(0).first;
  
  smap::const_iterator it; 
  it= xemit.find("c1"); // z
  for(int k = 0; k < it->second.size(); k++)
    {
      double value= atof( it->second.at(k).c_str() );
      j["zcoor"].push_back( value );
    }

  it= xemit.find("c4"); // xrms
  for(int k = 0; k < it->second.size(); k++)
    {
      double value= atof( it->second.at(k).c_str() );
      j["xsigma"].push_back( value );
    }

  it= yemit.find("c4"); // yrms
  for(int k = 0; k < it->second.size(); k++)
    {
      double value= atof( it->second.at(k).c_str() );
      j["ysigma"].push_back( value );
    }
  
  it= zemit.find("c4"); // zrms
  for(int k = 0; k < it->second.size(); k++)
    {
      double value= atof( it->second.at(k).c_str() );
      j["zsigma"].push_back( value );
    }

  it= xemit.find("c6"); // xemit
  for(int k = 0; k < it->second.size(); k++)
    {
      double value= atof( it->second.at(k).c_str() );
      j["xemit"].push_back( value );
    }

  it= yemit.find("c6"); // yemit
  for(int k = 0; k < it->second.size(); k++)
    {
      double value= atof( it->second.at(k).c_str() );
      j["yemit"].push_back( value );
    }

  it= zemit.find("c3"); // Ekin (MeV)
  for(int k = 0; k < it->second.size(); k++)
    {
      double xrms= j.find("xsigma")->at( k );
      double xeps= j.find("xemit")->at( k );
      double Ekin= atof( it->second.at(k).c_str() );
      double ratio= (EREST_MeV+Ekin)/EREST_MeV;
      double xbeta= ratio*xrms*xrms/xeps;
      j["xbeta"].push_back( xbeta );
    }

  for(int k = 0; k < it->second.size(); k++)
    {
      double yrms= j.find("ysigma")->at( k );
      double yeps= j.find("yemit")->at( k );
      double Ekin= atof( it->second.at(k).c_str() );
      double ratio= (EREST_MeV+Ekin)/EREST_MeV;
      double ybeta= ratio*yrms*yrms/yeps;
      j["ybeta"].push_back( ybeta );
    }

  smap::const_iterator jt= zemit.find("c5"); // DErms (KeV)
  for(int k = 0; k < it->second.size(); k++)
    {
      double Ekin= atof( it->second.at(k).c_str() );
      double DErms= atof( jt->second.at(k).c_str() );
      double ratio= (0.1*DErms)/Ekin; // %
      j["spread"].push_back( ratio );
    }

  it= xemit.find("c7"); // xx' (mm)
  for(int k = 0; k < it->second.size(); k++)
    {
      double xrms= j.find("xsigma")->at( k );
      double xbeta= j.find("xbeta")->at( k );
      double value= atof( it->second.at(k).c_str() );
      double xalfa= -1.0*value*xbeta/xrms;
      j["xalfa"].push_back( xalfa );
    }

  it= yemit.find("c7"); // yy' (mm)
  for(int k = 0; k < it->second.size(); k++)
    {
      double yrms= j.find("ysigma")->at( k );
      double ybeta= j.find("ybeta")->at( k );
      double value= atof( it->second.at(k).c_str() );
      double yalfa= -1.0*value*ybeta/yrms;
      j["yalfa"].push_back( yalfa );
    }
   
  // creates the output file => xxx-plot.txt
  string fname= userDir_ + getFileName("plot",false);
  ofstream outfile;
  outfile.open(fname.c_str(),ios::out);
  outfile << std::setw(4) << j << std::endl;
  outfile.close();

  return j;
}

smap softwareAstra::getTableContents(ifstream& infile)
{
  smap vmap;
  string buf;
  while (infile >> buf)
    {
      vmap[ "c1" ].push_back( buf );
      infile >> buf;
      vmap[ "c2" ].push_back( buf );
      infile >> buf;
      vmap[ "c3" ].push_back( buf );
      infile >> buf;
      vmap[ "c4" ].push_back( buf );
      infile >> buf;
      vmap[ "c5" ].push_back( buf );
      infile >> buf;
      vmap[ "c6" ].push_back( buf );
      infile >> buf;
      vmap[ "c7" ].push_back( buf );    
    }

  return vmap;
}