import pyvista as pv

import sys
import numpy as np
import time

filename = "position_buffer.txt"

all_worm_points = []
boundary_points = []

BOUNDARY_TYPE = 3

all_point_types = []


def load_points(filename):
    global all_worm_points
    global boundary_points
    global all_point_types

    worm_points = []
    types = []
    line_count = 0

    pcount = 0

    time_count = 0

    logStep = None

    include_boundary = True

    for l in open(filename):
        ws = l.split()
        # print(ws)
        if line_count == 6:
            numOfElasticP = int(ws[0])
        if line_count == 7:
            numOfLiquidP = int(ws[0])
        if line_count == 8:
            numOfBoundaryP = int(ws[0])
        if line_count == 9:
            timeStep = float(ws[0])
        if line_count == 10:
            logStep = int(ws[0])

        if len(ws) == 4:
            type = float(ws[3])

            if not type == BOUNDARY_TYPE:
                worm_points.append([float(ws[0]), float(ws[1]), float(ws[2])])
                types.append(type)
            else:
                if include_boundary:
                    boundary_points.append([float(ws[0]), float(ws[1]), float(ws[2])])

        if logStep is not None:
            pcount += 1

            if pcount == numOfBoundaryP + numOfElasticP + numOfLiquidP:
                print(
                    "End of one batch: %i worm points added, %i boundary points, %i total points at line %i, time: %i"
                    % (
                        len(worm_points),
                        len(boundary_points),
                        pcount,
                        line_count,
                        time_count,
                    )
                )
                all_worm_points.append(worm_points)

                all_point_types.append(types)

                worm_points = []
                types = []
                pcount = 0
                numOfBoundaryP = 0

                time_count += 1

        line_count += 1

    # all_points_np = np.array(all_points)

    print(
        f"Loaded positions with %i elastic, %i liquid and %i boundary points (%i total), %i lines"
        % (
            numOfElasticP,
            numOfLiquidP,
            numOfBoundaryP,
            numOfElasticP + numOfLiquidP + numOfBoundaryP,
            line_count,
        )
    )

    print("Num of time points found: %i" % len(all_worm_points))

    return all_worm_points, all_point_types


def toggle(flag):
    print("Setting to: %s" % flag)


def create_plotter():
    pl = pv.Plotter(window_size=[1200, 800])
    pl.set_background("#bbbbbb")

    pl.camera.position = (500, 300, 150.0)
    pl.camera.focal_point = (50, 0, 150)
    pl.camera.up = (0.0, 1.0, 0.0)
    pl.add_axes_at_origin(labels_off=True)
    # pl.add_mesh(pv.Cube(center=(50,0,50)))

    # pl.add_checkbox_button_widget(toggle, value=True)

    return pl


last_mesh = None


def create_mesh(step, plotter_=None):
    colours = {1.1: "lightblue", 2.1: "green", 2.2: "turquoise", 3: "#eeeeee"}
    colours = {1.1: "blue", 2.2: "#008080"}
    step_count = step
    value = step_count

    global last_mesh

    global all_worm_points
    global boundary_points
    global all_point_types

    render_here = False

    if plotter_ == None:
        global plotter
        render_here = True
    else:
        plotter = plotter_

    index = int(value)

    print("Changing to time point: %s (%s) " % (index, value))

    curr_points = all_worm_points[index]
    curr_types = all_point_types[index]

    if last_mesh is None:
        last_mesh = pv.PolyData(curr_points)
        last_mesh["types"] = curr_types
        print(last_mesh)
        print("Initial add of worm mesh")
        last_actor = plotter.add_mesh(
            last_mesh,
            render_points_as_spheres=False,
            cmap=[c for c in colours.values()],
            point_size=3,
            show_scalar_bar=False,
        )

        if len(boundary_points) > 0:
            boundary_mesh = pv.PolyData(boundary_points)
            # boundary_mesh["types"] = curr_t[ypes
            print(boundary_mesh)
            print("Initial add of boundary mesh")
            boundary_actor = plotter.add_mesh(
                boundary_mesh,
                render_points_as_spheres=False,
                point_size=3,
                show_scalar_bar=False,
            )
    else:
        print("Updating mesh...")
        last_mesh.points = curr_points

    plotter.render()

    time.sleep(0.1)

    return


if __name__ == "__main__":
    if len(sys.argv) == 2:
        filename = sys.argv[1]

    load_points(filename)

    plotter = create_plotter()
    create_mesh(0)

    max_time = len(all_worm_points) - 1

    plotter.add_slider_widget(
        create_mesh, rng=[0, max_time], value=max_time, title="Time point"
    )
    plotter.add_timer_event(
        max_steps=len(all_worm_points), duration=200, callback=create_mesh
    )

    plotter.show()