#%%
import torch
import os
from torch_geometric.data import Data
import matplotlib.pyplot as plt
from torch_geometric.utils import to_networkx
import plotly.graph_objects as go
import pyvista as pv
import numpy as np

def save_graph_to_paraview(data: Data, filename_prefix="graph"):

    # Convert PyTorch tensors to NumPy arrays
    points = data.pos.numpy()  # Node positions
    src, dst = data.edge_index.numpy()  # Edge indices
    edge_attr = data.edge_attr.numpy()  # Edge attributes

    # Create a PolyData object for nodes
    point_cloud = pv.PolyData(points)
    
    # If node features exist, store them as attributes
    if data.x is not None:
        node_features = data.x.numpy()
        for i in range(node_features.shape[1]):  # Add each feature separately
            point_cloud[f"feature_{i}"] = node_features[:, i]

    # Create line connectivity for edges
    num_edges = src.shape[0]
    lines = np.empty((num_edges, 3), dtype=np.int32)
    lines[:, 0] = 2  # Each edge has 2 points
    lines[:, 1] = src
    lines[:, 2] = dst
    lines = lines.flatten()

    # Create an UnstructuredGrid for edges
    edges = pv.PolyData()
    edges.points = points
    edges.lines = lines

    # Attach edge attributes
    if edge_attr is not None:
        for i in range(edge_attr.shape[1]):  # Store each edge attribute separately
            edges[f"edge_attr_{i}"] = edge_attr[:, i]

    # Save as VTK for ParaView
    nodes_file = f"{filename_prefix}_nodes.vtp"
    edges_file = f"{filename_prefix}_edges.vtp"
    point_cloud.save(nodes_file)
    edges.save(edges_file)

    print(f"Saved graph to '{nodes_file}' and '{edges_file}' for ParaView.")

def visualize_graph_plotly(data, feature_index=0, edge_threshold=None):

    print(f"Creating interactive visualization for {data.num_nodes} nodes and {data.num_edges} edges...")

    # Get node positions
    x_coords = data.pos[:, 0].tolist()
    y_coords = data.pos[:, 1].tolist()
    z_coords = data.pos[:, 2].tolist()

    # Get node colors based on the chosen feature (normalize for better visualization)
    feature_values = data.x[:, feature_index].numpy()  # Extract Ux or any other feature
    feature_values = (feature_values - feature_values.min()) / (feature_values.max() - feature_values.min())  # Normalize

    # Create node scatter plot with colors
    node_trace = go.Scatter3d(
        x=x_coords, y=y_coords, z=z_coords,
        mode='markers',
        marker=dict(size=3, color=feature_values, colorscale='Viridis', opacity=1.0, colorbar=dict(title="Feature Value")),
        name='Nodes'
    )

    # Create edge lines
    edge_x, edge_y, edge_z = [], [], []
    for i in range(data.edge_index.shape[1]):
        src, dst = data.edge_index[:, i].tolist()
        p1, p2 = data.pos[src], data.pos[dst]

        # Apply distance threshold (if given)
        if edge_threshold is None or torch.norm(p1 - p2).item() <= edge_threshold:
            edge_x.extend([p1[0].item(), p2[0].item(), None])
            edge_y.extend([p1[1].item(), p2[1].item(), None])
            edge_z.extend([p1[2].item(), p2[2].item(), None])

    edge_trace = go.Scatter3d(
        x=edge_x, y=edge_y, z=edge_z,
        mode='lines',
        line=dict(color='gray', width=1.0),
        name='Edges'
    )

    # Create figure
    fig = go.Figure(data=[edge_trace, node_trace])
    fig.update_layout(title="Interactive 3D Graph with Node Feature Coloring", showlegend=False)

    # Show interactive plot
    fig.show()





time = "0.533713"  # Your time value
graph_dir = "graph_data"

# Load the node features, edge indices, and node positions
node_features = torch.load(os.path.join(graph_dir, f"node_features_{time}.pt"))
edge_index = torch.load(os.path.join(graph_dir, f"edge_index_{time}.pt"))
node_positions = torch.load(os.path.join(graph_dir, f"node_positions_{time}.pt"))
src, dst = edge_index
edge_attr = node_positions[dst] - node_positions[src]

# Print shape information
print(f"Node features shape: {node_features.shape}")
print(f"Node positions shape: {node_positions.shape}")
print(f"Edge index shape: {edge_index.shape}")
print("Edge attributes shape:", edge_attr.shape)  # Should be (1920, 3)

data = Data(
    x=node_features,           # Node features (velocity)
    edge_index=edge_index,     # Edge indices
    edge_attr=edge_attr,
    pos=node_positions         # Node positions
)

visualize_graph_plotly(data, feature_index=0)
save_graph_to_paraview(data, filename_prefix="graph")

#%%