import numpy as np
import pyvista
import os

def transform(src, dest, phasenames):
    vtkfile = src
    grid = pyvista.read(vtkfile)

    phi_vars = [grid.get_array(each) for each in phasenames]
    # convert to numpy array
    phi_vars = [np.array(each) for each in phi_vars]

    # create a new array
    phi = np.zeros_like(phi_vars[0])
    phi[:] = -1.0

    for i, each in enumerate(phi_vars):
        phi[np.where(each > 0.6)] = i

    # add scalar data: phi to grid
    grid["phi"] = phi

    # save as vtu file
    grid.save(dest)

if __name__ == "__main__":

    src_folder = "threephase-4"
    dest_folder = "threephase-4-phi"
    phasenames = ["phi_liquid", "phi_pbpd3", "phi_fcc"]

    if not os.path.exists(dest_folder):
        os.makedirs(dest_folder)

    vtkfiles = [each for each in os.listdir(src_folder) if each.endswith(".vtk")]

    for vtkfile in vtkfiles:
        src = os.path.join(src_folder, vtkfile)
        dest = os.path.join(dest_folder, vtkfile)
        transform(src, dest, phasenames)