from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import torch
from torch_geometric.data import Data

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from onescience.utils.YParams import YParams


def resolve_path(path: str) -> Path:
    value = Path(path).expanduser()
    return value if value.is_absolute() else ROOT / value


def make_sample(num_nodes: int, input_dim: int, output_dim: int, surface_ratio: float, seed: int) -> Data:
    generator = torch.Generator().manual_seed(seed)
    pos = torch.rand((num_nodes, 2), generator=generator)
    freestream = torch.ones((num_nodes, 2))
    freestream[:, 1] = torch.linspace(-0.2, 0.2, num_nodes)
    sdf = torch.linalg.norm(pos - 0.5, dim=1, keepdim=True) - 0.25
    normals = torch.nn.functional.normalize(pos - 0.5, dim=1)
    x = torch.cat([pos, freestream, sdf, normals], dim=1)
    if x.size(1) != input_dim:
        raise ValueError(f"fake input_dim must be 7 for this AirfRANS feature layout, got {input_dim}")

    y0 = torch.sin(np.pi * pos[:, :1]) + 0.1 * freestream[:, :1]
    y1 = torch.cos(np.pi * pos[:, 1:2]) + 0.1 * freestream[:, 1:2]
    pressure = pos[:, :1] - pos[:, 1:2]
    viscosity = torch.abs(sdf) + 0.01
    y = torch.cat([y0, y1, pressure, viscosity], dim=1)
    if y.size(1) != output_dim:
        raise ValueError(f"fake output_dim must be 4 for this AirfRANS target layout, got {output_dim}")

    surf = torch.zeros(num_nodes, dtype=torch.bool)
    surf_count = max(1, int(num_nodes * surface_ratio))
    surf[:surf_count] = True
    perm = torch.randperm(num_nodes, generator=generator)
    return Data(pos=pos[perm], x=x[perm].float(), y=y[perm].float(), surf=surf[perm])


def main() -> int:
    cfg = YParams(str(ROOT / "conf" / "config.yaml"), "fake_data")
    cfg_data = YParams(str(ROOT / "conf" / "config.yaml"), "datapipe")
    cfg_train = YParams(str(ROOT / "conf" / "config.yaml"), "training")
    data_dir = resolve_path(cfg_data.source.data_dir)
    stats_dir = resolve_path(cfg_data.source.stats_dir)
    data_dir.mkdir(parents=True, exist_ok=True)
    stats_dir.mkdir(parents=True, exist_ok=True)

    split_counts = {
        cfg_data.data.splits.train_name: int(cfg.num_train),
        cfg_data.data.splits.val_name: int(cfg.num_val),
        cfg_data.data.splits.test_name: int(cfg.num_test),
    }
    manifest: dict[str, list[str]] = {}
    all_x, all_y = [], []
    seed = int(cfg_train.seed)

    for split, count in split_counts.items():
        names = []
        for idx in range(count):
            name = f"{split}_{idx:04d}"
            sample = make_sample(
                num_nodes=int(cfg.num_nodes),
                input_dim=int(cfg.input_dim),
                output_dim=int(cfg.output_dim),
                surface_ratio=float(cfg.surface_ratio),
                seed=seed + len(all_x),
            )
            torch.save(sample, data_dir / f"{name}.pt")
            names.append(name)
            if split == cfg_data.data.splits.train_name:
                all_x.append(sample.x.numpy())
                all_y.append(sample.y.numpy())
        manifest[split] = names

    (data_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    train_x = np.concatenate(all_x, axis=0)
    train_y = np.concatenate(all_y, axis=0)
    np.save(stats_dir / "mean_in.npy", train_x.mean(axis=0).astype(np.float32))
    np.save(stats_dir / "std_in.npy", (train_x.std(axis=0) + 1e-6).astype(np.float32))
    np.save(stats_dir / "mean_out.npy", train_y.mean(axis=0).astype(np.float32))
    np.save(stats_dir / "std_out.npy", (train_y.std(axis=0) + 1e-6).astype(np.float32))
    print(f"Fake AirfRANS graph data written to {data_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())