#!/usr/bin/env python

from vtkmodules.vtkIOFides import vtkFidesWriter, vtkFidesReader
from vtkmodules.vtkCommonDataModel import vtkPartitionedDataSetCollection
from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
from vtkmodules.vtkFiltersCore import vtkConvertToMultiBlockDataSet, vtkConvertToPartitionedDataSetCollection
from vtkmodules.vtkFiltersParallelDIY2 import vtkRedistributeDataSetFilter
from vtkmodules.vtkImagingCore import vtkRTAnalyticSource
from vtkmodules.util.misc import vtkGetTempDir
try:
    # Not directly used here, but if it's not imported when VTK_USE_MPI, we get following error:
    # Attempting to use an MPI routine before initializing MPICH
    from mpi4py import MPI
except ImportError:
    pass

import os.path
import shutil
VTK_TEMP_DIR = vtkGetTempDir()

from vtk.test import Testing

class TestFidesWriterBasic(Testing.vtkTest):
    def setUp(self):
        self._cleanup_artifacts()

    def tearDown(self):
        self._cleanup_artifacts()

    def _cleanup_artifacts(self):
        """Removes all ADIOS2 output directories and config files generated by these tests."""
        artifacts = [
            'testDataSet.bp',
            'testPDS.bp',
            'testPDSC.bp',
            'testMB.bp',
            'testExplicitBP5.bp',
            'testExplicitBP4.bp',
            'adios2_bp5.xml',
            'adios2_bp4.xml'
        ]

        for artifact in artifacts:
            path = os.path.join(VTK_TEMP_DIR, artifact)
            if os.path.isdir(path):
                shutil.rmtree(path, ignore_errors=True)
            elif os.path.exists(path):
                os.remove(path)

    def testDataSet(self):
        wavelet = vtkRTAnalyticSource()

        filename = os.path.join(VTK_TEMP_DIR, 'testDataSet.bp')
        writer = vtkFidesWriter()
        writer.SetFileName(filename)
        writer.SetInputConnection(wavelet.GetOutputPort())
        # setting to true but not turning on an array means no array will be written
        writer.SetChooseFieldsToWrite(True)
        writer.Write()

        reader = vtkFidesReader()
        reader.SetFileName(filename)
        reader.UpdateInformation()
        self.assertTrue(reader.GetOutputInformation(0).Has(
            vtkStreamingDemandDrivenPipeline.TIME_STEPS()))
        nsteps = reader.GetOutputInformation(0).Length(vtkStreamingDemandDrivenPipeline.TIME_STEPS())
        self.assertEqual(nsteps, 1)

        self.assertEqual(reader.GetNumberOfPointArrays(), 0)
        reader.Update()

        pdsc = reader.GetOutputDataObject(0)
        self.assertTrue(isinstance(pdsc, vtkPartitionedDataSetCollection))
        pds = pdsc.GetPartitionedDataSet(0)
        nParts = pds.GetNumberOfPartitions()
        self.assertEqual(nParts, 1)

        ds = pds.GetPartition(0)
        self.assertEqual(ds.GetNumberOfCells(), 8000)
        self.assertEqual(ds.GetNumberOfPoints(), 9261)


    def testPartitionedDataSet(self):
        wavelet = vtkRTAnalyticSource()
        # get the data into a vtkPartitionedDataSet
        rdsf = vtkRedistributeDataSetFilter()
        rdsf.SetInputConnection(wavelet.GetOutputPort())
        rdsf.SetNumberOfPartitions(4)

        filename = os.path.join(VTK_TEMP_DIR, 'testPDS.bp')
        writer = vtkFidesWriter()
        writer.SetFileName(filename)
        writer.SetInputConnection(rdsf.GetOutputPort())
        writer.Write()

        reader = vtkFidesReader()
        reader.SetFileName(filename)
        reader.UpdateInformation()
        self.assertTrue(reader.GetOutputInformation(0).Has(
            vtkStreamingDemandDrivenPipeline.TIME_STEPS()))
        nsteps = reader.GetOutputInformation(0).Length(vtkStreamingDemandDrivenPipeline.TIME_STEPS())
        self.assertEqual(nsteps, 1)

        self.assertEqual(reader.GetNumberOfPointArrays(), 1)
        reader.Update()

        pdsc = reader.GetOutputDataObject(0)
        self.assertTrue(isinstance(pdsc, vtkPartitionedDataSetCollection))
        pds = pdsc.GetPartitionedDataSet(0)
        nParts = pds.GetNumberOfPartitions()
        self.assertEqual(nParts, 1)

        ds = pds.GetPartition(0)
        self.assertEqual(ds.GetNumberOfCells(), 8000)
        self.assertEqual(ds.GetNumberOfPoints(), 9261)


    def testPartitionedDataSetCollection(self):
        wavelet = vtkRTAnalyticSource()

        pdsc = vtkConvertToPartitionedDataSetCollection()
        pdsc.SetInputConnection(wavelet.GetOutputPort())

        filename = os.path.join(VTK_TEMP_DIR, 'testPDSC.bp')
        writer = vtkFidesWriter()
        writer.SetFileName(filename)
        writer.SetInputConnection(pdsc.GetOutputPort())
        writer.Write()

        reader = vtkFidesReader()
        reader.SetFileName(filename)
        reader.UpdateInformation()
        self.assertTrue(reader.GetOutputInformation(0).Has(
            vtkStreamingDemandDrivenPipeline.TIME_STEPS()))
        nsteps = reader.GetOutputInformation(0).Length(vtkStreamingDemandDrivenPipeline.TIME_STEPS())
        self.assertEqual(nsteps, 1)

        self.assertEqual(reader.GetNumberOfPointArrays(), 1)
        reader.Update()

        pdsc = reader.GetOutputDataObject(0)
        self.assertTrue(isinstance(pdsc, vtkPartitionedDataSetCollection))
        pds = pdsc.GetPartitionedDataSet(0)
        nParts = pds.GetNumberOfPartitions()
        self.assertEqual(nParts, 1)

        ds = pds.GetPartition(0)
        self.assertEqual(ds.GetNumberOfCells(), 8000)
        self.assertEqual(ds.GetNumberOfPoints(), 9261)


    def testMultiBlockDataSet(self):
        wavelet = vtkRTAnalyticSource()

        mbds = vtkConvertToMultiBlockDataSet()
        mbds.SetInputConnection(wavelet.GetOutputPort())

        filename = os.path.join(VTK_TEMP_DIR, 'testMB.bp')
        writer = vtkFidesWriter()
        writer.SetFileName(filename)
        writer.SetInputConnection(mbds.GetOutputPort())
        writer.Write()

        reader = vtkFidesReader()
        reader.SetFileName(filename)
        reader.UpdateInformation()
        self.assertTrue(reader.GetOutputInformation(0).Has(
            vtkStreamingDemandDrivenPipeline.TIME_STEPS()))
        nsteps = reader.GetOutputInformation(0).Length(vtkStreamingDemandDrivenPipeline.TIME_STEPS())
        self.assertEqual(nsteps, 1)

        self.assertEqual(reader.GetNumberOfPointArrays(), 1)
        reader.Update()

        pdsc = reader.GetOutputDataObject(0)
        # vtkFidesReader always reads in a PDSC, as well as vtkFidesWriter always converts MB to PDSC
        self.assertTrue(isinstance(pdsc, vtkPartitionedDataSetCollection))
        pds = pdsc.GetPartitionedDataSet(0)
        nParts = pds.GetNumberOfPartitions()
        self.assertEqual(nParts, 1)

        ds = pds.GetPartition(0)
        self.assertEqual(ds.GetNumberOfCells(), 8000)
        self.assertEqual(ds.GetNumberOfPoints(), 9261)

    def testAdiosConfigFile(self):
        wavelet = vtkRTAnalyticSource()
        wavelet.Update()

        # First, explicitly request BP5
        bp5_config_path = os.path.join(VTK_TEMP_DIR, 'adios2_bp5.xml')
        with open(bp5_config_path, 'w') as f:
            f.write("""<?xml version="1.0"?>
<adios-config>
    <io name="fides-write-io">
        <engine type="BP5"/>
    </io>
</adios-config>""")

        bp5_filename = os.path.join(VTK_TEMP_DIR, 'testExplicitBP5.bp')

        writer5 = vtkFidesWriter()
        writer5.SetFileName(bp5_filename)
        writer5.SetInputConnection(wavelet.GetOutputPort())
        writer5.SetAdiosConfigFile(bp5_config_path)
        writer5.Write()

        # Verify BP5 files
        self.assertTrue(os.path.exists(os.path.join(bp5_filename, 'data.0')), "BP5 write failed completely.")
        self.assertTrue(os.path.exists(os.path.join(bp5_filename, 'mmd.0')), "BP5 config ignored: mmd.0 is missing.")

        # Next, explicitly request BP4
        bp4_config_path = os.path.join(VTK_TEMP_DIR, 'adios2_bp4.xml')
        with open(bp4_config_path, 'w') as f:
            f.write("""<?xml version="1.0"?>
<adios-config>
    <io name="fides-write-io">
        <engine type="BP4"/>
    </io>
</adios-config>""")

        bp4_filename = os.path.join(VTK_TEMP_DIR, 'testExplicitBP4.bp')

        writer4 = vtkFidesWriter()
        writer4.SetFileName(bp4_filename)
        writer4.SetInputConnection(wavelet.GetOutputPort())
        writer4.SetAdiosConfigFile(bp4_config_path)
        writer4.Write()

        # Verify BP4 files (ADIOS2 docs state that BP4 does not include mmd.0)
        self.assertTrue(os.path.exists(os.path.join(bp4_filename, 'data.0')), "BP4 write failed completely.")
        self.assertFalse(os.path.exists(os.path.join(bp4_filename, 'mmd.0')), "BP4 config ignored: mmd.0 was generated.")


if __name__ == "__main__":
    Testing.main([(TestFidesWriterBasic, 'test')])
