#include "NuTransODE.h"
#include "nuosc.h"
#include "earth.h"
#include "ode-steppers.h"

using namespace blitz;

static const complex<double> EYE(0,1);

NuTransODE::NuTransODE(ComplexMatrix u, double baseline,
                       double dm2_sol, double dm2_atm, double energy,
                       bool use_matter_effects, bool antineutrino)
    : OdeFunc(Range(1,3))
      ,_Uvac(GetRange(), GetRange())
      ,_baseline(baseline)
      ,_use_matter_effects(use_matter_effects)
      ,_antineutrino(antineutrino)
{
    this->SetStepper(runge_kutta_adaptive_stepper);

    ComplexMatrix msqrd = mass_squared_matrix(dm2_sol, dm2_atm);
    ComplexMatrix Udagger = hermitian_conjugate(u);

    // -i/2E (U)(M^2)(U^dagger)
    _Uvac = -EYE/(2.0*energy) * 
        matrix_product(u,matrix_product(msqrd,Udagger));
}

ComplexVector NuTransODE::operator()(double x, ComplexVector y)
{
    // hbar*c in eV*cm
    const double hbarc = 1.9732696e-05;

    ComplexMatrix u(this->GetRange(),this->GetRange());
    u = _Uvac;

    if (_use_matter_effects) {
        // sqrt(2) GF Ne in eV, when rho=g/cm^3
        double matter = 7.6e-14*earth_density(x,_baseline); 
        matter *= earth_electron_fraction(x,_baseline);
        if (_antineutrino) matter *= -1.0;
        u(1,1) += -EYE*matter;
    }

    using namespace blitz::tensor; // for i,j
    ComplexVector v(GetRange());
    v = sum(u(i,j)*y(j), j)/hbarc;
    return v;
}