import logging
import os
import sys
from pathlib import Path

import numpy as np
import torch

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


from model.meshgraphnet import MeshGraphNet
from onescience.utils.YParams import YParams
from onescience.launch.utils import load_checkpoint 
from fake_data import build_cylinder_flow_datapipe


def build_model(model_params, device):
    mlp_act = "silu" if model_params.recompute_activation else "relu"
    return MeshGraphNet(
        input_dim_nodes=model_params.num_input_features,
        input_dim_edges=model_params.num_edge_features,
        output_dim=model_params.num_output_features,
        processor_size=model_params.processor_size,
        hidden_dim_processor=model_params.hidden_dim_processor,
        num_layers_node_processor=model_params.num_layers_node_processor,
        num_layers_edge_processor=model_params.num_layers_edge_processor,
        hidden_dim_node_encoder=model_params.hidden_dim_node_encoder,
        hidden_dim_edge_encoder=model_params.hidden_dim_edge_encoder,
        hidden_dim_node_decoder=model_params.hidden_dim_node_decoder,
        mlp_activation_fn=mlp_act,
        do_concat_trick=model_params.do_concat_trick,
        num_processor_checkpoint_segments=model_params.num_processor_checkpoint_segments,
        recompute_activation=model_params.recompute_activation,
    ).to(device)


def resolve_device(device_name: str):
    if device_name == "cpu":
        return torch.device("cpu")
    if device_name in ("cuda", "gpu"):
        if not torch.cuda.is_available():
            raise RuntimeError("Config requested cuda device, but torch.cuda.is_available() is false.")
        return torch.device("cuda:0")
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def main():
    os.chdir(PROJECT_ROOT)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    logger = logging.getLogger("mesh_graph_net.inference")

    config_path = PROJECT_ROOT / "config" / "config.yaml"
    cfg_model = YParams(config_path, "model")
    cfg_data = YParams(config_path, "datapipe")
    cfg_train = YParams(config_path, "training")
    cfg_inference = YParams(config_path, "inference")
    model_params = cfg_model.specific_params[cfg_model.name]

    device = resolve_device(getattr(cfg_inference, "device", "auto"))
    logger.info("Using device: %s", device)
    datapipe = build_cylinder_flow_datapipe(
        params=cfg_data,
        distributed=False,
        project_root=PROJECT_ROOT,
    )
    loader = datapipe.test_dataloader()
    model = build_model(model_params, device)
    checkpoint_dir = PROJECT_ROOT / getattr(cfg_inference, "checkpoint_dir", cfg_train.checkpoint_dir)
    epoch = load_checkpoint(checkpoint_dir, models=model, device=device)
    if epoch == 0:
        logger.warning("No checkpoint found in %s; running with randomly initialized weights", checkpoint_dir)
    model.eval()

    predictions, targets = [], []
    with torch.no_grad():
        for batch in loader:
            graph = batch[0] if isinstance(batch, (tuple, list)) else batch
            graph = graph.to(device)
            pred = model(graph.ndata["x"], graph.edata["x"], graph)
            predictions.append(pred.cpu().numpy())
            targets.append(graph.ndata["y"].cpu().numpy())

    output_path = PROJECT_ROOT / cfg_inference.output_path
    output_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez(output_path, prediction=np.concatenate(predictions, axis=0), target=np.concatenate(targets, axis=0))
    logger.info("Saved inference results to %s", output_path)


if __name__ == "__main__":
    main()