///////////////////////////////////////////////////////////////////////////////
//
// Implementation of oscillations of neutrinos in matter in a
// three-neutrino framework with NSI.
//
// This  class inherits from the PMNS_Fast class
//
// jcoelho\@apc.in2p3.fr
///////////////////////////////////////////////////////////////////////////////

#include <iostream>

#include "PMNS_NSI.h"

using namespace OscProb;

using namespace std;

//.............................................................................
///
/// Constructor. \sa PMNS_Base::PMNS_Base
///
/// This class is restricted to 3 neutrino flavours.
///
PMNS_NSI::PMNS_NSI() : PMNS_Fast(), fEps()
{
  SetStdPath();
  SetNSI(0., 0., 0., 0., 0., 0., 0., 0., 0.);
  SetFermCoup(1, 0, 0);
}

//.............................................................................
///
/// Nothing to clean.
///
PMNS_NSI::~PMNS_NSI() {}

//.............................................................................
///
/// Set all NSI parameters at once.
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param eps_ee      - The real parameter eps_ee
/// @param eps_emu     - The absolute value of the complex parameter eps_emu
/// @param eps_etau    - The absolute value of the complex parameter eps_etau
/// @param eps_mumu    - The real parameter eps_mumu
/// @param eps_mutau   - The absolute value of the complex parameter eps_mutau
/// @param eps_tautau  - The real parameter eps_tautau
/// @param delta_emu   - The phase of the complex parameter eps_emu in radians
/// @param delta_etau  - The phase of the complex parameter eps_etau in radians
/// @param delta_mutau - The phase of the complex parameter eps_mutau in radians
///
void PMNS_NSI::SetNSI(double eps_ee, double eps_emu, double eps_etau,
                      double eps_mumu, double eps_mutau, double eps_tautau,
                      double delta_emu, double delta_etau, double delta_mutau)
{
  SetEps(0, 0, eps_ee, 0);
  SetEps(1, 1, eps_mumu, 0);
  SetEps(2, 2, eps_tautau, 0);

  SetEps(0, 1, eps_emu, delta_emu);
  SetEps(0, 2, eps_etau, delta_etau);
  SetEps(1, 2, eps_mutau, delta_mutau);
}

//.............................................................................
///
/// Set any given NSI parameter.
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// Flavours are:\n
/// - 0 = nue, 1 = numu, 2 = nutau
///
/// Requires that flvi < flvj. Will notify you if input is wrong.
/// If flvi > flvj, will assume reverse order and swap flvi and flvj.
///
/// @param flvi  - The first flavour index
/// @param flvj  - The second flavour index
/// @param val   - The absolute value of the parameter
/// @param phase - The complex phase of the parameter in radians
///
void PMNS_NSI::SetEps(int flvi, int flvj, double val, double phase)
{
  if (flvi > flvj) {
    cerr << "WARNING: First argument should be smaller or equal to second "
            "argument"
         << endl
         << "Setting reverse order (Eps_" << flvj << flvi << "). " << endl;
    int temp = flvi;
    flvi     = flvj;
    flvj     = temp;
  }
  if (flvi < 0 || flvi > 2 || flvj < flvi || flvj > 2) {
    cerr << "WARNING: Eps_" << flvi << flvj << " not valid for " << fNumNus
         << " neutrinos. Doing nothing." << endl;
    return;
  }

  complexD h = val;

  if (flvi != flvj) h *= complexD(cos(phase), sin(phase));

  bool isSame = (fEps[flvi][flvj] == h);

  if (!isSame) ClearCache();

  fGotES *= isSame;

  fEps[flvi][flvj] = h;
}

//.............................................................................
///
/// Get any given NSI parameter.
///
/// Flavours are:\n
/// - 0 = nue, 1 = numu, 2 = nutau
///
/// Requires that flvi < flvj. Will notify you if input is wrong.
/// If flvi > flvj, will assume reverse order and swap flvi and flvj.
///
/// @param flvi  - The first flavour index
/// @param flvj  - The second flavour index
///
complexD PMNS_NSI::GetEps(int flvi, int flvj)
{
  if (flvi > flvj) {
    cerr << "WARNING: First argument should be smaller or equal to second "
            "argument"
         << endl
         << "Setting reverse order (Eps_" << flvj << flvi << "). " << endl;
    int temp = flvi;
    flvi     = flvj;
    flvj     = temp;
  }
  if (flvi < 0 || flvi > 2 || flvj < flvi || flvj > 2) {
    cerr << "WARNING: Eps_" << flvi << flvj << " not valid for " << fNumNus
         << " neutrinos. Returning 0." << endl;
    return zero;
  }

  return fEps[flvi][flvj];
}

//.............................................................................
///
/// Set eps_ee parameter
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param a - The value of the real parameter eps_ee
///
void PMNS_NSI::SetEps_ee(double a) { SetEps(0, 0, a, 0); }

//.............................................................................
///
/// Set eps_mumu parameter
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param a - The value of the real parameter eps_mumu
///
void PMNS_NSI::SetEps_mumu(double a) { SetEps(1, 1, a, 0); }

//.............................................................................
///
/// Set eps_tautau parameter
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param a - The value of the real parameter eps_tautau
///
void PMNS_NSI::SetEps_tautau(double a) { SetEps(2, 2, a, 0); }

//.............................................................................
///
/// Set eps_emu parameter
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param a   - The absolute value of the parameter eps_emu
/// @param phi - The phase of the complex parameter eps_emu in radians
///
void PMNS_NSI::SetEps_emu(double a, double phi) { SetEps(0, 1, a, phi); }

//.............................................................................
///
/// Set eps_etau parameter
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param a   - The absolute value of the parameter eps_etau
/// @param phi - The phase of the complex parameter eps_etau in radians
///
void PMNS_NSI::SetEps_etau(double a, double phi) { SetEps(0, 2, a, phi); }

//.............................................................................
///
/// Set eps_mutau parameter
///
/// This will check if value is changing to keep track of whether
/// the eigensystem needs to be recomputed.
///
/// @param a   - The absolute value of the parameter eps_mutau
/// @param phi - The phase of the complex parameter eps_mutau in radians
///
void PMNS_NSI::SetEps_mutau(double a, double phi) { SetEps(1, 2, a, phi); }

//.............................................................................
///
/// Build the full Hamiltonian in matter
///
/// Here we divide the mass squared matrix Hms by the 2E
/// to obtain the vacuum Hamiltonian in eV. Then, the matter
/// potential is added to each flavour pair.
///
void PMNS_NSI::UpdateHam()
{
  double lv = 2 * kGeV2eV * fEnergy; // 2*E in eV

  double kr2GNe   = kK2 * M_SQRT2 * kGf * fPath.density;
  double kr2GNnsi = kr2GNe;

  kr2GNe *= fPath.zoa;      // Std matter potential in eV
  kr2GNnsi *= GetZoACoup(); // NSI matter potential in eV

  // Finish build Hamiltonian in matter with dimension of eV
  for (int i = 0; i < fNumNus; i++) {
    for (int j = i; j < fNumNus; j++) {
      if (!fIsNuBar)
        fHam[i][j] = fHms[i][j] / lv + kr2GNnsi * fEps[i][j];
      else
        fHam[i][j] = conj(fHms[i][j] / lv - kr2GNnsi * fEps[i][j]);
    }
  }

  if (!fIsNuBar)
    fHam[0][0] += kr2GNe;
  else
    fHam[0][0] -= kr2GNe;
}

//.............................................................................
///
/// Set the NSI relative coupling to a given fermion.
/// This factor represents what fraction of the fermion are seen
/// by the interaction.
///
/// @param c - fermion coupling strength
/// @param i - fermion index (0,1,2) -> (e,u,d)
///
void PMNS_NSI::SetCoupByIndex(double c, int i)
{
  bool isSame = (fNSIcoup[i] == c);

  if (!isSame) ClearCache();

  fGotES *= isSame;

  fNSIcoup[i] = c;
}

//.............................................................................
///
/// Set the NSI relative coupling to electrons.
/// This factor represents what fraction of the electrons are seen
/// by the interaction.
///
/// @param e - electron coupling strength
///
void PMNS_NSI::SetElecCoup(double e) { SetCoupByIndex(e, 0); }

//.............................................................................
///
/// Set the NSI relative coupling to u-quarks.
/// This factor represents what fraction of the u-quarks are seen
/// by the interaction.
///
/// @param u - u-quark coupling strength
///
void PMNS_NSI::SetUpCoup(double u) { SetCoupByIndex(u, 1); }

//.............................................................................
///
/// Set the NSI relative coupling to d-quarks.
/// This factor represents what fraction of the d-quarks are seen
/// by the interaction.
///
/// @param d - d-quark coupling strength
///
void PMNS_NSI::SetDownCoup(double d) { SetCoupByIndex(d, 2); }

//.............................................................................
///
/// Set the NSI relative couplings to each fermion.
/// These factors represent what fraction of each fermion are seen
/// by the interaction.
///
/// @param e - electron coupling strength
/// @param u - u-quark coupling strength
/// @param d - d-quark coupling strength
///
void PMNS_NSI::SetFermCoup(double e, double u, double d)
{
  SetElecCoup(e); // electron coupling
  SetUpCoup(u);   // u-quark coupling
  SetDownCoup(d); // d-quark coupling
}

//.............................................................................
///
/// Get the NSI relative coupling to electrons.
///
double PMNS_NSI::GetElecCoup() { return fNSIcoup[0]; }

//.............................................................................
///
/// Get the NSI relative coupling to u-quarks.
///
double PMNS_NSI::GetUpCoup() { return fNSIcoup[1]; }

//.............................................................................
///
/// Get the NSI relative coupling to d-quarks.
///
double PMNS_NSI::GetDownCoup() { return fNSIcoup[2]; }

//.............................................................................
///
/// Get the effective NSI Z/A dependence.
///
double PMNS_NSI::GetZoACoup()
{
  return fNSIcoup[0] * fPath.zoa          // electrons: Z
         + fNSIcoup[1] * (1 + fPath.zoa)  // u-quarks:  A + Z
         + fNSIcoup[2] * (2 - fPath.zoa); // d-quarks: 2A - Z
}

///////////////////////////////////////////////////////////////////////////////