#include <iomanip>
#include <cmath>
#include <string>

#include <armadillo>

#include <gsl/gsl_sf_coupling.h>

#include "Nilsson.hh"

namespace PRM {
  TensorOperator::TensorOperator(int dim, int rank) {
    Set(dim, rank);
  }

  void TensorOperator::Set(int dim, int rank) {
    values.clear();
    fRank = rank;
    for (int i=0; i<2*rank + 1; ++i) {
      values.emplace_back(dim, dim, arma::fill::zeros);
    }
  }  

  arma::mat &TensorOperator::operator()(int rank) {
    //indexes go from -r, -r+1, -r+2...0...r-2,r-1,r
    //index = rank + 2*maxrank
    return values[rank + fRank];
  }
  
  arma::mat &TensorOperator::GetValue(int rank) {
    //indexes go from -r, -r+1, -r+2...0...r-2,r-1,r
    //index = rank + 2*maxrank
    return values[rank + fRank];
  }
  
  Nilsson::Nilsson(int Z, int A) : fZ(Z),
                                           fA(A),
                                           fNNeut(10),
                                           fNProt(10),
                                           fGammaFunc(151,0) {
    fN = A - Z;
    if (fZ%2 == 0 && fN%2 == 1) {
      fOdd = 1; //odd neutron
      fNOdd = fN;
    }
    else if (fN%2 == 0 && fZ%2 == 1) {
      fOdd = 0; //odd proton
      fNOdd = fZ;
    }

    fAvL2.push_back(0); //N=0 shell has average l*l = 0

    //note that this is how gampn.f calculates them
    //l(l+1) averaged over all particles in the shell    
    for (int N=2; N<15; ++N) { 
      fAvL2.push_back(fAvL2[N-2]+N);
    }

    //gamma function
    //reimplemented from gampn.f
    //check if this is significantly faster than gsl gamma function
    fGammaFunc[52] = 1.0;
    fGammaFunc[51] = 1.77245385;
    for (int i=1; i<=28; ++i) {
      fGammaFunc[2*i+52] = i*fGammaFunc[2*i+50];
      fGammaFunc[2*i+51] = (i-0.5)*fGammaFunc[2*i+49];
    }
    fGammaFunc[49] = -2.0*fGammaFunc[51];
    for (int i=i; i<=20; ++i) {
      fGammaFunc[-2*i+49] = fGammaFunc[-2*i+51]/(-i-0.5);
    }
    fCoupledShells = 0;
  }

  double Nilsson::Gamma(int i) {
    return fGammaFunc[i];
  }

  void Nilsson::SetNMax(int nn, int np) {
    fNNeut = nn;
    fNProt = np;
  }

  void Nilsson::SetKappas(std::vector<double> kappas) {
    fKappa = kappas;
  }

  void Nilsson::SetMus(std::vector<double> mus) {
    fMu = mus;
  }

  void Nilsson::SetDeformation(double epsp, double gamma, double eps4, double eps6) {
    //see Sec. 3 of Physica Scripta 8, 17 (1973)
    fEpsP = epsp;
    fGamma = gamma * PI / 180.0;
    fEps2 = fEpsP * std::cos(fGamma);
    fEps22 = fEpsP * std::sin(fGamma) / std::sqrt(3./4.);
    fEps4 = eps4;
    fEps6 = eps6;

    VolumeConservation();
  }

  void Nilsson::VolumeConservation() {
    if (std::abs(fEps4) > 1e-4) {
      std::cerr << "|eps_4| > 0 not implemented!" << std::endl;
      exit(1);
    }
    fOmOm = std::pow((1. + fEps2/3. + fEps22/2.) * (1. + fEps2/3. - fEps22/2.) * (1. - 2.*fEps2/3.), -1./3.);
  }

  void Nilsson::SetCranking(double omrot) {
    fOmRot = 0;
  }

  void Nilsson::Solve() {
    MakeHamiltonian();
    Diagonalize();
  }

  void Nilsson::MakeBasisStates(int parity) {
    fBasisStates.clear();
    
    int maxShell = 0;
    if (fOdd) { maxShell = fNNeut; }
    else { maxShell = fNProt; }

    if (maxShell%2 == 1 && parity == 1) {
      maxShell = maxShell - 1;
    }
    else if (maxShell%2 == 0 && parity == -1) {
      maxShell = maxShell - 1;
    }
    
    for (int N = maxShell; N >= 0; N = N-2) {
      int omgmax = 2*N+1;
      if ((omgmax-1)%4 != 0) { omgmax = 2*N-1; }
      for (int omg2 = omgmax ; omg2 >= -(2*N+1); omg2 = omg2-4) {
        for (int l2 = 2*N; l2 >= omg2-1; l2 = l2 - 4) {
          if (l2 < 0) { continue; }
          int l = l2/2;
          for (int s2 = 1; s2 >= -1; s2 = s2 - 2) {
            int j2 = l2 + s2;            
            if (j2 < std::abs(omg2)) { continue; }
            fBasisStates.emplace_back(N, l, j2, omg2);
          }
        }
      }
    }

    fMaxOmg = 0;
    for (int i=0; i<fBasisStates.size(); ++i) {
      if (std::abs(fBasisStates[i].omega2) > fMaxOmg) {
        fMaxOmg = std::abs(fBasisStates[i].omega2);
      }
    }
  }

  void Nilsson::PrintBasisStates(std::ostream &out, int ncols/*=5*/) {
    std::vector<int> offsets(ncols,0); //columns
    int nrows = std::ceil((double)fBasisStates.size()/ncols);
    for (int i = 0; i < ncols; ++i) {
      offsets[i] = nrows*i;
    }
  
    for (int row = 0; row < nrows; ++row) {
      for (int col = 0; col < ncols; ++col) {
        int indx = col*nrows + row;
        if (indx >= fBasisStates.size()) { continue; }
        out << std::setw(4) << indx+1 << " ";
        fBasisStates[indx].Print(out);
        out << std::setw(8) << " ";
      }
      out << std::endl;
    }
  }

  void Nilsson::MakeHamiltonian() {
    //define and initialize hamiltonian of correct dimension
    if (fKappa.size() < std::min(fNNeut, fNProt)) { std::cerr << "Set Kappas before trying to diagonalize! " << std::endl; exit(1); }
    if (fMu.size() < std::min(fNNeut, fNProt)) { std::cerr << "Set Kappas before trying to diagonalize! " << std::endl; exit(1); }
    
    int dim = fBasisStates.size();
    fHamiltonian = arma::mat(dim, dim, arma::fill::zeros);

    //cranking things (?)
    double OmX = 1.0 - 2./3. * fEpsP * std::cos(fGamma + 2.*PI/3.);
    double OmY = 1.0 - 2./3. * fEpsP * std::cos(fGamma - 2.*PI/3.);
    double OmZ = 1.0 - 2./3. * fEpsP * std::cos(fGamma);

    fFac = 0.5 * (std::sqrt(OmY/OmX) + std::sqrt(OmX/OmY)) - 1;
    fHac = 0.5 * (std::sqrt(OmY/OmX) - std::sqrt(OmX/OmY));
    fOmRot = fOmRot / fOmOm;
    
    for (int i=0; i<fBasisStates.size(); ++i) {
      auto &bsi = fBasisStates[i];
      int Ni = bsi.N;
      int li = bsi.l;
      int j2i = bsi.j2;
      int omg2i = bsi.omega2;
      for (int j=i; j<fBasisStates.size(); ++j) {
        auto &bsj = fBasisStates[j];
        int Nj = bsj.N;
        int lj = bsj.l;
        int j2j = bsj.j2;
        int omg2j = bsj.omega2;

        if (Ni > Nj + fCoupledShells ||
            omg2j > omg2i+8 ||
            omg2j < omg2i-8) { continue; }

        double E20 = 0.0;
        double E22 = 0.0;
        double E40 = 0.0;
        double E42 = 0.0;
        double E44 = 0.0;
        double E60 = 0.0;
        double ELZ = 0.0;  //come up with more descriptive names
        double END = 0.0;  //these are both cranking things I think

        if (omg2i == omg2j + 8 || omg2i == omg2j - 8) {
          E44 = fEps4 * std::sqrt(70.0)/12. * std::pow(std::sin(fGamma),2) *
            ClebschGordan(j2j, 8, j2i, -1, 0, -1) * (ClebschGordan(j2j, 8, j2i, omg2j, 8, omg2i) +
                                                   ClebschGordan(j2j, 8, j2i, omg2j, -8, omg2i));
        }
        else if (omg2i == omg2j + 4 || omg2i == omg2j - 4) {
          if (Ni == Nj) {
            E22 = fEps22/std::sqrt(6.0) * ClebschGordan(j2j, 4, j2i, -1, 0, -1) *
              (ClebschGordan(j2j, 4, j2i, omg2j, 4, omg2i) +
               ClebschGordan(j2j, 4, j2i, omg2j, -4, omg2i));
          }
          if (Ni == Nj + 2) {
            END = fOmRot/std::sqrt(3./2.) * fHac * ClebschGordan(j2j, 4, j2i, -1, 0, -1) *
              (ClebschGordan(j2j, 4, j2i, omg2j, 4, omg2i) -
               ClebschGordan(j2j, 4, j2i, omg2j, -4, omg2i));    
          }
          E42 = -fEps4 * std::sqrt(30.0)/12. * std::sin(2*fGamma) *
            ClebschGordan(j2j, 8, j2i, -1, 0, -1) * (ClebschGordan(j2j, 8, j2i, omg2j, 4, omg2i) +
                                                    ClebschGordan(j2j, 8 ,j2i, omg2j, -4, omg2i));                                                
        }
        else {
          if (Ni == Nj) {
            E20 = -2.0 * fEps2 / 3.0 * ClebschGordan(j2j, 4, j2i, -1, 0, -1) * ClebschGordan(j2j, 4, j2i, omg2j, 0, omg2i);
          }
          
          E40 = fEps4 * ((5.0 * std::pow(std::cos(fGamma),2) + 1)/6.0) *
            ClebschGordan(j2j, 8, j2i, -1, 0, -1) * ClebschGordan(j2j, 8, j2i, omg2j, 0, omg2i);
          E60 = fEps6 * ClebschGordan(j2j, 12, j2i, -1, 0, -1) * ClebschGordan(j2j, 12, j2i, omg2j, 0, omg2i);
          if (Ni == Nj && li == lj) {
            ELZ = -fOmRot * fFac * ((omg2j - 1)/2.
                                    * ClebschGordan(2*lj, 1, j2i, omg2j-1, 1, omg2j)
                                    * ClebschGordan(2*lj, 1, j2j, omg2j-1, 1, omg2j) +
                                    (omg2j + 1)/2.
                                    * ClebschGordan(2*lj, 1, j2i, omg2j+1, -1, omg2j)
                                    * ClebschGordan(2*lj, 1, j2j, omg2j+1, -1, omg2j));
          }
        }

      double radialME = GetRadialME(Ni, li, Nj, lj, 2);        
        
        double matEl = radialME * std::sqrt((j2j + 1.0)/(j2i + 1.0)) *
          (E20 + E22 + E40 + E42 + E44 + E60 + END) + ELZ;

        if (i == j) {
          matEl += Nj + 3./2. - fKappa[Nj]/fOmOm*(j2j*(j2j+2.)/4. - lj*(lj+1) - 3./4.) -
            fKappa[Nj]/fOmOm * fMu[Nj] * lj*(lj+1) + fKappa[Nj]*fMu[Nj]/fOmOm * fAvL2[Nj] -
            fOmRot*omg2j/2.0;
        }

        //matEl = -matEl;

        fHamiltonian(i,j) = matEl;
        fHamiltonian(j,i) = matEl;
      }
    }
  }

  void Nilsson::Diagonalize() {
    arma::eig_sym(fEnergies, fCompositions, fHamiltonian);
    
    fOmegas = std::vector<int>(fEnergies.size(), 0);
    for (int i=0; i<fEnergies.size(); ++i) {
      arma::vec wf = fCompositions.col(i);
      for (int j=0; j<wf.size(); ++j) {
        if (std::abs(wf[j])>1e-4) {
          //the basis state tells us what the Omega is
          fOmegas[i] = std::abs(fBasisStates[j].omega2);
          break;
        }
      }
    }        
  }

  void Nilsson::SetDecoupleLevels(std::vector<int> levels) {
    fDecoupleLevels.clear();
    for (int i=0; i<levels.size(); ++i) {
      fDecoupleLevels.push_back(levels[i]-1);
    }
  }

  void Nilsson::Decouple(std::vector<int> levels) {
    SetDecoupleLevels(levels);
    if (levels.size() > 0) {
      Decouple();
    }
  }
  
  void Nilsson::Decouple() {
    //calculate "decoupling factors" i.e. operators needed for the particle-rotor
    //part of the calculation

    int dim = fEnergies.size();
    int nlevels = fDecoupleLevels.size();
    fABCDE.Set(dim, 2);  // -2,+2 = E, -1 = C, +1 = A+B, 0 = D.
    fNorm.Set(dim, 0);
    fJz.Set(dim, 0);
    fJz2.Set(dim, 0);
    fS.Set(dim, 1);

    fJr.Set(dim, 1);
    fSr.Set(dim, 1);
    fFr.Set(dim, 1);

    fR2J.Set(dim, 1);
    fR4J.Set(dim, 1);
    fR2S.Set(dim, 1);
    fR4S.Set(dim, 1);
    fR2D.Set(dim, 1);
    fR4D.Set(dim, 1);

    //loop through basis states
    for (int i=0; i<fBasisStates.size(); ++i) {
      auto &bsi = fBasisStates[i];
      int Ni = bsi.N;
      int li = bsi.l;
      int j2i = bsi.j2;
      int omg2i = bsi.omega2;
      for (int j=0; j<fBasisStates.size(); ++j) {
        auto &bsj = fBasisStates[j];
        int Nj = bsj.N;
        int lj = bsj.l;
        int j2j = bsj.j2;
        int omg2j = bsj.omega2;

        double rad0 = 1;        
        if (Nj != Ni || lj != li) { rad0 = 0.0; }

        double rad2 = GetRadialME(Nj,lj,Ni,li,2);
        double rad4 = GetRadialME(Nj,lj,Ni,li,4);

        double redmd = RedMED1(2*lj,1,j2j,2*li,1,j2i); //look this up in Edmonds

        if (omg2i == omg2j) { //rank zero only if omegas are equal
          //loop over nilsson states
          double CG = (ClebschGordan(2*lj,1,j2i,omg2i-1,1,omg2i) *
                       ClebschGordan(2*lj,1,j2j,omg2j-1,1,omg2j)) -
                      (ClebschGordan(2*lj,1,j2i,omg2i+1,-1,omg2i)*
                       ClebschGordan(2*lj,1,j2j,omg2j+1,-1,omg2j));

          for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
            for (int inj = ini; inj<nlevels; ++inj) {
              int nj = fDecoupleLevels[inj];
              double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //access is row, column

              if (Ni == Nj && li == lj && j2i == j2j) {
                fNorm(0)(ni,nj) += coeffs;
                //fABCDE(0)(ni,nj) += coeffs*(j2j*(j2j+2) - omg2j*omg2j)/4.0;
                //fJz2(0)(ni,nj) += coeffs*omg2j*omg2j/4.;                
              }
              if (li == lj) {                
                fS(0)(ni,nj) += coeffs * CG * 0.5 * rad0;
                fR2S(0)(ni,nj) += coeffs * CG * 0.5 * rad2;
                fR4S(0)(ni,nj) += coeffs * CG * 0.5 * rad4;
              }
              if (j2i == j2j) {
                fJz(0)(ni,nj) += coeffs * omg2j/2. * rad0;
                fR2J(0)(ni,nj) += coeffs * omg2j/2. * rad2;
                fR4J(0)(ni,nj) += coeffs * omg2j/2. * rad4;
              }
              double coeffD = std::pow(-1, (j2i - omg2i)/2.)/std::sqrt(3) *
                ClebschGordan(j2j, j2i, 2, omg2j,-omg2i, 0);
              fR2D(0)(ni,nj) += coeffs*coeffD*redmd*rad2;
              fR4D(0)(ni,nj) += coeffs*coeffD*redmd*rad4;
            }
          }

          //this is correction stuff which we always do
          for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
            if (fCompositions(i,ni) == 0) { continue; }
            for (int nj=0; nj<dim; ++nj) {
              double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //these could be the wrong way around

              if (li == lj) {
                fSr(0)(ni,nj) += coeffs * CG * 0.5 * rad0;
              }
              if (j2i == j2j) {
                fJr(0)(ni,nj) += coeffs * omg2j/2. * rad0;
              }

            }
          }   
        }


        //+1 components        
        if (omg2i == -omg2j+2) {
          double CG = ClebschGordan(2*lj,1,j2i,omg2i-1,1,omg2i) *
            ClebschGordan(2*lj,1,j2j,-(omg2j-1),-1,-omg2j) *
            std::pow(-1,(j2j-omg2j)/2.);

          double CGf = 0.0;
          if (std::abs(Nj-Ni) == 2 && std::abs(lj-li) <= 2) {
            CGf = std::pow(-1, (j2i-omg2i)/2.) * (Ni - Nj)/std::sqrt(3./2.) *
              std::sqrt((j2i+1.0)/(j2j+1.0)) * ClebschGordan(j2i, 4, j2j, -omg2i, 2, omg2j) *
              ClebschGordan(j2i, 4, j2j, -1, 0, -1);             
          }

          for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
            for (int inj = ini; inj<nlevels; ++inj) {
              int nj = fDecoupleLevels[inj];
              double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //these could be the wrong way around
              if (lj == li) {
                fS(1)(ni,nj) += coeffs * CG * rad0;
                fR2S(1)(ni,nj) += coeffs * CG * rad2;
                fR4S(1)(ni,nj) += coeffs * CG * rad4;
              }
              if (j2j == j2i) {
                double x = std::pow(-1, (j2j - omg2j)/2.) * 0.5 * std::sqrt((j2j+omg2j)*(j2j-omg2j+2.));
                //fABCDE(1)(ni,nj) += coeffs * x * rad0;
                fR2J(1)(ni,nj) += coeffs * x * rad2;
                fR4J(1)(ni,nj) += coeffs * x * rad4; 
              }

              double coeffD = -1/std::sqrt(3) * ClebschGordan(j2j,j2i,2,omg2j,omg2i,2);
              fR2D(1)(ni,nj) += coeffs * coeffD * redmd * rad2;
              fR4D(1)(ni,nj) += coeffs * coeffD * redmd * rad4;              
            }
          }

          //this is correction stuff which we always do
          for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
            if (fCompositions(i,ni) == 0) { continue; }
            for (int nj=0; nj<dim; ++nj) {
              double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //these could be the wrong way around
              if (li == lj) {
                fSr(1)(ni,nj) += coeffs * CG * rad0;
              }
              if (j2i == j2j) {
                double x = std::pow(-1, (j2j - omg2j)/2.) * 0.5 * std::sqrt((j2j+omg2j)*(j2j-omg2j+2.));
                fJr(1)(ni,nj) += coeffs * x * rad0;
              }              
              fFr(-1)(ni,nj) += coeffs * CGf * rad2;              
            }
          }
        }

        // -1 components
        if (omg2i == -omg2j - 2) {
          double CG = ClebschGordan(2*lj,1,j2i,omg2i+1,-1,omg2i) *
            ClebschGordan(2*lj,1,j2j,-(omg2j+1),1,-omg2j) *
            std::pow(-1,(j2j-omg2j)/2.);

          double CGf = 0.0;
          if (std::abs(Nj-Ni) == 2 && std::abs(lj-li) <= 2) {
            CGf = std::pow(-1, (j2i-omg2i)/2.) * (Ni - Nj)/std::sqrt(3./2.) *
              std::sqrt((j2i+1.0)/(j2j+1.0)) * ClebschGordan(j2i, 4, j2j, -omg2i, -2, omg2j) *
              ClebschGordan(j2i, 4, j2j, -1, 0, -1);             
          }

          for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
            for (int inj = ini; inj<nlevels; ++inj) {
              int nj = fDecoupleLevels[inj];
              double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //these could be the wrong way around
              if (lj == li) {
                fS(-1)(ni,nj) += coeffs * CG * rad0;
                fR2S(-1)(ni,nj) += coeffs * CG * rad2;
                fR4S(-1)(ni,nj) += coeffs * CG * rad4;
              }
              if (j2j == j2i) {
                double x = std::pow(-1, (j2j - omg2j)/2.) * 0.5 * std::sqrt((j2j-omg2j)*(j2j+omg2j+2.));
                //fABCDE(-1)(ni,nj) += coeffs * x * rad0;
                fR2J(-1)(ni,nj) += coeffs * x * rad2;
                fR4J(-1)(ni,nj) += coeffs * x * rad4; 
              }

              double coeffD = -1/std::sqrt(3) * ClebschGordan(j2j,j2i,2,omg2j,omg2i,-2);
              fR2D(-1)(ni,nj) += coeffs * coeffD * redmd * rad2;
              fR4D(-1)(ni,nj) += coeffs * coeffD * redmd * rad4;              
            }
          }

          for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
            if (fCompositions(i,ni) == 0) { continue; }
            for (int nj=0; nj<dim; ++nj) {
              double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //these could be the wrong way around
              if (li == lj) {
                fSr(-1)(ni,nj) += coeffs * CG * rad0;
              }
              if (j2i == j2j) {
                double x = std::pow(-1, (j2j - omg2j)/2.) * 0.5 * std::sqrt((j2j-omg2j)*(j2j+omg2j+2.));
                fJr(-1)(ni,nj) += coeffs * x * rad0;
              }
              fFr(1)(ni,nj) += coeffs * CGf * rad2;
            }
          }
        }

        //+2 and -2 components
        if (omg2i == omg2j-4 || false) { //we never do this
          //this is only for no corrections
          ;
        }

        //correction to Jz from stretching basis, nonaxiality
        if (std::abs(Ni-Nj) != 2) { continue; }
        int mu = omg2j-omg2i;
        if (std::abs(mu) != 4 || std::abs(lj-li) > 2) { continue; }
        double CGf = (Nj - Ni)/std::sqrt(6) * std::sqrt((j2i+1.)/(j2j+1)) *
          ClebschGordan(j2i, 4, j2j, -1, 0, -1) * ClebschGordan(j2j, 4, j2i, omg2j, mu, omg2i);
        if (mu == 4) { CGf = - CGf; }
        for (int ini = 0; ini<nlevels; ++ini) {
            int ni = fDecoupleLevels[ini];
          if (fCompositions(i,ni) == 0) { continue; }
          for (int nj=0; nj<dim; ++nj) {
            double coeffs = fCompositions(i,ni)*fCompositions(j,nj); //these could be the wrong way around
            fFr(0) += coeffs * CGf * rad2;
          }
        }        
      }
    }

    //corrections for stretched basis
    double OmX = 1.0 - 2./3. * fEpsP * std::cos(fGamma + 2.*PI/3.);
    double OmY = 1.0 - 2./3. * fEpsP * std::cos(fGamma - 2.*PI/3.);
    double OmZ = 1.0 - 2./3. * fEpsP * std::cos(fGamma);

    double f1 = std::sqrt(OmZ/OmY);
    double f2 = std::sqrt(OmZ/OmX);

    double fac1 = (f1 + 1./f1 + 1./f2 + f2)/4.;
    double fac2 = (f1 + 1./f1 - 1./f2 - f2)/4.;
    double fac3 = (f1 - 1./f1 + 1./f2 - f2)/4.;
    double fac4 = (f1 - 1./f1 - 1./f2 + f2)/4.;
    double fac5 = (f2/f1 + f1/f2)/2.;
    double fac6 = (f2/f1 - f1/f2)/2.;

    for (int ini = 0; ini<nlevels; ++ini) {
      int ni = fDecoupleLevels[ini];
      for (int nj= 0; nj<dim; ++nj) {
        double jtmp = fJr(1)(ni,nj);
        fJr(1)(ni,nj) =
          fac1 * jtmp                       + (1.-fac1)*fSr(1)(ni,nj)
          + fac2 * fJr(-1)(ni,nj)    -  fac2*fSr(-1)(ni,nj)
          + fac3 * fFr(1)(ni,nj)     +  fac4 * fFr(-1)(ni,nj);

        fJr(-1)(ni,nj) =
          fac1 * fJr(-1)(ni,nj)      + (1.-fac1)*fSr(-1)(ni,nj)
          +fac2 * jtmp                     - fac2 * fSr(1)(ni,nj)
          + fac3 *fFr(-1)(ni,nj)     + fac4 * fFr(1)(ni,nj);

        fJr(0)(ni,nj) =
          fac5 * fJr(0)(ni,nj)          + (1.-fac5)*fSr(0)(ni,nj)
          + fac6 * fFr(0)(ni,nj);
      }
    }

    for (int ini = 0; ini<nlevels; ++ini) {
      int ni = fDecoupleLevels[ini];
      for (int inj = 0; inj<nlevels; ++inj) {
        int nj = fDecoupleLevels[inj];
        fABCDE(1)(ni,nj) = fJr(1)(ni,nj);
        fABCDE(-1)(ni,nj) = fJr(-1)(ni,nj);
        fJz(0)(ni,nj) = fJr(0)(ni,nj);

        for (int i=0; i<dim; ++i) {
          fABCDE(0)(nj,ni) += (fJr(1)(ni,i)*fJr(1)(nj,i)
                            + fJr(-1)(ni,i)*fJr(-1)(nj,i))/2.0;
            
          fABCDE(2)(nj,ni) += (fJr(1)(ni,i)*fJr(-1)(nj,i)
                               + fJr(-1)(ni,i)*fJr(1)(nj,i));
          fJz2(0)(nj,ni) += fJr(0)(ni,i) * fJr(0)(nj,i);
        }
      }
    }
  }


  void Nilsson::DiagQMom() {
    double fac = (1. + fEps2/3. + fEps22/2.)*(1.+fEps2/3. - fEps22/2.)*(1.-2.*fEps2/3.);
    int dim = fEnergies.size();
    
    fQ20s = arma::vec(dim, arma::fill::zeros);
    fQ22s = arma::vec(dim, arma::fill::zeros);
    fQ20  = arma::vec(dim, arma::fill::zeros);
    fQ22  = arma::vec(dim, arma::fill::zeros);
    fRho2 = arma::vec(dim, arma::fill::zeros);
    fR2   = arma::vec(dim, arma::fill::zeros);
    fRJz  = arma::vec(dim, arma::fill::zeros);

    for (int i=0; i<dim; ++i) {
      auto &bsi = fBasisStates[i];
      int Ni = bsi.N;
      int li = bsi.l;
      int j2i = bsi.j2;
      int omg2i = bsi.omega2;
      for (int j=0; j<dim; ++j) {
        auto &bsj = fBasisStates[j];
        int Nj = bsj.N;
        int lj = bsj.l;
        int j2j = bsj.j2;
        int omg2j = bsj.omega2;

        if (std::abs(Ni-Nj) > 2 || std::abs(j2i - j2j) > 4 || std::abs(li - lj) > 2 || std::abs(omg2i - omg2j) > 4) { continue; }

        double CG1(0),CG2(0),CG3(0);
        double radME;
        double HLz(0),HJz(0);

        radME = GetRadialME(Ni,li,Nj,lj,2);

        if (omg2i == omg2j) { CG1 = ClebschGordan(j2j, 4, j2i, -1, 0, -1) * ClebschGordan(j2j, 4, j2i, omg2j, 0, omg2j); }
        if (omg2i == omg2j+4) { CG2 = ClebschGordan(j2j, 4, j2i, -1, 0, -1) * ClebschGordan(j2j, 4, j2i, omg2j, -4, omg2i); }
        if (omg2i == omg2j-4) { CG3 = ClebschGordan(j2j, 4, j2i, -1, 0, -1) * ClebschGordan(j2j, 4, j2i, omg2j, -4, omg2i); }
        if (Ni == Nj && li == lj && j2i == j2j && omg2i == omg2j) { HJz = omg2j/2.; }
        if (Ni == Nj && li == lj && omg2i == omg2j) {
          HLz = (omg2j - 1)/2.*ClebschGordan(2*lj,1,j2i,omg2j-1,1,omg2j) *
            ClebschGordan(2*lj,1,j2j,omg2j-1,1,omg2j) +
            (omg2j+1)/2.*ClebschGordan(2*lj,1,j2i,omg2j+1,-1,omg2j) *
            ClebschGordan(2*lj,1,j2j,omg2j+1,-1,omg2j);
        }

        for (int niq = 0; niq<dim; ++niq) {
          fQ20s(niq) += fCompositions(i,niq)*fCompositions(j,niq) * 2.0 * radME * std::sqrt((j2j+1.)/(j2i+1.))*CG1;          
          fQ22s(niq) += fCompositions(i,niq)*fCompositions(j,niq) * 2.0 * radME * std::sqrt((j2j+1.)/(j2i+1.))*(CG2+CG3);
          if (li==lj && j2i==j2j && omg2i==omg2j) {
            fRho2(niq) += fCompositions(i,niq)*fCompositions(j,niq) * radME;
          }
          fRJz(niq) += fCompositions(i,niq)*fCompositions(j,niq) *
            (HJz + fFac*HLz - fHac/std::sqrt(3./2.)*(CG2-CG3)*std::sqrt((j2j+1.)/(j2i+1.))* radME * (Ni - Nj)/2.);
        }
      }
    }

    //now transform from stretched into real coordinates
    for (int niq = 0; niq<dim; ++niq) {
      fQ20(niq) = ((1.+fEps2/3. - fEps22*fEps22/6.)*fQ20s(niq) +
                   fEps22/std::sqrt(6.)*(1.-2.*fEps2/3.)*fQ22s(niq)+
                   (2.*fEps2/3.*(1.+fEps2/3.)-fEps22*fEps22/6.)*fRho2(niq))/fac;

      fQ22(niq) = ((1.+fEps2/3.)*fQ22s(niq) +
                   fEps22/std::sqrt(24.)*fQ20s(niq) -
                   fEps22/std::sqrt(6.)*fRho2(niq))/fac * (1.-2.*fEps2/3.);

      fR2(niq) = (((1.+fEps2/3.)*(1.-fEps2/3) - 1/3.*std::pow(fEps22/2.,2))*fRho2(niq) +
                  ((1.+fEps2/3.)*fEps2/3. - 1/3.*std::pow(fEps22/3.,2))*fQ20s(niq) -
                  fEps22/2.*(1.-2.*fEps2/3.)*std::sqrt(2./3.)*fQ22s(niq))/fac;
    }
  }

  void Nilsson::PrintWavefunctions(std::ostream &out, int maxrow/*=5*/) {
    int width = 2;
    if (fBasisStates.size() > 99) { width = 3; }

    int rowct = 0;
    for (int i=0; i<fEnergies.size(); ++i) {
      rowct = 0;
      out << "E = " << std::setw(7) << fEnergies[i] << "   ";
      arma::vec wf = fCompositions.col(i);
      for (int j=0; j<wf.size(); ++j) {
        if (std::abs(wf[j])>1e-4) {
          rowct += 1;
          out << std::fixed << std::setprecision(5) << std::showpos
              << wf[j] << " |" << std::noshowpos << std::setw(width) << j+1 << "> ";
          if (rowct%maxrow == 0) { out << std::endl << std::setw(14) << " "; }
        }
      }
      out << std::endl;
    }
  }

  void Nilsson::PrintDecouple(std::ostream &out) {
    int nlevels = fDecoupleLevels.size();
    if (nlevels == 0) { return; }
    std::vector<std::string> colnames = {"Ei", "Ef", "A+B", "C", "D", "E", "Jz", "S0", "S+", "S-", "<R2S>", "<R4S>", "R2DZ", "R2D+", "R2D-"};
    for (int i=0; i<colnames.size(); ++i){
      out << std::setw(8) << colnames[i];
    }
    out << std::endl;
    int dim = fEnergies.size();

    for (int ini = 0; ini<nlevels; ++ini) {
      int i = fDecoupleLevels[ini];
      for (int inj = 0; inj<nlevels; ++inj) {
        int j = fDecoupleLevels[inj];
        double snorm = (std::abs(fS(0)(i,j)) +
                        std::abs(fS(-1)(i,j)) +
                        std::abs(fS(1)(i,j)));
        double avr2s = (std::abs(fR2S(0)(i,j)) +
                        std::abs(fR2S(-1)(i,j)) +
                        std::abs(fR2S(1)(i,j)))/snorm;
        double avr4s = (std::abs(fR4S(0)(i,j)) +
                        std::abs(fR4S(-1)(i,j)) +
                        std::abs(fR4S(1)(i,j)))/snorm;

        if (snorm == 0) { continue; }
        out << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fEnergies[j] << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fEnergies[i] << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fABCDE(1)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fABCDE(-1)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fABCDE(0)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fABCDE(2)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fJz(0)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fS(0)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fS(1)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fS(-1)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << avr2s << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << avr4s << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fR2D(0)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fR2D(1)(i,j) << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fR2D(-1)(i,j) << std::endl;
      }
      out << std::endl;
    }    
  }

  void Nilsson::PrintDiagQMom(std::ostream &out) {
    std::vector<std::string> colnames = {" ", "Energy", "<Q20>", "<Q22>", "<R2>", "<JZ>"};
    for (int i=0; i<colnames.size(); ++i){
      out << std::setw(8) << colnames[i];
    }
    out << std::endl;
    int dim = fEnergies.size();
    for (int ni=0; ni<dim; ++ni) {
      out << std::fixed << std::setw(7) << ni << " ";
      out << std::fixed << std::setprecision(3) << std::setw(7) << fEnergies[ni] << " ";
      out << std::fixed << std::setprecision(3) << std::setw(7) << fQ20(ni) << " ";
      out << std::fixed << std::setprecision(3) << std::setw(7) << fQ22(ni) << " ";
      out << std::fixed << std::setprecision(3) << std::setw(7) << fR2(ni) << " ";
      out << std::fixed << std::setprecision(3) << std::setw(7) << fRJz(ni) << std::endl;
    }
  }

  void Nilsson::WriteEnergies(std::ostream &out) {
    //group by omega, starting from 1/2 and moving up
    out << fEpsP << "    " << fGamma*180.0/PI << "    ";

    //for (int omg2 = 1; omg2 <= GetMaxOmega(); omg2 = omg2 + 2) {
      for (int i=0; i<fEnergies.size(); ++i) {
        //    if (fOmegas[i] == omg2) {
          out << fEnergies[i] << "   ";
          //     }
      }
      //    }
    out << std::endl;
  }

  int Nilsson::GetMaxOmega() {
    return fMaxOmg;
  }
  
  double Nilsson::ClebschGordan(int j1, int j2, int J, int m1, int m2, int M) {
    //arguments are all integer 2J or 2M
    int phase = std::pow(-1, -(double)j1/2. + (double)j2/2. - (double)M/2.);
    double scale = std::sqrt(J + 1);
    return phase * scale * gsl_sf_coupling_3j(j1, j2, J,
                                              m1, m2, -M);
  }

  double Nilsson::RedMED1(int l2, int s2, int j2, int l2p, int s2p, int j2p) {
    //arguments are 2*l, 2*s, 2*j, etc.
    if (l2 + l2p < 4) { return 0; }
    if (std::abs(l2 - l2p) > 4) { return 0; }
    if (j2 + j2p < 2) { return 0; }
    if (std::abs(j2-j2p) > 2) { return 0; }

    //9j symbol
    //look up edmonds 7.1.5, 6.4.17, 6.3.3, and 6.3.4
    //double ninej = gsl_sf_coupling_9j();

    double sixja;
    double sixjb;
    if (j2p == (l2p-1)) {
      double xt = (j2+l2p+5.)*(l2p+3.-j2)/4.;
      double xn = l2p*(l2p+1)*12.;
      sixja = std::pow(-1, (j2 + l2p + 3)/2.)*std::sqrt(xt/xn);
    }
    else if (j2p == (l2p+1)) {
      double xt = (j2+3.-l2p)*(j2+l2p-1.)/4.;
      double xn = (l2p+1.)*(l2p+2)*12.;
      sixja = std::pow(-1, (j2 + l2p + 3)/2.)*std::sqrt(xt/xn);
    }
    else {
      std::cerr << "Invalid arguments for 6j symbol" << std::endl;
      exit(1);
    }

    if (j2 == (l2-1)) {
      double xt = (l2p + l2 + 6.)*(l2 + 4. - l2p)/4.;
      double xn = (l2*(l2+1))*20.0;
      sixjb = std::pow(-1, (l2+l2p)/2.)*std::sqrt(xt/xn);
    }
    else if (j2 == (l2+1)) {
      double xt = (l2p + 4. - l2)*(l2p + l2 - 2.)/4.;
      double xn = (l2+1.)*(l2+2)*20.0;
      sixjb = std::pow(-1, (l2+l2p)/2.)*std::sqrt(xt/xn);
    }
    else {
      std::cerr << "Invalid arguments for 6j symbol" << std::endl;
    }
    double sixjc = std::sqrt(3.)/6.;

    double ninej = -sixja*sixjb/sixjc/3.;
    
    double redmC2 = std::sqrt(l2+1.)*ClebschGordan(l2, 4, l2p, 0, 0, 0);
    return -std::sqrt(10.)/2. * redmC2 * std::sqrt(3./2.) * std::sqrt((j2p+1.)*(j2+1.)*3.) * ninej;    
  }

  double Nilsson::GetRadialME(int Ni, int li, int Nj, int lj, int T) {
    if (T == 2) {
      Ni = 2*Ni;
      li = 2*li;
      Nj = 2*Nj;
      lj = 2*lj;
      //use closed formula
      if (Nj < Ni) {
        if (Nj != Ni-4) { return 0; }
        if (lj < li) {
          if (lj != li - 4) { return 0; }
          return std::sqrt((Ni+li+2.)*(Ni+li-2))/4.;
        }
        else if (li == lj) {
          return std::sqrt((Ni-li)*(Ni+li+2))/4.;
        }
        else {
          if (lj != li + 4) { return 0; }
          return std::sqrt((Ni-li)*(Ni-li-4))/4.;
        }        
      }
      else if (Nj == Ni) {
        if (lj < li) {
          if (lj != li - 4) { return 0; }
          return std::sqrt((Ni-li + 4)*(Ni+li+2))/2.;
        }
        else if (lj == li) {
          return (Ni+3)/2.;
        }
        else {
          if (lj != li + 4) { return 0; }
          return std::sqrt((Ni-li)*(Ni+li+6))/2.;
        }
      }
      else {
        if (Nj != Ni + 4) { return 0; }
        if (lj < li) {
          if (lj != li - 4 ) { return 0; }
          return std::sqrt((Ni-li+4)*(Ni-li+8))/4.;
        }
        else if (lj == li) {
          return std::sqrt((Ni-li+4)*(Ni+li+6))/4.;
        }
        else {
          if (lj != li+4) { return 0; }
          return std::sqrt((Ni+li+6)*(Ni+li+10))/4.;
        }
      }
    }
    else if (T==3) {      
      std::cerr << "Warning! T=3 not implemented!" << std::endl;
      return 0;
    }
    int Np = (Ni - li)/2 + 1;
    int N = (Nj - lj)/2 + 1;    
    
    if (Np < N)  {
      int tmpN, tmpL;
      tmpN = Ni;
      tmpL = li;
      Ni = Nj;
      li = lj;
      Nj = tmpN;
      lj = tmpL;
      Np = (Ni - li)/2 + 1;
      N = (Nj - lj)/2 + 1;
    }

    int C = (lj - li - T)/2;
    int D = (li - lj - T)/2;
    int E = D + Np - N - 1;
    int Nalt = 0;
    
    if (2*C != (lj - li - T)) {
      double me = 0;
      for (int k = 1; k<= N; ++k) {
        int k1 = 2*(N-k+1)+50;
        int k2 = lj - li - T + 2*(k-1)+50;
        int k3 = li - lj - T + 2*(k + Np - N - 1) + 50;
        int k4 = 2*(Np - N + k) + 50;
        int k5 = 2*(N - k + 1) + lj + li + 1 + T + 50;
        int k6 = lj - li - T + 50;
        int k7 = li - lj - T + 50;
        me += Gamma(2*N + 50) * Gamma(2*Np + 50) / Gamma(k1) * Gamma(k2)/Gamma(2*k + 50) *
          Gamma(k3)/Gamma(k4) * Gamma(k5)/( Gamma(k6) * Gamma(k7) );          
      }
      int k1 = 2*N + 2*lj + 51;
      int k2 = 2*Np + 2*li + 51;
      me = me/std::sqrt(Gamma(2*N + 50) * Gamma(k1) * Gamma(2*Np + 50) * Gamma(k2)) * std::pow(-1, Np-N);
      return me;      
    }
    else {
      if (C <= 0 && D <= 0) {
        if (E > -1) { Nalt = 0; }
        else if (C <= E+1) { Nalt = -E; }
        else { Nalt = 1 - C; }
      }
      else if (C <= 0 && D > 0) {
        Nalt = 1 - C; 
      }
      else if (C > 0 && D <= 0) {
        if (E > -1) { Nalt = 0; }
        else { Nalt = -E; }
      }
      else {
        std::cerr << "SEVERE ERROR IN GetRadialME!!!" << std::endl;
        std::cerr << Ni << "   " << li << "   " << Nj << "   " << lj << "    " << T << std::endl;
        std::cerr << C << "   " << D << "   " << E << std::endl;
        exit(1);
      }
    }

    if (Nalt < 0) {
      std::cerr << "SEVERE ERROR IN GetRadialME!!!" << std::endl;
      std::cerr << Ni << "   " << li << "   " << Nj << "   " << lj << "    " << T << std::endl;
      std::cerr << C << "   " << D << "   " << E << std::endl;
      std::cerr << Nalt << std::endl;
      exit(1);
    }

    if (Nalt == 0) { return 0; }

    double me = 0;
    double gli = 0;
    double gla = 0;
    if (Nalt > N) { Nalt = N; }
    for (int k=1; k<=Nalt; ++k) {
      if (C <= 0) {
        int ka = 2 - 2*C + 50;
        int kb = 2*(2-k-C)+50;
        gli = std::pow(-1, k+1)*Gamma(ka)/Gamma(kb);
      }
      else {
        int ka = 2*(C+k-1)+50;
        int kb = 2*C + 50;
        gli = Gamma(ka)/Gamma(kb);
      }
      if (D <= 0) {
        int ka = 2-2*D+50;
        int kb = 2*(2-D-k-Np+N)+50;
        gla = std::pow(-1, 1+k+Np-N)*Gamma(ka)/Gamma(kb);
      }
      else {
        int ka = 2*(D+k+Np-N-1)+50;
        int kb = 2*D+50;
        gla = Gamma(ka)/Gamma(kb);
      }
      int k1 = 2*(N-k+1)+50;
      int k2 = 2*(N-k+1)+lj+li+1+T+50;
      int k3 = 2*(Np-N+k)+50;

      me += Gamma(2*N+50)/Gamma(2*k+50) * Gamma(2*Np+50)/Gamma(k1) * Gamma(k2)/Gamma(k3) * gli * gla;      
    }
    int ka = 2*(N+lj)+51;
    int kb = 2*(Np+li)+51;
    me = me * std::pow(-1, N-Np)/std::sqrt(Gamma(2*N+50) * Gamma(2*Np+50) * Gamma(ka) * Gamma(kb));
    return me;
  }
}        
