#include "Wiggle.h"
#include "earth.h"
#include "ode-steppers.h"
#include "options.h"
#include <iostream>

using namespace blitz;

Wiggle::Wiggle()
{
}

Wiggle::Wiggle(NuoscParam& np, WiggleControl& wc)
{
    this->SetOptions(np,wc);
}

Wiggle::~Wiggle()
{
}

bool Wiggle::BuildLookups(void)
{

//    const double step_precision = 1.0e-14;
    const double step_precision = 1.0e-15;
//    const double step_precision = 1.0e-16;
    const int stepper_to_use = 3;

    ComplexMatrix u(Range(1,3),Range(1,3));
    bool antineutrino = false;
    if (nuosc_params.type < 0) {   // anti-nu
        u = mixing_matrix(nuosc_params.theta_12(),
                          nuosc_params.theta_23(),
                          nuosc_params.theta_13(),
                          -1.0*nuosc_params.cpphase);
        antineutrino = true;
    }
    else
        u = mixing_matrix(nuosc_params.theta_12(),
                          nuosc_params.theta_23(),
                          nuosc_params.theta_13(),
                          nuosc_params.cpphase);
    
//    cerr << "Using U= " << u;

    double x0[20], xf[20];
    int nregions = earth_get_slant_distances(x0,xf,
                                             nuosc_params.baseline);

    // Move integration bounds away from boundary slightly
//    for (int i=0; i<nregions; ++i) {
//        double red = 0.00001*(xf[i]-x0[i]);
//        x0[i] += red;
//        xf[i] -= red;
//    }

    double nominal_step_size[20];

    for (int n=0; n < nregions; ++n) {
//        nominal_step_size[n] = fabs(xf[n]-x0[n])*step_precision;
        nominal_step_size[n] = fabs(xf[n]-x0[n])*1.0e-4;
//        cerr << n << ": " << nominal_step_size[n] << endl;
    }

//    cerr << "Passing through " << nregions << " earth regions\n";


    ComplexVector initial_neutrino(Range(1,3));
    initial_neutrino = 0;
    int mag = abs(nuosc_params.type);
    int sign = nuosc_params.type/mag;
    initial_neutrino(mag) = sign;
        

//    cerr << "initial_neutrino: " << initial_neutrino << endl;

    for (double e = control.min_energy; e <= control.max_energy; e += control.energy_step) {
        double energy;
        if (control.use_log_energy) energy = pow(10.0,e);
        else energy = e;
        if (energy <= 0.0) continue;
        NuTransODE ode(u,
                       nuosc_params.baseline,
                       nuosc_params.dm2sol,
                       nuosc_params.dm2atm,
                       energy,
                       nuosc_params.matter,
                       antineutrino);
        switch (stepper_to_use) {
        case 1:
            ode.SetStepper(euler_stepper);
            break;
        case 2:
            ode.SetStepper(runge_kutta_4th_stepper);
            break;
        case 3: default:
            ode.SetStepper(runge_kutta_adaptive_stepper);
            break;
        }

        ComplexVector nu(Range(1,3));

        nu = ode.Solve(x0,xf,nominal_step_size,nregions,initial_neutrino, step_precision);

        E.push_back(energy);
        Pnue.push_back(abs(nu(1)*complex_conjugate(nu(1))));
        Pnumu.push_back(abs(nu(2)*complex_conjugate(nu(2))));
        Pnutau.push_back(abs(nu(3)*complex_conjugate(nu(3))));


        cerr << E.back() << " " 
             << Pnue.back() << " "
             << Pnumu.back() << " "
             << Pnutau.back() << endl;
    }

    return true;
}


WiggleControl::WiggleControl()
{
    this->init();
}
WiggleControl::WiggleControl(int& argc, const char** argv)
{
    this->init();
    this->parse_args(argc,argv);
}

WiggleControl::~WiggleControl()
{
    if (options) delete options;
}

Options* WiggleControl::parse_args(int& argc, const char** argv)
{
    const char * optv[] = {
        "I+lin <linear energy (eV): min max [step]>",
        "O+log <log10 energy (exp): min max [step], def=\"8 11 0.001\">",
        0
    };
    options = new Options(*argv, optv);
    OptArgvIter optitr(argc-1,argv+1);
    const char* optarg;

    vector<const char*> newargv;
    newargv.push_back(argv[0]);

    char optchar;
    options->ctrls(Options::PARSE_POS);
    while ((optchar = (*options)(optitr,optarg))) {
        switch (optchar) {
        case 'I': case 'O':
            if (optchar == 'I') use_log_energy = false;
            else use_log_energy = true;
            if (optarg) min_energy = atof(optarg);
            else usage("Bad min_energy value");
            optarg = optitr();
            if (optarg) max_energy = atof(optarg);
            else usage("Bad max_energy value");
            optarg = optitr();
            if (optarg) energy_step = atof(optarg);
            else energy_step = (use_log_energy ? 0.001 : 100e6);
            break;
        case Options::POSITIONAL:
            newargv.push_back(optarg);
            break;
        default:
            newargv.push_back(argv[optitr.index()]);
            break;
        }
    }

    argc =  newargv.size();
    for (int i = 0; i < argc; ++i) {
        argv[i] = newargv[i];
    }
    return options;
}

void WiggleControl::usage(const char* message)
{
    if (!options) return;
    options->usage(cerr,message);
}
void WiggleControl::init()
{
    options = 0;
    use_log_energy = false;
    min_energy = 0;
    max_energy = 10e9;
    energy_step = 100e6;
}