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

import itk
import numpy as np

from itkwasm import InterfaceTypes, TextStream, BinaryStream, PipelineInput, PipelineOutput, Pipeline, TextFile, BinaryFile, Image, Mesh

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([])

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)

    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

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)

        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

        # Currently required per https://github.com/bytecodealliance/wasmtime-py/issues/132
        del pipeline

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_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
