import argparse
import sys
import torch
from torch.utils.data import DataLoader, Subset
from omegaconf import OmegaConf

# Ensure local imports work when running from repo root
if "src" not in sys.path:
    sys.path.insert(0, "src")

from mixhub.data.dataset import MixtureTask
from mixhub.data.data import DATA_CATALOG
from mixhub.data.splits import SplitLoader
from mixhub.data.collate import custom_collate
from mixhub.model.model_builder import build_mixture_model


def main(config_path: str, checkpoint_path: str, split_num: int, index_in_split: int):
    config = OmegaConf.load(config_path)

    device = torch.device(
        "cuda" if torch.cuda.is_available() and config.device == "cuda" else "cpu"
    )

    mixture_task = MixtureTask(
        property=config.dataset.property,
        dataset=DATA_CATALOG[config.dataset.name](),
        featurization=config.dataset.featurization,
    )

    split_loader = SplitLoader(split_type="kfold")
    train_indices, _, _ = split_loader(
        property=mixture_task.property,
        cache_dir=mixture_task.dataset.data_dir,
        split_num=split_num,
    )

    target_idx = int(train_indices[index_in_split])
    sample_dataset = Subset(mixture_task, [target_idx])
    loader = DataLoader(
        sample_dataset,
        batch_size=1,
        collate_fn=custom_collate,
        num_workers=0,
    )
    batch = next(iter(loader))

    model = build_mixture_model(config=config.mixture_model)
    state = torch.load(checkpoint_path, weights_only=False)
    model.load_state_dict(state)
    model = model.to(device)
    model.eval()

    with torch.no_grad():
        pred = model(
            batch["features"].to(device),
            batch["ids"].to(device),
            batch["fractions"].to(device),
            batch["context"].to(device),
        )

    print(f"device={device}")
    print(f"sample_index={target_idx}")
    print(f"pred={pred.flatten().cpu().item()}")
    print(f"label={batch['label'].flatten().cpu().item()}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run single-sample inference")
    parser.add_argument("--config", type=str, default="outputs/hparams_test_run.yaml")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="outputs/best_model_dict_test_run.pt",
    )
    parser.add_argument("--split_num", type=int, default=0)
    parser.add_argument("--index_in_split", type=int, default=0)

    args = parser.parse_args()

    main(
        config_path=args.config,
        checkpoint_path=args.checkpoint,
        split_num=args.split_num,
        index_in_split=args.index_in_split,
    )