#include <iomanip>

#include "Pairing.hh"

namespace PRM {
  Pairing::Pairing(Nilsson &nilpos, Nilsson &nilneg) :
    fNilPos(nilpos),
    fNilNeg(nilneg),
    fNOdd(0),
    fTolerance(1e-6),
    fDelta(12./std::sqrt(nilpos.fA)),
    fDelta2(0),
    fEferm(0),
    fG0(0),
    fG2(0),
    fMaxIter(100),
    fEfac(0),
    fImin(0),
    fImax(120)
  {
    int posctr = 0;
    int negctr = 0;
    fImax = fNilPos.fEnergies.size()-1;
    while (true) {
      double posen = 99999;
      double negen = 99999;
      if (posctr < fNilPos.fEnergies.size()) {
        posen = fNilPos.fEnergies[posctr];
      }

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

      if (posen == 99999 && negen == 99999) {
        break;
      }
         
      if (posen < negen) {
        fQ.push_back(fNilPos.fQ20[posctr]);
        fEnergies.push_back(posen);
        ++posctr;
      }
      else {
        fQ.push_back(fNilNeg.fQ20[negctr]);
        fEnergies.push_back(negen);
        ++negctr;
      }
    }

    if (fNilPos.fOdd == 1) { //odd neutron
      fEfac = (1. + (fNilPos.fA - 2.*fNilPos.fZ)/fNilPos.fA/3.);
      fEferm = fEnergies[(int)fNilPos.fN/2];
      fNOdd = fNilPos.fA - fNilPos.fZ;
    }
    else { //odd proton
      fEfac = (1. - (fNilPos.fA - 2.*fNilPos.fZ)/fNilPos.fA/3.);
      fEferm = fEnergies[(int)fNilPos.fZ/2];
      fNOdd =  fNilPos.fZ;
    }
    fEfac *= 41.0/(std::pow(fNilPos.fA,1./3.)) * fNilPos.fOmOm;
    fEferm *= fEfac;

    for (int i=0; i<fEnergies.size(); ++i) {
      fEnergies[i] *= fEfac;
    }
  };

  void Pairing::SetParams(double G0, double G2) {
    fG0 = G0;
    fG2 = G2;
  }

  void Pairing::SetGNParams(double GN0, double GN1) {
    fGN0 = GN0;
    fGN1 = GN1;
    fG2 = 0.0;
    if (fNilPos.fOdd == 1) { //odd neutron
      fG0 = (GN0 + GN1*(fNilPos.fA-2.0*fNilPos.fZ)/fNilPos.fA)/(double)fNilPos.fA;
    }
    else { //odd proton
      fG0 = (GN0 - GN1*(fNilPos.fA-2.0*fNilPos.fZ)/fNilPos.fA)/(double)fNilPos.fA;
    }
  }

  void Pairing::SetIPair(int ipair) {
    fImin = (int)fNOdd/2 - (int)(std::sqrt((double)(ipair * fNOdd)))-1;
    fImax = (int)fNOdd/2 + (int)(std::sqrt((double)(ipair * fNOdd)))-1;
    if (fImin < 0) { fImin = 0; }
  }
  void Pairing::Solve(int ipair) {
    SetIPair(ipair);
    Solve();
  }

  void Pairing::Solve(int imin, int imax) {
    fImin = imin;
    fImax = imax;
    Solve();
  }
  
  void Pairing::Solve() {
    double ene = fNOdd - 2*fImin;

    if (std::abs(fG0) < 1e-10) { fG0 = 1e-10; }
    if (std::abs(fG2) < 1e-10) { fG2 = 1e-10; }

    for (int iter = 0; iter<fMaxIter; ++iter) {
      double y0(0);
      double yy0(0);
      double y2(0);

      double ak1(0);
      double ak1t(0);

      double a0(0);
      double a1(0);
      double a2(0);

      double b0(0);
      double b1(0);
      double b2(0);
      double b3(0);
      double b4(0);

      double a0t(0);
      double a1t(0);
      double a2t(0);

      double b0t(0);
      double b2t(0);
    
      for (int i=fImin; i<=fImax; ++i) {
        double delta = fDelta + fQ[i] * fDelta2; //monopole + quadrupole delta
        double sign = std::abs(delta)/delta;

        double Q = fQ[i];
        double Q2 = Q*Q;

        double DeltaEi = fEnergies[i] - fEferm;
        double root = std::sqrt(std::pow(DeltaEi,2) + std::pow(delta,2));
        double root3 = sign*std::pow(std::pow(DeltaEi,2) + std::pow(delta,2), 3./2.);

        y0 += 1./(sign*root);
        yy0 += DeltaEi/(sign*root);
        y2 += Q2/(sign*root);
      
        ak1 += Q/(sign*root);
        ak1t += Q/(root);
        a0 += DeltaEi/root3;
        a1 += DeltaEi*Q/root3;
        a2 += DeltaEi*Q2/root3;

        b0 += 1./root3;
        b1 += Q/root3;
        b2 += Q2/root3;
        b3 += std::pow(Q,3)/root3;
        b4 += std::pow(Q,4)/(root3*sign);

        /*
          a0t += DeltaEi/(root3*sign);
          a1t += DeltaEi*Q/(root3*sign);
          a2t += DeltaEi*Q2/(root3*sign);

          b0t += 1./(root3*sign);
          b2t += 1.*Q2/(root3*sign);
        */
      }

      double ak0 = 2./fG0 - y0;
      double ak2 = 2./fG2 - y2;

      arma::mat A(3,3,arma::fill::zeros);
      arma::vec b(3, arma::fill::zeros);
      arma::vec x(3, arma::fill::zeros);

      b(0) = ak0*ak2 - ak1 * ak1;
      b(1) = fDelta * ak0 - fDelta2 * ak1;
      b(2) = yy0 + ene - (fImax-fImin+1);

      A(0,0) = 2.*a1*ak1 + ak2*a0 + ak0*a2;
      A(0,1) = -2*ak1*(b1*fDelta+fDelta2*b2) - ak2*(fDelta*b0+fDelta2*b1)
        - ak0*(fDelta*b2 + fDelta2*b3);
      A(0,2) = -2.*ak1*(b2*fDelta+fDelta2*b3) - ak2*(fDelta*b1+fDelta2*b2)
        - ak0*(fDelta*b3 + fDelta2*b4);

      A(1,0) = fDelta*a0 + fDelta2*a1;
      A(1,1) = -ak0 - b0*fDelta*fDelta - 2.*b1*fDelta*fDelta2 - b2*fDelta2*fDelta2;
      A(1,2) =  ak1 - b1*fDelta*fDelta - 2.*b2*fDelta*fDelta2 - b3*fDelta2*fDelta2;

      A(2,0) = b0*fDelta*fDelta + 2.*b1*fDelta*fDelta2 + b2*fDelta2*fDelta2;
      A(2,1) = a0*fDelta + a1*fDelta2;
      A(2,2) = a1*fDelta + a2*fDelta2;

      //std::cout << A << std::endl;
      //std::cout << b << std::endl;

      if (iter < 3 ) {
        double norm = b0*b0*std::pow(fDelta,3) + a0*a0*fDelta;
        x(0) = (b(2)*fDelta*b0 + ak0*fDelta*a0)/norm;
        x(1) = -(b0*fDelta*fDelta*ak0 - a0*b(2))/norm;
        x(2) = 0.0;
        if (iter == 2) {
          fDelta2 = 0;
          for (int i=fImin; i<=fImax; ++i) {
            double DeltaEi = fEnergies[i] - fEferm - x(0);
            double root = std::sqrt(std::pow(DeltaEi, 2) + std::pow(fDelta + x(1), 2));
            double V = 0.5 * ( 1 - DeltaEi/root );
            fDelta2 += fG2 * fQ[i] * std::sqrt((1.-V)*V);
          }
        }
      }
      else {
        //we now solve Ax = b
        arma::mat invA = arma::inv(A);
        x = invA*b;
      }
      //this is our increment to delta and Eferm
      fEferm += x(0);
      fDelta += x(1);
      fDelta2 += x(2);
      if (fG2 < 1e-9) { fDelta2 = 0; }
      if (std::abs(x(0)) >= 2 ) {
        std::cerr << "Error in finding BCS solution" << std::endl;
        exit(1);
      }
      if (iter < 3) { continue; }

      double conv = std::max(std::max(std::abs(x(0)), std::abs(x(1))),
                             std::abs(x(2)));
      if (conv < fTolerance) { break; }                                   
    }
    return;
  }

  void Pairing::PrintLevels(std::ostream &out, int ncols/*=4*/) {
    std::vector<int> offsets;
    for (int col=0; col<ncols; ++col) {
      offsets.push_back(col*fEnergies.size()/ncols);
    }
    
    std::vector<std::string> colnames = {"#", "Energy", "Q20"};    
    for (int col=0; col<ncols; ++col) {
      for (int i=0; i<colnames.size(); ++i){
        out << std::setw(8) << colnames[i];
      }
      out << std::setw(8) << " ";
    }
    out << std::endl;

    for (int i=0; i<fEnergies.size()/ncols; ++i) {
      out << " ";
      for (int col=0; col<ncols; ++col) {
        out << std::fixed << std::setw(7) << i+offsets[col] << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fEnergies[i+offsets[col]] << " ";
        out << std::fixed << std::setprecision(3) << std::setw(7) << fQ[i+offsets[col]] << " ";
        out << std::setw(8) << " ";
      }
      out << std::endl;
    }    
  }

  void Pairing::PrintSummary(std::ostream &out) {
    out << "Pairing factors: GN0 = " << fGN0 << ", GN1 = " << fGN1 << std::endl;
    out << "                 G0 = " << fG0 << ",  " << "G2 = " << fG2 << std::endl;
    out << "   Levels from " << fImin << " to " << fImax << " considered" << std::endl;
    out << "   Delta = " << fDelta << ", lambda = " << fEferm << ", hw0 = " << fEfac << " MeV, lambda/hw0 = " << fEferm/fEfac << std::endl;
    out << std::endl;
  }
    
}
