/**
 * @file nubeam.h
 * @brief Nubeam C++ class
 * @version $Id: nubeam.cpp 864 2011-02-24 22:10:47Z pletzer $
 */

#include <nubeam.h>
#include <czspline_capi.h>
#include "cpp/multiarray.h"

#ifdef __MPI
#define MPICH_IGNORE_CXX_SEEK
#include <mpi.h>
#endif

#include <iostream>
#include <limits>
#include <algorithm>

Nubeam::Nubeam() {

  this->ps = 0;
  this->commF = 0;
  this->myid = 0;

  this->rngSeed = 1234567890;
  this->numberOfBeamParticles   = -1;
  this->numberOfFusionParticles = -1;

  this->plasmaStateFile = "";
  this->machDescr = "";
  this->geqdsk    = "";
  this->runid = "";
  this->perBeamSpecies = "";
  this->fileNbCtrlInit = "";

  int debug = 3;
  this->ps = new PlasmaState("nubeam", debug);

  this->minRadius = 0.0;

  // fill in default values in structs
  nbCtrlInitDefault(&this->nbCtrlInit);
  nbCtrlUpdateDefault(&this->nbCtrlUpdt);
}

Nubeam::~Nubeam() {
  delete this->ps;
}

void Nubeam::setFortranComm(int commF) {
  this->commF = commF;
  this->myid = 0;
  int numProcs = 1;
#ifdef __MPI
  MPI_Comm comm = MPI_Comm_f2c(commF);
  MPI_Comm_rank(comm, &this->myid);
  MPI_Comm_size(comm, &numProcs);
  nbiSetComm(&commF);
#endif
  nbiSetMyid(&this->myid);
  nbiSetNumprocs(&numProcs);
}

void Nubeam::setPreactdir(const std::string& path) {
  int overwrite = 1;
  setenv ("PREACTDIR", path.c_str(), overwrite);
}

void Nubeam::setAdasdir(const std::string& path) {
  int overwrite = 1;
  setenv ("ADASDIR", path.c_str(), overwrite);
}


void Nubeam::setPlasmaStateFile(const std::string& file) {
  this->plasmaStateFile = file;
}

void Nubeam::setRunid(const std::string& runid) {
  this->runid = runid;
}

void Nubeam::setTimeInterval(const double t0, const double t1) {
  this->ps->setData("t0", t0);
  this->ps->setData("t1", t1);
}

void Nubeam::setRandomNumberGeneratorSeed(const int n) {
  this->rngSeed = n;
}

void Nubeam::setGeqdskFile(const std::string& geqdsk, const double minRadius) {
  this->geqdsk = geqdsk;
  this->minRadius = minRadius;
}

void Nubeam::setMachineDescriptionFile(const std::string& file) {
  this->machDescr = file;
}

void Nubeam::setGrid(const std::string& name,
		     const std::vector<double>& values) {
  if (this->myid == 0) {
    // store values
    this->gridMap[name] = values;
  }
}

void Nubeam::setThermalSpecies(const int z, const int c, const int a) {
  if (this->myid == 0) {
    std::vector<int> zca(3);

    if (this->thermalSpecies.size() == 0) {
// electron species must be present and preceed any other thermal
// species
      zca[0] = -1; zca[1] = -1; zca[2] = 0;
      this->thermalSpecies.push_back(zca);
    }

    zca[0] = z; zca[1] = c; zca[2] = a;
    this->thermalSpecies.push_back(zca);
  }
}

void Nubeam::setNeutralBeamSpecies(const int z, const int c, const int a) {
  if (this->myid == 0) {
    std::vector<int> zca(3);
    zca[0] = z; zca[1] = c; zca[2] = a;
    this->neutralBeamSpecies.push_back(zca);
  }
}

void Nubeam::setImpuritySpecies(const int z, const int c, const int a) {
  if (this->myid == 0) {
    std::vector<int> zca(3);
    zca[0] = z; zca[1] = c; zca[2] = a;
    this->impuritySpecies.push_back(zca);
  }
}

void Nubeam::setNbCtrlInitFile(const std::string& file) {
  this->fileNbCtrlInit = file;
}

void Nubeam::initialize() {

  if (this->plasmaStateFile != "") {

    // load plasma state from netcdf file

    if (this->myid == 0) {
      this->ps->getState(this->plasmaStateFile);
    }

  }
  else {

    // build plasma state

    if (this->myid == 0) {

      // set machine description
      const std::string action = "new"; // this will erase plasma state
      this->ps->readMachDescr(this->machDescr, action, this->geqdsk);

      // set runid
      this->ps->setData("runid", this->runid);

      // set grid dimensions
      for (std::map<std::string, std::vector<double> >::const_iterator it =
	     this->gridMap.begin(); it != this->gridMap.end(); ++it) {
	std::string name = it->first;
	int nsize = (int) it->second.size();
	this->ps->setData("n" + name, nsize);
      }            

      // allocate
      this->ps->alloc();

      // set axes
      for (std::map<std::string, std::vector<double> >::const_iterator it =
	     this->gridMap.begin(); it != this->gridMap.end(); ++it) {
	this->ps->setData(it->first, it->second);
      }

      // set species
      for (size_t i = 0; i < this->thermalSpecies.size(); ++i) {
	std::vector<int> zca = this->thermalSpecies[i];
	this->ps->setThermalSpecies(zca[0], zca[1], zca[2]);
      }
      for (size_t i = 0; i < this->impuritySpecies.size(); ++i) {
	std::vector<int> zca = this->impuritySpecies[i];
	this->ps->setImpuritySpecies(zca[0], zca[1], zca[2]);
      }
      // only a single neutral beam species supported at present
      if (this->neutralBeamSpecies.size() != 1) {
	this->catchError(100, __LINE__, __FILE__);
      }
      for (size_t i = 0; i < this->neutralBeamSpecies.size(); ++i) {
	std::vector<int> zca = this->neutralBeamSpecies[i];
	this->ps->setNeutralBeamSpecies(zca[0], zca[1], zca[2]);
	// this will set this->perBeamSpecies
	// which must be one of "H", "D", "T", "He3", or "He4"
	this->setSpeciesNameFromZA(zca[0], zca[2]);
      }

      // Set the ion beam species, must be consistent with setNeutralBeamSpecies
      // call. Also assume that there is only one call to setNeutralBeamSpecies!
      int nbeam = this->ps->getData<int>("nbeam"); // number of beams
      // only support one type of beam species for the time being
      std::vector<std::string> nbion(nbeam, this->perBeamSpecies);
      this->ps->setData("nbion", nbion);

      this->ps->finishSpecies();

      // set equilibrium
      const int kcur_option = 0; // take q(psi) from file
      const double rho_curbrk = 0.9; // rho beyond which interpolate to match
                                     // total current
      this->ps->updateEquilibrium(this->geqdsk, this->minRadius,
				  kcur_option, rho_curbrk);
      std::vector<double> null;
      this->ps->verifyMhdEq(null);
      this->ps->deriveMhdEq("Everything", null);
    }
  }

  // initialize nb ctrl init struct
  int ier = 0;
  nbCtrlInitRead(this->fileNbCtrlInit.c_str(), &this->nbCtrlInit, &ier,
		 this->fileNbCtrlInit.size());
  this->catchError(ier, __LINE__, __FILE__);

  // set the seed
  this->nbCtrlInit.nseed = this->rngSeed;

  // initialize nb ctrl struct
  ier = 0;
  nbCtrlInitialize(this->ps->getHandle(), &this->nbCtrlInit, &ier);
  this->catchError(ier, __LINE__, __FILE__);

  nbCtrlNostep(this->ps->getHandle(), &ier);
  this->catchError(ier, __LINE__, __FILE__);

  // set working directory to be local ("")
  std::string workpath = "";
  nbWorkpathReset(workpath.c_str(), workpath.size());

  // set particle numbers, override
  if (this->numberOfBeamParticles >= 0) {
    this->nbCtrlUpdt.nptcls = this->numberOfBeamParticles;
    // suggested by D. McCune
    this->nbCtrlInit.ndep0 = this->numberOfBeamParticles / 2;
  }
  if (this->numberOfFusionParticles >= 0) {
    this->nbCtrlUpdt.nptclf = this->numberOfFusionParticles;
  }
}

void Nubeam::setNumberOfBeamParticles(const int n) {
  this->numberOfBeamParticles = n;
}

void Nubeam::setNumberOfFusionParticles(const int n) {
  this->numberOfFusionParticles = n;
}

void Nubeam::setFractionsOfBeamCurrents(const std::vector<double>& atHalfVoltage,
					const std::vector<double>& atFullVoltage) {
  if (this->myid == 0) {
    this->ps->setData("frac_half", atHalfVoltage);
    this->ps->setData("frac_full", atFullVoltage);
  }
}

void Nubeam::setZEffective(const std::vector<double>& zeff) {
  if (this->myid == 0) {
    this->ps->setData("zeff", zeff);
  }
}

void Nubeam::setTemperatures(const std::vector<double>& te,
			    const std::vector<double>& ti) {

  if (this->myid == 0) {
    // note that the dimensions are reversed with respect to C/C++
    std::vector<size_t> dims = this->ps->getDims("ts");
    MultiArray<double> temperatures;
    temperatures.setDims(dims);
    size_t ndims = dims.size();
    std::vector<size_t> indices(ndims);
    std::vector<double> kevTemp;
    const double ev2kev = 1.e-3;
    indices[1] = 0;
    for (size_t i = 0; i < dims[0]; ++i) {
      indices[0] = i;
      temperatures[indices] = te[i] * ev2kev;
      double val = ti[i] * ev2kev;
      kevTemp.push_back(val);
    }

    for (size_t j = 1; j < dims[1]; ++j) {
      indices[1] = j;
      for (size_t i = 0; i < dims[0]; ++i) {
	indices[0] = i;
	// set the temperature for all ion species to this value
	// (as suggested by D McCune)
	temperatures[indices] = kevTemp[i];
      }
    }
    this->ps->setData("ts", temperatures);
    // also set Ti, which is a weighted average of ion temperatures
    this->ps->setData("ti", kevTemp);
  }
}

void Nubeam::setDensities(const std::vector<double>& ne,
			    const std::vector<double>& ni) {

  if (this->myid == 0) {
    // note that the dimensions are reversed with respect to C/C++
    std::vector<size_t> dims = this->ps->getDims("ns");
    MultiArray<double> densities;
    densities.setDims(dims);
    size_t ndims = dims.size();
    std::vector<size_t> indices(ndims);
    for (size_t i = 0; i < dims[0]; ++i) {
      indices[1] = 0; indices[0] = i;
      densities[indices] = ne[i];
      indices[1] = 1; indices[0] = i;
      densities[indices] = ni[i]; // only leading density is set
      for (size_t j = 2; j < dims[1]; ++j) {
	indices[1] = j; indices[0] = i;
	// all other densities set to zero
	densities[indices] = 0.0;
      }
    }
    this->ps->setData("ns", densities);
    this->ps->setData("ni", ni);
  }
}

int Nubeam::getNumBeams() {
  return this->ps->getData<int>("nbeam");
}

void Nubeam::setBeamEnergies(const std::vector<double>& evolt_nbi) {

  if (this->myid == 0) {
    const double ev2kev = 1.e-3;
    size_t n = evolt_nbi.size();
    std::vector<double> kvolt_nbi(n);
    for (size_t i = 0; i < n; ++i) {
      // unit conversion from eV to keV
      kvolt_nbi[i] = evolt_nbi[i] * ev2kev;
    }
    this->ps->setData("kvolt_nbi", kvolt_nbi);
  }
}

void Nubeam::setBeamPowers(const std::vector<double>& power_nbi) {

  if (this->myid == 0) {
    this->ps->setData("power_nbi", power_nbi);
  }
}

void Nubeam::haloThermalNeutralRecaptureOn() {

  this->nbCtrlInit.levmod_halo = 1;
  this->nbCtrlUpdt.levmod_halo = 1;
}

void Nubeam::haloThermalNeutralRecaptureOff() {

  this->nbCtrlInit.levmod_halo = 0;
  this->nbCtrlUpdt.levmod_halo = 0;
}

void Nubeam::debugOn() {
  
  this->nbCtrlInit.nltest_output = 1;
  this->nbCtrlUpdt.nltest_output = 1;
}

void Nubeam::debugOff() {
  
  this->nbCtrlInit.nltest_output = 0;
  this->nbCtrlUpdt.nltest_output = 0;
}

void Nubeam::updateNbCtrl(const std::string& file) {

  if (this->myid == 0) {
    int ier = 0;
    if (file.find(".cdf") != std::string::npos ||
	file.find(".nc") != std::string::npos) {
      // Netcdf
      nbCtrlUpdateCdfread(file.c_str(), &this->nbCtrlUpdt, &ier, file.size());
      this->catchError(ier, __LINE__, __FILE__);
    }
    else {
      // Fortran namelist
      nbCtrlUpdateRead(file.c_str(), &this->nbCtrlUpdt, &ier, file.size());
      this->catchError(ier, __LINE__, __FILE__);
    }
  }
}

void Nubeam::checkAndFix() {
  int ier = 0;
  psStateMemoryUpdate(this->ps->getHandle(), &ier);
  this->catchError(ier, __LINE__, __FILE__);
}

void Nubeam::savePlasmaState(const std::string& file) {
  if (this->myid == 0) {
    this->ps->store(file);
  }
}

void Nubeam::enforceQuasiNeutrality() {
  
  if (this->myid != 0) return;

  std::vector<size_t> dims;
  dims = this->ps->getDims("nbeami");
  size_t nspecBeam = dims[1];
  size_t nrho_nbi = dims[0] + 1;
  std::vector<double> rho_nbi = this->ps->getData<std::vector<double> >("rho_nbi");
  std::vector<double> rho = this->ps->getData<std::vector<double> >("rho");

  std::vector<double> rhoCell(rho.size() - 1);
  for (size_t i = 0; i < rhoCell.size(); ++i) {
    rhoCell[i] = 0.5 * (rho[i] + rho[i + 1]);
  }

// beam ion density 
  std::vector<double> nbeami =  this->ps->getData<std::vector<double> >("nbeami");

  std::vector<std::vector<double> > nbeamiInterp;
  for (size_t j = 0; j < nspecBeam; ++j) {
// project onto nodes
    std::vector<double> data(nrho_nbi);
    data[0] = nbeami[0];
    for (size_t i = 1; i < nrho_nbi; ++i) {
      data[i] = 0.5 * (nbeami[i - 1]  + nbeami[i]);
    }
    data[nrho_nbi - 1] = nbeami[nrho_nbi - 2];

// interpolate ont rho cells
    nbeamiInterp.push_back( this->interp1d(rho_nbi, data, rhoCell) );
  }

// all species density list (incl. electrons)
  dims = this->ps->getDims("ns");
  std::vector<double> ns = this->ps->getData<std::vector<double> >("ns");
  size_t nrhoM1 = dims[0];
  size_t nspecThermP1 = dims[1];
  for (size_t i = 0; i < nrhoM1; ++i) {
    double sumIonDens = 0.0;
    for (size_t j = 0; j < nspecBeam; ++j) {
      sumIonDens += nbeamiInterp[j][i];
    }
// iterate over ions only
    for (size_t j = 1; j < nspecThermP1; ++j) {
      size_t k = j + nspecThermP1 * i;
      sumIonDens += ns[k];
    }
// 1st index is electron
    size_t j = 0;
    size_t k = j + nspecThermP1 * i;
    ns[k] = sumIonDens; // set electron density
  }
  
// need to cast into MultiArray when setting
  MultiArray<double> nsMul;
  dims = this->ps->getDims("ns");
  nsMul.setDims(dims);
  nsMul.setData(ns);
  this->ps->setData("ns", nsMul);
}

void Nubeam::run() {
  int ier = 0;
  if (this->myid == 0) {
    nbCtrlStep(this->ps->getHandle(),
	       this->ps->getHandle(),
	       &this->nbCtrlUpdt, &ier);
    this->catchError(ier, __LINE__, __FILE__);
  }
  else {
    nbCtrlStepChild(&ier);
    this->catchError(ier, __LINE__, __FILE__);
  }
}

void Nubeam::loadEnergySources(const std::vector<double>& rho,
			std::vector<double>& pe,
			std::vector<double>& pi) const {
  
  size_t n = rho.size();

  if (this->myid == 0 && n > 0) {
// enable interpolation services
    // this->ps->xLoad();

// d Vol / drho
    std::vector<double> dV_drho = this->ps->interp1d("vol", 1, rho);

// pbe and pbi are the deposited powers per zone in W. we want d P/drho
// 1. create a mesh that starts at rho=0 and ends at rho=1
// 2. extrapolate cell values of pbe and pbi to first and last node
// 3. interpolate d P /drho onto rho 

// nodal values
    std::vector<double> rho_nbi = this->ps->getData< std::vector<double> >("rho_nbi");
    size_t n = rho_nbi.size();    
    if (n < 3) {
      this->catchError(1, __LINE__, __FILE__);
    }
    double rhoMin = rho_nbi[0];
    double rhoMax = rho_nbi[n-1];
// cell
    std::vector<double> pbe = this->ps->getData< std::vector<double> >("pbe");
    std::vector<double> pbi = this->ps->getData< std::vector<double> >("pbi");

// cell centered mesh extended with 2 node values at rho=rhoMin and rhoMax
// this is required to ensure that target interpolation points are always 
// contained within the original mesh

    std::vector<double> rho_nbiCellX(n + 1);
    std::vector<double> pbeCellX(n + 1);
    std::vector<double> pbiCellX(n + 1);

    rho_nbiCellX[0] = 0.0;
    pbeCellX[0] = pbe[0] - 0.5 * (pbe[1] - pbe[0]); // linear extrapolation
    pbiCellX[0] = pbi[0] - 0.5 * (pbi[1] - pbi[0]);
    for (size_t i = 0; i < n - 1; ++i) {
      size_t i1 = i + 1;
      rho_nbiCellX[i1] = 0.5 * (rho_nbi[i] + rho_nbi[i+1]); // mid cell value
      pbeCellX[i1] = pbe[i];
      pbiCellX[i1] = pbi[i];
    }
    rho_nbiCellX[n] = rho_nbi[n-1];
    pbeCellX[n] = pbe[n-2] + 0.5 * (pbe[n-2] - pbe[n-3]); // linear extrapolation
    pbiCellX[n] = pbi[n-2] + 0.5 * (pbi[n-2] - pbi[n-3]); // linear extrapolation
// interpolate onto rho
    std::vector<double> pelc = this->interp1d(rho_nbiCellX, pbeCellX, rho);
    std::vector<double> pion = this->interp1d(rho_nbiCellX, pbiCellX, rho);
   
// compute d Power / d Volume
    double dRho = (rhoMax - rhoMin) / double(n - 1);
    for (size_t i = 0; i < rho.size(); ++i) {
      double pe_drho = pelc[i] / dRho;
      double pi_drho = pion[i] / dRho;
      pe[i] = pe_drho / dV_drho[i];
      pi[i] = pi_drho / dV_drho[i];
    }
  }
}

void Nubeam::loadParticleSources(const std::vector<double>& rho,
			std::vector<double>& se,
			std::vector<double>& si) const {
  size_t n = rho.size();
  if (this->myid == 0 && n > 0) {
    for (size_t i = 0; i < n; ++i) {
// to implement
      se[i] = 0.0;
      si[i] = 0.0;
    }
  }
}

std::vector<double> Nubeam::interp1d(const std::vector<double>& rhoOri,
				     const std::vector<double>& datOri,
				     const std::vector<double>& rho) const {
  int ier = 0;
  int n1 = rhoOri.size();
  if (n1 != datOri.size()) {
    this->catchError(100, __LINE__, __FILE__);
  }
  int bcs1[2];
  int m1 = rho.size();
  double rhoNew[n1];
  double datNew[datOri.size()];
  double rhon[m1];
  double res[m1];
  std::copy(rhoOri.begin(), rhoOri.end(), rhoNew);
  std::copy(datOri.begin(), datOri.end(), datNew);
  std::copy(rho.begin(), rho.end(), rhon);
  // not-a-knot
  bcs1[0] = 0; bcs1[1] = 0;
  czspline_init1(&n1, bcs1, &ier);
  this->catchError(ier, __LINE__, __FILE__);
  czspline_set_axes1(&n1, rhoNew, &ier);
  this->catchError(ier, __LINE__, __FILE__);
  czspline_setup1(&n1, datNew, &ier);
  this->catchError(ier, __LINE__, __FILE__);
  czspline_interp1_cloud(&m1, rhon, res, &ier);
  this->catchError(ier, __LINE__, __FILE__);
  czspline_free1(&ier);
  this->catchError(ier, __LINE__, __FILE__);
  std::vector<double> res2;
  res2.insert(res2.end(), &res[0], &res[m1]);
  return res2;
}

std::vector<double> Nubeam::getDensityFromPlasmaState(int specieIndex) {
  const std::string PROF = "ns";
  std::vector<size_t> dims = this->ps->getDims(PROF);
// note that dims are returned in the order dictated by fortran
  size_t nspecies = dims[1];
  size_t nrhoCell = dims[0];
// array of size nspecies * nrhoCell
  std::vector<double> flatArray = this->ps->getData< std::vector<double> >(PROF);
  std::vector<double> res(nrhoCell);
  for (size_t i = 0; i < nrhoCell; ++i) {
    int k = nrhoCell * specieIndex + i;
    res[i] = flatArray[k];
  }
  return res;
}

std::vector<double> Nubeam::getTemperatureFromPlasmaState(int specieIndex) {
  const std::string PROF = "ts";
  std::vector<size_t> dims = this->ps->getDims(PROF);
// note that dims are returned in the order dictated by fortran
  size_t nspecies = dims[1];
  size_t nrhoCell = dims[0];
// array of size nspecies * nrhoCell
  std::vector<double> flatArray = this->ps->getData< std::vector<double> >(PROF);
  std::vector<double> res(nrhoCell);
  for (size_t i = 0; i < nrhoCell; ++i) {
    int k = nrhoCell * specieIndex  + i;
    res[i] = flatArray[k];
  }
  return res;
}

void Nubeam::setSpeciesNameFromZA(const int zi, const int ai) {

  // map (z,a) to string (eg "He3")
  std::map<int, std::string> zca2name;
  const int base = 1000;

  // Hydrogen: Z=1, A=1
  int z = 1; int a = 1;
  zca2name[base*z + a] = "H";

  // Deuterium: Z=1, A=2
  z = 1; a = 2;
  zca2name[base*z + a] = "D";

  // Tritium: Z=1, A=3
  z = 1; a = 3;
  zca2name[base*z + a] = "T";

  // He3: Z=2, A=3
  z = 2; a = 3;
  zca2name[base*z + a] = "He3";

  // He4: Z=2 A=4
  z = 2; a = 4;
  zca2name[base*z + a] = "He4";

  // Carbon
  z = 6; a = 12;
  zca2name[base*z + a] = "C";

  // Nitrogen
  z = 7; a = 14;
  zca2name[base*z + a] = "N";

  // Oxygen
  z = 8; a = 16;
  zca2name[base*z + a] = "O";

  this->perBeamSpecies = zca2name[base*zi + ai];
}

void Nubeam::catchError(const int ier, const int lineNumber,
			const char* fileName) const {
  if (ier != 0) {
    NubeamError err(fileName, lineNumber, ier);
    std::cerr << "[" << this->myid << "] ERROR in file "
	      << fileName << " at line " << lineNumber
	      << " error code=" << ier << '\n';
    throw err;
  }
}
