"""ParaView-side renderer. Run only with pvpython."""

import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import paraview.simple as pv

os.environ.setdefault("PV_DEBUG_SKIP_OPENGL_VERSION_CHECK", "1")


def npz_source(config, files, name):
    source = pv.ProgrammableSource(registrationName=name)
    source.OutputDataSetType = "vtkImageData"
    import numpy as np

    with np.load(files[0], mmap_mode="r", allow_pickle=False) as archive:
        u_name, v_name = config["fields"]
        if u_name not in archive or v_name not in archive:
            raise ValueError("NPZ does not contain configured velocity fields")
        if archive[u_name].shape != archive[v_name].shape:
            raise ValueError("velocity field shapes differ")
        shape = archive[u_name].shape
    order = config["axis_order"]
    frames = shape[0] if order in ("tyx", "txy") and len(shape) == 3 else 1
    if (order in ("tyx", "txy") and len(shape) != 3) or (order in ("yx", "xy") and len(shape) != 2):
        raise ValueError("NPZ shape is incompatible with axis_order")
    source.ScriptRequestInformation = """
from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
info = self.GetExecutive().GetOutputInformation(0)
times = list(range(%d))
info.Set(vtkStreamingDemandDrivenPipeline.TIME_STEPS(), times, len(times))
info.Set(vtkStreamingDemandDrivenPipeline.TIME_RANGE(), [times[0], times[-1]], 2)
""" % frames
    payload = json.dumps({"path": files[0], "fields": config["fields"], "order": order, "frames": frames})
    source.Script = """
import json
import numpy as np
from vtkmodules.numpy_interface import dataset_adapter as dsa
from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
cfg = json.loads(%r)
info = self.GetExecutive().GetOutputInformation(0)
time = info.Get(vtkStreamingDemandDrivenPipeline.UPDATE_TIME_STEP())
index = max(0, min(cfg['frames'] - 1, int(round(time or 0))))
with np.load(cfg['path'], allow_pickle=False) as archive:
    def read(name):
        array = np.asarray(archive[name])
        if cfg['order'] == 'tyx': array = array[index]
        elif cfg['order'] == 'txy': array = array[index].T
        elif cfg['order'] == 'xy': array = array.T
        return np.ascontiguousarray(array)
    u, v = read(cfg['fields'][0]), read(cfg['fields'][1])
    output.SetDimensions(u.shape[1], u.shape[0], 1)
    wrapped = dsa.WrapDataObject(output)
    wrapped.PointData.append(u.ravel(order='C'), 'U')
    wrapped.PointData.append(v.ravel(order='C'), 'V')
""" % payload
    return source


def source(config, files, name):
    if config["format"] == "npz":
        return npz_source(config, files, name)
    reader = pv.TecplotReader(registrationName=name, FileNames=files)
    reader.DataArrayStatus = ["flag", "U", "V"]
    return reader


def vector_field(input_proxy, name="Velocity", u="U", v="V"):
    zero = pv.Calculator(registrationName=name + "ZeroZ", Input=input_proxy)
    zero.ResultArrayName = name + "Z"
    zero.Function = "0"
    vector = pv.MergeVectorComponents(registrationName=name, Input=zero)
    vector.XArray, vector.YArray, vector.ZArray = u, v, name + "Z"
    return vector


def scalar_field(input_proxy, preset, reference=None):
    if preset == "comparison":
        difference = pv.ProgrammableFilter(registrationName="Difference", Input=[input_proxy, reference])
        difference.Script = """import numpy as np
u = inputs[0].PointData['U'] - inputs[1].PointData['U']
v = inputs[0].PointData['V'] - inputs[1].PointData['V']
output.PointData.append(np.sqrt(u*u + v*v), 'VelocityError')
"""
        return difference, "VelocityError", False
    vector = vector_field(input_proxy)
    if preset == "speed":
        speed = pv.Calculator(registrationName="Speed", Input=vector)
        speed.ResultArrayName, speed.Function = "Speed", "mag(Velocity)"
        return speed, "Speed", False
    gradient = pv.Gradient(registrationName="Vorticity", Input=vector)
    gradient.ScalarArray = ["POINTS", "Velocity"]
    gradient.ComputeGradient = 0
    gradient.ComputeVorticity = 1
    vorticity = pv.Calculator(registrationName="ZVorticity", Input=gradient)
    vorticity.ResultArrayName, vorticity.Function = "Z_Vort", "Vorticity_Z"
    return vorticity, "Z_Vort", True


def show(proxy, view, array_name, value_range, diverging, colorbar):
    low, high = value_range
    middle = 0.0 if diverging and low < 0 < high else (low + high) / 2
    if diverging:
        points = [low, 0.231, 0.298, 0.753, middle, 1, 1, 1, high, 0.706, 0.016, 0.149]
    else:
        points = [low, 0.267, 0.005, 0.329, middle, 0.128, 0.567, 0.551, high, 0.993, 0.906, 0.144]
    lut = pv.GetColorTransferFunction(array_name)
    lut.AutomaticRescaleRangeMode, lut.RGBPoints = "Never", points
    display = pv.Show(proxy, view)
    display.ColorArrayName, display.LookupTable = ["POINTS", array_name], lut
    display.SetScalarBarVisibility(view, bool(colorbar))
    if colorbar:
        pv.GetScalarBar(lut, view).Title = colorbar


def main(config_path):
    config = json.loads(Path(config_path).read_text())
    primary = source(config, config["source"], "Flow")
    reference = source(config, config["reference"], "Reference") if config["reference"] else None
    if reference is not None:
        primary.UpdatePipelineInformation()
        reference.UpdatePipelineInformation()
        if len(times(primary)) != len(times(reference)):
            raise ValueError("source and reference frame counts differ")
    scalar, array_name, diverging = scalar_field(primary, config["preset"], reference)
    view = pv.CreateView("RenderView")
    view.ViewSize = [1200, 480]
    view.InteractionMode, view.OrientationAxesVisibility = "2D", 0
    view.CameraParallelProjection, view.Background, view.UseLight = 1, [1, 1, 1], 0
    show(scalar, view, array_name, config["range"], diverging, config["colorbar"])
    values = times(primary)
    output = Path(config["output"])
    if output.suffix.lower() != ".mp4":
        index = config["frame"]
        if index < 0 or index >= len(values):
            raise IndexError("frame is outside the available range")
        primary.UpdatePipeline(time=values[index])
        view.ResetCamera()
        pv.Render(view)
        pv.SaveScreenshot(str(output), view=view)
        return
    with tempfile.TemporaryDirectory(prefix=output.stem + "-", dir=output.parent) as directory:
        directory = Path(directory)
        for index, value in enumerate(values):
            primary.UpdatePipeline(time=value)
            if index == 0:
                view.ResetCamera()
            pv.Render(view)
            pv.SaveScreenshot(str(directory / ("frame.%06d.png" % index)), view=view)
        temporary = output.with_name(output.stem + ".tmp" + output.suffix)
        subprocess.run([config["ffmpeg"], "-y", "-framerate", str(config["fps"]), "-i", str(directory / "frame.%06d.png"), "-c:v", "libx264", "-pix_fmt", "yuv420p", str(temporary)], check=True)
        os.replace(temporary, output)


def times(proxy):
    proxy.UpdatePipelineInformation()
    return list(getattr(proxy, "TimestepValues", []) or [0.0])


if __name__ == "__main__":
    main(sys.argv[1])
