from pathlib import Path, PurePosixPath
import tempfile
from dataclasses import asdict
import sys

import pytest

import itk
import numpy as np

from itkwasm import (
    InterfaceTypes,
    TextStream,
    BinaryStream,
    PipelineInput,
    PipelineOutput,
    Pipeline,
    TextFile,
    BinaryFile,
    Image,
    Mesh,
    Transform,
    TransformList,
    TransformParameterizations,
    FloatTypes,
)

test_input_dir = Path(__file__).resolve().parent / "input"
test_baseline_dir = Path(__file__).resolve().parent / "baseline"


def test_stdout_stderr():
    pipeline = Pipeline(test_input_dir / "stdout-stderr-test.wasi.wasm")
    pipeline.run([])

    # Test re-run
    pipeline.run([])


def test_pipeline_bytes():
    pipeline_path = test_input_dir / "stdout-stderr-test.wasi.wasm"
    with open(pipeline_path, "rb") as fp:
        wasm_bytes = fp.read()
    pipeline = Pipeline(wasm_bytes)
    pipeline.run([])


def test_pipeline_input_output_streams():
    pipeline = Pipeline(test_input_dir / "input-output-files-test.wasi.wasm")

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.TextStream, TextStream("The answer is 42.")),
        PipelineInput(InterfaceTypes.BinaryStream, BinaryStream(bytes([222, 173, 190, 239]))),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.TextStream),
        PipelineOutput(InterfaceTypes.BinaryStream),
    ]

    args = [
        "--memory-io",
        "--input-text-stream",
        "0",
        "--input-binary-stream",
        "1",
        "--output-text-stream",
        "0",
        "--output-binary-stream",
        "1",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

    # Test re-run
    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

    assert outputs[0].type == InterfaceTypes.TextStream
    assert outputs[0].data.data == "The answer is 42."
    assert outputs[1].type, InterfaceTypes.BinaryStream
    assert outputs[1].data.data[0], 222
    assert outputs[1].data.data[1], 173
    assert outputs[1].data.data[2], 190
    assert outputs[1].data.data[3], 239


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="Windows tempfile resource, https://github.com/bytecodealliance/wasmtime-py/issues/132",
)
def test_pipeline_input_output_files():
    pipeline = Pipeline(test_input_dir / "input-output-files-test.wasi.wasm")
    input_text_file = PurePosixPath(test_input_dir / "input.txt")
    input_binary_file = PurePosixPath(test_input_dir / "input.bin")

    with tempfile.TemporaryDirectory() as tmpdirname:
        output_text_file = PurePosixPath(Path(tmpdirname) / "output.txt")
        output_binary_file = PurePosixPath(Path(tmpdirname) / "output.bin")

        pipeline_inputs = [
            PipelineInput(InterfaceTypes.TextFile, TextFile(input_text_file)),
            PipelineInput(InterfaceTypes.BinaryFile, BinaryFile(input_binary_file)),
        ]

        pipeline_outputs = [
            PipelineOutput(InterfaceTypes.TextFile, TextFile(output_text_file)),
            PipelineOutput(InterfaceTypes.BinaryFile, BinaryFile(output_binary_file)),
        ]

        args = [
            "--memory-io",
            "--use-files",
            "--input-text-file",
            str(input_text_file),
            "--input-binary-file",
            str(input_binary_file),
            "--output-text-file",
            str(output_text_file),
            "--output-binary-file",
            str(output_binary_file),
        ]

        outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

        # Test re-run
        outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

        assert outputs[0].type == InterfaceTypes.TextFile
        with open(outputs[0].data.path, "r") as fp:
            content = fp.read()
            assert content == "The answer is 42."
        assert outputs[1].type, InterfaceTypes.BinaryFile
        with open(outputs[1].data.path, "rb") as fp:
            content = fp.read()
            assert content[0] == 222
            assert content[1] == 173
            assert content[2] == 190
            assert content[3] == 239


def test_pipeline_write_read_image():
    pipeline = Pipeline(test_input_dir / "median-filter-test.wasi.wasm")

    data = test_input_dir / "cthead1.png"
    itk_image = itk.imread(data, itk.UC)
    itk_image_dict = itk.dict_from_image(itk_image)
    itkwasm_image = Image(**itk_image_dict)

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.Image, itkwasm_image),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.Image),
    ]

    args = [
        "--memory-io",
        "0",
        "0",
        "--radius",
        "2",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

    out_image = itk.image_from_dict(asdict(outputs[0].data))
    # To be addressed in itk-5.3.1
    out_image.SetRegions([256, 256])

    baseline = itk.imread(test_baseline_dir / "test_pipeline_write_read_image.png")

    difference = np.sum(itk.comparison_image_filter(out_image, baseline))
    assert difference == 0.0


def test_pipeline_dask_array_input():
    pipeline = Pipeline(test_input_dir / "median-filter-test.wasi.wasm")

    data = test_input_dir / "cthead1.png"
    itk_image = itk.imread(data, itk.UC)
    itk_image_dict = itk.dict_from_image(itk_image)
    itkwasm_image = Image(**itk_image_dict)

    from dask.array import from_array

    itkwasm_image.data = from_array(itkwasm_image.data)

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.Image, itkwasm_image),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.Image),
    ]

    args = [
        "--memory-io",
        "0",
        "0",
        "--radius",
        "2",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

    out_image = itk.image_from_dict(asdict(outputs[0].data))
    # To be addressed in itk-5.3.1
    out_image.SetRegions([256, 256])

    baseline = itk.imread(test_baseline_dir / "test_pipeline_write_read_image.png")

    difference = np.sum(itk.comparison_image_filter(out_image, baseline))
    assert difference == 0.0


def test_pipeline_write_read_mesh():
    pipeline = Pipeline(test_input_dir / "mesh-read-write-test.wasi.wasm")

    data = test_input_dir / "cow.vtk"
    itk_mesh = itk.meshread(data)
    itk_mesh_dict = itk.dict_from_mesh(itk_mesh)
    itkwasm_mesh = Mesh(**itk_mesh_dict)

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.Mesh, itkwasm_mesh),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.Mesh),
    ]

    args = [
        "--memory-io",
        "0",
        "0",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

    out_mesh_dict = asdict(outputs[0].data)
    # Native ITK Python binaries require uint64
    out_mesh_dict["cells"] = out_mesh_dict["cells"].astype(np.uint64)
    out_mesh_dict["meshType"]["cellComponentType"] = "uint64"
    out_mesh = itk.mesh_from_dict(out_mesh_dict)

    assert out_mesh.GetNumberOfPoints() == 2903
    assert out_mesh.GetNumberOfCells() == 3263


def test_pipeline_write_read_polydata():
    pipeline = Pipeline(test_input_dir / "mesh-to-poly-data.wasi.wasm")

    data = test_input_dir / "cow.vtk"
    itk_mesh = itk.meshread(data)
    itk_mesh_dict = itk.dict_from_mesh(itk_mesh)
    itkwasm_mesh = Mesh(**itk_mesh_dict)

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.Mesh, itkwasm_mesh),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.PolyData),
    ]

    args = [
        "--memory-io",
        "0",
        "0",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)
    polydata = outputs[0].data

    pipeline = Pipeline(test_input_dir / "poly-data-to-mesh.wasi.wasm")

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.PolyData, polydata),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.Mesh),
    ]

    args = [
        "--memory-io",
        "0",
        "0",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)

    out_mesh_dict = asdict(outputs[0].data)

    # native itk python binaries require uint64
    out_mesh_dict["cells"] = out_mesh_dict["cells"].astype(np.uint64)
    out_mesh_dict["meshType"]["cellComponentType"] = "uint64"
    assert np.isclose(out_mesh_dict["points"][0], 3.71636)
    out_mesh = itk.mesh_from_dict(out_mesh_dict)

    assert out_mesh.GetNumberOfPoints() == 2903
    assert out_mesh.GetNumberOfCells() == 3263

def test_pipeline_write_read_transform():
    pipeline = Pipeline(test_input_dir / "transform-read-write-test.wasi.wasm")

    data = test_input_dir / "LinearTransform.h5"
    itk_transform = itk.transformread(data)
    itk_transform_list = itk.dict_from_transform(itk_transform)
    itkwasm_transform_list = [Transform(**t) for t in itk_transform_list]

    pipeline_inputs = [
        PipelineInput(InterfaceTypes.TransformList, itkwasm_transform_list),
    ]

    pipeline_outputs = [
        PipelineOutput(InterfaceTypes.TransformList),
    ]

    args = [
        "--memory-io",
        "0",
        "0",
    ]

    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)
    transform_list = outputs[0].data

    assert len(transform_list) == 1
    transform = transform_list[0]
    assert transform.transformType.transformParameterization == TransformParameterizations.Affine
    assert transform.transformType.parametersValueType == FloatTypes.Float64
    assert transform.numberOfParameters == 12
    assert transform.numberOfFixedParameters == 3
    np.testing.assert_allclose(transform.fixedParameters, np.array([0.0, 0.0, 0.0]))
    np.testing.assert_allclose(transform.parameters, np.array([
      0.65631490118447, 0.5806583745824385, -0.4817536741017158,
      -0.7407986817430222, 0.37486398378429736, -0.5573995934598175,
      -0.14306664045479867, 0.7227121458012518, 0.676179776908723,
      -65.99999999999997, 69.00000000000004, 32.000000000000036]))