import os
import re
import vtkmodules.all as vtk
from vtkmodules.util import numpy_support
import numpy as np
import imageio


def read_vtk(filename):
    reader = vtk.vtkUnstructuredGridReader()
    reader.SetFileName(filename)
    reader.Update()
    return reader.GetOutput()


def create_actor(dataset, color):
    mapper = vtk.vtkDataSetMapper()
    mapper.SetInputData(dataset)

    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    actor.GetProperty().SetColor(color)
    actor.GetProperty().SetDiffuse(0.8)
    actor.GetProperty().SetSpecular(0.3)
    actor.GetProperty().SetSpecularPower(20)
    return actor


def create_glyph_actor(dataset, radius, color):
    # Sphere to use as glyph
    sphere_source = vtk.vtkSphereSource()
    sphere_source.SetRadius(radius)
    sphere_source.SetThetaResolution(5)
    sphere_source.SetPhiResolution(5)
    sphere_source.Update()

    glyph = vtk.vtkGlyph3D()
    glyph.SetSourceConnection(sphere_source.GetOutputPort())
    glyph.SetInputData(dataset)
    glyph.SetScaleModeToDataScalingOff()  # same size spheres
    glyph.Update()

    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(glyph.GetOutputPort())

    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    actor.GetProperty().SetColor(color)
    actor.GetProperty().SetDiffuse(0.8)
    actor.GetProperty().SetSpecular(0.3)
    actor.GetProperty().SetSpecularPower(20)
    actor.GetProperty().SetInterpolationToPhong()
    return actor


def list_folders(directory):
    try:
        # List all entries in the directory
        entries = os.listdir(directory)
        # Filter only directories
        folders = [
            entry for entry in entries
            if os.path.isdir(os.path.join(directory, entry))
        ]
        return folders
    except FileNotFoundError:
        print(f"Error: Directory '{directory}' does not exist.")
        return []
    except PermissionError:
        print(f"Error: Permission denied to access '{directory}'.")
        return []


def numeric_key(filename):
    # Extract number between last underscore and .vtk
    match = re.search(r'_(\d+)\.vtk$', filename)
    if match:
        return int(match.group(1))
    else:
        return -1  # or something to put files without numbers first or last


def render_to_image(render_window):
    window_to_image = vtk.vtkWindowToImageFilter()
    window_to_image.SetInput(render_window)
    window_to_image.Update()

    vtk_image = window_to_image.GetOutput()
    width, height, _ = vtk_image.GetDimensions()

    vtk_array = vtk_image.GetPointData().GetScalars()
    components = vtk_array.GetNumberOfComponents()
    arr = numpy_support.vtk_to_numpy(vtk_array).reshape(
        height, width, components)
    arr = np.flip(arr, 0)  # Flip vertical

    return arr


def main(folder_path, output_gif):
    files = sorted([f for f in os.listdir(folder_path) if f.endswith('.vtk')],
                   key=numeric_key)

    # Create renderer and render window
    renderer = vtk.vtkRenderer()
    render_window = vtk.vtkRenderWindow()
    render_window.SetOffScreenRendering(1)
    render_window.AddRenderer(renderer)
    render_window.SetSize(600, 600)

    # Create render window interactor (not really needed for offscreen)
    interactor = vtk.vtkRenderWindowInteractor()
    interactor.SetRenderWindow(render_window)

    images = []

    for i, f in enumerate(files):
        renderer.RemoveAllViewProps()  # Clear previous actors

        polydata = read_vtk(os.path.join(folder_path, f))
        color = (
            0.1, 0.4, 0.8
        )
        actor = create_glyph_actor(polydata, radius=0.015, color=color)
        renderer.AddActor(actor)

        # Set background white
        renderer.SetBackground(1, 1, 1)

        # Reset camera so we see the object well
        renderer.ResetCamera()
        camera = renderer.GetActiveCamera()
        camera.SetPosition(4, 1, 4)
        camera.SetFocalPoint(0, 0, 0)
        camera.SetViewUp(0, 1, 0)
        renderer.ResetCameraClippingRange()
        render_window.Render()

        img = render_to_image(render_window)
        images.append(img)

    # Save GIF
    imageio.mimsave(output_gif, images, loop=0, duration=0.2)
    print(f"GIF saved to {output_gif}")


if __name__ == "__main__":

    folder = "inductiva_output"

    sims_folder = list_folders(folder)
    for i, sim_folder in enumerate(sims_folder):
        main(folder + "/" + sim_folder + "/outputs/vtk",
             f"gifs/output_{i}.gif")