#!/usr/bin/python3

# This script is intended to be invoked by the ReviveData cactus thorn [1].
# [1] https://git.khirnov.net/thorns/ReviveData.git/

import argparse
import sys

import numpy as np
from scipy import interpolate

from cactus_utils import simdata
from nr_analysis_axi import interp

metric_str = {
    'gxx' : '11', 'gyy' : '22', 'gzz' : '33',
    'gxy' : '12', 'gxz' : '13', 'gyz' : '23',
}
curv_str = {
    'kxx' : '11', 'kyy' : '22', 'kzz' : '33',
    'kxy' : '12', 'kxz' : '13', 'kyz' : '23',
}
shift_str = {
    'betax' : 'beta1', 'betay' : 'beta2', 'betaz' : 'beta3',
}

interp_order = 8

parser = argparse.ArgumentParser(description = 'Revive cactus simulation data')

parser.add_argument('simdata_path')
parser.add_argument('load_time', type = float)
parser.add_argument('load_iteration', type = int)
parser.add_argument('varname')
parser.add_argument('x0', type = float)
parser.add_argument('z0', type = float)
parser.add_argument('dx', type = float)
parser.add_argument('dz', type = float)
parser.add_argument('gridsize_x', type = int)
parser.add_argument('gridsize_z', type = int)

args = parser.parse_args(sys.argv[1:])

x = np.linspace(args.x0, args.x0 + args.dx * (args.gridsize_x - 1), args.gridsize_x)
z = np.linspace(args.z0, args.z0 + args.dz * (args.gridsize_z - 1), args.gridsize_z)

X, Z = np.meshgrid(x, z)

sd = simdata.SimulationData(args.simdata_path)

if args.load_iteration < 0:
    s = sd.rl[-1].slice(time = args.load_time)
    load_iteration = s.it
else:
    load_iteration = args.load_iteration

varname = args.varname
if '::' in varname:
    varname = varname.split('::', 1)[1]

load_rl = None
for rl in sd.rl[::-1]:
    s = rl.slice(load_iteration)
    if (s.x[0] <= args.x0 and s.z[0] <= args.z0 and
        s.x[-1] >= x[-1] and s.z[-1] >= z[-1]):
        load_rl = rl
        break

if load_rl is None:
    sys.stderr.write('No refinement level contains all the data required\n')
    sys.exit(1)

s = load_rl.slice(load_iteration)
src = None
try:
    if varname in metric_str:
        src = s['gt' + metric_str[varname]][:] / (s['phi'][:] ** 2)
    elif varname in curv_str:
        src = (s['At' + curv_str[varname]][:] + (1.0 / 3.0) * s['gt' + curv_str[varname]][:] * s['trK'][:]) / (s['phi'][:] ** 2)
    elif varname in shift_str:
        src = s[shift_str[varname]][:]
    elif varname == 'alp' or varname == 'alpha':
        src = s['alpha'][:]
    else:
        sys.stderr.write('Unknown gridfunction: %s\n' % varname)
        sys.exit(1)
except IndexError:
    sys.stderr.write('Gridfunction "%s" not in data, assuming zero\n' % varname)
    np.zeros(z.shape + x.shape).tofile(sys.stdout.buffer.raw)
    sys.exit(0)

src = src.squeeze()

src_interp = interp.Interp2D_C([s.z[0], s.x[0]], [s.z[1] - s.z[0], s.x[1] - s.x[0]], src, interp_order)

res = src_interp(Z.flatten(), X.flatten())
res.tofile(sys.stdout.buffer.raw)