#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt

# -----------------------------
# LOAD DATA
# -----------------------------
M = np.load("mach_field.npy")
mask = np.load("mask.npy")

# Mask invalid
M_plot = np.where(mask, M, np.nan)

# -----------------------------
# BASIC MAP
# -----------------------------
plt.figure(figsize=(16, 8))
plt.imshow(M_plot, cmap="viridis")
plt.colorbar(label="Mach")
plt.title("Global Mach Field (Stage 1)")
plt.tight_layout()
plt.show()

# -----------------------------
# HISTOGRAM
# -----------------------------
valid = M[mask]

plt.figure()
plt.hist(valid, bins=200)
plt.xlabel("Mach")
plt.ylabel("Count")
plt.title("Mach Distribution")
plt.tight_layout()
plt.show()

# -----------------------------
# CLIPPED MAP (BETTER CONTRAST)
# -----------------------------
vmin = np.nanpercentile(M_plot, 5)
vmax = np.nanpercentile(M_plot, 95)

plt.figure(figsize=(16, 8))
plt.imshow(M_plot, vmin=vmin, vmax=vmax, cmap="viridis")
plt.colorbar(label="Mach")
plt.title("Clipped Mach Field (5–95%)")
plt.tight_layout()
plt.show()

# -----------------------------
# OCEAN vs LAND HISTOGRAM
# -----------------------------
import rasterio

with rasterio.open("SRTM15+V2.tiff") as src:
    h = src.read(1)

ocean = (h <= -200)
land = (h >= 200)

M_ocean = M[np.logical_and(mask, ocean)]
M_land  = M[np.logical_and(mask, land)]

plt.figure()
plt.hist(M_ocean, bins=200, alpha=0.5, label="Ocean")
plt.hist(M_land, bins=200, alpha=0.5, label="Land")
plt.legend()
plt.xlabel("Mach")
plt.title("Ocean vs Land Mach Distribution")
plt.tight_layout()
plt.show()

# -----------------------------
# EXPORT TO VTK (FIXED)
# -----------------------------
try:
    import pyvista as pv

    # Use ImageData instead of deprecated UniformGrid
    grid = pv.ImageData()

    # Dimensions must be +1 for cell data
    ny, nx = M.shape
    grid.dimensions = (nx + 1, ny + 1, 1)

    # Flatten in Fortran order to match VTK expectations
    grid.cell_data["Mach"] = M.flatten(order="F")

    grid.save("mach_field.vtk")
    print("Saved mach_field.vtk for ParaView.")

except Exception as e:
    print("VTK export skipped:", e)

# -----------------------------
# CARTOPY MAP (OPTIONAL)
# -----------------------------
try:
    import cartopy.crs as ccrs

    plt.figure(figsize=(16, 8))
    ax = plt.axes(projection=ccrs.Robinson())
    ax.coastlines()

    ax.imshow(
        M_plot,
        transform=ccrs.PlateCarree(),
        cmap="viridis",
        vmin=vmin,
        vmax=vmax,
        origin="upper"
    )

    plt.title("Mach Field (Robinson Projection)")
    plt.tight_layout()
    plt.show()

except Exception as e:
    print("Cartopy plot skipped:", e)