"""Run chunked inference on the full 251 x 251 x 44 x 3 paper-resolution domain."""

import importlib.util
from pathlib import Path

import numpy as np
import torch
import yaml


ROOT = Path(__file__).resolve().parents[1]


def load_model_module():
    spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def coordinates_for_indices(indices, y_axis, x_axis, times, pressures):
    x_count, p_count, t_count = len(x_axis), len(pressures), len(times)
    y_index = indices // (x_count * p_count * t_count)
    remainder = indices % (x_count * p_count * t_count)
    x_index = remainder // (p_count * t_count)
    remainder %= p_count * t_count
    p_index = remainder // t_count
    t_index = remainder % t_count
    return np.column_stack((y_axis[y_index], x_axis[x_index], times[t_index], pressures[p_index])).astype(np.float32)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    module = load_model_module()
    requested = config["runtime"]["device"]
    device = torch.device("cuda" if requested == "auto" and torch.cuda.is_available() else ("cpu" if requested == "auto" else requested))
    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
    if checkpoint.get("format_version") != config["data"]["format_version"]:
        raise ValueError("checkpoint format_version mismatch")
    if checkpoint["input_order"] != list(module.INPUT_ORDER) or checkpoint["output_order"] != list(module.OUTPUT_ORDER):
        raise ValueError("checkpoint coordinate or variable protocol mismatch")
    model = module.PINNTC(**checkpoint["model_config"]).to(device)
    model.load_state_dict(checkpoint["model_state_dict"]); model.eval()
    data = config["data"]
    y_axis = np.linspace(-data["horizontal_extent_m"], data["horizontal_extent_m"], data["grid_points"], dtype=np.float32)
    x_axis = y_axis.copy()
    times = np.asarray(data["observation_times_s"], dtype=np.float32)
    pressures = np.linspace(data["pressure_min_pa"], data["pressure_max_pa"], data["pressure_levels"], dtype=np.float32)
    shape = (len(y_axis), len(x_axis), len(pressures), len(times), 4)
    expected = (251, 251, 44, 3, 4)
    if shape != expected:
        raise ValueError(f"dense grid must remain {expected}, got {shape}")
    output = ROOT / config["paths"]["prediction"]
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(".work.npy")
    prediction = np.lib.format.open_memmap(temporary, mode="w+", dtype=np.float32, shape=shape)
    flat = prediction.reshape(-1, 4)
    chunk_size = int(config["runtime"]["inference_chunk_size"])
    with torch.inference_mode():
        for start in range(0, len(flat), chunk_size):
            stop = min(start + chunk_size, len(flat))
            coordinates = coordinates_for_indices(np.arange(start, stop, dtype=np.int64), y_axis, x_axis, times, pressures)
            values = model(torch.from_numpy(coordinates).to(device)).cpu().numpy().astype(np.float32)
            if not np.isfinite(values).all():
                raise FloatingPointError(f"non-finite prediction in rows {start}:{stop}")
            flat[start:stop] = values
    prediction.flush()
    np.savez(output, predictions=prediction, y=y_axis, x=x_axis, time=times, pressure=pressures,
             input_order=np.asarray(module.INPUT_ORDER), output_order=np.asarray(module.OUTPUT_ORDER),
             format_version=np.asarray(config["data"]["format_version"]), layout=np.asarray("YXPTV"))
    del prediction
    temporary.unlink()
    print(f"predictions={output.relative_to(ROOT)} shape={shape} bytes={output.stat().st_size}")


if __name__ == "__main__":
    main()