#!/usr/bin/env python3
# Project : From geodynamic to Seismic observations in the Earth's inner core
# Author : Marine Lasbleis
""" Define classes for geodynnamical models combining growth, translation and rotation. """

from __future__ import division
from __future__ import absolute_import


import numpy as np
import matplotlib.pyplot as plt  # for figures
import math
from scipy.integrate import ode
from scipy.optimize import fsolve
import warnings

# personal routines
from . import positions
from . import intersection
from . import geodyn
from . import mineral_phys


# def exact_translation(point, velocity, direction=positions.CartesianPoint(1,0,0)):
#     # TODO : is this function used anywhere??
#     RICB = 1221.
#     x_0, y_0, z_0 = point.x, point.y, point.z
#     mean_direction = np.sqrt(direction.x**2+ direction.y**2+ direction.z**2)
#     a, b, c = direction.x/mean_direction, direction.y/mean_direction, direction.z/mean_direction
#
#     solution_1 = x_0*a+y_0*b+z_0*c + np.sqrt((x_0*a+y_0*b+z_0*c)**2-(x_0**2+y_0**2+z_0**2-RICB**2))
#     solution_2 = x_0*a+y_0*b+z_0*c - np.sqrt((x_0*a+y_0*b+z_0*c)**2-(x_0**2+y_0**2+z_0**2-RICB**    2))
#     ## TO DO : verify that we can remove solution_2 ? Is solution_1 always max?
#     return max(solution_1, solution_2)


class ModelTRG(geodyn.Model):

    def __init__(self):
        pass

    def velocity(self, t, position):
        """ Velocity at the given position and given time.

            Need to be implemented in derived classes.
            Needed in cartesian coordinates.

            Example :
            velocity, in cartesian coordinates

            position is a cartesian position [x, y, z]
            vt is the translation velocity (numpy array (1,3))
            rotation velocity is $\omega \vec{e}_z \times \vec{r}$.

            return vt + np.array([-omega * r[1], omega * r[0], 0.] )
            """
        raise NotImplementedError(
            "need to implement velocity() in derived class!")

    def radius_ic(self, t):
        """ radius of the inner core with time.

            Need to be implemented in derived classes.
            """
        raise NotImplementedError(
            "need to implement radius_ic() in derived class!")

    def growth_ic(self, t):
        """ growth rate of the inner core with time.

            Need to be implemented in derived classes.
            """
        raise NotImplementedError(
            "need to implement growth_ic() in derived class!")

    def effective_growth_rate(self, t, point):
        """ Effective growth rate at the point r.

        v_{g_eff} = || v_growth - v_geodynamic*e_r ||
        v_geodynamic is already in cartesian coordinates.
        v_growth = ||v_growth|| * vec{e}_r (the unit vector for the radial direction)
        point.er() gives the cartesian coordinates of the vector e_r
        point.proj_er(vect) gives the value of the vector projected on the vector e_r
        r is the position, described as x,y,z
        This function is used for points that are at the surface: r(t) is a point at the surface of the inner core at the time t.
        """
        r = np.array([point.x, point.y, point.z])
        vitesse = point.proj_er(self.velocity(t, r))  # projected on e_r
        growth = self.growth_ic(t) - vitesse
        return growth

    def effective_growth_rate2(self, t, point):
        """ Effective growth rate at the point r defined by the radial derivative of the field age

        \dot{r} = -1/( d(age)/dr )

        Args:
            t: time
            point: position (Point instance)
        Return:
            scalar
        """
        r = np.array([point.x, point.y, point.z])
        dadr = radial_derivative(
            self.crystallisation_time, r, self.rICB * 0.001, self.tau_ic)
        # if the age derivative is too close to 0, then we set it to 1.e-6.
        # Avoid having too large growth rates.
        if abs(dadr) < 1.e-2:
            message = "Age growth may be divergent. Value of da/dr = {}, modified for {}, corresponding to age growth of {} m/years".format(
                dadr, np.sign(dadr) * 1.e-2, 1e2 * self.length_unit / self.time_unit)
            warnings.warn(message)
            dadr = np.sign(dadr) * 1.e-2
        return 1. / (dadr * self.time_unit / self.length_unit)

    def find_time_beforex0(self, point, t0, t1):
        """ find the intersection between the trajectory and the radius of the IC
        if needed, can be re defined in derived class!

        point : [x, y, z]
        """
        return intersection.zero_brentq(
            self.distance_to_radius, point, t0, a=0., b=t1)

    def proxy_singlepoint(self, point, proxy_type):
        """ evaluate the proxy on a single positions.Point instance.

        return a dictionnary {}
        """
        proxy = {}  # empty dictionnary
        x, y, z = point.x, point.y, point.z
        # be careful, it's the hemisphere at the time t_end (not at time of
        # freezing)
        if proxy_type == "hemispheres":
            proxy["hemispheres"] = self.hemispheres(point)
            return proxy  # no calculation of time, etc.
        time = self.crystallisation_time([x, y, z], self.tau_ic)
        position_crys = self.crystallisation_position([x, y, z], time)
        # calculate the proxy needed (proxy_type value)
        if proxy_type == "age":
            proxy["age"] = (self.tau_ic - time) * 1e-6 * \
                self.time_unit  # age is in Myears!
        elif proxy_type == "domain_size":
            proxy["domain_size"] = self.domain_size(self.tau_ic - time)
        elif proxy_type == "dV_V":
            proxy["dV_V"] = self.dV_V(self.tau_ic - time)
        elif proxy_type == "theta" or proxy_type == "phi":
            point = self.integration_trajectory(time, [x, y, z], self.tau_ic)
            proxy["position"] = positions.CartesianPoint(
                point[0], point[1], point[2])
            proxy["phi"] = proxy["position"].phi