run_clay_embeddings.py

"""
Clay Foundation Model - GeoAI Embeddings with Cosine Similarity Validation

Fetches Sentinel-2 imagery for a small region, generates embeddings using
Clay v1.5, and validates via cosine similarity between patches.
"""

import os
import sys
from pathlib import Path

import numpy as np
import torch
import yaml
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.metrics.pairwise import cosine_similarity

try:
    import pystac_client
    import planetary_computer
    import stackstac
except ImportError:
    print("Install dependencies: pip install pystac-client planetary-computer stackstac")
    sys.exit(1)

try:
    from claymodel.module import ClayMAEModule
except ImportError:
    print("Install Clay model: pip install git+https://github.com/Clay-foundation/model.git")
    sys.exit(1)


CHECKPOINT_PATH = os.environ.get("CLAY_CHECKPOINT", "v1.5/clay-v1.5.ckpt")
METADATA_PATH = "configs/metadata.yaml"

# Small region: Bangalore, India (urban + vegetation mix)
BBOX = [77.55, 12.93, 77.65, 13.03]  # [west, south, east, north]
LATLON_CENTER = [12.98, 77.60]
TIME_RANGE = "2024-01-01/2024-03-31"
MAX_CLOUD_COVER = 10
PATCH_SIZE = 256

S2_BANDS = [
    "B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B11", "B12"
]
PLATFORM = "sentinel-2-l2a"


def load_metadata():
    with open(METADATA_PATH) as f:
        return yaml.safe_load(f)


def fetch_sentinel2_data():
    """Fetch Sentinel-2 data for the configured region from Planetary Computer."""
    print(f"Fetching Sentinel-2 data for bbox={BBOX}, time={TIME_RANGE}...")

    catalog = pystac_client.Client.open(
        "https://planetarycomputer.microsoft.com/api/stac/v1",
        modifier=planetary_computer.sign_inplace,
    )

    search = catalog.search(
        collections=["sentinel-2-l2a"],
        bbox=BBOX,
        datetime=TIME_RANGE,
        query={"eo:cloud_cover": {"lt": MAX_CLOUD_COVER}},
    )

    items = list(search.items())
    if not items:
        print("No items found. Try expanding the time range or cloud cover threshold.")
        sys.exit(1)

    print(f"Found {len(items)} scenes. Using the first one.")
    item = items[0]
    print(f"  Scene: {item.id}")
    print(f"  Date: {item.datetime}")

    stack = stackstac.stack(
        [item],
        assets=S2_BANDS,
        resolution=10,
        bounds_latlon=BBOX,
        dtype="float64",
        rescale=False,
        epsg=32643,
    )

    data = stack.compute()
    print(f"  Data shape: {data.shape}")

    image = data.values[0].astype(np.float32)
    return image, item.datetime


def create_patches(image, patch_size=PATCH_SIZE):
    """Split image into non-overlapping patches of patch_size x patch_size."""
    bands, h, w = image.shape
    patches = []
    coords = []

    n_rows = h // patch_size
    n_cols = w // patch_size

    for i in range(n_rows):
        for j in range(n_cols):
            patch = image[
                :,
                i * patch_size : (i + 1) * patch_size,
                j * patch_size : (j + 1) * patch_size,
            ]
            if not np.isnan(patch).any() and patch.std() > 0:
                patches.append(patch)
                coords.append((i, j))

    print(f"Created {len(patches)} valid patches of size {patch_size}x{patch_size}")
    return patches, coords


def normalize_patches(patches, metadata):
    """Normalize patches using Clay's Sentinel-2 statistics from metadata."""
    sensor = metadata[PLATFORM]
    means = [sensor["bands"]["mean"][b] for b in sensor["band_order"]]
    stds = [sensor["bands"]["std"][b] for b in sensor["band_order"]]

    means = np.array(means, dtype=np.float32).reshape(1, -1, 1, 1)
    stds = np.array(stds, dtype=np.float32).reshape(1, -1, 1, 1)

    patches_array = np.stack(patches, axis=0).astype(np.float32)
    normalized = (patches_array - means) / stds
    return normalized


def generate_embeddings(patches_normalized, metadata, scene_datetime, device="cpu"):
    """Generate Clay embeddings for normalized patches."""
    if not Path(CHECKPOINT_PATH).exists():
        print(f"\nCheckpoint not found at: {CHECKPOINT_PATH}")
        print("Download it with:")
        print("  python -c \"from huggingface_hub import hf_hub_download; "
              "hf_hub_download('made-with-clay/Clay', 'v1.5/clay-v1.5.ckpt', local_dir='.')\"")
        sys.exit(1)

    print(f"\nLoading Clay v1.5 model from {CHECKPOINT_PATH}...")
    model = ClayMAEModule.load_from_checkpoint(
        CHECKPOINT_PATH,
        metadata_path=METADATA_PATH,
        map_location=device,
    )
    model.eval()
    model.to(device)

    # Disable masking for inference
    model.model.encoder.mask_ratio = 0.0
    model.model.encoder.shuffle = False

    sensor = metadata[PLATFORM]
    wavelengths = torch.tensor(
        [sensor["bands"]["wavelength"][b] for b in sensor["band_order"]],
        dtype=torch.float32,
    )
    gsd = torch.tensor(sensor["gsd"], dtype=torch.float32)

    batch_size = patches_normalized.shape[0]

    # Encode time as [week_sin, week_cos, hour_sin, hour_cos]
    if scene_datetime:
        week = scene_datetime.isocalendar()[1]
        hour = scene_datetime.hour
    else:
        week, hour = 0, 0
    week_norm = 2 * np.pi * week / 52
    hour_norm = 2 * np.pi * hour / 24
    time_vec = [np.sin(week_norm), np.cos(week_norm), np.sin(hour_norm), np.cos(hour_norm)]
    time_tensor = torch.tensor([time_vec] * batch_size, dtype=torch.float32, device=device)

    # Lat/lon encoding [lat_sin, lat_cos, lon_sin, lon_cos]
    lat, lon = LATLON_CENTER
    lat_rad = np.radians(lat)
    lon_rad = np.radians(lon)
    latlon_vec = [np.sin(lat_rad), np.cos(lat_rad), np.sin(lon_rad), np.cos(lon_rad)]
    latlon_tensor = torch.tensor(
        [latlon_vec] * batch_size, dtype=torch.float32, device=device
    )

    print(f"Generating embeddings for {batch_size} patches (one at a time)...")
    all_embeddings = []

    for idx in range(batch_size):
        chip = torch.tensor(
            patches_normalized[idx : idx + 1], dtype=torch.float32, device=device
        )
        datacube = {
            "pixels": chip,
            "time": time_tensor[idx : idx + 1],
            "latlon": latlon_tensor[idx : idx + 1],
            "gsd": gsd.to(device),
            "waves": wavelengths.to(device),
        }
        with torch.no_grad():
            encoded_patches, _, _, _ = model.model.encoder(datacube)
        cls_emb = encoded_patches[:, 0, :].cpu()
        all_embeddings.append(cls_emb)
        print(f"  Patch {idx + 1}/{batch_size} done")

    embeddings_np = torch.cat(all_embeddings, dim=0).numpy()
    print(f"Embeddings shape: {embeddings_np.shape}")
    return embeddings_np


def compute_and_validate_similarity(embeddings, coords, patches):
    """Compute cosine similarity and validate spatial coherence."""
    print("\n--- Cosine Similarity Analysis ---")

    sim_matrix = cosine_similarity(embeddings)
    n = len(embeddings)

    upper_tri_vals = sim_matrix[np.triu_indices(n, k=1)]
    print(f"Similarity matrix shape: {sim_matrix.shape}")
    print(f"Min similarity: {upper_tri_vals.min():.4f}")
    print(f"Max similarity: {upper_tri_vals.max():.4f}")
    print(f"Mean similarity: {upper_tri_vals.mean():.4f}")
    print(f"Std similarity:  {upper_tri_vals.std():.4f}")

    adjacent_sims = []
    distant_sims = []

    for i in range(n):
        for j in range(i + 1, n):
            r1, c1 = coords[i]
            r2, c2 = coords[j]
            dist = abs(r1 - r2) + abs(c1 - c2)

            if dist == 1:
                adjacent_sims.append(sim_matrix[i, j])
            elif dist >= 4:
                distant_sims.append(sim_matrix[i, j])

    print(f"\n--- Spatial Coherence Validation ---")
    if adjacent_sims and distant_sims:
        mean_adjacent = np.mean(adjacent_sims)
        mean_distant = np.mean(distant_sims)

        print(f"Adjacent patches (dist=1): mean={mean_adjacent:.4f}, n={len(adjacent_sims)}")
        print(f"Distant patches (dist>=4): mean={mean_distant:.4f}, n={len(distant_sims)}")
        print(f"Difference: {mean_adjacent - mean_distant:+.4f}")

        if mean_adjacent > mean_distant:
            print("PASS: Adjacent patches are more similar than distant patches.")
        else:
            print("NOTE: No clear spatial gradient — scene may be homogeneous.")
    else:
        print("Not enough patch pairs for adjacency validation.")

    diag = np.diag(sim_matrix)
    print(f"\nSelf-similarity (should be 1.0): min={diag.min():.6f}, max={diag.max():.6f}")

    upper_tri = np.triu_indices(n, k=1)
    sims_flat = sim_matrix[upper_tri]
    most_idx = np.argmax(sims_flat)
    least_idx = np.argmin(sims_flat)

    i_m, j_m = upper_tri[0][most_idx], upper_tri[1][most_idx]
    i_l, j_l = upper_tri[0][least_idx], upper_tri[1][least_idx]

    print(f"\nMost similar:  patches {i_m} & {j_m} "
          f"(coords {coords[i_m]}, {coords[j_m]}) = {sims_flat[most_idx]:.4f}")
    print(f"Least similar: patches {i_l} & {j_l} "
          f"(coords {coords[i_l]}, {coords[j_l]}) = {sims_flat[least_idx]:.4f}")

    return sim_matrix


def plot_results(sim_matrix, coords, patches, embeddings):
    """Visualize similarity matrix, patches, and embedding space."""
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    im = axes[0].imshow(sim_matrix, cmap="RdYlGn", vmin=-1, vmax=1)
    axes[0].set_title("Pairwise Cosine Similarity")
    axes[0].set_xlabel("Patch Index")
    axes[0].set_ylabel("Patch Index")
    plt.colorbar(im, ax=axes[0])

    n_show = min(8, len(patches))
    rgb_strip = []
    for i in range(n_show):
        rgb = patches[i][[2, 1, 0]]
        rgb = np.clip(rgb / 3000, 0, 1)
        rgb_strip.append(rgb.transpose(1, 2, 0))

    if rgb_strip:
        combined = np.concatenate(rgb_strip, axis=1)
        axes[1].imshow(combined)
        axes[1].set_title(f"First {n_show} patches (RGB)")
        axes[1].axis("off")

    from sklearn.decomposition import PCA

    if len(embeddings) >= 3:
        pca = PCA(n_components=2)
        emb_2d = pca.fit_transform(embeddings)
        scatter = axes[2].scatter(
            emb_2d[:, 0],
            emb_2d[:, 1],
            c=range(len(emb_2d)),
            cmap="viridis",
            s=60,
            edgecolors="k",
            linewidths=0.5,
        )
        axes[2].set_title("Embeddings (PCA 2D)")
        axes[2].set_xlabel("PC1")
        axes[2].set_ylabel("PC2")
        plt.colorbar(scatter, ax=axes[2], label="Patch index")

    plt.tight_layout()
    output_path = "clay_similarity_results.png"
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    print(f"\nVisualization saved to: {output_path}")
    plt.close()


def main():
    print("=" * 60)
    print("Clay Foundation Model - GeoAI Cosine Similarity Validation")
    print("=" * 60)

    device = "cpu"
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    print(f"Device: {device}")

    metadata = load_metadata()

    image, scene_datetime = fetch_sentinel2_data()

    patches, coords = create_patches(image)
    if len(patches) < 2:
        print("Need at least 2 valid patches. Try a larger region.")
        sys.exit(1)

    max_patches = 16
    if len(patches) > max_patches:
        print(f"Using first {max_patches} patches for efficiency.")
        patches = patches[:max_patches]
        coords = coords[:max_patches]

    patches_normalized = normalize_patches(patches, metadata)

    embeddings = generate_embeddings(
        patches_normalized, metadata, scene_datetime, device=device
    )

    sim_matrix = compute_and_validate_similarity(embeddings, coords, patches)

    plot_results(sim_matrix, coords, patches, embeddings)

    print("\nDone!")


if __name__ == "__main__":
    main()