#!/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 matplotlib.pyplot import subplots, 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_mdk)
from pylib.numerical.ode_model import disk

PRINT_RESULTS = False
PLOT_RESULTS = False

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])
  show()

p = .01, .002, .015                  # parameter of the model
f = disk(*p)
fnm = disk(*p, method='nm')
fnmmdk = disk(*p, method='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

  def test_e1(self):
    xe1 = e1(f, self.x0, self.t)
    #plotx1rphi(xe1, self.t, ("Euler (Runge-Kutta 1th order),
    #  increments: %d" % self.nt))
    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)
    if PLOT_RESULTS:
      plotx1rphi(xe1_2, self.t_2, ("Euler (Runge-Kutta 1th order), \
                                   increments: %d" % self.nt_2))
    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]
    ))

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

  def test_e2(self):
    xe2 = e2(f, self.x0, self.t)
    if PLOT_RESULTS:
     plotx1rphi(xe2, self.t, ("Runge-Kutta 2th order, \
                              increments: %d" % self.nt))
    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)
    if PLOT_RESULTS:
      plotx1rphi(xe4, self.t, ("Runge-Kutta 4th order, \
                               increments: %d" % self.nt))
    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,
                         max_iterations=self.max_iterations,
                         tol=self.tol)
    if PLOT_RESULTS:
     plotx1rphi(xi1, self.t,
                ("Backward Euler, increments: %d, \
                 iterations: %d" % (self.nt, sum(iterations))))
    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, gamma=self.gamma,
      beta=self.beta, max_iterations=1, tol=self.tol)
    if PRINT_RESULTS:
      print('\nNewmark Newton-Raphson x and xp. Max. Iterations:',
            1)
      print(xnm)
      print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    if PLOT_RESULTS:
     plotx1rphi(x, self.t,
                ("Newmark, increments: %d, \
                 iterations: %d" % (self.nt, sum(iterations))))
    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, gamma=self.gamma,
      beta=self.beta, max_iterations=self.max_iterations,
      tol=self.tol)
    if PRINT_RESULTS:
      print('\nNewmark Newton-Raphson x and xp. Max. Iterations:',
            self.max_iterations)
      print(xnm)
      print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    if PLOT_RESULTS:
      plotx1rphi(x, self.t,
                 ("Newmark, increments: %d, \
                  iterations: %d" % (self.nt, sum(iterations))))
    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_mdk(
      fnmmdk, (0,0,0), (0,0,0), (0,0,0), self.t, gamma=self.gamma,
      beta=self.beta, max_iterations=1, tol=self.tol)
    if PRINT_RESULTS:
      print('\nNewmark Newton-Raphson mdk x and xp. Max. Iterations:',
            1)
      print(xnm)
      print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    if PLOT_RESULTS:
      plotx1rphi(x, self.t,
                 ("Newmark RDK, increments: %d, \
                  iterations: %d" % (self.nt, sum(iterations))))
    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_mdk(
      fnmmdk, (0,0,0), (0,0,0), (0,0,0), self.t, gamma=self.gamma,
      beta=self.beta, max_iterations=100, tol=self.tol)
    if PRINT_RESULTS:
      print('\nNewmark Newton-Raphson mdk x and xp. Max. Iterations:',
            100)
      print(xnm)
      print(xpnm)
    x = concatenate((xnm, xpnm), axis=1)
    if PLOT_RESULTS:
      plotx1rphi(x, self.t,
                 ("Newmark RDK, increments: %d, \
                  iterations: %d" % (self.nt, sum(iterations))))
    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))


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