from __future__ import division

import matplotlib.pyplot as plt
import numpy as np


def evolve(state, t_step, gamma, w_squared, Ampl, Omega):
    '''Evolution function for the problem. Takes state, produces the evalution of the system at this point.'''
    position_prime = state[2];
    speed_prime = - 2 * gamma * state[2] - w_squared**2 * state[1] - Ampl * np.cos(Omega * state[0]);
    return np.array([1, position_prime, speed_prime]);


def rk_coefs(state, t_step, gamma, w_squared, Ampl, Omega):
    '''RK-4 coeficents. Takes state, returns increment of the state.'''
    coef = evolve(state, t_step, gamma, w_squared, Ampl, Omega) * t_step;

    temp_state = state + coef / 2;
    temp_evolution = evolve(temp_state, t_step, gamma, w_squared, Ampl, Omega) * t_step;
    coef = coef + 2 * temp_evolution;

    temp_state = state + temp_evolution / 2;
    temp_evolution = evolve(temp_state, t_step, gamma, w_squared, Ampl, Omega) * t_step;
    coef = coef + 2 * temp_evolution;

    temp_state = state + temp_evolution;
    temp_evolution = evolve(temp_state, t_step, gamma, w_squared, Ampl, Omega) * t_step;
    coef = coef + temp_evolution;

    coef = coef/ 6;
    return coef;


def main():
    Time = np.arange(100);
    X = np.zeros(100);
    state = np.array([0,0,1]);
    for time in Time:
        X[time] = state[1];
        temp = rk_coefs(state, 0.1, 0, 1, 0, 0)
        state = state + temp;
    plt.plot(Time, X)
    plt.plot(Time, np.sin(Time*0.1))
    plt.show()


if __name__ == '__main__':
    main();
