/**
 * Quick program to solve some differential equations
 *
 *  Copyright 2013 by David Zaslavsky <diazona@ellipsix.net>
 *
 * This is free software: you can redistribute  it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This file is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * To view the license, see <http://www.gnu.org/licenses/>.
 *
 * @license GPL-3.0+ <http://spdx.org/licenses/GPL-3.0+>
 */

// compile with
//  gcc -std=c99 -o dfeqsolver dfeqsolver.c -lgsl -lconfuse

#include <assert.h>
#include <stdio.h>
#include <confuse.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_odeiv2.h>

static const size_t dim = 4;
#define E_INWATER 56
#define E_OUTOFWATER 57
#define E_SINK 58

struct parameters {
    double density;
    double width;
    double radius;
    double mass; // of bike
    double coefficient; // -2*density*width/(3*mass*radius^2)
    double g;
    double weight; // m*g
};

void report_cfg_error(cfg_t* cfg, const char* fmt, va_list ap) {
    fprintf(stderr, fmt, ap);
}

void read_options(struct parameters* p, double* y, char* filename, int argc, char** argv) {
    cfg_opt_t opts[] = {
        CFG_SIMPLE_FLOAT("density", &(p->density)),
        CFG_SIMPLE_FLOAT("width", &(p->width)),
        CFG_SIMPLE_FLOAT("radius", &(p->radius)),
        CFG_SIMPLE_FLOAT("mass", &(p->mass)),
        CFG_SIMPLE_FLOAT("g", &(p->g)),
        CFG_SIMPLE_FLOAT("depth", &(y[0])),
        CFG_SIMPLE_FLOAT("fall", &(y[1])),
        CFG_SIMPLE_FLOAT("speed", &(y[3])),
        CFG_END()
    };
    p->density = 1000; // kg/m^3
    p->width = 0.1; // m
    p->radius = 0.5; // m
    p->mass = 140; // kg
    p->g = 9.80665; // m/s^2
    cfg_t* cfg = cfg_init(opts, 0);
    cfg_set_error_function(cfg, report_cfg_error);
    cfg_parse(cfg, filename);
    for (size_t i = 1; i < argc; i++) {
        cfg_parse_buf(cfg, argv[i]);
    }
    cfg_free(cfg);
    p->coefficient = -2 * p->density * p->width / (3 * p->mass * gsl_pow_2(p->radius));
    p->weight = p->mass * p->g;
}

int bounce_equations_implementation(double t, const double y[], double dydt[], void* params) {
    double c = ((struct parameters*)params)->coefficient;
    double R = ((struct parameters*)params)->radius;
    double g = ((struct parameters*)params)->g;
    double  h = y[0], hp = y[1],
            x = y[2], xp = y[3]; // p for prime, i.e. time derivative
    if (h < 0) {
        return E_OUTOFWATER;
    }
    else if (h > R) {
        return E_SINK;
    }
    assert(!gsl_isnan(h));
    assert(!gsl_isnan(hp));
    assert(!gsl_isnan(x));
    assert(!gsl_isnan(xp));
    // dh/dt = h'
    dydt[0] = hp;
    // dh'/dt = -2Dwx'^2/3mR^2 (h(2R-h))^(3/2) + g
    dydt[1] = c * gsl_pow_2(xp) * sqrt(gsl_pow_3(h * (2 * R - h))) + g;
    // dx/dt = x'
    dydt[2] = xp;
    // dx'/dt = -2Dwx'^2/3mR^2 h^2(3R-h)
    dydt[3] = c * gsl_pow_2(xp) * gsl_pow_2(h) * (3 * R - h);
    return GSL_SUCCESS;
}

int bounce_jacobian_implementation(double t, const double y[], double* dfdy, double dfdt[], void* params) {
    double c = ((struct parameters*)params)->coefficient;
    double R = ((struct parameters*)params)->radius;
    double  h = y[0], hp = y[1],
            x = y[2], xp = y[3]; // p for prime, i.e. time derivative
    if (h < 0) {
        return E_OUTOFWATER;
    }
    else if (h > R) {
        return E_SINK;
    }
    assert(!gsl_isnan(h));
    assert(!gsl_isnan(hp));
    assert(!gsl_isnan(x));
    assert(!gsl_isnan(xp));
    // explicit time derivatives are all zero
    dfdt[0] = dfdt[1] = dfdt[2] = dfdt[3] = 0;
    // Jacobian elements
    dfdy[0 * dim + 0] = 0;
    dfdy[0 * dim + 1] = 1;
    dfdy[0 * dim + 2] = 0;
    dfdy[0 * dim + 3] = 0;
    dfdy[1 * dim + 0] = c * gsl_pow_2(xp) * 3 * (R - h) * sqrt(h * (2 * R - h));
    dfdy[1 * dim + 1] = 0;
    dfdy[1 * dim + 2] = 0;
    dfdy[1 * dim + 3] = c * 2 * xp * sqrt(gsl_pow_3(h * (2 * R - h)));
    dfdy[2 * dim + 0] = 0;
    dfdy[2 * dim + 1] = 0;
    dfdy[2 * dim + 2] = 0;
    dfdy[2 * dim + 3] = 1;
    dfdy[3 * dim + 0] = c * gsl_pow_2(xp) * 3 * h * (2 * R - h);
    dfdy[3 * dim + 1] = 0;
    dfdy[3 * dim + 2] = 0;
    dfdy[3 * dim + 3] = c * 2 * xp * gsl_pow_2(h) * (3 * R - h);
    return GSL_SUCCESS;
}

int projectile_equations_implementation(double t, const double y[], double dydt[], void* params) {
    double g = ((struct parameters*)params)->g;
    double  h = y[0], hp = y[1],
            x = y[2], xp = y[3]; // p for prime, i.e. time derivative
    assert(!gsl_isnan(h));
    assert(!gsl_isnan(hp));
    assert(!gsl_isnan(x));
    assert(!gsl_isnan(xp));
    // dh/dt = h'
    dydt[0] = hp;
    // dh'/dt = g
    dydt[1] = g;
    // dx/dt = x'
    dydt[2] = xp;
    // dx'/dt = 0
    dydt[3] = 0;
    return GSL_SUCCESS;
}

int projectile_jacobian_implementation(double t, const double y[], double* dfdy, double dfdt[], void* params) {
    double  h = y[0], hp = y[1],
            x = y[2], xp = y[3]; // p for prime, i.e. time derivative
    assert(!gsl_isnan(h));
    assert(!gsl_isnan(hp));
    assert(!gsl_isnan(x));
    assert(!gsl_isnan(xp));
    // explicit time derivatives are all zero
    dfdt[0] = dfdt[1] = dfdt[2] = dfdt[3] = 0;
    // Jacobian elements
    dfdy[0 * dim + 0] = 0;
    dfdy[0 * dim + 1] = 1;
    dfdy[0 * dim + 2] = 0;
    dfdy[0 * dim + 3] = 0;
    dfdy[1 * dim + 0] = 0;
    dfdy[1 * dim + 1] = 0;
    dfdy[1 * dim + 2] = 0;
    dfdy[1 * dim + 3] = 0;
    dfdy[2 * dim + 0] = 0;
    dfdy[2 * dim + 1] = 0;
    dfdy[2 * dim + 2] = 0;
    dfdy[2 * dim + 3] = 1;
    dfdy[3 * dim + 0] = 0;
    dfdy[3 * dim + 1] = 0;
    dfdy[3 * dim + 2] = 0;
    dfdy[3 * dim + 3] = 0;
    return GSL_SUCCESS;
}

static const size_t maxiter = 10000;
static const size_t maxphase = 3;

int solve_bounce(struct parameters* p, double y[], double* t, double dt) {
    int status;
    
    gsl_odeiv2_system sys;
    sys.function = bounce_equations_implementation;
    sys.jacobian = bounce_jacobian_implementation;
    sys.dimension = dim;
    sys.params = p;
    
    gsl_odeiv2_driver* drv = gsl_odeiv2_driver_alloc_y_new(&sys, gsl_odeiv2_step_rk8pd, 1e-6, 1e-12, 1e-3);
    
    size_t i;
    for (i = 0; i < maxiter; i++) {
        status = gsl_odeiv2_driver_apply(drv, t, *t + dt, y);
        printf("%.5g %.5g %.5g %.5g %.5g\n", *t, -y[0], -y[1], y[2], y[3]);
        switch (status) {
            case GSL_SUCCESS:
                break;
            case E_OUTOFWATER:
                status = 0;
                goto endloop; // height is back to zero
            case E_SINK:
                goto endloop;
            default:
                fprintf(stderr, "error! return value %d\n", status);
                goto endloop;
        }
    }
    if (i == maxiter && status == 0) {
        status = 1;
    }
    endloop:
    gsl_odeiv2_driver_free(drv);
    
    return status;
}

int solve_projectile(struct parameters* p, double y[], double *t, double dt) {
    int status;
    
    gsl_odeiv2_system sys;
    sys.function = projectile_equations_implementation;
    sys.jacobian = projectile_jacobian_implementation;
    sys.dimension = dim;
    sys.params = p;
    
    gsl_odeiv2_driver* drv = gsl_odeiv2_driver_alloc_y_new(&sys, gsl_odeiv2_step_rk8pd, 1e-6, 1e-12, 1e-3);
    
    size_t i;
    for (i = 0; i < maxiter; i++) {
        status = gsl_odeiv2_driver_apply(drv, t, *t + dt, y);
        printf("%.5g %.5g %.5g %.5g %.5g\n", *t, -y[0], -y[1], y[2], y[3]);
        switch (status) {
            case GSL_SUCCESS:
                if (i > 2 && y[0] > 0) {
                    goto endloop;
                }
                break;
            default:
                fprintf(stderr, "error! return value %d\n", status);
                goto endloop;
        }
    }
    if (i == maxiter && status == 0) {
        status = 1;
    }
    endloop:
    gsl_odeiv2_driver_free(drv);
    
    return status;
}
    
int main(int argc, char** argv) {
    static const double dt = 0.001;
    
    struct parameters p;
    double y[4] = {0, 0.3, 0, 26.8};
    read_options(&p, y, "motorcycle.cfg", argc, argv);
    
    double t = 0;
    
    int status = 0;

#ifndef FIRSTBOUNCEONLY
    for (size_t i = 0; i < maxphase; i++) {
#endif
        status = solve_bounce(&p, y, &t, dt);
        fprintf(stderr, "finished bounce\n");
        switch (status) {
            case E_SINK:
            case 0:
                break;
            default:
                return 1;
        }
#ifndef FIRSTBOUNCEONLY  
        status = solve_projectile(&p, y, &t, dt);
        fprintf(stderr, "finished projectile\n");
        switch (status) {
            case 0:
                break;
            default:
                return 1;
        }
    }
#endif
    
    return 0;
}