//
// Program: osc_many.cxx
//
// Description: Solve the the armonic oscillator equation using various methods:
//              - Euler;
//              - Euler Cromer;
//              - Euler Richardson;
//              - Mid Point;
//              - Verlet.
//
// Author: A. Garfagnini
//
// Date: 9 October 2012
//

#include <iostream>
#include <fstream>
#include <unistd.h>
#include <libgen.h>
#include <stdlib.h>

enum algorithm {EULER, EULER_CROMER, MID_POINT, EULER_RICHARDSON, VERLET};

struct oscillator {
    int    curr_step;
    double x;
    double v;
    double v_nhalf;
    double t;
    double energy;
};

using namespace std;

void usage(const char * pname)
{
    cerr << "Usage: " << pname << " [scheme] [-o file_name]" << endl;
    cerr << "\nSolve the simple armonic oscillator using any of the "
         << "following methods:" << endl;
    cerr << " -e (default) : Euler\n";
    cerr << " -c           : Euler-Cromer\n";
    cerr << " -m           : Mid Point\n";
    cerr << " -r           : Euler-Richardson\n";
    cerr << " -v           : Verlet\n";
    cerr << "\nOther options:\n";
    cerr << " -o file_name : save output in a user definable file name\n";
}

void print_scheme(ostream & out, algorithm scheme)
{
    switch (scheme)
    {
        case EULER :
            out << "# Scheme: EULER" << endl;
            break;

        case EULER_CROMER :
            out << "# Scheme: EULER CROMER" << endl;
            break;

        case MID_POINT :
            out << "# Scheme: MID POINT" << endl;
            break;

        case EULER_RICHARDSON :
            out << "# Scheme: EULER RICHARDSON" << endl;
            break;

        case VERLET :
            out << "# Scheme: VERLET" << endl;
            break;
    }
}

void print_step(ostream & out, oscillator & z, bool print_header=false)
{
    if (print_header) {
        out << "# Time  x  v  Energy" << endl;
    }
    out << z.t << " " << z.x << " " << z.v << " " << z.energy << endl;
}

void step_Euler(oscillator &z, double h)
{
    double v_temp = z.v;
    z.v -=  h * z.x;
    z.x +=  h * v_temp;
    z.t +=  h;

    z.energy = z.x * z.x + z.v * z.v;
}

void step_EulerCromer(oscillator &z, double h)
{
    z.v -=  h * z.x;
    z.x +=  h * z.v;
    z.t +=  h;

    z.energy = z.x * z.x + z.v * z.v;
}

void step_MidPoint(oscillator &z, double h)
{
    double v_temp = z.v;
    z.v -=  h * z.x;
    z.x +=  h * (z.v + v_temp)/2.0;
    z.t +=  h;

    z.energy = z.x * z.x + z.v * z.v;
}

void step_EulerRichardson(oscillator &z, double h)
{
    if (z.curr_step == 0) {
        z.v_nhalf = z.v - h * z.x / 2.0;
    } else {
        z.v_nhalf -= h * z.x;
    }

    z.x += h * z.v_nhalf;
    z.v -= h * z.x;
    z.t += h;
    z.curr_step ++;

    z.energy = z.x * z.x + z.v * z.v;
}

// - Main heart of the program: evolve the system integrating the equation
// - using the selected scheme
void evolve(void method(oscillator &, double), oscillator & z,
            double t_f, double h, ostream & out)
{
    while (z.t < t_f) {

        method(z, h);
        print_step(out, z);

    }
}

int main(int argc, char * argv[])
{
   int opt;
   ofstream out_file;
   algorithm scheme = EULER;
   

    // Handle command line options
    while((opt = getopt(argc, argv, "o:ecmrvh?")) != -1) {

        switch (opt) {

        case 'o':
            out_file.open(optarg);
            break;

        case 'e':
            scheme = EULER;
            break;

        case 'c':
            scheme = EULER_CROMER;
            break;

        case 'm':
            scheme = MID_POINT;
            break;

        case 'r':
            scheme = EULER_RICHARDSON;
            break;

        case 'v':
            scheme = VERLET;
            break;

        case 'h':
        case '?':
        default:
            usage( basename( argv[0] ) );
            return 1;
        }
    }

    cout << "Time step h: ";
    double h;
    cin >> h;
    if (out_file.is_open()) {
        out_file << "# Time step h: " << h << endl;
    }

    cout << "T_final: ";
    double t_f;
    cin >> t_f;
    if (out_file.is_open()) {
        out_file << "# T_final: " << t_f << endl;
    }

    ostream & myout = (out_file.is_open()) ? out_file : cout;
    print_scheme(myout, scheme);

    // - Init the system: set the starting conditions
    oscillator z;
    z.curr_step = 0;
    z.x = 1.0;
    z.v = 0.0;
    z.v_nhalf = 0.0;
    z.t = 0.0;
    z.energy = z.x * z.x + z.v * z.v;

    print_step(myout, z, true);

    switch(scheme) {

        case EULER:
            evolve(step_Euler, z, t_f, h, myout);
            break;

        case EULER_CROMER:
            evolve(step_EulerCromer, z, t_f, h, myout);
            break;

        case MID_POINT:
            evolve(step_MidPoint, z, t_f, h, myout);
            break;

        case EULER_RICHARDSON:
            evolve(step_EulerRichardson, z, t_f, h, myout);
            break;
    }

    if (out_file.is_open()) {
        out_file.close();
    }

    return 0;
}