#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
import rasterio

# -----------------------------
# CONFIG
# -----------------------------
DEM_FILE = "SRTM15+V2.tiff"

# -----------------------------
# LOAD SHAPE
# -----------------------------
with rasterio.open(DEM_FILE) as src:
    h = src.read(1)
    ny, nx = h.shape

# -----------------------------
# LOAD MEMMAP DATA
# -----------------------------
M_plus = np.memmap("mach_plus.dat", dtype="float32", mode="r", shape=(ny, nx))
M_minus = np.memmap("mach_minus.dat", dtype="float32", mode="r", shape=(ny, nx))

valid_plus = np.memmap("valid_plus.dat", dtype="bool", mode="r", shape=(ny, nx))
valid_minus = np.memmap("valid_minus.dat", dtype="bool", mode="r", shape=(ny, nx))

# Mask invalid
Mp = np.where(valid_plus, M_plus, np.nan)
Mm = np.where(valid_minus, M_minus, np.nan)

# -----------------------------
# SUMMARY STATS
# -----------------------------
def summarize(name, M, valid):
    vals = M[valid]
    print(f"{name}:")
    print(f"  valid fraction: {valid.mean():.4f}")
    print(f"  mean: {vals.mean():.4f}")
    print(f"  std: {vals.std():.4f}")

summarize("M_plus", M_plus, valid_plus)
summarize("M_minus", M_minus, valid_minus)

# -----------------------------
# CLIPPING RANGE (shared)
# -----------------------------
combined = np.concatenate([
    M_plus[valid_plus],
    M_minus[valid_minus]
])

vmin = np.percentile(combined, 5)
vmax = np.percentile(combined, 95)

# -----------------------------
# SIDE-BY-SIDE MAPS
# -----------------------------
plt.figure(figsize=(18, 8))

plt.subplot(1, 2, 1)
plt.imshow(Mp, cmap="viridis", vmin=vmin, vmax=vmax)
plt.title("Mach Field M+")
plt.colorbar()

plt.subplot(1, 2, 2)
plt.imshow(Mm, cmap="viridis", vmin=vmin, vmax=vmax)
plt.title("Mach Field M-")
plt.colorbar()

plt.tight_layout()
plt.show()

# -----------------------------
# DIFFERENCE MAP
# -----------------------------
diff = Mp - Mm

plt.figure(figsize=(16, 8))
plt.imshow(diff, cmap="coolwarm", vmin=-0.5, vmax=0.5)
plt.colorbar(label="M+ - M-")
plt.title("Directional Difference (M+ - M-)")
plt.tight_layout()
plt.show()

# -----------------------------
# HISTOGRAM COMPARISON
# -----------------------------
plt.figure()

plt.hist(M_plus[valid_plus], bins=200, alpha=0.5, label="M+")
plt.hist(M_minus[valid_minus], bins=200, alpha=0.5, label="M-")

plt.legend()
plt.xlabel("Mach")
plt.ylabel("Count")
plt.title("Mach Distribution (M+ vs M-)")
plt.tight_layout()
plt.show()

# -----------------------------
# OCEAN vs LAND (for each direction)
# -----------------------------
ocean = (h <= -200)
land = (h >= 200)

plt.figure()

plt.hist(M_plus[np.logical_and(valid_plus, ocean)], bins=200, alpha=0.5, label="Ocean M+")
plt.hist(M_plus[np.logical_and(valid_plus, land)], bins=200, alpha=0.5, label="Land M+")

plt.legend()
plt.title("Ocean vs Land (M+)")
plt.tight_layout()
plt.show()

plt.figure()

plt.hist(M_minus[np.logical_and(valid_minus, ocean)], bins=200, alpha=0.5, label="Ocean M-")
plt.hist(M_minus[np.logical_and(valid_minus, land)], bins=200, alpha=0.5, label="Land M-")

plt.legend()
plt.title("Ocean vs Land (M-)")
plt.tight_layout()
plt.show()