import argparse
import json
import importlib.util
import os
import sys
from pathlib import Path

import torch
from tqdm import tqdm


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

from model import build_model, infer_task_type
import onescience
from onescience.distributed.manager import DistributedManager
from onescience.utils.YParams import YParams


def resolve_path(path_value):
    path = Path(path_value)
    return path if path.is_absolute() else PROJECT_ROOT / path


def parse_args():
    parser = argparse.ArgumentParser(description="Run CFDBench inference.")
    parser.add_argument("--model", default=None, help="Override root.model.name and read weight/<model>.pt by default.")
    return parser.parse_args()


def load_config(model_name=None):
    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
    if model_name:
        cfg.model.name = model_name
    elif os.environ.get("CFDBENCH_MODEL_NAME"):
        cfg.model.name = os.environ["CFDBENCH_MODEL_NAME"]
    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
    cfg.inference.output_dir = str(resolve_path(cfg.inference.output_dir))
    return cfg


def checkpoint_path(cfg):
    path_value = cfg.inference.get("checkpoint_path", "auto")
    if path_value == "auto":
        return PROJECT_ROOT / "weight" / f"{cfg.model.name}.pt"
    return resolve_path(path_value)


def output_path(cfg, task_type):
    output_dir = Path(cfg.inference.output_dir)
    if cfg.inference.get("group_by_model", False):
        output_dir = output_dir / task_type / cfg.model.name
    return output_dir


def select_device(requested, dist):
    if requested == "auto":
        return dist.device
    if requested.startswith("cuda") and not torch.cuda.is_available():
        raise RuntimeError(f"Requested device {requested!r}, but CUDA is not available.")
    return torch.device(requested)


def load_cfdbench_datapipe_class():
    runtime_root = Path(onescience.__file__).resolve().parent
    datapipe_file = runtime_root / "datapipes" / "cfd" / "cfdbench.py"
    spec = importlib.util.spec_from_file_location("_onescience_cfdbench_datapipe", datapipe_file)
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot load CFDBench datapipe from {datapipe_file}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module.CFDBenchDatapipe


def predict_full_frame(model, batch, task_type):
    if task_type == "auto":
        return model.generate(
            inputs=batch["inputs"],
            case_params=batch["case_params"],
            mask=batch.get("mask"),
        )

    labels = batch["label"]
    return model.generate_one(
        case_params=batch["case_params"],
        t=batch["t"],
        height=labels.shape[-2],
        width=labels.shape[-1],
    )


def main():
    args = parse_args()
    DistributedManager.initialize()
    dist = DistributedManager()
    cfg = load_config(args.model)
    task_type = infer_task_type(cfg.model.name)
    cfg.datapipe.data.task_type = task_type
    device = select_device(cfg.inference.get("device", "auto"), dist)

    ckpt_path = checkpoint_path(cfg)
    if not ckpt_path.is_file():
        raise FileNotFoundError(f"Missing checkpoint: {ckpt_path}. Run scripts/train.py first or update root.inference.checkpoint_path.")

    torch.set_num_threads(1)
    CFDBenchDatapipe = load_cfdbench_datapipe_class()
    datapipe = CFDBenchDatapipe(cfg.datapipe, distributed=False)
    model = build_model(cfg).to(device)
    model.load_state_dict(torch.load(ckpt_path, map_location=device, weights_only=True))
    model.eval()

    output_dir = output_path(cfg, task_type)
    output_dir.mkdir(parents=True, exist_ok=True)

    preds = []
    score_lists = {name: [] for name in model.loss_fn.get_score_names()}
    with torch.no_grad():
        for step, batch in enumerate(tqdm(datapipe.test_dataloader(), desc="Inferring")):
            if step >= cfg.inference.max_batches:
                break
            batch = {key: value.to(device) for key, value in batch.items()}
            pred = predict_full_frame(model, batch, task_type)
            label = batch["label"][:, : pred.shape[1]]
            loss = model.loss_fn(preds=pred, labels=label)
            preds.append(pred.detach().cpu())
            for key, value in loss.items():
                score_lists[key].append(float(value.detach().cpu()))

    if not preds:
        raise RuntimeError("Inference loader is empty. Check dataset split and data directory.")

    torch.save(torch.cat(preds, dim=0), output_dir / "preds.pt")
    means = {key: sum(values) / len(values) for key, values in score_lists.items() if values}
    with open(output_dir / "scores.json", "w", encoding="utf-8") as f:
        json.dump({"mean": means, "all": score_lists}, f, indent=2)
    print(f"Saved predictions to {output_dir / 'preds.pt'}")
    print(f"Scores: {means}")

    DistributedManager.cleanup()


if __name__ == "__main__":
    main()