#!/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('..'))
from pylib.numerical.ode import (e1, e2, e4, i1,
  newmark_newtonraphson, newmark_newtonraphson_rdk)
from pylib.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()