"""Generate sparse structured observations and collocation coordinates lazily."""

import argparse
import importlib.util
from pathlib import Path

import numpy as np
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 structured_observations(config, mode):
    data = config["data"]
    extent = float(data["horizontal_extent_m"])
    levels = np.linspace(float(data["pressure_min_pa"]), float(data["pressure_max_pa"]), int(data["pressure_levels"]))
    times = np.asarray(data["observation_times_s"], dtype=np.float64)
    if times.shape != (int(data["time_steps"]),) or not np.array_equal(times, [-10800.0, 0.0, 10800.0]):
        raise ValueError("paper observation times must be [-3h, 0h, +3h]")
    line = np.linspace(-extent, extent, int(data["observation_line_points"]))
    coordinates = []
    for time_index, time in enumerate(times):
        active_mode = ("Cross" if time_index % 2 == 0 else "Plus") if mode == "Switch" else mode
        for pressure in levels[:: int(data["observation_pressure_stride"])]:
            if active_mode == "Plus":
                coordinates.extend((y, x, time, pressure) for y, x in zip(np.zeros_like(line), line))
                coordinates.extend((y, x, time, pressure) for y, x in zip(line, np.zeros_like(line)))
            elif active_mode == "Cross":
                coordinates.extend((y, x, time, pressure) for y, x in zip(line, line))
                coordinates.extend((y, x, time, pressure) for y, x in zip(line, -line))
            else:
                raise ValueError("observation mode must be Cross, Plus, or Switch")
            edge = np.linspace(-extent, extent, int(data["boundary_points_per_edge"]))
            coordinates.extend((-extent, value, time, pressure) for value in edge)
            coordinates.extend((extent, value, time, pressure) for value in edge)
            coordinates.extend((value, -extent, time, pressure) for value in edge[1:-1])
            coordinates.extend((value, extent, time, pressure) for value in edge[1:-1])
    return np.unique(np.asarray(coordinates, dtype=np.float32), axis=0)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=("Cross", "Plus", "Switch"))
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    mode = args.mode or config["data"]["observation_mode"]
    output = ROOT / config["data"]["root"]
    output.mkdir(parents=True, exist_ok=True)
    path = output / "training_points.npz"
    if path.exists() and not args.force:
        print(f"exists={path.relative_to(ROOT)} (use --force to regenerate)")
        return
    module = load_model_module()
    observations = structured_observations(config, mode)
    targets = module.analytic_vortex_numpy(observations)[:, :3]
    rng = np.random.default_rng(int(config["seed"]))
    count = int(config["data"]["collocation_points"])
    collocation = np.column_stack((
        rng.uniform(-config["data"]["horizontal_extent_m"], config["data"]["horizontal_extent_m"], (count, 2)),
        rng.uniform(-config["data"]["time_extent_s"], config["data"]["time_extent_s"], count),
        rng.uniform(config["data"]["pressure_min_pa"], config["data"]["pressure_max_pa"], count),
    )).astype(np.float32)
    # The first two random columns already follow the required [y, x] order.
    np.savez(path, observation_coordinates=observations, observation_targets=targets,
             collocation_coordinates=collocation, input_order=np.asarray(module.INPUT_ORDER),
             supervised_order=np.asarray(module.OUTPUT_ORDER[:3]), observation_mode=np.asarray(mode),
             dense_shape=np.asarray([data_dim := int(config["data"]["grid_points"]), data_dim,
                                     int(config["data"]["pressure_levels"]), int(config["data"]["time_steps"])]))
    print(f"generated={path.relative_to(ROOT)} observations={len(observations)} collocation={len(collocation)} mode={mode}")


if __name__ == "__main__":
    main()