##################################################
#              Postprocessing code               #
#              Author: Tim Mulder                #
#               Date: 24.03.2025                 #
##################################################

import  numpy as np
import  io
from    functools import lru_cache
from    datetime import datetime
from    PIL import Image as im

def plot_ABCD(ABCD, output_path):
    ABCD = np.where(ABCD > 0, 155, ABCD)
    ABCD = np.where(ABCD < 0, -100, ABCD)
    ABCD += 100

    # creating image object of
    # above array
    data = im.fromarray(ABCD)
    if data.mode != 'RGB':
        data = data.convert('RGB')

    # saving the final output
    # as a jpg file
    data.save(output_path / 'ABCD_Matrix.jpg')

@lru_cache(maxsize=1)
def _segments_to_vtk(all_verts: tuple[tuple[int]], all_elements: tuple[tuple[int]]) -> str:
    with io.StringIO() as s:
        s.write('# vtk DataFile Version 2.0\n')
        s.write(f'HETS {datetime.now().strftime("%Y-%m-%d")}\n')
        s.write('ASCII\n')
        s.write('DATASET UNSTRUCTURED_GRID\n')
        # Write the elements to the file
        s.write(f'POINTS {len(all_verts)} double\n')
        for vert in all_verts:
            s.write(f'{vert[0]} {vert[1]} {vert[2]}\n')
        s.write(f'CELLS {len(all_elements)} {len(all_elements) * 9}\n')
        for cell in all_elements:
            s.write('8 ' + ' '.join([str(i) for i in cell]) + '\n')
        s.write(f'CELL_TYPES {len(all_elements)}\n')
        # VTK_HEXAHEDRON = cell type 12
        for _ in range(len(all_elements)):
            s.write('12\n')
        return s.getvalue()

# To install.
import pyvista as pv
from vtk import VTK_HEXAHEDRON

def ensure_numpy_array(arr):
    """Ensure input is a NumPy array."""
    if not isinstance(arr, np.ndarray):
        arr = np.array(arr)
    return arr

def data_to_vtu(r_centers, z_centers, dr, dz, data_fields, n_theta=36, filename="revolved.vtu"):
    # --- Inputs ---
    # 1D arrays: cell-centered
    r_centers = ensure_numpy_array(r_centers)   # len N_r
    z_centers = ensure_numpy_array(z_centers)   # len N_z
    dr = ensure_numpy_array(dr)                 # len N_r
    dz = ensure_numpy_array(dz)                 # len N_z

    # Check if all lengths are ok.
    len_r_centers = len(r_centers)
    len_z_centers = len(z_centers)
    len_dr = len(dr)
    len_dz = len(dz)

    if not (len_r_centers == len_z_centers == len_dr == len_dz):
        raise ValueError(
            f"Array length mismatch in data_to_vtu:\n"
            f"  r_centers: {len_r_centers}\n"
            f"  z_centers: {len_z_centers}\n"
            f"  dr:        {len_dr}\n"
            f"  dz:        {len_dz}\n"
            f"All arrays must have the same length."
        )

    # The amount of angular points.
    theta = np.linspace(0, 2 * np.pi, n_theta, endpoint=False)

    # Creating the empty arrays.
    points = []
    cells = []
    cell_types = []
    cell_data_dict = {name: [] for name in data_fields.keys()}

    # Filling the arrays.
    for idx in range(len(r_centers)):
        r_c = r_centers[idx]
        z_c = z_centers[idx]
        dr_i = dr[idx]
        dz_i = dz[idx]

        r1 = r_c - dr_i / 2
        r2 = r_c + dr_i / 2
        z1 = z_c - dz_i / 2
        z2 = z_c + dz_i / 2

        for k in range(n_theta):
            th1 = theta[k]
            th2 = theta[(k + 1) % n_theta]

            def p(r, z, th): return [r * np.cos(th), r * np.sin(th), z]

            pts = [
                p(r1, z1, th1),  # 0
                p(r2, z1, th1),  # 1
                p(r2, z2, th1),  # 2
                p(r1, z2, th1),  # 3
                p(r1, z1, th2),  # 4
                p(r2, z1, th2),  # 5
                p(r2, z2, th2),  # 6
                p(r1, z2, th2),  # 7
            ]

            base_idx = len(points)
            points.extend(pts)

            cell = [8] + [base_idx + i for i in range(8)]
            cells.extend(cell)
            cell_types.append(VTK_HEXAHEDRON)

    # Validate all data_fields arrays have consistent lengths
    expected_length = len(r_centers)  # Should match number of cells
    field_lengths = {name: len(array) for name, array in data_fields.items()}

    # Check if any field has incorrect length
    mismatched_fields = {name: length for name, length in field_lengths.items()
                         if length != expected_length}

    if mismatched_fields:
        error_msg = (
            f"Data field length mismatch in data_to_vtu:\n"
            f"  Expected length (number of cells): {expected_length}\n"
            f"  Mismatched fields:\n"
        )
        for name, length in mismatched_fields.items():
            error_msg += f"    {name}: {length}\n"

        # Also show correctly sized fields for reference
        correct_fields = {name: length for name, length in field_lengths.items()
                         if length == expected_length}
        if correct_fields:
            error_msg += "  Correctly sized fields:\n"
            for name, length in correct_fields.items():
                error_msg += f"    {name}: {length}\n"

        raise ValueError(error_msg)

    # Process data fields with error handling
    for idx in range(expected_length):  # loop over each 2D cell
        for k in range(n_theta):  # revolve into n_theta 3D cells
            for name, array in data_fields.items():
                try:
                    cell_data_dict[name].append(array[idx])
                except IndexError as e:
                    raise IndexError(
                        f"Index error when processing data field '{name}' at idx={idx}:\n"
                        f"  Array length: {len(array)}\n"
                        f"  Requested index: {idx}\n"
                        f"  Expected length: {expected_length}\n"
                        f"  Original error: {e}"
                    ) from e

    # Create VTK grid
    points = np.array(points)
    cells = np.array(cells)

    grid = pv.UnstructuredGrid(cells, cell_types, points)

    # Adding the data
    for name, values in cell_data_dict.items():
        grid.cell_data[name] = np.array(values)

    grid.save(filename)

def write_pvd_file(filenames, timesteps, output_path="results.pvd"):
    from xml.etree.ElementTree import Element, SubElement, ElementTree

    assert len(filenames) == len(timesteps), "Mismatch between file and time lists"

    vtkfile = Element("VTKFile", type="Collection", version="0.1", byte_order="LittleEndian")
    collection = SubElement(vtkfile, "Collection")

    for filename, t in zip(filenames, timesteps):
        SubElement(collection, "DataSet", {
            "timestep": f"{t:.16g}",  # full precision for floating-point
            "group": "",
            "part": "0",
            "file": filename
        })

    tree = ElementTree(vtkfile)
    tree.write(output_path, encoding="utf-8", xml_declaration=True)

def vector_to_vtu(
    r,
    z,
    vec_r,
    vec_z,
    n_theta=36,
    filename='field_vector.vts',
    seed_points = None,
    n_stream_seeds: int = 100,
    vector_name: str = "V",
    step_length: float = 0.5,
):
    """
    Save a single timestep of a 2D (r, z) vector field revolved around the z-axis to a 3D VTU file.

    Parameters:
    - r: 1D array of radial positions, shape (nr,)
    - z: 1D array of axial positions, shape (nz,)
    - vec_r: 1D array of radial vector component, shape (nr * nz,)
    - vec_z: 1D array of axial vector component, shape (nr * nz,)
    - n_theta: number of segments for revolving around the z-axis
    - filename: full path to save the .vtu file
    - n_stream_seeds: how many seed points to randomly place for integration
    - vector_name: name of the vector field
    - step_length: integration step size
    """

    nr, nz = len(r), len(z)
    R, Z = np.meshgrid(r, z, indexing="ij")  # Shape: (nr, nz)
    R_flat = R.ravel()
    Z_flat = Z.ravel()

    # Generate 3D points and vectors by rotating around the z-axis
    theta = np.linspace(0, 2 * np.pi, n_theta, endpoint=False)
    points = []
    vectors = []

    for i in range(nr * nz):
        r_i = R_flat[i]
        z_i = Z_flat[i]
        v_r = vec_r[i]
        v_z = vec_z[i]

        for t in theta:
            x = r_i * np.cos(t)
            y = r_i * np.sin(t)
            points.append([x, y, z_i])
            vx = v_r * np.cos(t)
            vy = v_r * np.sin(t)
            vectors.append([vx, vy, v_z])

    points = np.array(points)
    vectors = np.array(vectors)

    # Create structured grid (theta fastest, then z, then r)
    dims = (n_theta, nz, nr)
    structured = pv.StructuredGrid()
    structured.points = points
    structured.dimensions = dims
    structured.point_data[vector_name] = vectors

    # Convert to unstructured grid (required for stream tracing)
    ugrid = structured.cast_to_unstructured_grid()

    # Generate random seed points within bounding box
    bounds = ugrid.bounds
    seed = pv.PolyData()
    if seed_points is not None:
        seed.points = seed_points
    else:
        seed.points = np.random.uniform(
            low=[bounds[0], bounds[2], bounds[4]],
            high=[bounds[1], bounds[3], bounds[5]],
            size=(n_stream_seeds, 3)
        )

    # Generate streamlines
    streamlines = ugrid.streamlines_from_source(
        vectors=vector_name,
        source=seed,
        integration_direction="both",
        initial_step_length=step_length,
        terminal_speed=1e-5
    )

    # Save as .vtp
    streamlines.save(filename)

    # Returning the seed points so make sure that are always the same for all vtp's.
    return seed.points


