#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Test of numerical solver.

:Date: 2019-05-25

.. module:: test_ode
  :platform: *nix, Windows
  :synopsis: Test of numerical solver.

.. moduleauthor:: Daniel Weschke <daniel.weschke@directbox.de>
"""
from __future__ import division, print_function
import unittest
import math
from numpy import linspace, allclose, sqrt, square, concatenate
from numpy.linalg import inv
from matplotlib.pyplot import figure, subplots, plot, show

import os
import sys
sys.path.insert(0, os.path.abspath('../src'))
from numerical.ode import (e1, e2, e4, i1, newmark_newtonraphson,
  newmark_newtonraphson_rdk)
from numerical.ode_model import disk, disk_nm, disk_nmmdk

def plotx1rphi(x, t, title):
  """Plot plane rotation (x, y, phi)
  """
  x1 = x[:, 0]
  x2 = x[:, 1]
  x3 = x[:, 2]
  xp1 = x[:, 3]
  xp2 = x[:, 4]
  xp3 = x[:, 5]
  fig, ax1 = subplots()
  ax1.set_title(title)
  ax1.plot(t, x1, 'b-')
  #ax1.plot(t, x3, 'b-')
  ax1.set_xlabel('time t in s')
  # Make the y-axis label and tick labels match the line color.
  ax1.set_ylabel('x_1', color='b')
  for tl in ax1.get_yticklabels():
    tl.set_color('b')
  envelope = sqrt(square(x1)+square(x2))
  ax1.plot(t, envelope, 'b--', t, -envelope, 'b--')
  ax2 = ax1.twinx()
  ax2.plot(t, xp3, 'r')
  ax2.set_ylabel('phi', color='r')
  for tl in ax2.get_yticklabels():
    tl.set_color('r')
  if max(xp3) < 2:
    ax2.set_ylim([0,2])

f = disk
fnm = disk_nm
fnmmdk = disk_nmmdk


class TestODE(unittest.TestCase):

  Dt = .01                           # Delta t
  tend = .1                          # t_end
  #tend = 300
  nt = math.ceil(tend/Dt)            # number of timesteps
  t = linspace(0., tend, nt+1)       # time vector
  x0 = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  # initial conditions
  p = .01, .002, .015                # parameter of the model

  def test_e1(self):
    xe1 = e1(f, self.x0, self.t, *self.p)
    #plotx1rphi(xe1, self.t, ("Euler (Runge-Kutta 1th order),
    #  increments: %d" % self.nt))
    #show()
    self.assertTrue(allclose(xe1[:, 5],
      [0., 0.00015, 0.0003, 0.00045, 0.0006,  0.00075, 0.0009,
       0.00105, 0.0012, 0.00135, 0.0015 ]))

  Dt_2 = Dt/5.
  nt_2 = math.ceil(tend/Dt_2)
  t_2 = linspace(0., tend, nt_2+1)

  def test_e1_2(self):
    xe1_2 = e1(f, self.x0, self.t_2, *self.p)
    #plotx1rphi(xe1_2, self.t_2, ("Euler (Runge-Kutta 1th order),
    #  increments: %d" % self.nt_2))
    #show()
    self.assertTrue(allclose(xe1_2[:, 5],
[0.00000000e+00, 3.00000000e-05, 6.00000000e-05, 8.99999998e-05,
 1.19999999e-04, 1.49999998e-04, 1.79999995e-04, 2.09999992e-04,
 2.39999987e-04, 2.69999980e-04, 2.99999971e-04, 3.29999960e-04,
 3.59999947e-04, 3.89999931e-04, 4.19999913e-04, 4.49999891e-04,
 4.79999866e-04, 5.09999837e-04, 5.39999804e-04, 5.69999767e-04,
 5.99999726e-04, 6.29999681e-04, 6.59999630e-04, 6.89999575e-04,
 7.19999514e-04, 7.49999448e-04, 7.79999376e-04, 8.09999298e-04,
 8.39999214e-04, 8.69999123e-04, 8.99999026e-04, 9.29998921e-04,
 9.59998810e-04, 9.89998691e-04, 1.01999856e-03, 1.04999843e-03,
 1.07999829e-03, 1.10999814e-03, 1.13999798e-03, 1.16999781e-03,
 1.19999763e-03, 1.22999744e-03, 1.25999725e-03, 1.28999704e-03,
 1.31999682e-03, 1.34999660e-03, 1.37999636e-03, 1.40999611e-03,
 1.43999585e-03, 1.46999558e-03, 1.49999530e-03]
    ))

  maxIterations = 1000
  tol = 1e-9  # against residuum

  def test_e2(self):
    xe2 = e2(f, self.x0, self.t, *self.p)
    #plotx1rphi(xe2, self.t, ("Runge-Kutta 2th order,
    #  increments: %d" % self.nt))
    #show()
    self.assertTrue(allclose(xe2[:, 5],
      [0., 0.00015, 0.0003, 0.00045, 0.0006,  0.00075, 0.0009,
       0.00105, 0.0012, 0.00135, 0.0015 ]))

  def test_e4(self):
    xe4 = e4(f, self.x0, self.t, *self.p)
    #plotx1rphi(xe4, self.t, ("Runge-Kutta 4th order,
    #  increments: %d" % self.nt))
    #show()
    self.assertTrue(allclose(xe4[:, 5],
      [0., 0.00015, 0.0003, 0.00045, 0.0006,  0.00075, 0.0009,
       0.00105, 0.0012, 0.00135, 0.0015 ]))

  def test_i1(self):
    xi1, iterations = i1(f, self.x0, self.t, *self.p,
      self.maxIterations, self.tol)
    #plotx1rphi(xi1, self.t, ("Backward Euler, increments: %d,
    #  iterations: %d" % (self.nt, sum(iterations))))
    #show()
    self.assertTrue(allclose(xi1[:, 5],
      [0., 0.00015, 0.0003, 0.00045, 0.0006,  0.00075, 0.0009,
       0.00105, 0.0012, 0.00135, 0.0015 ]))

  gamma = .5  # newmark parameter for velocity
  beta = .25  # newmark parameter for displacement

  def test_nm(self):
    xnm, xpnm, xppnm, iterations = newmark_newtonraphson(fnm,
      (0,0,0), (0,0,0), (0,0,0), self.t, *self.p,
      self.gamma, self.beta, 1, self.tol)
    #print(xnm)
    #print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    #plotx1rphi(x, self.t, ("Newmark, increments: %d,
    #  iterations: %d" % (self.nt, sum(iterations))))
    #show()
    self.assertTrue(allclose(x[:, 5],
[0.00000000e+00, 7.49999925e-05, 2.24999951e-04, 3.74999839e-04,
 5.24999621e-04, 6.74999269e-04, 8.24998752e-04, 9.74998039e-04,
 1.12499710e-03, 1.27499591e-03, 1.42499444e-03]
    ))

  def test_nm_2(self):
    xnm, xpnm, xppnm, iterations = newmark_newtonraphson(fnm,
      (0,0,0), (0,0,0), (0,0,0), self.t, *self.p,
      self.gamma, self.beta, self.maxIterations, self.tol)
    #print(xnm)
    #print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    #plotx1rphi(x, self.t, ("Newmark, increments: %d,
    #  iterations: %d" % (self.nt, sum(iterations))))
    #show()
    self.assertTrue(allclose(x[:, 5],
[0.00000000e+00, 7.49999925e-05, 2.24999951e-04, 3.74999839e-04,
 5.24999621e-04, 6.74999269e-04, 8.24998752e-04, 9.74998039e-04,
 1.12499710e-03, 1.27499591e-03, 1.42499444e-03]
    ))

  def test_nmmdk(self):
    xnm, xpnm, xppnm, iterations = newmark_newtonraphson_rdk(
      fnmmdk, (0,0,0), (0,0,0), (0,0,0), self.t, *self.p,
      self.gamma, self.beta, 1, self.tol)
    #print(xnm)
    #print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    #plotx1rphi(x, self.t, ("Newmark, increments: %d,
    #  iterations: %d" % (self.nt, sum(iterations))))
    #show()
    self.assertTrue(allclose(x[:, 5],
[0.00000000e+00, 7.49999963e-05, 2.24999974e-04, 3.74999906e-04,
 5.24999764e-04, 6.74999516e-04, 8.24999134e-04, 9.74998586e-04,
 1.12499784e-03, 1.27499688e-03, 1.42499566e-03]
    ))

  def test_nmmdk_2(self):
    xnm, xpnm, xppnm, iterations = newmark_newtonraphson_rdk(
      fnmmdk, (0,0,0), (0,0,0), (0,0,0), self.t, *self.p,
      self.gamma, self.beta, 100, self.tol)
    #print(xnm)
    #print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    #plotx1rphi(x, self.t, ("Newmark, increments: %d,
    #  iterations: %d" % (self.nt, sum(iterations))))
    #show()
    self.assertTrue(allclose(x[:, 5],
[0.00000000e+00, 7.49999963e-05, 2.24999974e-04, 3.74999906e-04,
 5.24999764e-04, 6.74999516e-04, 8.24999134e-04, 9.74998586e-04,
 1.12499784e-03, 1.27499688e-03, 1.42499566e-03]
    ))

  #from scipy.integrate import odeint
  #xode = odeint(disk, x0, t, p, printmessg=True)
  #plotx1rphi(xode, t, ("ODE, Inkremente: %d" % nt))
  #show()


if __name__ == '__main__':
  unittest.main()