from __future__ import print_function
from __future__ import division
from builtins import range
from past.utils import old_div
from proteus import Comm, Profiling
import numpy as np
import numpy.testing as npt
import unittest
import random
from math import cos,sin,cosh,sinh,pi,tanh,log,atan2,asin,acos,radians
import sys,os
import logging
import pytest

import cython

comm = Comm.init()
Profiling.procID = comm.rank()
def getpath():
    path = sys.path[0]+'/'
    return path

def remove_files(filenames):
    ''' delete files in filenames list '''
    for f in filenames:
        if os.path.exists(f):
            try:
                os.remove(f)
            except OSError as e:
                p
rint ("Error: %s - %s" %(e.filename,e.strerror))
        else:
            pass

Profiling.logEvent("Testing Bodydynamics")

class TestAuxFunctions(unittest.TestCase):
    def testForwardEler(self):
        from proteus.mprans.BodyDynamics import forward_euler
     
   p0 = 5.0
        v0 = 2.0
        a  = 0.5
        dt = 0.01
        v = v0 + a*dt
        p = p0 + v*dt
        p1, v1 = forward_euler(p0, v0, a, dt)
        self.assertTrue(p==p1)
        self.assertTrue(v==v1)

    def testRungeKutta(self):
        from proteus.mprans.BodyDynamics import runge_kutta
        u0 = 5.0
        v0 = 2.0
        dt = 0.01
        substeps = 20
        F = 150.0
        K = 100000.0
        C = 1000.0
        m = 50.0
        a0 = old_div((F - C*v0 - K*u0), m)
        velCheck = True
        # Loop through substeps
        for i in range(substeps):
            # 1st step
            u1 = u0
            v1 = v0
            a1 = old_div((F - C*v1 - K*u1), m)
            # 2nd step
            u2 = u0 + v1*dt/2.
            v2 = v0 + a1*dt/2.
            a2 = old_div((F - C*v2 - K*u2), m)
            # 3rd step
            u3 = u0 + v2*dt/2.
            v3 = v0 + a2*dt/2.
            a3 = old_div((F - C*v3 - K*u3), m)
            # 4th step
            u4 = u0 + v3*dt
            v4 = v0 + a3*dt
            a4 = old_div((F - C*v4 - K*u4), m)
            # Final step
            u = u0 + dt/6.*(v1+2*v2+2*v3+v4)
            v = v0 + dt/6.*(a1+2*a2+2*a3+a4)
            a = old_div((F - C*v - K*u), m)
            # Check for friction module, when velocity is reversed
            if velCheck:
                if v0*v < 0.0:
                    break
        uc, vc, ac = runge_kutta(u0, v0, a0, dt, substeps, F, K, C, m, velCheck)
        self.assertTrue(u==uc)
        self.assertTrue(v==vc)
        self.assertTrue(a==ac)

    def testEulerAngles(self):
        from proteus.mprans.BodyDynamics import getEulerAngles
        coords = np.array([ [1.,0.,0.],
                            [0.,1.,0.],
                            [0.,0.,1.] ])     
        anglex = radians(5.0)
        angley = radians(10.)
        anglez = radians(15.)
        # rotation around x-axis 
        rotx = np.array([  [  1.0,   0.0,          0.0 ],
                           [  0.0,   cos(anglex),  sin(anglex) ],
                           [  0.0,  -sin(anglex),  cos(anglex) ] ])       
        # rotation around y-axis 
        roty = np.array([  [  cos(angley),  0.0,  -sin(angley) ],
                           [  0.0,          1.0,   0.0 ],
                           [  sin(angley),  0.0,   cos(angley) ] ])
        # rotation around z-axis 
        rotz = np.array([  [  cos(anglez),  sin(anglez),  0.0 ],
                           [ -sin(anglez),  cos(anglez),  0.0 ],
                           [  0.0,         0.0,           1.0 ] ])
        # total 3D rotation
        rot3D = np.dot(np.dot(rotx,roty),rotz)
        coordinates = np.dot( rot3D, coords) 
        ax = atan2(coordinates[1,2], coordinates[2,2])
        ay = asin(-coordinates[0,2])
        az = atan2(coordinates[0,1], coordinates[0,0])
        angles = [ax, ay, az]
        rotation = getEulerAngles(coordinates)
        self.assertTrue(angles==rotation)

class TestRigidBody(unittest.TestCase):

    def testCalculateInit(self):
        from proteus import Domain  
        from proteus.mprans import SpatialTools as st
        from proteus.mprans import BodyDynamics as bd   
        domain = Domain.PlanarStraightLineGraphDomain()
        pos = np.array([1.0, 1.0, 0.0])
        rot = np.array([ [1.,0.,0.], [0.,1.,0.], [0.,0.,1.] ])  
        dim=(0.3, 0.385)
        caisson = st.Rectangle(domain, dim=dim, coords=[pos[0], pos[1]])
        caisson2D = bd.RigidBody(shape=caisson)
        st.assembleDomain(domain)      
        c2d = caisson2D
        nd = 2
        angle = bd.getEulerAngles(rot)

        c2d.calculate_init()

        npt.assert_equal(c2d.position, pos)
        npt.assert_equal(c2d.last_position, pos)
        npt.assert_equal(c2d.rotation, rot)
        npt.assert_equal(c2d.last_rotation, rot)
        npt.assert_equal(c2d.rotation_euler, angle)
        npt.assert_equal(c2d.last_rotation_euler, angle)

    def testStoreLastValues(self):
        from proteus import Domain  
        from proteus.mprans import SpatialTools as st
        from proteus.mprans import BodyDynamics as bd   
        domain = Domain.PlanarStraightLineGraphDomain()
        pos = np.array([1.0, 1.0, 0.0])
        rot = np.array([ [1.,0.,0.], [0.,1.,0.], [0.,0.,1.] ])  
        dim=(0.3, 0.385)
        caisson = st.Rectangle(domain, dim=dim, coords=[pos[0], pos[1]])
        caisson2D = bd.RigidBody(shape=caisson)
        st.assembleDomain(domain)      
        c2d = caisson2D
        nd = 2
        anglez = radians(5.0)
        rotz = np.array([  [  cos(anglez),  sin(anglez),  0.0 ],
                           [ -sin(anglez),  cos(anglez),  0.0 ],
                           [  0.0,         0.0,           1.0 ] ])

        c2d.calculate_init()

        # Imposing values on rigid body's parameters
        c2d.position = np.array([2.0, 3.0, 0.0])
        c2d.velocity = np.array([0.5, -0.3, 0.0])
        c2d.acceleration = np.array([0.1, 0.2, 0.0])
        c2d.rotation = rotz
        c2d.rotation_euler =  bd.getEulerAngles(rotz)
        c2d.ang_disp = np.array([0.0, 0.0, 0.001])
        c2d.ang_vel = np.array([0.0, 0.0, -0.003])
        c2d.ang_acc = np.array([0.0, 0.0, 0.005])
        c2d.F = np.array([100.0, -300.0, 0.0])
        c2d.M = np.array([0.0, 0.0, 50.0])
        c2d.pivot = c2d.barycenter
        c2d.ux = 0.5
        c2d.uy = 0.05
        c2d.uz = 0.0

        c2d._store_last_values()
        
        npt.assert_equal(c2d.position, c2d.last_position)
        npt.assert_equal(c2d.velocity, c2d.last_velocity)
        npt.assert_equal(c2d.acceleration, c2d.last_acceleration)
        npt.assert_equal(c2d.rotation, c2d.last_rotation)
        npt.assert_equal(c2d.rotation_euler, c2d.last_rotation_euler)
        npt.assert_equal(c2d.ang_disp, c2d.last_ang_disp)
        npt.assert_equal(c2d.ang_vel, c2d.last_ang_vel)
        npt.assert_equal(c2d.ang_acc, c2d.last_ang_acc)
        npt.assert_equal(c2d.F, c2d.last_F)
        npt.assert_equal(c2d.M, c2d.last_M)
        npt.assert_equal(c2d.pivot, c2d.last_pivot)
        npt.assert_equal(c2d.ux, c2d.last_ux)
        npt.assert_equal(c2d.uy, c2d.last_uy)
        npt.assert_equal(c2d.uz, c2d.last_uz)

    def testGetInertia(self):
        from proteus import Domain  
        from proteus import SpatialTools as st    
        from proteus.mprans import SpatialTools as mst
        from proteus.mprans import BodyDynamics as bd   

        # 2d case
        domain2D = Domain.PlanarStraightLineGraphDomain()
        pos = np.array([1.0, 1.0, 0.0])
        dim = (0.3, 0.385)
        caisson = mst.Rectangle(domain=domain2D, dim=dim, coords=[pos[0], pos[1]])
        caisson2D = bd.RigidBody(shape=caisson)
        mst.assembleDomain(domain2D)      
        c2d = caisson2D
        c2d.nd = 2
        Ix = (old_div((dim[0]**2),12.))
        Iy = (old_div((dim[1]**2),12.))
        It = Ix + Iy 
        mass = 50.0
        I1 = It*mass
        c2d.mass = mass
        c2d.It = It
        I2 = c2d.getInertia()
        npt.assert_equal(I1, I2)

        # 3d case
        pos, dim, caisson = [], [], []
        domain3D = Domain.PiecewiseLinearComplexDomain()
        pos = np.array([0.0, 0.0, 0.0])
        dim = (0.3, 0.38