import pandas as pd
import numpy as np
import plotly.graph_objects as go
import os
import subprocess
from PIL import Image


"""
Edit the CSV to read so that you can view different hours of the day, e.g. change 
generated_data/satellite_positions_hour_0.csv to generated_data/satellite_positions_hour_7.csv

"""

# Load satellite data
sat_df = pd.read_csv("generated_data/satellite_positions_hour_6.csv")
sun_vec = np.array([
    sat_df["sun_x_km"].iloc[0],
    sat_df["sun_y_km"].iloc[0],
    sat_df["sun_z_km"].iloc[0],
])
sun_vec[0] *= -1  # Flip the x-axis component only
sun_vec_unit = sun_vec / np.linalg.norm(sun_vec)


# Load solar farms
solar_df = pd.read_csv("solar_data/solar_sites.csv").sort_values("Capacity (MW)", ascending=False).head(500)

# Load image from local disk (ensure this file exists!)
img_path = "bluemarble_texture.png"  # <-- put a local Blue Marble RGB image here
img = Image.open(img_path).resize((360, 180))
texture = np.array(img.transpose(Image.FLIP_TOP_BOTTOM))  # flips vertically


# Earth geometry
R = 6371
lons = np.linspace(-180, 180, texture.shape[1])
lats = np.linspace(-90, 90, texture.shape[0])
lon_grid, lat_grid = np.meshgrid(lons, lats)
lon_rad = np.radians(lon_grid)
lat_rad = np.radians(lat_grid)

x = R * np.cos(lat_rad) * np.cos(lon_rad)
y = R * np.cos(lat_rad) * np.sin(lon_rad)
z = R * np.sin(lat_rad)

# Normals and lighting
surface_vectors = np.stack([x, y, z], axis=-1)
norms = np.linalg.norm(surface_vectors, axis=-1, keepdims=True)
unit_vectors = surface_vectors / norms
# Directional lighting from sun
cos_angles = np.dot(unit_vectors, sun_vec_unit)
cos_angles = np.clip(cos_angles, 0, 1)

# Final brightness = ambient + diffuse
lighting = 0.0 + 3 * cos_angles  # tweakable
# Normalize texture to [0, 1] float
texture_float = texture.astype(np.float32) / 255.0

# Apply brightness scaling uniformly to each pixel
shaded_texture = (texture_float * lighting[..., np.newaxis])  # shape: (H, W, 3)

# Clip and convert back to 0–255
shaded_texture = np.clip(shaded_texture * 255, 0, 255).astype(np.uint8)

# Flatten RGB for plotly vertexcolor
rgb_colors = [
    f"rgb({r},{g},{b})"
    for r, g, b in shaded_texture.reshape(-1, 3)
]

# Triangulate globe for Mesh3d
nx, ny = texture.shape[1], texture.shape[0]
vertices = np.column_stack((x.flatten(), y.flatten(), z.flatten()))
triangles = []
for i in range(ny - 1):
    for j in range(nx - 1):
        idx = i * nx + j
        triangles.append([idx, idx + 1, idx + nx])
        triangles.append([idx + 1, idx + nx + 1, idx + nx])
i, j, k = np.array(triangles).T

earth_mesh = go.Mesh3d(
    x=vertices[:, 0],
    y=vertices[:, 1],
    z=vertices[:, 2],
    i=i, j=j, k=k,
    vertexcolor=rgb_colors,
    name="Earth",
    lighting=dict(ambient=1),
    hoverinfo="skip",
    showscale=False,
    opacity=1
)

# Solar farms
lat_sf = np.radians(solar_df["Latitude"])
lon_sf = np.radians(solar_df["Longitude"])
lift = 100
x_sf = (R + lift) * np.cos(lat_sf) * np.cos(lon_sf)
y_sf = (R + lift) * np.cos(lat_sf) * np.sin(lon_sf)
z_sf = (R + lift) * np.sin(lat_sf)

solar_trace = go.Scatter3d(
    x=x_sf, y=y_sf, z=z_sf,
    mode="markers",
    marker=dict(size=3, color="yellow"),
    text=solar_df["Project Name"] + " (" + solar_df["Capacity (MW)"].astype(str) + " MW)",
    hoverinfo="text",
    name="Solar Farms"
)

# Satellites
in_range = sat_df[sat_df["near"]]
out_of_range = sat_df[~sat_df["near"]]

sat_in = go.Scatter3d(
    x=in_range["x_km"], y=in_range["y_km"], z=in_range["z_km"],
    mode="markers",
    marker=dict(size=3, color="green"),
    name="In Range",
    text=in_range.apply(lambda row: (
        f"🛰 Sat ID: {row['satellite_id']}<br>"
        f"⚡ Power: {row['power_transmitted_W']:.6e} W<br>"
        f"☀️ Max Irr: {row["max_irradiance_W_m2"]:.6e} W/m²<br>"
        f"☀️ Spot Radius: {row['r_spot_km']:.3f} km<br>"
        f"🌐 Farm: {row['closest_farm']}<br>"
        f"📐 Phi: {row["phi_deg"]:.2f}°<br>"
        f"📏 Distance: {row['distance_km']:.2f} km"
    ), axis=1),
    hoverinfo="text"
)


sat_out = go.Scatter3d(
    x=out_of_range["x_km"], y=out_of_range["y_km"], z=out_of_range["z_km"],
    mode="markers",
    marker=dict(size=3, color="red"),
    name="Out of Range",
    text=out_of_range.apply(lambda row: f"Sat ID: {row['satellite_id']}", axis=1),
    hoverinfo="text"
)

# Final plot
fig = go.Figure(data=[earth_mesh, sat_in, sat_out, solar_trace])
fig.update_layout(
    title="3D Earth (Sunlit) with Satellites and Solar Farms",
    scene=dict(
        xaxis=dict(visible=False),
        yaxis=dict(visible=False),
        zaxis=dict(visible=False),
        aspectmode="data",
        bgcolor="#000000"
    ),
    margin=dict(l=0, r=0, t=60, b=0),
    paper_bgcolor="#000000",
    font=dict(color="white")
)

# Save and open
html_path = os.path.abspath("satellite_3d_scene.html")
fig.write_html(html_path)

chrome_path = r"C:\Program Files\Google\Chrome\Application\chrome.exe"
subprocess.run([chrome_path, html_path], shell=True)
