import argparse
from pathlib import Path
import importlib.util

import numpy as np
import onnxruntime as ort


# -----------------------------
# Helpers: import modules by file path (avoid "utils" name conflicts)
# -----------------------------
ROOT = Path(__file__).resolve().parents[1]  # inference_example/

def _import_from_path(module_name: str, file_path: Path):
    spec = importlib.util.spec_from_file_location(module_name, str(file_path))
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot import {module_name} from {file_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


# Load our local helper files explicitly
io = _import_from_path("fuxicfd_io", ROOT / "utils" / "io.py")
pre = _import_from_path("fuxicfd_pre", ROOT / "utils" / "preprocessing.py")
post = _import_from_path("fuxicfd_post", ROOT / "utils" / "postprocessing.py")

load_example_input = io.load_example_input
save_prediction_npz = io.save_prediction_npz
build_model_input = pre.build_model_input
denormalize_and_split = post.denormalize_and_split


def _resolve_path(p: str) -> Path:
    """Resolve a path relative to inference_example/ if it's not absolute."""
    path = Path(p)
    if path.is_absolute():
        return path
    return (ROOT / path).resolve()


def main():
    parser = argparse.ArgumentParser(description="FuXi-CFD ONNX inference example (publish-ready).")
    parser.add_argument("--model", type=str, default="../model/fuxicfd_model.onnx",
                        help="Path to ONNX model. Default: ../model/fuxicfd_model.onnx (relative to inference_example/)")
    parser.add_argument("--input", type=str, default="data/inputs.npz",
                        help="Input file (.npy dict or .npz). Default: data/inputs.npz")
    parser.add_argument("--output", type=str, default="data/prediction.npz",
                        help="Output prediction .npz. Default: data/prediction.npz")
    parser.add_argument("--norm_in", type=str, default="normalization/scaler_input.npy",
                        help="Input normalization stats (.npy dict). Default: normalization/scaler_input.npy")
    parser.add_argument("--norm_out", type=str, default="normalization/scaler_output.npy",
                        help="Output normalization stats (.npy dict). Default: normalization/scaler_output.npy")
    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda"],
                        help="Execution device for onnxruntime. Default: cpu")
    args = parser.parse_args()

    # Resolve paths relative to inference_example/
    model_path = _resolve_path(args.model)
    input_path = _resolve_path(args.input)
    output_path = _resolve_path(args.output)
    norm_in_path = _resolve_path(args.norm_in)
    norm_out_path = _resolve_path(args.norm_out)

    # --- sanity checks ---
    if not model_path.exists():
        raise FileNotFoundError(
            f"ONNX model file not found:\n  {model_path}\n\n"
            f"Please place your model at:\n  {ROOT.parent / 'model' / 'fuxicfd_model.onnx'}\n"
            f"or pass --model with the correct path."
        )
    if not norm_in_path.exists():
        raise FileNotFoundError(f"Input normalization file not found: {norm_in_path}")
    if not norm_out_path.exists():
        raise FileNotFoundError(f"Output normalization file not found: {norm_out_path}")
    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_path}")

    # Load normalization dicts
    in_stats = np.load(norm_in_path, allow_pickle=True).item()
    out_stats = np.load(norm_out_path, allow_pickle=True).item()

    # Select providers
    providers = ["CPUExecutionProvider"]
    if args.device == "cuda":
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]

    # Create session
    sess = ort.InferenceSession(str(model_path), providers=providers)
    input_name = sess.get_inputs()[0].name

    # Load and preprocess inputs
    x_dict = load_example_input(str(input_path))
    x = build_model_input(x_dict, in_stats)          # (4,300,300) float32
    ort_inputs = x[None, ...]                        # (1,4,300,300)

    # Run inference
    pred = sess.run(None, {input_name: ort_inputs})[0]  # (1,27,4,300,300) expected

    # Postprocess
    u, v, w, k = denormalize_and_split(pred, out_stats)

    # Save output
    save_prediction_npz(str(output_path), u=u, v=v, w=w, k=k)
    print(f"[OK] Saved prediction to: {output_path}")


if __name__ == "__main__":
    main()