#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Numerical solver of ordinary differential equations.

Solves the initial value problem for systems of first order ordinary differential
equations.

:Date: 2015-09-21

.. module:: solver
  :platform: *nix, Windows
  :synopsis: Numerical solver.

.. moduleauthor:: Daniel Weschke <daniel.weschke@directbox.de>
"""

from __future__ import division, print_function
from numpy import array, isnan, sum, zeros, dot
from numpy.linalg import norm, inv

def e1(f, x0, t, *p, verbose=False):
  r"""Explicit first-order method /
  (standard, or forward) Euler method /
  Runge-Kutta 1st order method.

  de:
  Euler'sche Polygonzugverfahren / explizite Euler-Verfahren /
  Euler-Cauchy-Verfahren / Euler-vorwärts-Verfahren

  :param f: the function to solve
  :type f: function
  :param x0: initial condition
  :type x0: list
  :param t: time
  :type t: list
  :param `*p`: parameters of the function (thickness, diameter, ...)
  :param verbose: print information (default = False)
  :type verbose: bool

  Approximate the solution of the initial value problem

  .. math ::
    \dot{x} &= f(t,x) \\
    x(t_0) &= x_0

  Choose a value h for the size of every step and set

  .. math ::
    t_i = t_0 + i h ~,\quad i=1,2,\ldots,n

  One step :math:`h` of the Euler method from :math:`t_i` to  :math:`t_{i+1}` is

  .. math ::
    x_{i+1} &= x_i + (t_{i+1}-t_i) f(t_i, x_i) \\
    x_{i+1} &= x_i + h f(t_i, x_i) \\

  Example 1:

  .. math ::
    m\ddot{u} + d\dot{u} + ku = f(t) \\
    \ddot{u} = m^{-1}(f(t) - d\dot{u} - ku) \\

  with

  .. math ::
    x_1 &= u        &\quad \dot{x}_1 = \dot{u} = x_2 \\
    x_2 &= \dot{u}  &\quad \dot{x}_2 = \ddot{u} \\

  becomes

  .. math ::
    \dot{x}_1 &= x_2 \\
    \dot{x}_2 &= m^{-1}(f(t) - d x_2 - k x_1) \\

  or

  .. math ::
    \dot{x} &= f(t,x) \\
    \begin{bmatrix} \dot{x}_1 \\ \dot{x}_1 \end{bmatrix} &=
    \begin{bmatrix} x_2 \\ m^{-1}(f(t) - d x_2 - k x_1) \end{bmatrix}  \\
    &=
    \begin{bmatrix} 0 \\ m^{-1} f(t) \end{bmatrix} +
    \begin{bmatrix} 0 & 1 \\ -m^{-1} k & -m^{-1} d \end{bmatrix}
    \begin{bmatrix} x_1 \\ x_2 \end{bmatrix}

  Example 2:

  .. math ::
    m(u)\ddot{u} + d(u,\dot{u})\dot{u} + k(u)u = f(t) \\
    \ddot{u} = m^{-1}(u)(f(t) - d(u,\dot{u})\dot{u} - k(u)u) \\

  with

  .. math ::
    x_1 &= u        &\quad \dot{x}_1 = \dot{u} = x_2 \\
    x_2 &= \dot{u}  &\quad \dot{x}_2 = \ddot{u} \\

  becomes

  .. math ::
    \dot{x}_1 &= x_2 \\
    \dot{x}_2 &= m^{-1}(x_1)(f(t) - d(x_1,x_2) x_2 - k(x_1) x_1) \\

  or

  .. math ::
    \dot{x} &= f(t,x) \\
    \begin{bmatrix} \dot{x}_1 \\ \dot{x}_2 \end{bmatrix} &=
    \begin{bmatrix} x_2 \\ m^{-1}(x_1)(f(t) - d(x_1,x_2) x_2 - k(x_1) x_1) \end{bmatrix}  \\
    &=
    \begin{bmatrix} 0 \\ m^{-1}(x_1) f(t) \end{bmatrix} +
    \begin{bmatrix} 0 & 1 \\ -m^{-1}(x_1) k(x_1) & -m^{-1} d(x_1,x_2) \end{bmatrix}
    \begin{bmatrix} x_1 \\ x_2 \end{bmatrix}

  The Euler method is a first-order method,
  which means that the local error (error per step) is proportional to the
  square of the step size, and the global error (error at a given time) is
  proportional to the step size.
  """
  x = zeros((len(t), len(x0)))  # Preallocate array
  x[0,:] = x0            # Initial condition gives solution at t=0.
  for i in range(len(t)-1):
    Dt = t[i+1]-t[i]
    dxdt = array(f(x[i,:], t[i], *p))
    x[i+1,:] = x[i,:] + dxdt*Dt # Approximate solution at next value of x
  if verbose:
    print('Numerical integration of ODE using explicit first-order method (Euler / Runge-Kutta) was successful.')
  return x

def e2(f, x0, t, *p, verbose=False):
  r"""Explicit second-order method / Runge-Kutta 2nd order method.

  :param f: the function to solve
  :type f: function
  :param x0: initial condition
  :type x0: list
  :param t: time
  :type t: list
  :param `*p`: parameters of the function (thickness, diameter, ...)
  :param verbose: print information (default = False)
  :type verbose: bool
  """
  x = zeros((len(t), len(x0))) 
  x[0,:] = x0                            # initial condition
  for i in range(len(t)-1):                              # calculation loop
    Dt = t[i+1]-t[i]
    k_1 = array(f(x[i,:], t[i], *p))
    k_2 = array(f(x[i,:]+0.5*Dt*k_1, t[i]+0.5*Dt, *p))
    x[i+1,:] = x[i,:] + k_2*Dt  # main equation
  if verbose:
    print('Numerical integration of ODE using explicit 2th-order method (Runge-Kutta) was successful.')
  return x

def e4(f, x0, t, *p, verbose=False):
  r"""Explicit fourth-order method / Runge-Kutta 4th order method.

  :param f: the function to solve
  :type f: function
  :param x0: initial condition
  :type x0: list
  :param t: time
  :type t: list
  :param `*p`: parameters of the function (thickness, diameter, ...)
  :param verbose: print information (default = False)
  :type verbose: bool
  """
  x = zeros((len(t), len(x0)))  # Preallocate array
  x[0,:] = x0            # Initial condition gives solution at t=0.
  for i in range(len(t)-1):                              # calculation loop
    Dt = t[i+1]-t[i]
    k_1 = array(f(x[i,:], t[i], *p))
    k_2 = array(f(x[i,:]+0.5*Dt*k_1, t[i]+0.5*Dt, *p))
    k_3 = array(f(x[i,:]+0.5*Dt*k_2, t[i]+0.5*Dt, *p))
    k_4 = array(f(x[i,:]+k_3*Dt, t[i]+Dt, *p))
    x[i+1,:] = x[i,:] + 1./6*(k_1+2*k_2+2*k_3+k_4)*Dt  # main equation
  if verbose:
    print('Numerical integration of ODE using explicit 4th-order method (Runge-Kutta) was successful.')
  return x

def i1(f, x0, t, *p, max_iterations=1000, tol=1e-9, verbose=False):
  r"""Implicite first-order method / backward Euler method.

  :param f: the function to solve
  :type f: function
  :param x0: initial condition
  :type x0: list
  :param t: time
  :type t: list
  :param `*p`: parameters of the function (thickness, diameter, ...)
  :param max_iterations: maximum number of iterations
  :type max_iterations: int
  :param tol: tolerance against residuum (default = 1e-9)
  :type tol: float
  :param verbose: print information (default = False)
  :type verbose: bool
  
  The backward Euler method has order one and is A-stable.
  """
  iterations = zeros((len(t), 1))
  x = zeros((len(t), len(x0)))  # Preallocate array
  x[0,:] = x0            # Initial condition gives solution at t=0.
  for i in range(len(t)-1):
    Dt = t[i+1]-t[i]
    x1 = x[i,:]
    for j in range(max_iterations):        # fixed-point iteration
      dxdt = array(f(x1, t[i+1], *p))
      x11 = x[i,:] + dxdt*Dt  # Approximate solution at next value of x
      residuum = norm(x11-x1)/norm(x11)
      x1 = x11
      if residuum < tol:
        break
    iterations[i] = j+1
    x[i+1,:] = x1
  if verbose:
    print('Numerical integration of ODE using implicite first-order method (Euler) was successful.')
  return x, iterations

def newmark_newtonraphson(f, x0, xp0, xpp0, t, *p, gamma=.5, beta=.25, max_iterations=1000, tol=1e-9, verbose=False):
  r"""Newmark method.

  :param f: the function to solve
  :type f: function
  :param x0: initial condition
  :type x0: list
  :param xp0: initial condition
  :type xp0: list
  :param xpp0: initial condition
  :type xpp0: list
  :param t: time
  :type t: list
  :param `*p`: parameters of the function (thickness, diameter, ...)
  :param gamma: newmark parameter for velocity (default = 0.5)
  :type gamma: float
  :param beta: newmark parameter for displacement (default = 0.25)
  :type beta: float
  :param max_iterations: maximum number of iterations
  :type max_iterations: int
  :param tol: tolerance against residuum (default = 1e-9)
  :type tol: float
  :param verbose: print information (default = False)
  :type verbose: bool
  """
  iterations = zeros((len(t), 1))
  x = zeros((len(t), len(x0)))      # Preallocate array
  xp = zeros((len(t), len(xp0)))    # Preallocate array
  xpp = zeros((len(t), len(xpp0)))  # Preallocate array
  x[0,:] = x0               # Initial condition gives solution at t=0.
  xp[0,:] = xp0             # Initial condition gives solution at t=0.
  xpp[0,:] = xpp0           # Initial condition gives solution at t=0.
  for i in range(len(t)-1):
    Dt = t[i+1]-t[i]
        
    xi = x[i,:].reshape(3,1)
    xpi = xp[i,:].reshape(3,1)
    xppi = xpp[i,:].reshape(3,1)
    x1 = xi
    xp1 = xpi
    xpp1 = xppi
    j = 0
    for j in range(max_iterations):        # fixed-point iteration
      #dxdt = array(f(t[i+1], x1, p))
      #x11 = x[i,:] + dxdt*Dt  # Approximate solution at next value of x
            
      N, dN, dNp, dNpp = f(x1.reshape(-1,).tolist(),
        xp1.reshape(-1,).tolist(), xpp1.reshape(-1,).tolist(),
        t[i], *p)
      if isnan(sum(dN)) or isnan(sum(dNp)) or isnan(sum(dNpp)):
        print('divergiert')
        break

      xpp11 = xpp1 - dot(inv(dNpp), (N + dot(dN, (x1-xi)) + dot(dNp, (xp1-xpi))))
      xp1 = xpi + Dt*( (1-gamma)*xppi + gamma*xpp11 )
      x1 = xi + Dt*xpi + Dt**2*( (.5-beta)*xppi + beta*xpp11 )
            
      residuum = norm(xpp11-xpp1)/norm(xpp11)
      xpp1 = xpp11
      if residuum < tol:
        break
    iterations[i] = j+1
        
    xpp[i+1,:] = xpp1.reshape(-1,).tolist()
    xp[i+1,:] = xp1.reshape(-1,).tolist()
    x[i+1,:] = x1.reshape(-1,).tolist()
  if verbose:
    print('Numerical integration of ODE using explicite newmark method was successful.')
  return x, xp, xpp, iterations
  # x = concatenate((x, xp, xpp), axis=1)

def newmark_newtonraphson_rdk(fnm, x0, xp0, xpp0, t, *p, gamma=.5, beta=.25, maxIterations=1000, tol=1e-9, verbose=False):
  r"""Newmark method.

  :param f: the function to solve
  :type f: function
  :param x0: initial condition
  :type x0: list
  :param xp0: initial condition
  :type xp0: list
  :param xpp0: initial condition
  :type xpp0: list
  :param t: time
  :type t: list
  :param `*p`: parameters of the function (thickness, diameter, ...)
  :param gamma: newmark parameter for velocity (default = 0.5)
  :type gamma: float
  :param beta: newmark parameter for displacement (default = 0.25)
  :type beta: float
  :param max_iterations: maximum number of iterations
  :type max_iterations: int
  :param tol: tolerance against residuum (default = 1e-9)
  :type tol: float
  :param verbose: print information (default = False)
  :type verbose: bool
  """
  iterations = zeros((len(t), 1))
  x = zeros((len(t), len(x0)))      # Preallocate array
  xp = zeros((len(t), len(xp0)))    # Preallocate array
  xpp = zeros((len(t), len(xpp0)))  # Preallocate array
  x[0,:] = x0               # Initial condition gives solution at t=0.
  xp[0,:] = xp0             # Initial condition gives solution at t=0.
  xpp[0,:] = xpp0           # Initial condition gives solution at t=0.
  for i in range(len(t)-1):
    Dt = t[i+1]-t[i]

    rm, rmx, rmxpp, rd, rdx, rdxp, rk, rkx, f = fnm(x[i,:], xp[i,:], xpp[i,:], t[i], *p)
        
    xi = x[i,:].reshape(3,1)
    xpi = xp[i,:].reshape(3,1)
    xppi = xpp[i,:].reshape(3,1)
    x1 = xi
    xp1 = xpi
    xpp1 = xppi
    j = 0
    for j in range(maxIterations):        # fixed-point iteration
      #dxdt = array(f(t[i+1], x1, p))
      #x11 = x[i,:] + dxdt*Dt  # Approximate solution at next value of x

      r = (rmx+rdx+rkx)*Dt**2./4 + rdxp*Dt/2 + rmxpp 
      rp = f - (rm + dot(rmx, (Dt*xpi+Dt**2./4*xppi)) - dot(rmxpp, xppi) + \
                rd + dot(rdx, (Dt*xpi+Dt**2./4*xppi)) + dot(rdxp, Dt/2*xppi) + \
                rk + dot(rkx, (Dt*xpi+Dt**2./4*xppi)) )
      xpp11 = dot(inv(r), rp)
      xp1 = xpi + Dt*( (1-gamma)*xppi + gamma*xpp11 )
      x1 = xi + Dt*xpi + Dt**2*( (.5-beta)*xppi + beta*xpp11 )

      residuum = norm(xpp11-xpp1)/norm(xpp11)
      xpp1 = xpp11
      if residuum < tol:
        break
    iterations[i] = j+1

    xpp[i+1,:] = xpp1.reshape(-1,).tolist()
    xp[i+1,:] = xp1.reshape(-1,).tolist()
    x[i+1,:] = x1.reshape(-1,).tolist()
  if verbose:
    print('Numerical integration of ODE using explicite newmark method was successful.')
  return x, xp, xpp, iterations
  # x = concatenate((x, xp, xpp), axis=1)