#include <iomanip>

#include "LevelScheme.hh"

namespace PRM {
  Level::Level() {};
  Level::Level(int spin, int indx, int par, double energy) :
    fSpin(spin),
    fIndx(indx),
    fEnergy(energy),
    fParity(par),
    fSPDom(-1),
    fKDom(-1),
    fMaxSP(-1),
    fMaxK(-1),
    fSPDomGlob(-1) {}

  void Level::Print(std::ostream &out) {
    out << std::fixed << std::setprecision(1) << std::setw(8) << fEnergy*1000.0 << "   " <<
      std::fixed << std::setw(3) << fSpin << "/2   " <<
      std::fixed << std::setw(3) << fIndx << "   ";
    if (fParity == 1) {
      out << "+   ";
    }
    else if (fParity == -1) {
      out << "-   ";
    }        
    out << std::fixed << std::setw(3) << fSPDom+1 << "  ";
    out << std::fixed << std::setw(3) << fKDom << "/2  "
        << std::fixed << std::setw(6) << fMaxSP+1 << "  ";
    out << std::fixed << std::setw(6) << fMaxK << "/2  " << 
      std::fixed << std::setw(10) << fSPDomGlob+1 << std::endl;
  }

  void Level::WriteDraw(std::ostream &out, std::string name) {
    out << "L    " << name << "     " <<
      std::fixed << std::setprecision(1) << std::setw(8) << fEnergy*1000.0 << "   " <<
      (double)fSPDomGlob+1-0.3 << "   " << (double)fSPDomGlob+1+0.3 <<
      std::fixed << std::setprecision(1) << std::setw(8) << fEnergy*1000.0 << "   "
        << fSpin << "/2_{"<< fIndx << "}^{";
    if (fParity == 1) {
      out << "+}";
    }
    else if (fParity == -1) {
      out << "-}";
    }        

    out << std::endl;
  }

  LevelScheme::LevelScheme(ParticleRotor &protpos, ParticleRotor &protneg) :
    fPartRotorPos(protpos),
    fPartRotorNeg(protneg)
  {
    //set up pos_splevel -> glob_splevel
    //and    neg_splevel -> glob_splevel
    fPosSPLevel.clear();
    fNegSPLevel.clear();
    int posctr = 0;
    int negctr = 0;
    while (true) {
      double posen = 99999;
      double negen = 99999;
      if (posctr < fPartRotorPos.fNilsson.fEnergies.size()) {
        posen = fPartRotorPos.fNilsson.fEnergies[posctr];
      }

      if (negctr < fPartRotorNeg.fNilsson.fEnergies.size()) {
        negen = fPartRotorNeg.fNilsson.fEnergies[negctr];
      }

      if (posen == 99999 && negen == 99999) {
        break;
      }
         
      if (posen < negen) {
        fPosSPLevel.push_back(posctr+negctr);
        ++posctr;
      }
      else {
        fNegSPLevel.push_back(posctr+negctr);
        ++negctr;
      }
    }

    double minen = 9e9;
    
    if (fPartRotorPos.fNilsson.fDecoupleLevels.size() > 0) {
    for (int spin = fPartRotorPos.fMinSpin; spin<fPartRotorPos.fMaxSpin; spin+=2) {
      if (fPartRotorPos.fEnergies[spin][0] < minen) {
        minen = fPartRotorPos.fEnergies[spin][0];
      }
    }
    }
    
    if (fPartRotorNeg.fNilsson.fDecoupleLevels.size() > 0) {
    for (int spin = fPartRotorNeg.fMinSpin; spin<fPartRotorNeg.fMaxSpin; spin+=2) {
      if (fPartRotorNeg.fEnergies[spin][0] < minen) {
        minen = fPartRotorNeg.fEnergies[spin][0];
      }
    }
    }
    
    if (fPartRotorPos.fNilsson.fDecoupleLevels.size() > 0) {
      for (int spin = fPartRotorPos.fMinSpin; spin<fPartRotorPos.fMaxSpin; spin+=2) {
        for (int indx = 0; indx < fPartRotorPos.fEnergies[spin].size(); ++indx) {
          fLevels[{spin,indx+1, 1}] = Level(spin,indx+1, 1, fPartRotorPos.fEnergies[spin][indx] - minen);
          fLevels[{spin,indx+1, 1}].fSPDom = fPartRotorPos.GetSPDom(spin,indx+1);
          fLevels[{spin,indx+1, 1}].fKDom = fPartRotorPos.GetKDom(spin,indx+1);
          fLevels[{spin,indx+1, 1}].fMaxK = fPartRotorPos.GetMaxKSP(spin,indx+1).first;
          fLevels[{spin,indx+1, 1}].fMaxSP = fPartRotorPos.GetMaxKSP(spin,indx+1).second;  
          fLevels[{spin,indx+1, 1}].fSPDomGlob = fPosSPLevel[fLevels[{spin,indx+1, 1}].fSPDom];
        }
      }
    }

    if (fPartRotorNeg.fNilsson.fDecoupleLevels.size() > 0) {
      for (int spin = fPartRotorNeg.fMinSpin; spin<fPartRotorNeg.fMaxSpin; spin+=2) {
        for (int indx = 0; indx < fPartRotorNeg.fEnergies[spin].size(); ++indx) {
          fLevels[{spin,indx+1, -1}] = Level(spin,indx+1, -1, fPartRotorNeg.fEnergies[spin][indx] - minen);
          fLevels[{spin,indx+1, -1}].fSPDom = fPartRotorNeg.GetSPDom(spin,indx+1);
          fLevels[{spin,indx+1, -1}].fKDom = fPartRotorNeg.GetKDom(spin,indx+1);
          fLevels[{spin,indx+1, -1}].fMaxK = fPartRotorNeg.GetMaxKSP(spin,indx+1).first;
          fLevels[{spin,indx+1, -1}].fMaxSP = fPartRotorNeg.GetMaxKSP(spin,indx+1).second;  
          fLevels[{spin,indx+1, -1}].fSPDomGlob = fNegSPLevel[fLevels[{spin,indx+1, -1}].fSPDom];
        }
      }
    }
  }

  void LevelScheme::SetPrintLevels(int nlevels) {
    fPrintLevels.clear();
    fPosPrintLevels.clear();
    fNegPrintLevels.clear();

    for (int lev = 0; lev < nlevels; ++lev) {
      double thisen = 9e9;
      int thisspin = 0;
      int thisindx = 0;
      int thispar = 0;

      std::map<std::tuple<int, int, int>, Level>::iterator it;
      for (it = fLevels.begin(); it != fLevels.end(); it++) {
        int done = 0;
        for (int i=0; i<fPrintLevels.size(); ++i) {
          if (it->second.fSpin == std::get<0>(fPrintLevels[i]) &&
              it->second.fIndx == std::get<1>(fPrintLevels[i]) &&
              it->second.fParity == std::get<2>(fPrintLevels[i])) {
            done = 1;
            break;
          }          
        }
        if (done) { continue; }
        if (it->second.fEnergy < thisen) {
          thisen = it->second.fEnergy;
          thisspin = it->second.fSpin;
          thisindx = it->second.fIndx;
          thispar = it->second.fParity;
        }
      }      
      fPrintLevels.push_back({thisspin,thisindx,thispar});
      if (thispar==1) {
        fPosPrintLevels.push_back({thisspin,thisindx});
      }
      else if (thispar==-1) {
        fNegPrintLevels.push_back({thisspin,thisindx});
      }
    }
  }
    
  void LevelScheme::Print(std::ostream &out, int nlevels) {
    SetPrintLevels(nlevels);

    int lev = 1;
    out << std::fixed << std::setw(3) << "#"
        << std::fixed << std::setw(12) << "Energy"
        << std::fixed << std::setw(8) << "J"
        << std::fixed << std::setw(6) << "i"
        << std::fixed << std::setw(10) << "Dom SP"
        << std::fixed << std::setw(7) << "Dom K"
        << std::fixed << std::setw(8) << "Max SP"
        << std::fixed << std::setw(10) << "Max K"
        << std::fixed << std::setw(12) << "Dom SP (G)"
        << std::endl;
    for (auto &state : fPrintLevels) {      
      out << std::fixed << std::setw(3) << lev << "    ";
      fLevels[state].Print(out);
      if ((lev)%5 == 0) { out << std::endl; }
      ++lev;
    }
  }

  void LevelScheme::WriteDraw(std::ostream &out, int nlevels) {
    SetPrintLevels(nlevels);

    int lev = 1;
    for (auto &state : fPrintLevels) {      
      fLevels[state].WriteDraw(out, std::to_string(lev));
      ++lev;
    }
  }
}
