import adios2
import numpy
import sys

from contextlib import contextmanager
from vtkmodules.vtkIOIOSS import vtkIOSSReader
from vtkmodules.vtkCommonDataModel import (
    vtkCompositeDataSet,
    vtkPartitionedDataSetCollection,
    vtkUnstructuredGrid,
)
from vtkmodules.vtkIOParallelXML import vtkXMLPartitionedDataSetCollectionWriter
from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
from vtkmodules.numpy_interface import dataset_adapter as dsa


@contextmanager
def managed_adios2_open(io, *args, **kwargs):
    engine = io.Open(*args, **kwargs)
    try:
        yield engine
    finally:
        engine.Close()


# the word 'block' can be confusied with ADIOS2 blocks. Replace all occurrences of block_ with mesh_
def convert_block_name(name: str) -> str:
    return name.replace('block', 'mesh')


def write_vtpc(datasetCollection: vtkPartitionedDataSetCollection, filename: str):
    """
        Writes a vtkPartitionedDataSetCollection to a .vtpc file.
    Args:
        datasetCollection (vtkPartitionedDataSetCollection): dataset collection
        filename (str): output file name
    """
    writer = vtkXMLPartitionedDataSetCollectionWriter()
    writer.SetFileName(filename)
    writer.SetInputData(datasetCollection)
    writer.Write()


def write_vtk_ugrid_to_bp_file(
    mesh: vtkUnstructuredGrid, step: int, group_name: str, bpIO, fh
):
    cells = mesh.GetCells()
    connectivity = dsa.vtkDataArrayToVTKArray(cells.GetConnectivityArray())
    offsets = dsa.vtkDataArrayToVTKArray(cells.GetOffsetsArray())
    # cellCounts can't be int64_t because vtkFidesReader tries to cast from long long to int.
    # ParaView Error -
    #   terminate called after throwing an instance of 'vtkm::cont::ErrorBadType'
    #   what():  Cast failed: vtkm::cont::ArrayHandle<long long, vtkm::cont::StorageTagBasic> --> vtkm::cont::ArrayHandle<int, vtkm::cont::StorageTagBasic>
    cellCounts = numpy.array(offsets[1:] - offsets[:-1], dtype=numpy.int32)
    cellTypes = dsa.vtkDataArrayToVTKArray(mesh.GetCellTypesArray())
    points = dsa.vtkDataArrayToVTKArray(mesh.GetPoints().GetData())
    all_fields = {"points": [], "cell_set": []}
    # vtk data model permits array with same name in both point and cell data.
    # In the CGNS dataset, both PointData and CellData have an array named 'ids'.
    # However, this is not allowable in the ADIOS file.
    processed_arrays = set()
    all_field_names = []
    all_field_associations = []

    for association, data_attributes in zip(["points", "cell_set"], [mesh.GetPointData(), mesh.GetCellData()]):
        for i in range(data_attributes.GetNumberOfArrays()):
            arr_name = data_attributes.GetArrayName(i)
            if data_attributes.GetArray(i).GetNumberOfComponents() > 3:
                # fides ReadVariableInternal only allocates arrays for variables with upto 3 components
                continue
            if arr_name in processed_arrays:
                # Prepend association to distinguish it from the other.
                arr_name = association + '_' + arr_name
            all_fields[association].append(
                {
                    "name": arr_name,
                    "data": dsa.vtkDataArrayToVTKArray(data_attributes.GetArray(i)),
                    "association": association
                }
            )
            processed_arrays.add(arr_name)
            all_field_names.append(arr_name)
            all_field_associations.append(association)

    bpIO.DefineAttribute("Fides_Data_Model", "unstructured")
    bpIO.DefineAttribute("Fides_Variable_List", all_field_names)
    bpIO.DefineAttribute("Fides_Variable_Associations", all_field_associations)

    if step == 0:
        # define variables for the first time
        pointsVar = bpIO.DefineVariable(
            group_name + "/" + "points",
            points,
            list(points.shape),
            [0, 0],
            list(points.shape),
        )
        conVar = bpIO.DefineVariable(
            group_name + "/" + "connectivity",
            connectivity,
            [connectivity.size],
            [0],
            [connectivity.size],
        )
        cellTypesVar = bpIO.DefineVariable(
            group_name + "/" + "cell_types",
            cellTypes,
            [cellTypes.size],
            [0],
            [cellTypes.size],
        )
        numVertsVar = bpIO.DefineVariable(
            group_name + "/" + "num_verts",
            cellCounts,
            [cellCounts.size],
            [0],
            [cellCounts.size],
        )
        vars = []
        for association, fields in all_fields.items():
            for field in fields:
                vars.append(
                    bpIO.DefineVariable(
                    group_name + "/" + field["name"],
                    field["data"],
                    list(field["data"].shape),
                    [0] * len(field["data"].shape),
                    list(field["data"].shape))
                )
    else:
        # reuse variables
        pointsVar = bpIO.InquireVariable(group_name + "/" + "points")
        conVar = bpIO.InquireVariable(group_name + "/" + "connectivity")
        cellTypesVar = bpIO.InquireVariable(group_name + "/" + "cell_types")
        numVertsVar = bpIO.InquireVariable(group_name + "/" + "num_verts")
        vars = []
        for association, fields in all_fields.items():
            for field in fields:
                vars.append(bpIO.InquireVariable(group_name + "/" + field["name"]))


    fh.Put(conVar, connectivity)
    fh.Put(numVertsVar, cellCounts)
    fh.Put(cellTypesVar, cellTypes)
    fh.Put(pointsVar, points)
    var_iter = iter(vars)
    for association, fields in all_fields.items():
        for field in fields:
            fh.Put(next(var_iter), field["data"])


def write_step(vtpc: vtkPartitionedDataSetCollection, step: int, assembly_selector: str, bpIO, fh):
    """
        Writes a vtkPartitionedDataSetCollection to a .bp (binary-pack) directory/file
    Args:
        datasetCollection (vtkPartitionedDataSetCollection): dataset collection
        filename (str): output file name
    """
    dasm = vtpc.GetDataAssembly()
    pivot_node = dasm.GetFirstNodeByPath(assembly_selector)
    if not assembly_selector.startswith('/'):
        pivot_node = dasm.GetFirstNodeByPath('/' + assembly_selector)
    if pivot_node < 0:
        print(f"Unable to find a node at {assembly_selector}")
        return
    children = dasm.GetChildNodes(pivot_node)

    fh.BeginStep()

    for child in children:
        for pd_idx in dasm.GetDataSetIndices(child):
            if vtpc.HasMetaData(pd_idx):
                name = assembly_selector + '/' + convert_block_name(vtpc.GetMetaData(pd_idx).Get(vtkCompositeDataSet.NAME()))
            else:
                name = f"{assembly_selector} + /mesh_{child}"
            for partition in range(vtpc.GetNumberOfPartitions(pd_idx)):
                write_vtk_ugrid_to_bp_file(vtpc.GetPartition(pd_idx, partition), step, name, bpIO, fh)
                # Write only the first partition
                break

    fh.EndStep()


def convert_ioss_file_to_bp_file(infile: str, outfile: str, assembly_selector: str):
    """Convert an ioss file to vtkPartitionedDataSetCollection and write it to a bp file

    Args:
        infile (str): input filename
        outfile (str): output filename
        assembly_selector (str): path to a node in the data assembly. all datasets under this node will be written out to bp file
    """
    ioss_reader = vtkIOSSReader()
    ioss_reader.SetFileName(infile)
    ioss_reader.UpdateInformation()

    adios = adios2.ADIOS()  # type: ignore
    bpIO = adios.DeclareIO("BPFile")
    bpIO.SetEngine("bp5")

    with managed_adios2_open(bpIO, outfile, adios2.Mode.Write) as fh:
        info = ioss_reader.GetOutputInformation(0)
        timesteps = info.Get(vtkStreamingDemandDrivenPipeline.TIME_STEPS())

        # for each timestep
        for step, t in enumerate(timesteps):
            ioss_reader.UpdateTimeStep(t)
            vtpc = ioss_reader.GetOutputDataObject(0)
            print(f"Write t[{step}]: {t}")
            write_step(vtpc, step, assembly_selector, bpIO, fh)


if __name__ == "__main__":
    # The cgns file has many meshes that are shared by different assemblies.
    # If all of them were written out, meshes would be duplicated under different groups
    # So, write only the meshes which are inside the selected assembly
    # "IOSS/assemblies/Conglomerate/Top/Prime"
    convert_ioss_file_to_bp_file(sys.argv[1], sys.argv[2], assembly_selector="IOSS/assemblies/Conglomerate/Top/Prime")
