import argparse
import os
import torch
import sys
import copy
import traceback

from torch.utils.data import DataLoader
from torch.utils.data import Subset
from omegaconf import OmegaConf

if "src" not in sys.path:
    sys.path.insert(0, "src")

VERBOSE = False


def debug(message: str) -> None:
    if VERBOSE:
        print(f"[run_model] {message}", flush=True)


def load_runtime():
    debug("Importing mixhub.data.dataset")
    from mixhub.data.dataset import MixtureTask
    debug("Importing mixhub.data.data")
    from mixhub.data.data import DATA_CATALOG
    debug("Importing mixhub.data.featurization_types")
    from mixhub.data.featurization_types import FEATURIZATION_TYPE
    debug("Importing mixhub.data.collate")
    from mixhub.data.collate import custom_collate
    debug("Importing mixhub.data.splits")
    from mixhub.data.splits import SplitLoader
    debug("Importing mixhub.model.train")
    from mixhub.model.train import train
    debug("Importing mixhub.model.model_builder")
    from mixhub.model.model_builder import build_mixture_model

    return {
        "MixtureTask": MixtureTask,
        "DATA_CATALOG": DATA_CATALOG,
        "FEATURIZATION_TYPE": FEATURIZATION_TYPE,
        "custom_collate": custom_collate,
        "SplitLoader": SplitLoader,
        "train": train,
        "build_mixture_model": build_mixture_model,
    }


def main(
    config,
    experiment_name,
    wandb_logger=None,
):
    config = copy.deepcopy(config)
    runtime = load_runtime()
    MixtureTask = runtime["MixtureTask"]
    DATA_CATALOG = runtime["DATA_CATALOG"]
    FEATURIZATION_TYPE = runtime["FEATURIZATION_TYPE"]
    custom_collate = runtime["custom_collate"]
    SplitLoader = runtime["SplitLoader"]
    train = runtime["train"]
    build_mixture_model = runtime["build_mixture_model"]

    torch.manual_seed(config.seed)
    device = torch.device(config.device)
    print(f"Running on: {device}", flush=True)

    root_dir = config.root_dir
    os.makedirs(root_dir, exist_ok=True)

    featurization = config.dataset.featurization

    if FEATURIZATION_TYPE[featurization] == "graphs" and config.mixture_model.mol_encoder.type != "gnn":
        raise ValueError(f"featurization is:{FEATURIZATION_TYPE[featurization]} but molecule encoder is: {config.mol_encoder.type}")

    if FEATURIZATION_TYPE[featurization] == "tensors" and config.mixture_model.mol_encoder.type == "gnn":
        raise ValueError(f"featurization is:{FEATURIZATION_TYPE[featurization]} but molecule encoder is: {config.mol_encoder.type}")

    # Dataset
    debug(f"Loading dataset: {config.dataset.name}")
    dataset = DATA_CATALOG[config.dataset.name]()
    property = config.dataset.property

    debug(
        "Building MixtureTask "
        f"(property={property}, featurization={featurization})"
    )
    mixture_task = MixtureTask(
        property=property,
        dataset=dataset,
        featurization=featurization,
    )
    debug(f"MixtureTask ready with {len(mixture_task)} samples")

    # Split Loader
    split_loader = SplitLoader(split_type="kfold")
    debug("Loading split 0 indices")
    train_indices, val_indices, _ = split_loader(
        property=mixture_task.property,
        cache_dir=mixture_task.dataset.data_dir,
        split_num=0,
    )
    debug(f"Split 0 sizes: train={len(train_indices)}, val={len(val_indices)}")

    # Data Loader
    debug("Creating DataLoaders")
    train_dataset = Subset(mixture_task, train_indices.tolist())
    val_dataset = Subset(mixture_task, val_indices.tolist())

    train_loader = DataLoader(
        train_dataset,
        batch_size=config.batch_size,
        shuffle=True,
        collate_fn=custom_collate,
        num_workers=config.num_workers,
        pin_memory=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=config.batch_size,
        collate_fn=custom_collate,
        num_workers=config.num_workers,
    )

    debug("Building model")
    model = build_mixture_model(config=config.mixture_model)
    debug(f"Moving model to {device}")
    model = model.to(device)

    # Save hyper parameters    
    OmegaConf.save(config, f"{root_dir}/hparams_{experiment_name}.yaml")

    # Training
    debug("Starting training")
    train(
        root_dir=root_dir,
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        loss_type=config.loss_type,
        lr_mol_encoder=config.lr_mol_encoder,
        lr_other=config.lr_other,
        device=device,
        weight_decay=config.weight_decay,
        max_epochs=config.max_epochs,
        patience=config.patience,
        experiment_name=experiment_name,
        wandb_logger=wandb_logger,
    )
    debug("Finished training")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run model training and evaluation on a random split")
    parser.add_argument("config", type=str, help="Path to the YAML configuration file")
    parser.add_argument("--experiment_name", type=str, default="test_run", help="Name of the experiment")
    parser.add_argument("--wandb_project", type=str, default=None, help="Name of the wandb project (optional)")
    parser.add_argument("--verbose", action="store_true", help="Print detailed diagnostic logs")

    args = parser.parse_args()
    VERBOSE = args.verbose
    if VERBOSE:
        os.environ["CHEMIXHUB_VERBOSE"] = "1"

    try:
        debug(f"Loading config from {args.config}")
        config = OmegaConf.load(args.config)
        experiment_name = args.experiment_name

        if args.wandb_project is not None:
            try:
                from torchtune.training.metric_logging import WandBLogger
            except ImportError as exc:
                raise ImportError(
                    "torchtune is required for WandB logging. Install torchtune/torchao or run without --wandb_project."
                ) from exc
            wandb_logger = WandBLogger(project=args.wandb_project)
        else:
            wandb_logger = None

        main(
            config=config,
            experiment_name=experiment_name,
            wandb_logger=wandb_logger,
        )
    except Exception as exc:
        print(f"Execution failed: {exc}", flush=True)
        traceback.print_exc()
        raise