"""Create structurally realistic full-grid samples with irregular HadEX-style masks."""

from pathlib import Path
import json
import numpy as np


INDICES = np.array(["TX90p", "TN90p", "TX10p", "TN10p"])
H, W = 144, 192


def europe_mask(lat, lon):
    yy, xx = np.meshgrid(lat, lon, indexing="ij")
    broad = (yy >= 30) & (yy <= 72) & (xx >= -25) & (xx <= 45)
    # A coarse geographic silhouette keeps the scientific global-grid contract.
    atlantic_cut = (xx < -10) & (yy < 44)
    southeast_cut = (xx > 30) & (yy < 40)
    north_cut = (yy > 68) & ((xx < 5) | (xx > 30))
    return (broad & ~atlantic_cut & ~southeast_cut & ~north_cut).astype(np.float32)


def main():
    rng = np.random.default_rng(42)
    root = Path(__file__).resolve().parents[1]
    output = root / "data"
    output.mkdir(exist_ok=True)
    lat = np.linspace(-89.375, 89.375, H, dtype=np.float32)
    lon = np.linspace(-179.0625, 179.0625, W, dtype=np.float32)
    land = europe_mask(lat, lon)
    yy, xx = np.meshgrid(lat, lon, indexing="ij")
    n = 8
    target = np.zeros((n, 1, H, W), dtype=np.float32)
    valid = np.zeros_like(target)
    index_ids = np.arange(n, dtype=np.int64) % 4
    for sample in range(n):
        phase = 0.55 * sample
        field = 50 + 21 * np.sin(np.deg2rad(2.3 * xx) + phase)
        field += 16 * np.cos(np.deg2rad(3.2 * yy) - 0.4 * phase)
        field += 5 * np.sin(np.deg2rad(xx + yy) * 4 + phase)
        field += rng.normal(0, 1.2, (H, W))
        if index_ids[sample] >= 2:
            field = 100 - field
        target[sample, 0] = np.clip(field, 0, 100) * land
        observed = land.copy()
        observed[rng.random((H, W)) < (0.35 + 0.04 * (sample % 3))] = 0
        for _ in range(5):
            cy, cx = rng.integers(45, 99), rng.integers(78, 121)
            ry, rx = rng.integers(3, 11), rng.integers(4, 15)
            hole = ((np.arange(H)[:, None] - cy) / ry) ** 2
            hole = hole + ((np.arange(W)[None, :] - cx) / rx) ** 2
            observed[hole < 1] = 0
        valid[sample, 0] = observed
    observed_values = target * valid
    np.savez_compressed(
        output / "crai_fake.npz", target=target, observed=observed_values,
        valid_mask=valid, europe_mask=land, index_ids=index_ids,
        index_names=INDICES, latitude=lat, longitude=lon,
    )
    metadata = {
        "kind": "structured_synthetic",
        "shape": [n, 1, H, W],
        "grid_resolution": {"longitude_degrees": 1.875, "latitude_degrees": 1.25},
        "indices": INDICES.tolist(),
        "mask": "global grid with coarse Europe land support and irregular missing regions",
    }
    (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
    print(f"wrote {output / 'crai_fake.npz'} with shape {target.shape}")


if __name__ == "__main__":
    main()