import numpy
from numpy import sin, cos, tan

from .base import BaseMetric

class Kerr(BaseMetric):
    def __init__(self, mass, a):
        super().__init__(mass, valid_coordinate="BoyerLindquist", kwargs={"a": a})
        self.a = a * self.R_s / 2

    def _metric_dimensionless(self, xs):
        """
        Returns the Kerr metric tensor in Boyer-Lindquist coordinates.

        Parameters
        ----------
        xs : list
            List of coordinates [t, r, theta, phi] of boyer-lindquist coordinates.
        mass : float
            Mass of the black hole in kg.
        a : float
            Spin parameter of the black hole.
        """
        xs = numpy.asarray(xs, dtype=float)

        if xs.ndim == 1:
            x1   = xs[1]
            x1_2 = x1 * x1
            s_x2 = numpy.sin(xs[2]) ** 2
            c_x2 = numpy.cos(xs[2]) ** 2
            a    = self.a
            a_2  = a * a
            R_s  = self.R_s

            Sigma = x1_2 + a_2 * c_x2
            Delta = x1_2 - R_s * x1 + a_2

            A = 1 - R_s * x1 / Sigma
            B = -Sigma / Delta
            C = -Sigma
            D = -(x1_2 + a_2 + R_s * x1 * a_2 / Sigma * s_x2) * s_x2
            E = 2 * R_s * x1 * a / Sigma * s_x2

            metric = numpy.diag([A, B, C, D])
            metric[0, 3] = metric[3, 0] = E / 2
            return metric

        # Multiple points: shape (N, 4)
        r    = xs[:, 1]
        r2   = r * r
        s_x2 = numpy.sin(xs[:, 2]) ** 2
        c_x2 = numpy.cos(xs[:, 2]) ** 2
        a    = self.a
        a_2  = a * a
        R_s  = self.R_s

        Sigma = r2 + a_2 * c_x2
        Delta = r2 - R_s * r + a_2

        A = 1 - R_s * r / Sigma
        B = -Sigma / Delta
        C = -Sigma
        D = -(r2 + a_2 + R_s * r * a_2 / Sigma * s_x2) * s_x2
        E = 2 * R_s * r * a / Sigma * s_x2

        N = len(xs)
        metrics = numpy.zeros((N, 4, 4))
        metrics[:, 0, 0] = A
        metrics[:, 1, 1] = B
        metrics[:, 2, 2] = C
        metrics[:, 3, 3] = D
        metrics[:, 0, 3] = metrics[:, 3, 0] = E / 2

        return metrics

    # Generated by utils/symbolic_to_numeric.py
    def _get_christoffel_symbols(self, xs):
        """
        Returns the Christoffel symbols of the metric.

        Parameters
        ----------
        xs : array of shape (4,) or (N, 4)
            Coordinates [x0, x1, x2, x3].
        """
        xs = numpy.asarray(xs, dtype=float)

        if xs.ndim == 2:
            N = len(xs)
            Gammas = numpy.zeros((N, 4, 4, 4))
            for i, x in enumerate(xs):
                Gammas[i] = self._get_christoffel_symbols(x)
            return Gammas

        x0, x1, x2, x3 = xs

        ############### Auxiliary variables ##############
        a = self.a
        R_s = self.R_s

        cos_I2_x2I_ = cos(2 * x2)
        cos_I4_x2I_ = cos(4 * x2)
        cos_Ix2I_ = cos(x2)
        sin_I2_x2I_ = sin(2 * x2)
        sin_Ix2I_ = sin(x2)
        tan_Ix2I_ = tan(x2)

        R_s_pow2 = R_s**2
        _1_over_tan_Ix2I_ = 1 / tan_Ix2I_
        a_pow2 = a**2
        a_pow4 = a**4
        a_pow6 = a**6
        cos_Ix2I__pow2 = cos_Ix2I_**2
        cos_Ix2I__pow4 = cos_Ix2I_**4
        sin_Ix2I__pow2 = sin_Ix2I_**2
        sin_Ix2I__pow4 = sin_Ix2I_**4
        x1_pow2 = x1**2
        x1_pow3 = x1**3
        x1_pow4 = x1**4
        x1_pow5 = x1**5
        x1_pow6 = x1**6

        # Modified manually for optimization
        a_pow2_cos_Ix2I__pow2 = a_pow2 * cos_Ix2I__pow2
        R_s_a_pow2_cos_Ix2I__pow2_x1 = R_s * a_pow2_cos_Ix2I__pow2 * x1
        R_s_x1_pow3 = R_s * x1_pow3
        _2_x1_pow2 = 2 * x1_pow2
        _R_s_x1 = -R_s * x1
        _a_pow2_cos_Ix2I__pow2_x1_pow2 = -a_pow2_cos_Ix2I__pow2 * x1_pow2
        a_pow2_cos_I2_x2I_ = a_pow2 * cos_I2_x2I_
        _a_pow2_sin_I2_x2I__over__Ia_pow2_cos_I2_x2I___plus__a_pow2__plus__2_x1_pow2I_ = (
            -a_pow2 * sin_I2_x2I_ / (a_pow2_cos_I2_x2I_ + a_pow2 + _2_x1_pow2)
        )
        a_pow2_x1_pow2 = a_pow2 * x1_pow2
        _a_pow2_x1_pow2 = -a_pow2_x1_pow2
        _cos_I4_x2I_ = -cos_I4_x2I_
        _x1_pow4 = -x1_pow4
        a_pow4_cos_Ix2I__pow2 = a_pow4 * cos_Ix2I__pow2

        ##################################################

        ############### Christoffel symbols ##############
        Gamma = numpy.zeros((4, 4, 4))

        # Only non-zero components and symmetry properties

        # mu = 0
        Gamma[0, 0, 1] = (
            R_s
            * (
                a_pow2_x1_pow2 * sin_Ix2I__pow2
                + a_pow4 * sin_Ix2I__pow2
                - a_pow4
                + x1_pow4
            )
            / (
                2
                * (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 2
                * (_R_s_x1 + a_pow2 + x1_pow2)
            )
        )
        Gamma[0, 0, 2] = (
            2
            * R_s
            * _a_pow2_sin_I2_x2I__over__Ia_pow2_cos_I2_x2I___plus__a_pow2__plus__2_x1_pow2I_
            * x1
            / (_2_x1_pow2 + a_pow2 + a_pow2_cos_I2_x2I_)
        )
        Gamma[0, 1, 3] = (
            R_s
            * a
            * sin_Ix2I__pow2
            * (
                _a_pow2_cos_Ix2I__pow2_x1_pow2
                + _a_pow2_x1_pow2
                + 3 * _x1_pow4
                + a_pow4_cos_Ix2I__pow2
            )
            / (
                2
                * (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 2
                * (_R_s_x1 + a_pow2 + x1_pow2)
            )
        )
        Gamma[0, 2, 3] = (
            R_s
            * a**3
            * cos_Ix2I_
            * sin_Ix2I_**3
            * x1
            / (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 2
        )

        Gamma[0, 1, 0] = Gamma[0, 0, 1]
        Gamma[0, 2, 0] = Gamma[0, 0, 2]
        Gamma[0, 3, 1] = Gamma[0, 1, 3]
        Gamma[0, 3, 2] = Gamma[0, 2, 3]

        # mu = 1
        Gamma[1, 0, 0] = (
            R_s
            * (
                R_s_a_pow2_cos_Ix2I__pow2_x1
                - R_s_x1_pow3
                + _a_pow2_cos_Ix2I__pow2_x1_pow2
                + _a_pow2_x1_pow2
                + a_pow4_cos_Ix2I__pow2
                + x1_pow4
            )
            / (2 * (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 3)
        )
        Gamma[1, 0, 3] = (
            R_s
            * a
            * sin_Ix2I__pow2
            * (
                -R_s_a_pow2_cos_Ix2I__pow2_x1
                + R_s_x1_pow3
                + _x1_pow4
                + a_pow2_cos_Ix2I__pow2 * x1_pow2
                + a_pow2_x1_pow2
                - a_pow4_cos_Ix2I__pow2
            )
            / (2 * (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 3)
        )
        Gamma[1, 1, 1] = (
            -R_s * a_pow2_cos_Ix2I__pow2 / 2
            + R_s * x1_pow2 / 2
            + a_pow2 * x1
            + a_pow2_cos_Ix2I__pow2 * x1
        ) / (
            R_s_a_pow2_cos_Ix2I__pow2_x1
            + R_s_x1_pow3
            + _a_pow2_cos_Ix2I__pow2_x1_pow2
            + _x1_pow4
            + a_pow2_x1_pow2
            + a_pow4_cos_Ix2I__pow2
        )
        Gamma[1, 1, 2] = (
            _a_pow2_sin_I2_x2I__over__Ia_pow2_cos_I2_x2I___plus__a_pow2__plus__2_x1_pow2I_
        )
        Gamma[1, 2, 2] = (
            x1 * (R_s * x1 + a_pow2 - x1_pow2) / (a_pow2_cos_Ix2I__pow2 + x1_pow2)
        )
        Gamma[1, 3, 3] = (
            sin_Ix2I__pow2
            * (
                R_s * a_pow2 * sin_Ix2I__pow2 * x1_pow4 / 2
                + 2 * R_s * a_pow2_cos_Ix2I__pow2 * x1_pow4
                + R_s * a_pow4 * cos_Ix2I__pow4 * x1_pow2
                - R_s * a_pow4 * sin_Ix2I__pow2 * x1_pow2 / 2
                - R_s * a_pow4 * x1_pow2 * (_cos_I4_x2I_ + 1) / 16
                + R_s * a_pow6 * (_cos_I4_x2I_ + 1) / 16
                + R_s * x1_pow6
                - R_s_pow2 * a_pow2 * sin_Ix2I__pow2 * x1_pow3 / 2
                + R_s_pow2 * a_pow4 * x1 * (_cos_I4_x2I_ + 1) / 16
                + a_pow2 * x1_pow5
                - 2 * a_pow2_cos_Ix2I__pow2 * x1_pow5
                - a_pow4 * cos_Ix2I__pow4 * x1_pow3
                + 2 * a_pow4_cos_Ix2I__pow2 * x1_pow3
                + a_pow6 * cos_Ix2I__pow4 * x1
                - x1**7
            )
            / (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 3
        )

        Gamma[1, 2, 1] = Gamma[1, 1, 2]
        Gamma[1, 3, 0] = Gamma[1, 0, 3]

        # mu = 2
        Gamma[2, 0, 0] = (
            4
            * R_s
            * _a_pow2_sin_I2_x2I__over__Ia_pow2_cos_I2_x2I___plus__a_pow2__plus__2_x1_pow2I_
            * x1
            / (_2_x1_pow2 + a_pow2 + a_pow2_cos_I2_x2I_) ** 2
        )
        Gamma[2, 0, 3] = (
            4
            * R_s
            * a
            * sin_I2_x2I_
            * x1
            * (a_pow2 + x1_pow2)
            / (_2_x1_pow2 + a_pow2 + a_pow2_cos_I2_x2I_) ** 3
        )
        Gamma[2, 1, 1] = (
            -a_pow2
            * cos_Ix2I_
            * sin_Ix2I_
            / (
                R_s_a_pow2_cos_Ix2I__pow2_x1
                + R_s_x1_pow3
                + _a_pow2_cos_Ix2I__pow2_x1_pow2
                - _a_pow2_x1_pow2
                + _x1_pow4
                + a_pow4_cos_Ix2I__pow2
            )
        )
        Gamma[2, 1, 2] = x1 / (a_pow2_cos_Ix2I__pow2 + x1_pow2)
        Gamma[2, 2, 2] = (
            _a_pow2_sin_I2_x2I__over__Ia_pow2_cos_I2_x2I___plus__a_pow2__plus__2_x1_pow2I_
        )
        Gamma[2, 3, 3] = (
            cos_Ix2I_
            * sin_Ix2I_
            * (
                -2 * R_s_x1_pow3 * a_pow2 * sin_Ix2I__pow2
                - _2_x1_pow2 * a_pow4_cos_Ix2I__pow2
                + _R_s_x1 * a_pow4 * sin_Ix2I__pow4
                + _R_s_x1 * a_pow4 * (_cos_I4_x2I_ + 1) / 4
                + _x1_pow4 * a_pow2
                + 2 * _x1_pow4 * a_pow2_cos_Ix2I__pow2
                - a_pow4 * cos_Ix2I__pow4 * x1_pow2
                - a_pow6 * cos_Ix2I__pow4
                - x1_pow6
            )
            / (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 3
        )

        Gamma[2, 2, 1] = Gamma[2, 1, 2]
        Gamma[2, 3, 0] = Gamma[2, 0, 3]

        # mu = 3
        Gamma[3, 0, 1] = (
            R_s
            * a
            * (-a_pow2_cos_Ix2I__pow2 + x1_pow2)
            / (
                2
                * (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 2
                * (_R_s_x1 + a_pow2 + x1_pow2)
            )
        )
        Gamma[3, 0, 2] = (
            _1_over_tan_Ix2I_ * _R_s_x1 * a / (a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 2
        )
        Gamma[3, 1, 3] = (
            R_s * _a_pow2_cos_Ix2I__pow2_x1_pow2
            + R_s * _a_pow2_x1_pow2 * sin_Ix2I__pow2 / 2
            + R_s * _x1_pow4
            + R_s * a_pow4 * (_cos_I4_x2I_ + 1) / 16
            + 2 * a_pow2_cos_Ix2I__pow2 * x1_pow3
            + a_pow4 * cos_Ix2I__pow4 * x1
            + x1_pow5
        ) / ((a_pow2_cos_Ix2I__pow2 + x1_pow2) ** 2 * (_R_s_x1 + a_pow2 + x1_pow2))
        Gamma[3, 2, 3] = (
            _1_over_tan_Ix2I_
            * (
                R_s * a_pow2 * sin_Ix2I__pow2 * x1
                + _2_x1_pow2 * a_pow2
                + 2 * _a_pow2_x1_pow2 * sin_Ix2I__pow2
                - 2 * a_pow4 * sin_Ix2I__pow2
                + a_pow4 * sin_Ix2I__pow4
                + a_pow4
                + x1_pow4
            )
            / (_2_x1_pow2 * a_pow2_cos_Ix2I__pow2 + a_pow4 * cos_Ix2I__pow4 + x1_pow4)
        )

        Gamma[3, 1, 0] = Gamma[3, 0, 1]
        Gamma[3, 2, 0] = Gamma[3, 0, 2]
        Gamma[3, 3, 1] = Gamma[3, 1, 3]
        Gamma[3, 3, 2] = Gamma[3, 2, 3]

        ##################################################

        return Gamma
