#include <iostream>
#include <cmath>

#include "ivp.h"
#include "vec.h"

// Function definitions
int main();
vec<4> f_2_1( double t, vec<4> x );
vec<4> f_2_16( double t, vec<4> x );
vec<4> f_3_1( double t, vec<4> x );
vec<4> f_4_1( double t, vec<4> x );
double f1( double t );
double f2( double t );

// Function declarations

// These are implementations of the equations found in the
// paper "Coupled spring equations" by
// Temple H. Fay and Sarah Duncan Graham, 
// Int. J. Math. Educ. Sci. Technol., 2003, Vol. 34, No. 1, pp.65–79.
//
// See: https://www.researchgate.net/publication/242300860_Coupled_spring_equations

// Function definitions

// Equations (2.1)
vec<4> f_2_1( double t, vec<4> x ) {
    double k1{1.0}; // First string coefficient
    double k2{1.0}; // Second string coefficient

    return vec<4>{
         x[1],
        -k1*x[0] - k2*(x[0] - x[2]),
         x[3],
        -k2*(x[2] - x[0])
    };
}

// Equations (2.16)
//  - adding damping
vec<4> f_2_16( double t, vec<4> x ) {
    double k1{1.0}; // First string coefficient
    double k2{1.0}; // Second string coefficient
    double d1{1.0}; // First damping coefficient
    double d2{1.0}; // Second damping coefficient

    return vec<4>{
         x[1],
        -d1*x[1] - k1*x[0] - k2*(x[0] - x[2]),
         x[3],
        -d2*x[3] - k2*(x[2] - x[0])
    };
}

// Equations (3.1)
//  - adding non-linearlity
vec<4> f_3_1( double t, vec<4> x ) {
    double k1{1.0}; // First string linear coefficient
    double k2{1.0}; // Second string linear coefficient
    double mu1{1.0}; // First string non-linear coefficient (cubic)
    double mu2{1.0}; // Second string non-linear coefficient (cubic)
    double d1{1.0}; // First damping coefficient
    double d2{1.0}; // Second damping coefficient

    double delta = x[0] - x[2];

    return vec<4>{
         x[1],
        -d1*x[1] - k1*x[0]  + mu1*x[0]*x[0]*x[0]
	         - k2*delta + mu2*delta*delta*delta,
         x[3],
        -d2*x[3] + k2*delta - mu2*delta*delta*delta
    };
}

// The forcing functions:

double amp1{0.1};
double omega1{0.1};

double amp2{0.2};
double omega2{0.2};

double f1( double t ) {
    return amp1*std::cos( omega1*t );
}

double f2( double t ) {
    return amp2*std::cos( omega2*t );
}

// Equations (4.1)
//  - adding a forcing function
vec<4> f_4_1( double t, vec<4> x ) {
    double k1{1.0}; // First string linear coefficient
    double k2{1.0}; // Second string linear coefficient
    double mu1{1.0}; // First string non-linear coefficient (cubic)
    double mu2{1.0}; // Second string non-linear coefficient (cubic)
    double d1{1.0}; // First damping coefficient
    double d2{1.0}; // Second damping coefficient

    double delta = x[0] - x[2];

    return vec<4>{
         x[1],
        -d1*x[1] - k1*x[0]  + mu1*x[0]*x[0]*x[0]
	         - k2*delta + mu2*delta*delta*delta + f1( t ),
         x[3],
        -d2*x[3] + k2*delta - mu2*delta*delta*delta + f2( t )
    };
}

int main() {
    std::cout.precision( 6 );
    std::clog.precision( 6 );

    ivp<vec<4>> x{ f_4_1, 0.0, vec<4>{-1.0, 0.0, -2.0, 0.0},
	           0.1, std::make_pair( 1e-5, 0.2 ), 1e-5, vec<4>::norm };

    double step{0.01};
    double t_max{10.0};

    for ( double t{0.0}; t < t_max + step/2.0; t += step ) {
       std::cout << "x1(" << t << ") = " << x( t )[0] << std::endl;
    }

    for ( double t{0.0}; t < t_max + step/2.0; t += step ) {
       std::cout << "x2(" << t << ") = " << x( t )[2] << std::endl;
    }

    return 0;
}