# -*- coding: utf-8 -*-

import itertools as it
import numpy as np

from . import diff
from . import utils

def _calc_dmetric(X, Z, metric, diff_op):
    dmetric = np.zeros((3,) + metric.shape)

    res = diff_op(metric, 3, X[0, 1] - X[0, 0])
    dmetric[0] = diff_op(metric, 3, X[0, 1] - X[0, 0])
    dmetric[2] = diff_op(metric, 2, Z[1, 0] - Z[0, 0])

    dmetric[1, 0, 0] = 0.0
    dmetric[1, 1, 1] = 0.0
    dmetric[1, 2, 2] = 0.0
    dmetric[1, 0, 1] = np.where(np.abs(X) > 1e-8, (metric[0, 0] - metric[1, 1]) / X, dmetric[0, 0, 0] - dmetric[0, 1, 1])
    dmetric[1, 1, 0] = dmetric[1, 0, 1]
    dmetric[1, 0, 2] = 0.0
    dmetric[1, 2, 0] = 0.0
    dmetric[1, 1, 2] = np.where(np.abs(X) > 1e-8, metric[0, 2] / X, dmetric[0, 0, 2])
    dmetric[1, 2, 1] = dmetric[1, 1, 2]

    return dmetric

def _calc_d2metric(X, Z, metric, dmetric, diff_op, diff2_op):
    d2metric = np.empty((3,) + dmetric.shape)

    dx = X[0, 1] - X[0, 0]
    dz = Z[1, 0] - Z[0, 0]
    X2 = X * X

    d2metric[0, 0] = diff2_op(metric, 3, dx)
    d2metric[2, 2] = diff2_op(metric, 2, dz)

    d2metric[0, 2] = diff_op(dmetric[0], 2, dz)
    d2metric[2, 0] = d2metric[0, 2]

    d2metric[0, 1, 0, 0] = 0
    d2metric[1, 0, 0, 0] = 0
    d2metric[1, 2, 0, 0] = 0
    d2metric[2, 1, 0, 0] = 0
    d2metric[1, 1, 0, 0] = np.where(np.abs(X) > 1e-8, dmetric[0, 0, 0] / X - 2.0 * (metric[0, 0] - metric[1, 1]) / X2, d2metric[0, 0, 1, 1])

    d2metric[1, 1, 1, 1] = np.where(np.abs(X) > 1e-8, dmetric[0, 1, 1] / X + 2.0 * (metric[0, 0] - metric[1, 1]) / X2, d2metric[0, 0, 0, 0])
    d2metric[0, 1, 1, 1] = 0
    d2metric[1, 0, 1, 1] = 0
    d2metric[2, 1, 1, 1] = 0
    d2metric[1, 2, 1, 1] = 0

    d2metric[1, 1, 2, 2] = np.where(np.abs(X) > 1e-8, dmetric[0, 2, 2] / X, d2metric[0, 0, 2, 2])
    d2metric[0, 1, 2, 2] = 0
    d2metric[1, 0, 2, 2] = 0
    d2metric[2, 1, 2, 2] = 0
    d2metric[1, 2, 2, 2] = 0

    d2metric[1, 1, 0, 1] = 0
    d2metric[0, 1, 0, 1] = np.where(np.abs(X) > 1e-8, (dmetric[0, 0, 0] - dmetric[0, 1, 1]) / X - (metric[0, 0] - metric[1, 1]) / X2, 0.5 * (d2metric[0, 0, 0, 0] - d2metric[0, 0, 1, 1]))
    d2metric[1, 0, 0, 1] = d2metric[0, 1, 0, 1]
    d2metric[1, 2, 0, 1] = np.where(np.abs(X) > 1e-8, (dmetric[2, 0, 0] - dmetric[2, 1, 1]) / X, d2metric[0, 2, 0, 0] - d2metric[0, 2, 1, 1])
    d2metric[2, 1, 0, 1] = d2metric[1, 2, 0, 1]


    d2metric[1, 1, 0, 2] = np.where(np.abs(X) > 1e-8, dmetric[0, 0, 2] / X - metric[0, 2] / X2, 0.5 * d2metric[0, 0, 0, 2])
    d2metric[0, 1, 0, 2] = 0
    d2metric[1, 0, 0, 2] = 0
    d2metric[1, 2, 0, 2] = 0
    d2metric[2, 1, 0, 2] = 0

    d2metric[1, 1, 1, 2] = 0
    d2metric[0, 1, 1, 2] = np.where(np.abs(X) > 1e-8, dmetric[0, 0, 2] / X - metric[0, 2] / X2, 0.5 * d2metric[0, 0, 0, 2])
    d2metric[1, 0, 1, 2] = d2metric[0, 1, 1, 2]
    d2metric[1, 2, 1, 2] = np.where(np.abs(X) > 1e-8, dmetric[2, 0, 2] / X, d2metric[0, 2, 0, 2])
    d2metric[2, 1, 1, 2] = d2metric[1, 2, 1, 2]

    d2metric[:, :, 1, 0] = d2metric[:, :, 0, 1]
    d2metric[:, :, 2, 0] = d2metric[:, :, 0, 2]
    d2metric[:, :, 2, 1] = d2metric[:, :, 1, 2]

    return d2metric

def _calc_christoffel(metric, metric_u, dmetric):
    Gamma = np.empty_like(dmetric)
    for i, j, k in it.product(range(3), repeat = 3):
        Gamma[i, j, k] = 0.5 * np.einsum('k...,k...', metric_u[i], dmetric[j, k] + dmetric[k, j] - dmetric[:, k, j])

    return Gamma

def calc_christoffel(x, z, metric, diff_op = diff.fd4):
    """
    Calculate Christoffel symbols

     i    1 il /                       \
    Γ   = -γ   | ∂ γ   + ∂ γ   - ∂ γ   |
     jk   2    \  j kl    k jl    l jk /

    using finite differences.

    :param array_like x: 1D array of x coordinates.
    :param array_like z: 1D array of z-coordinates.
    :param array_like metric: 4D array of spatial metric values at the grid
                              formed by x and z. metric[i, j, k, l] is the ijth
                              component of the metric at the point (X=x[l],
                              Z=z[k]).
    :rtype: array_like, shape (3, 3, 3, z.shape[0], x.shape[0])
    :return: Christoffel symbols, first axis is the upper index, following two
             axes are the two lower indices, final two axes correspond to the z
             and x grid points respectively.
    """
    X, Z  = np.meshgrid(x, z)

    metric_u = utils.matrix_invert(metric)
    dmetric  = _calc_dmetric(X, Z, metric, diff_op)

    return _calc_christoffel(metric, metric_u, dmetric)

def _calc_dchristoffel(metric, metric_u, dmetric, dmetric_u, d2metric):
    dGamma = np.empty_like(d2metric)
    for i, j, k, l in it.product(range(3), repeat = 4):
        dGamma[i, j, k, l] = 0.5 * (np.einsum('k...,k...', dmetric_u[i, j], dmetric[l, k] + dmetric[k, l] - dmetric[:, k, l]) +
                                    np.einsum('k...,k...',  metric_u[j],    d2metric[i, l, k] + d2metric[i, k, l] - d2metric[i, :, l, k]))
    return dGamma

def calc_riemann(metric, metric_u, dmetric, d2metric):
    dmetric_u = -np.einsum('ij...,km...,ljk...->lim...', metric_u, metric_u, dmetric)
    Gamma     = _calc_christoffel(metric, metric_u, dmetric)
    dGamma    = _calc_dchristoffel(metric, metric_u, dmetric, dmetric_u, d2metric)

    Riemann_uddd = (np.einsum('cabd...->abcd...', dGamma) - np.einsum('dabc...->abcd...', dGamma) +
                    np.einsum('akc...,kbd...->abcd...', Gamma, Gamma) - np.einsum('akd...,kbc...->abcd...', Gamma, Gamma))
    Riemann = np.einsum('ik...,klmn...->ilmn...', metric, Riemann_uddd)

    return Riemann

def calc_ricci(metric, metric_u, dmetric, d2metric):
    dmetric_u = -np.einsum('ij...,km...,ljk...->lim...', metric_u, metric_u, dmetric)
    Gamma     = _calc_christoffel(metric, metric_u, dmetric)
    dGamma    = _calc_dchristoffel(metric, metric_u, dmetric, dmetric_u, d2metric)

    Ricci  = (np.einsum('mmjk...->jk...', dGamma) - np.einsum('kmjm...->jk...', dGamma) +
              np.einsum('llm...,mjk...->jk...', Gamma, Gamma) - np.einsum('lkm...,mjl...->jk...', Gamma, Gamma))

    return Ricci

def calc_constraint_ham(metric, metric_u, dmetric, d2metric, curv):
    Ricci      = calc_ricci(metric, metric_u, dmetric, d2metric)
    R_scalar   = np.einsum('ij...,ij...', metric_u, Ricci)

    curv_ud    = np.einsum('ik...,kj...->ij...', metric_u, curv)
    curv_trace = curv_ud[0, 0] + curv_ud[1, 1] + curv_ud[2, 2]

    k2 = np.einsum('ij...,ji...', curv_ud, curv_ud)

    return R_scalar + curv_trace ** 2 - k2

def calc_constraint_mom(metric, metric_u, dmetric, curv_m, dcurv_m):
    Gamma      = _calc_christoffel(metric, metric_u, dmetric)
    curv_trace = curv_m[0, 0] + curv_m[1, 1] + curv_m[2, 2]

    M = np.empty((3,) + curv_trace.shape)
    for j in range(3):
        M[j] = np.einsum('ii...', dcurv_m[:, :, j]) + np.einsum('iik...,k...', Gamma, curv_m[:, j]) - np.einsum('ki...,ik...', Gamma[:, :, j], curv_m)

    return M
