#include <iomanip>

#include "InputFileReader.hh"

#include "ParticleRotor.hh"
#include "Nilsson.hh"
#include "LevelScheme.hh"

namespace PRM {
  Log::Log(std::ostream &stream) : ofile(), os(&stream) { }

  Log::Log(std::string fname) : ofile(fname) {
    if (!ofile.is_open()) {
      std::cerr << "Error! Log file not opened correctly" << std::endl;
      exit(1);
    }
    os = &ofile;
  }

  Log::~Log() {
    if (ofile.is_open()) {
      ofile.close();
    }
  }

  void Log::Open(std::string filename) {
    ofile.open(filename.c_str());
    if (!ofile.is_open()) {  
      std::cerr << "Error! Log file not opened correctly" << std::endl;
      exit(1);
    }
    os = &ofile;
  }

  void Log::Close() {
    if (ofile.is_open()) {
      ofile.close();
      os = &std::cout;
    }
  }

  std::ostream& Log::GetLog() {
    return *os;
  }
  
  InputFileReader::InputFileReader(std::string filename) :
    fLog(std::cout),
    fOpened(true)
  {
    FILE *file = fopen(filename.c_str(), "ra");
    if (!file) { std::cerr << "Error! file " << filename << " not opened" << std::endl; fOpened= false; return;}    
    std::stringstream ss;
    char cline[2048];

    int ind = 0;
    while(std::fgets(cline, sizeof cline, file)!=NULL) {
      std::string line(cline);
      if (line.size() == 0) { continue; }
      if (line[0] == '#') { continue; }
      if (line[0] == ';') { continue; }

      ss.clear();
      ss.str(line);

      std::string key;
      ss >> key;
      std::vector<std::string> values;
      std::string val;
      while (ss >> val) {
        if (!val.compare(0,1,"#")) { break; }
        values.push_back(val);
      }
      fLines[key]=values;
    }

    if (fLines.find("LOGFILE") != fLines.end()) {
      fLog.Open(fLines["LOGFILE"][0]);
    }
  }

  double InputFileReader::GetDoubleValue(std::string key) {
    if (fLines.find(key) == fLines.end()) { std::cerr << "Warning! Key " << key << " in input file not found!" << std::endl; return 0.0; }

    if (fLines[key].size() == 0) {  std::cerr << "Warning! No value for key " << key << " !" << std::endl; return 0.0; }

    return std::stod(fLines[key][0]);
  }

  std::vector<double> InputFileReader::GetDoubleValues(std::string key) {
    if (fLines.find(key) == fLines.end()) { std::cerr << "Warning! Key " << key << " in input file not found!" << std::endl; return {}; }
    if (fLines[key].size() == 0) {  std::cerr << "Warning! No value for key " << key << " !" << std::endl; return {}; }

    std::vector<double> retvals;
    for (std::string &val : fLines[key]) {
      retvals.push_back(std::stod(val));
    }
    return retvals;
  }

  int InputFileReader::GetIntegerValue(std::string key) {
    if (fLines.find(key) == fLines.end()) { std::cerr << "Warning! Key " << key << " in input file not found!" << std::endl; return 0.0; }

    if (fLines[key].size() == 0) {  std::cerr << "Warning! No value for key " << key << " !" << std::endl; return 0.0; }

    return std::stoi(fLines[key][0]);
  }

  std::vector<int> InputFileReader::GetIntegerValues(std::string key) {
    if (fLines.find(key) == fLines.end()) { std::cerr << "Warning! Key " << key << " in input file not found!" << std::endl; return {}; }
    if (fLines[key].size() == 0) {  std::cerr << "Warning! No value for key " << key << " !" << std::endl; return {}; }

    std::vector<int> retvals;
    for (std::string &val : fLines[key]) {
      retvals.push_back(std::stoi(val));
    }
    return retvals;
  }

  void InputFileReader::RunCalculation() {
    if (!fOpened) { std::cerr << "Cannot run calculation if file is not loaded properly" << std::endl; return; }
    int verbosity = GetIntegerValue("VERBOSITY");
        
    poscalc = new PRM::Nilsson(GetIntegerValue("Z"),GetIntegerValue("A"));
    negcalc = new PRM::Nilsson(GetIntegerValue("Z"),GetIntegerValue("A"));

    double eps = GetDoubleValue("EPS2");
    double gamma = GetDoubleValue("GAMMA");

    int nmax = GetIntegerValue("NMAX");
    poscalc->SetNMax(nmax,nmax);
    poscalc->MakeBasisStates(1);
    negcalc->SetNMax(nmax,nmax);
    negcalc->MakeBasisStates(-1);

    if (verbosity>0) {
      fLog.GetLog() << "=============================================" << std::endl;
      fLog.GetLog() << "        Particle-Rotor Calculations" << std::endl;
      fLog.GetLog() << " Based on Ragnarsson + Semmes codes" << std::endl;
      fLog.GetLog() << " Described in Nucl Phys A307 (1978) 189" << std::endl;
      fLog.GetLog() << "              Physica Scripta 8 (1973) 17" << std::endl;
      fLog.GetLog() << " Converted to C++ by Tim Gray (ORNL) 2023" << std::endl;
      fLog.GetLog() << "=============================================" << std::endl;
        
        
      fLog.GetLog() << "============= POSITIVE PARITY ===============" << std::endl;
      poscalc->PrintBasisStates(fLog.GetLog(), 5);

      fLog.GetLog() << "============= NEGATIVE PARITY ===============" << std::endl;
      negcalc->PrintBasisStates(fLog.GetLog(), 5);
    }

    poscalc->SetKappas(GetDoubleValues("KAPPAS"));
    poscalc->SetMus(GetDoubleValues("MUS"));

    negcalc->SetKappas(GetDoubleValues("KAPPAS"));
    negcalc->SetMus(GetDoubleValues("MUS"));

    if (fLines.find("NILPLOT") != fLines.end()) {
      if (verbosity > 0) {
        fLog.GetLog() << "Making Nilsson plot..." << std::endl;
      }
      std::ofstream diag_pos("nilsson_pos.dat");
      std::ofstream diag_neg("nilsson_neg.dat");
      for (double ep = GetDoubleValues("NILPLOT")[0]; ep <= GetDoubleValues("NILPLOT")[1]; ep += GetDoubleValues("NILPLOT")[2]) {

        if (verbosity > 0) {
          fLog.GetLog() << "[ eps = " << std::fixed << std::setprecision(2) << std::setw(3) << ep << " ] ";
        }
        poscalc->SetDeformation(ep, gamma, 0.0, 0.0);
        negcalc->SetDeformation(ep, gamma, 0.0, 0.0);

        poscalc->Solve();
        negcalc->Solve();

        poscalc->WriteEnergies(diag_pos);
        negcalc->WriteEnergies(diag_neg);

      }
      if (verbosity > 0) {
        fLog.GetLog() << std::endl;
      }
      diag_pos.close();
      diag_neg.close();

    }
  
    poscalc->SetDeformation(eps, gamma, 0.0, 0.0);
    negcalc->SetDeformation(eps, gamma, 0.0, 0.0);

    poscalc->Solve();
    negcalc->Solve();

    if (verbosity>0) {
      fLog.GetLog() << " ================== Positive-parity wavefunctions =================" << std::endl;
      poscalc->PrintWavefunctions(fLog.GetLog());
      fLog.GetLog() << " ================== Negative-parity wavefunctions =================" << std::endl;
      negcalc->PrintWavefunctions(fLog.GetLog());
    }

    if (fLines.find("NEG_DEC") == fLines.end() &&
        fLines.find("POS_DEC") == fLines.end()) {
      return;
    }

    if (fLines.find("NEG_DEC") != fLines.end()) {
      negcalc->Decouple(GetIntegerValues("NEG_DEC"));
    }
    
    if (fLines.find("POS_DEC") != fLines.end()) {
      poscalc->Decouple(GetIntegerValues("POS_DEC"));
    }
                     
    if (verbosity>0) {
      fLog.GetLog() << " ================== Positive-parity decoupling factors =================" << std::endl;
      poscalc->PrintDecouple(fLog.GetLog());
      fLog.GetLog() << " ================== Negative-parity decoupling factors =================" << std::endl;
      negcalc->PrintDecouple(fLog.GetLog());

    }

    poscalc->DiagQMom();
    negcalc->DiagQMom();
                     
    if (verbosity>0) {
      fLog.GetLog() << " ================== Positive-parity Q moments =================" << std::endl;
      poscalc->PrintDiagQMom(fLog.GetLog());
      fLog.GetLog() << " ================== Negative-parity Q moments =================" << std::endl;
      negcalc->PrintDiagQMom(fLog.GetLog());
    }

    paircalc = new PRM::Pairing(*poscalc,*negcalc);
    
    if (verbosity>1) {
      fLog.GetLog() << " ================== Energy levels for pairing calculation =================" << std::endl;
      paircalc->PrintLevels(fLog.GetLog());
    }

    if (fLines.find("IPAIR") != fLines.end()) {
      if (fLines.find("GN0") != fLines.end() &&
          fLines.find("GN1") != fLines.end() ) {
        double GN0 = GetDoubleValue("GN0");
        double GN1 = GetDoubleValue("GN1");
        paircalc->SetGNParams(GN0,GN1);
        paircalc->Solve(GetIntegerValue("IPAIR"));   
      }
      else {
        std::cerr << "IPAIR given but GN0 and GN1 not specified" << std::endl;
        exit(0);
      }
    }
    else if (fLines.find("DELTA") != fLines.end() &&
             fLines.find("LAMBDA") != fLines.end()) {
      paircalc->fDelta = GetDoubleValue("DELTA");
      paircalc->fEferm = GetDoubleValue("LAMBDA");
      paircalc->SetIPair(99);
    }
    else if (fLines.find("GN1") != fLines.end()) {
      paircalc->fDelta = GetDoubleValue("GN1")/std::sqrt(poscalc->fA);
      if (fLines.find("GN0") != fLines.end()) {
        paircalc->fEferm = GetDoubleValue("GN0");        
      }
      else {
        paircalc->fEferm = paircalc->fEnergies[(paircalc->fNOdd-1)/2];
      }
    }

    if (verbosity > 0) {
      fLog.GetLog() << " ================== Pairing calcultion results =================" << std::endl;
      paircalc->PrintSummary(fLog.GetLog());
    }
    
    partrotorpos = new PRM::ParticleRotor(*poscalc, *paircalc);  
    partrotorneg = new PRM::ParticleRotor(*negcalc, *paircalc);
    
    if (fLines.find("MOMENTS") != fLines.end()) {
      std::vector<double> As = GetDoubleValues("MOMENTS");
      partrotorpos->SetMomentsOfInertia(As[0], As[1], As[2]);
      partrotorneg->SetMomentsOfInertia(As[0], As[1], As[2]);
    }
    else if (fLines.find("B") != fLines.end()) {
      partrotorpos->SetB(GetDoubleValue("B"));
      partrotorpos->SetMomentsOfInertia();

      partrotorneg->SetB(GetDoubleValue("B"));
      partrotorneg->SetMomentsOfInertia();
    }
    else if (fLines.find("BSIM") != fLines.end()) {
      partrotorpos->SetBSimple(GetDoubleValue("BSIM"));
      partrotorpos->SetMomentsOfInertia();

      partrotorneg->SetBSimple(GetDoubleValue("BSIM"));
      partrotorneg->SetMomentsOfInertia();
    }
    else if (fLines.find("BGROD") != fLines.end()) {
      partrotorpos->SetBGrodzin();
      partrotorpos->SetMomentsOfInertia();

      partrotorneg->SetBGrodzin();
      partrotorneg->SetMomentsOfInertia();
    }

    if (verbosity>0) {
      fLog.GetLog() << " ================== Particle-rotor calculation parameters =================" << std::endl;
      partrotorpos->PrintMomentsOfInertia(fLog.GetLog());
      partrotorpos->PrintSPLevels(fLog.GetLog());
    }

    partrotorpos->MakeHamiltonians(fLog.GetLog());
    partrotorpos->Diagonalize();

    partrotorneg->MakeHamiltonians(fLog.GetLog());
    partrotorneg->Diagonalize();
    
    if (verbosity>0) {
      fLog.GetLog() << " ================== Positive-parity particle-rotor results =================" << std::endl;
      partrotorpos->PrintResults(fLog.GetLog(), 10.0);
      fLog.GetLog() << " ================== Negative-parity particle-rotor results =================" << std::endl;
      partrotorneg->PrintResults(fLog.GetLog(), 10.0);
    }

    int nstates = 15;
    if (fLines.find("NSTATES") != fLines.end()) {
      nstates = GetIntegerValue("NSTATES");
    }
    
    levelscheme = new PRM::LevelScheme(*partrotorpos, *partrotorneg);
    if (verbosity>0) {
      fLog.GetLog() << " ================== Particle-rotor level scheme =================" << std::endl;
      levelscheme->Print(fLog.GetLog(), nstates);
    }

    std::ofstream ampfilepos("prm_amplitudes_pos.dat");
    std::ofstream ampfileneg("prm_amplitudes_neg.dat");
    levelscheme->SetPrintLevels(nstates);

    partrotorpos->PrintAmplitudes(ampfilepos, levelscheme->fPosPrintLevels);
    partrotorneg->PrintAmplitudes(ampfileneg, levelscheme->fNegPrintLevels);
    
    ampfilepos.close();
    ampfileneg.close();
    
    std::ofstream lsfile("prm_levels.scheme");
    levelscheme->WriteDraw(lsfile, nstates);
    lsfile.close();

  }
}
