#!/usr/bin/python

import argparse
import sys

import numpy as np
from scipy import interpolate

from cactus_utils import simdata

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',
}

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:])

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] >= (args.x0 + args.dx * (args.gridsize_x - 1)) and
        s.z[-1] >= (args.z0 + args.dz * (args.gridsize_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)
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':
    src = s['alpha'][:]

src = src.squeeze()

src_interp = interpolate.RectBivariateSpline(s.z, s.x, src)

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)

res = src_interp(z, x)
res.tofile(sys.stdout)