#!/usr/bin/env python3
"""Generate the FAST-DeepONet Darcy inclusion benchmark from a fixed seed."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from scipy import sparse
from scipy.sparse.linalg import spsolve


GRID = 64
SEED = 33
SPLITS = (("train", 60), ("tune", 10), ("calib", 20), ("test", 10))
OFFICIAL_COMMIT = "9541c112176c42470cd39138dacbab8af6892232"


def make_basis(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    modes = (
        (1, 0, 0, 0), (0, 1, 0, 0), (1, 1, 0, 0), (1, 0, 1, 0),
        (0, 1, 0, 1), (1, 1, 1, 1), (2, 0, 0, 0), (0, 2, 0, 0),
    )
    fields = []
    for kx, ky, px, py in modes:
        x_wave = np.cos(2 * np.pi * kx * x + px * np.pi / 2) if kx else 1.0
        y_wave = np.cos(2 * np.pi * ky * y + py * np.pi / 2) if ky else 1.0
        field = x_wave * y_wave
        fields.append(field / np.max(np.abs(field)))
    return np.stack(fields)


def site_locations(rng: np.random.Generator) -> np.ndarray:
    axis = np.linspace(0.2, 0.8, 4)
    candidates = np.array([(a, b) for a in axis for b in axis])
    sites = candidates[rng.permutation(len(candidates))[:10]]
    return sites + rng.uniform(-0.04, 0.04, size=sites.shape)


def solve_pressure(log_coefficient: np.ndarray) -> np.ndarray:
    coefficient = np.exp(log_coefficient)
    inner = coefficient[1:-1, 1:-1]
    harmonic = lambda neighbor: 2 * inner * neighbor / (inner + neighbor)
    east = harmonic(coefficient[2:, 1:-1])
    west = harmonic(coefficient[:-2, 1:-1])
    north = harmonic(coefficient[1:-1, 2:])
    south = harmonic(coefficient[1:-1, :-2])
    side = GRID - 2
    nodes = np.arange(side * side).reshape(side, side)
    spacing_squared = (1.0 / GRID) ** 2

    connections = (
        (nodes[:-1, :], nodes[1:, :], east[:-1, :]),
        (nodes[1:, :], nodes[:-1, :], west[1:, :]),
        (nodes[:, :-1], nodes[:, 1:], north[:, :-1]),
        (nodes[:, 1:], nodes[:, :-1], south[:, 1:]),
    )
    row = np.concatenate([nodes.ravel()] + [src.ravel() for src, _, _ in connections])
    col = np.concatenate([nodes.ravel()] + [dst.ravel() for _, dst, _ in connections])
    diagonal = (east + west + north + south).ravel() / spacing_squared
    values = np.concatenate(
        [diagonal] + [-edge.ravel() / spacing_squared for _, _, edge in connections]
    )
    operator = sparse.csr_matrix(
        (values, (row, col)), shape=(side * side, side * side)
    )
    pressure = np.zeros((GRID, GRID))
    pressure[1:-1, 1:-1] = spsolve(
        operator, np.ones(side * side)
    ).reshape(side, side)
    return pressure


def generate_sample(
    rng: np.random.Generator,
    basis: np.ndarray,
    sites: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    factors = rng.normal(size=len(basis))
    weights = np.sqrt(np.exp(-0.5 * np.arange(len(basis))))
    log_coefficient = np.tensordot(factors * weights, basis, axes=1) * 0.8

    count = rng.integers(3, 7)
    chosen = sites[rng.permutation(len(sites))[:count]]
    mask = np.zeros((GRID, GRID))
    for center in chosen:
        sign = rng.choice([-1, 1])
        magnitude = rng.uniform(1.2, 1.8)
        inclusion = (
            (x - center[0]) ** 2 + (y - center[1]) ** 2 <= 0.045 ** 2
        ).astype(np.float64)
        log_coefficient += sign * magnitude * inclusion
        mask = np.maximum(mask, inclusion)

    log_coefficient = np.clip(log_coefficient, -3, 3)
    pressure = solve_pressure(log_coefficient)
    return tuple(
        array.astype(np.float32).ravel()
        for array in (log_coefficient, pressure, mask)
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "data",
    )
    parser.add_argument(
        "--overwrite", action="store_true", help="Replace existing NPZ samples"
    )
    args = parser.parse_args()
    root = args.output_dir.resolve()
    existing = [path for split, _ in SPLITS for path in (root / split).glob("*.npz")]
    if existing and not args.overwrite:
        parser.error(f"{root} already contains NPZ samples; use --overwrite to replace")

    axis = (np.arange(GRID) + 0.5) / GRID
    x, y = np.meshgrid(axis, axis, indexing="ij")
    rng = np.random.default_rng(SEED)
    basis = make_basis(x, y)
    sites = site_locations(rng)

    for split, count in SPLITS:
        directory = root / split
        directory.mkdir(parents=True, exist_ok=True)
        for index in range(count):
            permeability, pressure, mask = generate_sample(rng, basis, sites, x, y)
            np.savez(
                directory / f"darcy_inclusion_{split}_{index:03d}.npz",
                nx=np.int64(GRID),
                ny=np.int64(GRID),
                pde_name="darcy_inclusion",
                init_value=permeability,
                last_value=pressure,
                sparse_mask=mask,
            )
        print(f"{split}: generated {count} samples in {directory}", flush=True)

    manifest = {
        "source": "https://github.com/kwonji39/FAST-DeepONet",
        "source_commit": OFFICIAL_COMMIT,
        "experiment": "Table 2 Darcy inclusion",
        "seed": SEED,
        "grid_size": [GRID, GRID],
        "splits": dict(SPLITS),
        "arrays": ["init_value", "last_value", "sparse_mask"],
    }
    root.mkdir(parents=True, exist_ok=True)
    (root / "DATA_MANIFEST.json").write_text(
        json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
    )


if __name__ == "__main__":
    main()