#include <iomanip>

#include "CLI.hh"

#include <iostream>
#include <sstream>
#include <stdio.h>

#include "InputFileReader.hh"

namespace PRM {
  CLI::CLI() : fFormat( std::cout.flags()) { };

  void CLI::RunCalculation(std::string infile) {
    PRM::InputFileReader reader(infile);
    for (auto &line : reader.fLines) {
      fAnswers[line.first] = line.second;
    }
    
    reader.RunCalculation();

    poscalc = reader.poscalc;
    negcalc = reader.negcalc;

    paircalc = reader.paircalc;

    partrotorpos = reader.partrotorpos;
    partrotorneg = reader.partrotorneg;

    levelscheme = reader.levelscheme;

    fEps2 = poscalc->fEps2;
    fGamma = poscalc->fGamma;
  }
  
  void CLI::RunCalculation() {
    std::vector<int> iretvals;
    iretvals = GetIntegersAnswer("Enter nucleus (Z,A)", {22,47}, 2);
    int Z = iretvals[0];
    int A = iretvals[1];
    fAnswers["Z"] = {std::to_string(Z)};
    fAnswers["A"] = {std::to_string(A)};

    poscalc = new PRM::Nilsson(Z,A);
    negcalc = new PRM::Nilsson(Z,A);    
    
    fEps2 = GetDoubleAnswer("Enter deformation parameter (epsilon)", 0.3);
    fAnswers["EPS2"] = {std::to_string(fEps2)};
    fGamma = GetDoubleAnswer("Enter asymmetry parameter in degrees (gamma)", 0.0);
    fAnswers["GAMMA"] = {std::to_string(fGamma)};

    int nmax = GetIntegerAnswer("Enter max oscillator shells", 7);
    fAnswers["NMAX"] = {std::to_string(nmax)};

    poscalc->SetNMax(nmax,nmax);
    poscalc->MakeBasisStates(1);
    negcalc->SetNMax(nmax,nmax);
    negcalc->MakeBasisStates(-1);

    std::cout << "============= POSITIVE PARITY ===============" << std::endl;
    poscalc->PrintBasisStates(std::cout, 5);

    std::cout<< "============= NEGATIVE PARITY ===============" << std::endl;
    negcalc->PrintBasisStates(std::cout, 5);


    std::vector<double> kappas_p = {0.12,  0.12,  0.105,  0.090,  0.065,  0.060, 0.054, 0.054, 0.054};
    std::vector<double> kappas_n = {0.12,  0.12,  0.105,  0.095,  0.080,  0.062, 0.062, 0.062, 0.062};
    std::vector<double> mus_n = {0.0, 0.0, 0.0, 0.18, 0.22, 0.43, 0.34,  0.26,  0.26};
    std::vector<double> mus_p = {0.0, 0.0, 0.0, 0.30, 0.57, 0.65, 0.69,  0.69,  0.69};
    std::vector<double> kappas;
    if (Z%2==1) { kappas = GetDoublesAnswer("Enter kappas", kappas_p, nmax); }
    else { kappas = GetDoublesAnswer("Enter kappas", kappas_n, nmax); }
    std::vector<std::string> kapstr;
    std::vector<std::string> musstr;
    std::vector<double> mus;
    if (Z%2==1) { mus = GetDoublesAnswer("Enter mus", mus_p, nmax); }
    else { mus = GetDoublesAnswer("Enter mus", mus_n, nmax); }
    for (int i=0; i<kappas.size(); ++i) {
      kapstr.push_back(std::to_string(kappas[i]));
    }
    for (int i=0; i<mus.size(); ++i) {
      musstr.push_back(std::to_string(mus[i]));
    }
    fAnswers["KAPPAS"] = kapstr;
    fAnswers["MUS"] = musstr;

    poscalc->SetKappas(kappas);
    poscalc->SetMus(mus);

    negcalc->SetKappas(kappas);
    negcalc->SetMus(mus);

    int epsplot = GetIntegerAnswer("Would you like to make a epsilon plot? (1/0)", 0);    
    if (epsplot) {
      std::vector<double> params = GetDoublesAnswer("Enter epsilon low, epsilon high, increment", {-0.4,0.4,0.02}, 3);
      EpsPlot(params[0], params[1], params[2]);
    }

    int gamplot = GetIntegerAnswer("Would you like to make a gamma plot? (1/0)", 0);    
    if (gamplot) {
      std::vector<double> params = GetDoublesAnswer("Enter gamma low, gamma high, increment", {0,60,1}, 3);
      GamPlot(params[0], params[1], params[2]);
    }
    
    poscalc->SetDeformation(fEps2, fGamma, 0.0, 0.0);
    negcalc->SetDeformation(fEps2, fGamma, 0.0, 0.0);

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

    int yn = GetIntegerAnswer("Print wavefunctions? (1/0)", 0);
    if (yn) {
      poscalc->PrintWavefunctions(std::cout);
      negcalc->PrintWavefunctions(std::cout);      
    }
    
    std::vector<int> pos_dec = GetIntegersAnswer("Enter positive-parity levels for particle-rotor calculation", {}, 0);
    std::vector<int> neg_dec = GetIntegersAnswer("Enter negative-parity levels for particle-rotor calculation", {}, 0);
    if (pos_dec.size() > 0) {
      std::vector<std::string> posdecstr;
      for (int i=0; i<pos_dec.size(); ++i) {
        posdecstr.push_back(std::to_string(pos_dec[i]));
      }
      fAnswers["POS_DEC"] = posdecstr;
    }
    
    if (neg_dec.size() > 0) {
      std::vector<std::string> negdecstr;
      for (int i=0; i<neg_dec.size(); ++i) {
        negdecstr.push_back(std::to_string(neg_dec[i]));
      }
      fAnswers["NEG_DEC"] = negdecstr;
    }

    poscalc->Decouple(pos_dec);
    negcalc->Decouple(neg_dec);

    yn = GetIntegerAnswer("Print decoupling factors? (1/0)", 0);
    if (yn) {
      poscalc->PrintDecouple(std::cout);
      negcalc->PrintDecouple(std::cout);      
    }

    poscalc->DiagQMom();
    negcalc->DiagQMom();
    
    yn = GetIntegerAnswer("Print Q factors? (1/0)", 0);
    if (yn) {
      poscalc->PrintDiagQMom(std::cout);
      negcalc->PrintDiagQMom(std::cout);      
    }

    paircalc = new PRM::Pairing(*poscalc,*negcalc);
    yn = GetIntegerAnswer("Print pairing calculation input levels? (1/0)", 0);
    if (yn) {
      paircalc->PrintLevels(std::cout);
    }

    std::vector<double> gnpars = GetDoublesAnswer("Enter pairing parameters (GN0, GN1)", {21.0, 0.0}, 2);

    paircalc->SetGNParams(gnpars[0],gnpars[1]);
    if (gnpars[0] > 0.0) {
      fAnswers["GN0"] = {std::to_string(gnpars[0])};
    }
    fAnswers["GN1"] = {std::to_string(gnpars[1])};
    int ipair = GetIntegerAnswer("Enter ipair, controls how many levels above/below the Fermi surface are included", 10);
    if (ipair < 99) {
      paircalc->Solve(ipair);
      fAnswers["IPAIR"] = {std::to_string(ipair)};
    }
    else {
      paircalc->fDelta = gnpars[1]/std::sqrt(poscalc->fA);
      if (gnpars[0] > 0) {
        paircalc->fEferm = gnpars[0];
      }
      else {
        paircalc->fEferm = paircalc->fEnergies[(paircalc->fNOdd-1)/2];
      }
    }

    paircalc->PrintSummary(std::cout);

    partrotorpos = new PRM::ParticleRotor(*poscalc, *paircalc);
    partrotorneg = new PRM::ParticleRotor(*negcalc, *paircalc);

    while (true) {
      int opt = GetIntegerAnswer("Method for specifying moments of inertia:\n    1: Directly\n    2: B value fixed to the E(2+) of the core\n    3: hbar^2/B = E(2+)/A^2\n    4: Grodzin's relation\nEnter option", 2);
      if (opt == 1) {
        std::vector<double> As = GetDoublesAnswer("Enter moments of inertia (A1, A2, A3)", {0.15,0.15,150.0}, 3);
        partrotorpos->SetMomentsOfInertia(As[0], As[1], As[2]);
        partrotorneg->SetMomentsOfInertia(As[0], As[1], As[2]);
        fAnswers["MOMENTS"] = {std::to_string(As[0]),std::to_string(As[1]),std::to_string(As[2])};
        break;
      }
      else if (opt == 2) {
        double E2p = GetDoubleAnswer("Enter E(2+) in MeV", 0.9);
        partrotorpos->SetB(E2p);
        partrotorpos->SetMomentsOfInertia();

        partrotorneg->SetB(E2p);
        partrotorneg->SetMomentsOfInertia();
        fAnswers["B"] = {std::to_string(E2p)};                 
        break;
      }
      else if (opt == 3) {
        double E2p = GetDoubleAnswer("Enter E(2+) in MeV", 0.9);
        partrotorpos->SetBSimple(E2p);
        partrotorpos->SetMomentsOfInertia();

        partrotorneg->SetBSimple(E2p);
        partrotorneg->SetMomentsOfInertia();
        fAnswers["BSIM"] = {std::to_string(E2p)};
        break;
      }
      else if (opt == 4) {
        partrotorpos->SetBGrodzin();
        partrotorpos->SetMomentsOfInertia();

        partrotorneg->SetBGrodzin();
        partrotorneg->SetMomentsOfInertia();
        fAnswers["BGROD"] = {std::to_string(1)};
        break;
      }
      else { continue; }
    }
    
    partrotorpos->PrintMomentsOfInertia(std::cout);
    
    partrotorpos->PrintSPLevels(std::cout);
    partrotorneg->PrintSPLevels(std::cout);

    
    partrotorpos->MakeHamiltonians(std::cout);
    partrotorpos->Diagonalize();

    partrotorneg->MakeHamiltonians(std::cout);
    partrotorneg->Diagonalize();

    yn = GetIntegerAnswer("Print detailed results for each spin? (1/0)", 0);
    if (yn) {
      partrotorpos->PrintResults(std::cout, 10.0);
      partrotorneg->PrintResults(std::cout, 10.0);
    }

    levelscheme = new PRM::LevelScheme(*partrotorpos, *partrotorneg);
    int nstates = GetIntegerAnswer("Number of states to print level scheme for", 15);
    fAnswers["NSTATES"] = {std::to_string(nstates)};
    if (nstates > 0) {
      levelscheme->Print(std::cout, nstates);      
    }
  }

  void CLI::Prompt() {
    HelpMessage();
    
    while (true) {      
      std::cout << "PRM >> ";
      std::string answer = GetAnswer();
      
      if (answer.size() == 0) { continue; }
      if (!answer.compare(0,4,"quit") || !answer.compare(0,4,"exit")) { exit(0); }
      if (!answer.compare(0,4,"help")) {
        HelpMessage();
        continue;
      }
      if (!answer.compare(0,3,"run")) { RunCalculation(); continue; }
      else if (!answer.compare(0,4,"load")) {
        std::stringstream ss(answer);
        std::string val;
        ss >> val; //load
        if (!(ss >> val)) { continue; }        
        RunCalculation(val); continue;
      }
      else if (!answer.compare(0,4,"save")) {
        std::stringstream ss(answer);
        std::string val;
        ss >> val; //load
        if (!(ss >> val)) { continue; }
        Save(val);
      }
      else if(!answer.compare(0,7,"epsplot")) {
        char *end;
        for (int i=0; i<answer.size(); ++i) {
          if (!answer.compare(i,1,",")) {
            answer.replace(i,1," ");
          }
        }
      
        std::stringstream ss(answer);
        std::string val;
        std::vector<double> retvals;
        ss >> val;
        while (ss >> val) {
          double retval = std::strtod(val.c_str(), &end);
          if (end[0] == val[0]) { continue; }
          retvals.push_back(retval);
        }
        if (retvals.size() < 3) { continue; }

        EpsPlot(retvals[0], retvals[1], retvals[2]);        
      }
      else if(!answer.compare(0,7,"gamplot")) {
        char *end;
        for (int i=0; i<answer.size(); ++i) {
          if (!answer.compare(i,1,",")) {
            answer.replace(i,1," ");
          }
        }
        
        std::stringstream ss(answer);
        std::string val;
        std::vector<double> retvals;
        ss >> val;
        while (ss >> val) {
          double retval = std::strtod(val.c_str(), &end);
          if (end[0] == val[0]) { continue; }
          retvals.push_back(retval);
        }
        if (retvals.size() < 3) { continue; }

        GamPlot(retvals[0], retvals[1], retvals[2]);        
      }
      else if(!answer.compare(0,5,"print")) {

        std::stringstream ss(answer);
        std::string val;
        std::vector<double> retvals;
        ss >> val; //print
        if (ss >> val) {
          char *end;
          int retval = std::strtoll(val.c_str(), &end, 10);
          if (end[0] == val[0]) { continue; }

          if (retval > 0) {
            levelscheme->Print(std::cout, retval);
          }
        }
        continue;
      }
      else if (!answer.compare(0,4,"draw")) {
        std::stringstream ss(answer);
        std::string val;
        std::vector<double> retvals;
        ss >> val; //print
        if (!ss) { continue; }
        ss >> val;
        char *end;
        int retval = std::strtoll(val.c_str(), &end, 10);
        if (end[0] == val[0]) { continue; }

        if (retval > 0) {
          std::ofstream lsfile("prm_levels.scheme");
          levelscheme->WriteDraw(lsfile, retval);
          lsfile.close();
          popen("./plot_levelscheme.py", "r");
        }
        continue;
      }
      else if (!answer.compare(0,3,"amp")) {
        
        for (int i=0; i<answer.size(); ++i) {
          if (!answer.compare(i,1,",")) {
            answer.replace(i,1," ");
          }
        }
        
        std::stringstream ss(answer);
        std::string val;
        ss >> val; //print
        if (!ss) { continue; }

        std::vector<int> ints;
        while (ss >> val) {
          char *end;
          int retval = std::strtoll(val.c_str(), &end, 10);
          if (end[0] == val[0]) { continue; }
          ints.push_back(retval);
        }

        if (ints.size() == 1) {
          std::ofstream ampfilepos("prm_amplitudes_pos.dat");
          std::ofstream ampfileneg("prm_amplitudes_neg.dat");
          levelscheme->SetPrintLevels(ints[0]);
          partrotorpos->PrintAmplitudes(ampfilepos, levelscheme->fPosPrintLevels);
          partrotorneg->PrintAmplitudes(ampfileneg, levelscheme->fNegPrintLevels);
          ampfilepos.close();
          ampfileneg.close();
          popen("./plot_amplitudes.py", "r");
        }
        else {
          ss.clear();
          ss.str(answer);

          ss >> val; //amp
          std::vector<std::pair<int,int> > posstates;
          std::vector<std::pair<int,int> > negstates;
          while (ss >> val) {
            if (val.back()=='+') {
              for (int i=0; i<val.size(); ++i) {
                if (!val.compare(i,1,"_")) {
                  val.replace(i,1," ");
                }
              }
              std::stringstream valss(val.substr(0,val.size()-1));
              int spin;
              int indx;
              if (!(valss >> spin)) { continue; }
              if (!(valss >> indx)) { continue; }
              posstates.push_back({spin,indx});
            }
            else if (val.back()=='-') {
              for (int i=0; i<val.size(); ++i) {
                if (!val.compare(i,1,"_")) {
                  val.replace(i,1," ");
                }
              }
              std::stringstream valss(val.substr(0,val.size()-1));
              int spin;
              int indx;
              if (!(valss >> spin)) { continue; }
              if (!(valss >> indx)) { continue; }
              negstates.push_back({spin,indx});
            }
            else { continue; }
          }

          std::ofstream ampfilepos("prm_amplitudes_pos.dat");
          std::ofstream ampfileneg("prm_amplitudes_neg.dat");
          partrotorpos->PrintAmplitudes(ampfilepos, posstates);
          partrotorneg->PrintAmplitudes(ampfileneg, negstates);
          ampfilepos.close();
          ampfileneg.close();
          popen("./plot_amplitudes.py", "r");
        }
      }
    }
  }
  
  double CLI::GetDoubleAnswer(std::string question, double def) {
    std::cout.flags( fFormat );
    while (true) {
      std::cout << question << " [ " << def << " ]: ";
      std::string answer = GetAnswer();

      if (answer.size() == 0) { return def; }
      char *end;
      double retval = std::strtod(answer.c_str(), &end);
      if (end[0] == answer[0]) { continue; }
      return retval;
    }
  }

  int CLI::GetIntegerAnswer(std::string question, int def) {
    std::cout.flags( fFormat );
    while (true) {
      std::cout << question << " [ " << def << " ]: ";
      std::string answer = GetAnswer();

      if (answer.size() == 0) { return def; }
      char *end;
      int retval = std::strtoll(answer.c_str(), &end,10);
      if (end[0] == answer[0]) { continue; }
      return retval;
    }
  }

  std::vector<double> CLI::GetDoublesAnswer(std::string question, std::vector<double> defs, int minlen) {
    std::cout.flags( fFormat );
    while (true) {
      std::cout << question << " [ ";
      for (int i=0; i<defs.size(); ++i) {
        std::cout << defs[i] << " ";
      }
      std::cout << " ] : ";
      std::string answer = GetAnswer();

      if (answer.size() == 0) { return defs; }
      char *end;
      for (int i=0; i<answer.size(); ++i) {
        if (!answer.compare(i,1,",")) {
          answer.replace(i,1," ");
        }
      }
      
      std::stringstream ss(answer);
      std::string val;
      std::vector<double> retvals;
      while (ss >> val) {
        double retval = std::strtod(val.c_str(), &end);
        if (end[0] == val[0]) { continue; }
        retvals.push_back(retval);
      }
      if (retvals.size() < minlen) { continue; }
      return retvals;        
    }
  }

  std::vector<int> CLI::GetIntegersAnswer(std::string question, std::vector<int> defs, int minlen) {
    std::cout.flags( fFormat );
    while (true) {
      std::cout << question << " [ ";
      for (int i=0; i<defs.size(); ++i) {
        std::cout << defs[i] << " ";
      }
      std::cout << " ] : ";
      std::string answer = GetAnswer();
      
      if (answer.size() == 0) { return defs; }
      char *end;
      for (int i=0; i<answer.size(); ++i) {
        if (!answer.compare(i,1,",")) {
          answer.replace(i,1," ");
        }
      }
      
      std::stringstream ss(answer);
      std::string val;
      std::vector<int> retvals;
      while (ss >> val) {
        int retval = std::strtoll(val.c_str(), &end,10);
        if (end[0] == val[0]) { continue; }
        retvals.push_back(retval);
      }
      if (retvals.size() < minlen) { continue; }
      return retvals;        
    }
  }

  void CLI::HelpMessage() {
    std::cout << "run                            - question/answer prompt to run calculation" << std::endl;
    std::cout << "load filename                  - load parameters from input file and run the calculation" << std::endl;
    std::cout << "save filename                  - save current calculation parameters to input file" << std::endl;
    std::cout << "print nlevels                  - print lowes nlevels in current level scheme" << std::endl;
    std::cout << "draw nlevels                   - draw lowest nlevels with plot_levelscheme.py" << std::endl;
    std::cout << "amp nlevels                    - plot wavefunction amplitudes of the lowest nlevels using plot_amplitudes.py" << std::endl;
    std::cout << "amp list                       - plot wavefunction amplitudes of specified levels, format 7_1 5_1 7_2 3_1..." << std::endl;
    std::cout << "epsplot low high increment     - make epsilon plot from eps=low to eps=high" << std::endl;
    std::cout << "gamplot low high increment     - make gamma plot from gamma=low to gamma=high" << std::endl;
    std::cout << "quit                           - exit" << std::endl;
    std::cout << "help                           - print this message" << std::endl;
  }

  std::string CLI::GetAnswer() {
    std::string answer;
    std::getline(std::cin, answer);
    if (std::cin.eof()==1) {
      clearerr(stdin);
      std::cin.clear();
      std::cin.ignore();
    }
    return answer;
  }

  void CLI::Save(std::string filename) {
    std::ofstream ofile(filename);
    ofile << std::fixed << std::setw(15) << "VERBOSITY" << "     " << "1";
    for (auto &ans : fAnswers) {
      ofile << std::fixed << std::setw(15) << ans.first << "     ";
      std::cout.flags( fFormat );
      for (auto &s : ans.second) {
        ofile << s << "  ";
      }
      ofile << std::endl;
    }
    ofile.close();
  }

  void CLI::EpsPlot(double eplo, double ephi, double incr) {

      fAnswers["EPSPLOT"] = {std::to_string(eplo),std::to_string(ephi),std::to_string(incr)};
      std::cout << "Making Epsilon plot..." << std::endl;
      std::ofstream diag_pos("epsilon_pos.dat");
      std::ofstream diag_neg("epsilon_neg.dat");

      for (double ep = eplo; ep <= ephi; ep += incr) {

        std::cout << "[ eps = " << std::fixed << std::setprecision(2) << std::setw(3) << ep << " ] " << std::flush;
        
        poscalc->SetDeformation(ep, fGamma, 0.0, 0.0);
        negcalc->SetDeformation(ep, fGamma, 0.0, 0.0);

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

        poscalc->WriteEnergies(diag_pos);
        negcalc->WriteEnergies(diag_neg);        
      }
      std::cout << std::endl;

      diag_pos.close();
      diag_neg.close();
      popen(("./plot.py 0 "+std::to_string(poscalc->fNOdd)+" "+std::to_string(fEps2)).c_str(), "r");
  }

  void CLI::GamPlot(double gamlo, double gamhi, double incr) {

    fAnswers["GAMPLOT"] = {std::to_string(gamlo),std::to_string(gamhi),std::to_string(incr)};
    std::cout << "Making Gamma plot..." << std::endl;
    std::ofstream diag_pos("gamma_pos.dat");
    std::ofstream diag_neg("gamma_neg.dat");

    for (double gam = gamlo; gam <= gamhi; gam += incr) {

      std::cout << "[ gam = " << std::fixed << std::setprecision(2) << std::setw(3) << gam << " ] " << std::flush;
        
      poscalc->SetDeformation(fEps2, gam, 0.0, 0.0);
      negcalc->SetDeformation(fEps2, gam, 0.0, 0.0);

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

      poscalc->WriteEnergies(diag_pos);
      negcalc->WriteEnergies(diag_neg);        
    }
    std::cout << std::endl;

    diag_pos.close();
    diag_neg.close();
    popen(("./plot.py 1 "+std::to_string(poscalc->fNOdd)+" "+std::to_string(fGamma)).c_str(), "r");
  }
}
