#!/usr/bin/python3

import argparse
import sys

import numpy as np

import cactus_utils.simdata as simdata

import math_utils.basis as basis
import math_utils.nonlin_ode as nlo
import math_utils.series_expansion as se

from nr_analysis_axi import interp
from nr_analysis_axi import horizon
from nr_analysis_axi import utils

np.seterr(divide = 'ignore')

def extract_gamma_k(s, idx0, idx1):
    idx_y = np.where(abs(s.y) < 1e-10)[0][0]
    metric = np.array([[s['gt11'][:, idx_y], s['gt12'][:, idx_y], s['gt13'][:, idx_y]],
                       [s['gt12'][:, idx_y], s['gt22'][:, idx_y], s['gt23'][:, idx_y]],
                       [s['gt13'][:, idx_y], s['gt23'][:, idx_y], s['gt33'][:, idx_y]]]) / (s['phi'][:, idx_y] ** 2)

    curv = np.array([[s['At11'][:, idx_y], s['At12'][:, idx_y], s['At13'][:, idx_y]],
                     [s['At12'][:, idx_y], s['At22'][:, idx_y], s['At23'][:, idx_y]],
                     [s['At13'][:, idx_y], s['At23'][:, idx_y], s['At33'][:, idx_y]]]) / (s['phi'][:, idx_y] ** 2)
    curv += s['trK'][:, idx_y] * metric / 3.0
    return metric[:, :, idx0:idx1, idx0:idx1], curv[:, :, idx0:idx1, idx0:idx1]

def load_hor_vars(s):
    X, Z = np.meshgrid(s.x[5:-5], s.z[5:-5])
    X2 = utils.array_reflect(utils.array_reflect(X, 1.0, 0), -1.0, 1)
    Z2 = utils.array_reflect(utils.array_reflect(Z, -1.0, 0) , 1.0, 1)
    metric, curv = extract_gamma_k(s, 5, -5)
    metric2 = np.empty((3, 3, metric.shape[2] * 2 - 1, metric.shape[3] * 2 - 1))
    curv2   = np.empty((3, 3, curv.shape[2] * 2 - 1, curv.shape[3] * 2 - 1))
    for i in range(3):
        for j in range(3):
            parity_x = 1.0
            if i == 0:
                parity_x *= -1.0
            if j == 0:
                parity_x *= -1.0

            parity_z = 1.0
            if i == 2:
                parity_z *= -1.0
            if j == 2:
                parity_z *= -1.0

            metric2[i, j] = utils.array_reflect(utils.array_reflect(metric[i, j], parity_z, 0), parity_x, 1)
            curv2[i, j]   = utils.array_reflect(utils.array_reflect(curv[i, j], parity_z, 0),   parity_x, 1)
    return X2, Z2, metric2, curv2

parser = argparse.ArgumentParser(description = 'Find apparent horizons in Cactus simulation output.')
parser.add_argument('datadir', type = str, help = 'Path to the directory with the data files')
parser.add_argument('--rl', type = int, default = -1, help = 'refinement level to use')
parser.add_argument('--start', type = float, default = 0.0, help = 'first time level on which the horizon will be located')
parser.add_argument('--end', type = float, default = -1.0,
                    help = 'last time level on which the horizon will be located (default - equal to first)')
parser.add_argument('--step', type = float, default = -1.0,
                    help = 'time step between successive searches (default - as small as possible)')
parser.add_argument('--guess', type = str, default = '1.0',
                    help = 'Initial guesses for the spectral coefficients, separated by ":". Unspecified ones are taken to be zero.')
parser.add_argument('--order', type = int, default = 40,
                    help = 'Order of the solver')
parser.add_argument('--coeffs', dest = 'output_coeffs', action = 'store_true',
                    help = 'Print the spectral coefficients of the horizon (separated by ":")')
parser.add_argument('-v', '--verbose', dest = 'verbose', action = 'store_true',
                    help = 'Print extra debugging information')
parser.add_argument('--maxiter', type = int, default = 15,
                    help ='Maximum number of Newton iterations')
parser.set_defaults(verbose = False)

args = parser.parse_args()

if args.verbose:
    sys.stderr.write('Opening data dir: %s\n' % args.datadir)
sd = simdata.SimulationData(args.datadir)
if args.verbose:
    sys.stderr.write('Opened the data dir\n')

rl = args.rl
if rl < 0:
    rl = len(sd.rl) - 1
end = args.end
if end < 0:
    end = args.start

C  = basis.CosBasis(basis.CosBasis.PARITY_EVEN)

guess = list(map(float, args.guess.split(':')))
ig = se.SeriesExpansion(np.array(guess + [0.0] * (args.order - len(guess))), C)

last_time = None
for s1 in sd.df['trK'].rl[rl]:
    t  = s1.time
    it = s1.it

    if t < args.start:
        continue

    if t > end:
        break

    if args.step > 0 and last_time is not None and t - last_time < args.step:
        continue

    last_time = t

    if args.verbose:
        sys.stderr.write('Searching at time %g\n' % t)

    s = sd.rl[rl].slice(iteration = it)

    X, Z, metric, curv = load_hor_vars(s)
    ahc = horizon.AHCalc(X, Z, metric, curv)

    try:
        hor = nlo.nonlin_solve_1d(ig, ahc.Fs, maxiter = args.maxiter, verbose = args.verbose)
    except RuntimeError:
        sys.stdout.write('time %g: no horizon\n' % t)
        continue

    # calculate maxima of out- and in-going expansion on the horizon
    # outgoing should be zero
    interp_origin = (Z[0, 0], X[0, 0])
    interp_step   = (Z[1, 0] - Z[0, 0], X[0, 1] - X[0, 0])

    t = np.linspace(0, np.pi / 2, 255)
    interp_x = hor.eval(t) * np.sin(t)
    interp_z = hor.eval(t) * np.cos(t)

    expansion_in_2d  = ahc.calc_expansion(hor, -1)
    expansion_in_hor = interp.Interp2D_C(interp_origin, interp_step, expansion_in_2d, 6)(interp_z, interp_x)
    expansion_in_max = np.max(expansion_in_hor)
    expansion_in_min = np.min(expansion_in_hor)

    expansion_out_2d  = ahc.calc_expansion(hor, 1)
    expansion_out_hor = interp.Interp2D_C(interp_origin, interp_step, expansion_out_2d, 6)(interp_z, interp_x)
    expansion_out_max = np.max(expansion_out_hor)
    expansion_out_min = np.min(expansion_out_hor)

    sys.stdout.write('time %g: horizon found\t' % t)
    sys.stdout.write('x-axis value: %g\t' % hor.eval(np.array([np.pi / 2])))
    sys.stdout.write('z-axis value: %g\t' % hor.eval(np.array([0.0])))
    sys.stdout.write('mass: %g\t' % ahc.hor_mass(hor))
    sys.stdout.write('expansion in: min %g\tmax %g\n' % (expansion_in_min, expansion_in_max))
    sys.stdout.write('expansion out: min %g\tmax %g\n' % (expansion_out_min, expansion_out_max))
    if args.output_coeffs:
        sys.stdout.write('coeffs: ' + ':'.join(map(str, hor.coeffs)))
    sys.stdout.write('\n')

    ig = hor

sd.close()