#!/usr/bin/python3

import brill_data
import numpy as np

import sys

def _invert_matrix(mat):
    oldshape = mat.shape
    newshape = oldshape[:2] + (np.product(mat.shape[2:]),)

    mat = mat.reshape(newshape).transpose((2, 0, 1))
    inv = np.linalg.inv(mat)
    return inv.transpose((1, 2, 0)).reshape(oldshape)

def ham_calc(bd, rho, z):
    g   = np.empty((3, 3) + x.shape + z.shape)
    dg  = np.empty((3,) + g.shape)
    d2g = np.zeros((3,) + dg.shape)

    for i in range(3):
        for j in range(3):
            g[i, j] = bd.eval_metric(rho, z, i * 3 + j)

            dg[0, i, j] = bd.eval_metric(rho, z, i * 3 + j, [1, 0])
            dg[1, i, j] = bd.eval_metric(rho, z, i * 3 + j, [0, 1])
            dg[2, i, j] = 0.0

            d2g[0, 0, i, j] = bd.eval_metric(rho, z, i * 3 + j, [2, 0])
            d2g[0, 1, i, j] = bd.eval_metric(rho, z, i * 3 + j, [1, 1])
            d2g[1, 0, i, j] = d2g[1, 0, i, j]
            d2g[1, 1, i, j] = bd.eval_metric(rho, z, i * 3 + j, [0, 2])

            d2g[2, 2, i, j] = 0.0
            d2g[0, 2, i, j] = 0.0
            d2g[2, 0, i, j] = 0.0
            d2g[1, 2, i, j] = 0.0
            d2g[2, 1, i, j] = 0.0

    gu = _invert_matrix(g)

    Gl = 0.5 * (np.einsum('cab...->abc...', dg) + np.einsum('bac...->abc...', dg) - dg)
    G  = np.einsum('ij...,jkl...->ikl...', gu, Gl)
    R = (0.5 * (np.einsum('adbc...->abcd...', d2g) + np.einsum('bcad...->abcd...', d2g) - np.einsum('bdac...->abcd...', d2g) - np.einsum('acbd...->abcd...', d2g)) +
         np.einsum('ead...,ebc...->abcd...', Gl, G) - np.einsum('eac...,ebd...->abcd...', Gl, G))

    R = np.einsum('ijkl...,ik...->jl...', R, gu)
    R = np.einsum('ij...,ij...', R, gu)
    return R

tol = 5e-13

x = np.linspace(0.1, 256, 50)
z = np.linspace(0.1, 256, 50)

coeffs = (80, 30)
amplitudes = np.linspace(1, 12, 11)
hmax = np.empty_like(amplitudes)
for i, a in enumerate(amplitudes):
    bd = brill_data.BrillData(amplitude = a, nb_coeffs = coeffs)
    hmax[i] = np.max(np.abs(ham_calc(bd, x, z)))

if np.any(hmax > tol):
    sys.stderr.write('Large constraint violation: amplitudes %s, H %s\n'
                     % (amplitudes, hmax))
    sys.exit(1)