#!/usr/bin/env python3
import pandas as pd
import numpy as np
import plotly.express as px
from umap import UMAP
import plotly.io as pio

# Set plotly to render in browser for better interaction
pio.renderers.default = "browser"

# Load the data
print("Loading data...")
df = pd.read_csv("magnetic_features_with_density.csv")

# Separate features and target
target = "magnetic_density"
# total_magnetic_moment
feature_cols = [col for col in df.columns if col.startswith("orb_feat_")]
X = df[feature_cols].values
y = df[target].values

print(f"Performing UMAP on {len(feature_cols)} features...")
# Initialize and fit UMAP
umap_3d = UMAP(
    n_components=3,
    n_neighbors=15,
    min_dist=0.1,
    metric="euclidean",
    target_metric="euclidean",
    target_weight=0.5,
    random_state=42,
)

# Fit and transform the data
embedding = umap_3d.fit_transform(
    X,
)

# Create a new dataframe with the embeddings
viz_df = pd.DataFrame(
    {
        "UMAP1": embedding[:, 0],
        "UMAP2": embedding[:, 1],
        "UMAP3": embedding[:, 2],
        target: y,
        "MP ID": df["mp_id"],
        "formula": df["cif_path"].str.replace("cif_path/", ""),
    }
)

print("Creating 3D visualization...")
# Create an interactive 3D scatter plot
fig = px.scatter_3d(
    viz_df,
    x="UMAP1",
    y="UMAP2",
    z="UMAP3",
    color=target,
    hover_data=["formula"],
    color_continuous_scale="Viridis",
    title=f"UMAP visualization colored by {target}",
)

# Update the layout for better visualization
fig.update_layout(
    scene=dict(xaxis_title="UMAP1", yaxis_title="UMAP2", zaxis_title="UMAP3"),
    width=1000,
    height=800,
)

# Show the plot
fig.show()

# Save the plot as HTML for later viewing
fig.write_html("magnetic_umap_visualization.html")
print("Visualization saved as 'magnetic_umap_visualization.html'")