import argparse
import os
import torch
import sys
import copy
import traceback
import pandas as pd
import numpy as np

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_cv] {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.predict")
    from mixhub.model.predict import predict
    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,
        "predict": predict,
        "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"]
    predict = runtime["predict"]
    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")

    for i in range(5):

        run_name = f"cv_{i}"

        print(f"Training/validating on split {i}", flush=True)
        debug(f"Loading split {i} indices")
        train_indices, val_indices, test_indices = split_loader(
            property=mixture_task.property,
            cache_dir=mixture_task.dataset.data_dir,
            split_num=i,
        )
        debug(
            f"Split {i} sizes: "
            f"train={len(train_indices)}, val={len(val_indices)}, test={len(test_indices)}"
        )

        # Data Loader
        debug(f"Creating DataLoaders for split {i}")
        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(f"Building model for split {i}")
        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(f"Starting training for split {i}")
        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=run_name,
            wandb_logger=wandb_logger,
        )
        debug(f"Finished training for split {i}")

        print(f"Testing on split {i}", flush=True)

        # Data Loader (one big batch)
        test_dataset = Subset(mixture_task, test_indices.tolist())
        test_loader = DataLoader(
            test_dataset,
            batch_size=test_dataset.__len__(),
            collate_fn=custom_collate,
            num_workers=config.num_workers,
        )

        metric_dict, y_pred, y_test = predict(
            model=model,
            test_loader=test_loader,
            device=device,
        )

        print(metric_dict, flush=True)
        test_metrics = pd.DataFrame(metric_dict, index=["metrics"]).transpose()
        test_metrics.to_csv(os.path.join(config.root_dir, f"{run_name}_test_metrics.csv"))
    
        y_pred = y_pred.detach().cpu().numpy().flatten()
        y_test = y_test.detach().cpu().numpy().flatten()
        test_predictions = pd.DataFrame(
            {
                "Predicted_Experimental_Values": y_pred,
                "Ground_Truth": y_test,
                "MAE": np.abs(y_pred - y_test),
            },
            index=range(len(y_pred)),
        )
        test_predictions.to_csv(os.path.join(config.root_dir, f"{run_name}_test_predictions.csv"), index=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run cross validation experiment")
    parser.add_argument("config", type=str, help="Path to the YAML configuration file")
    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 = f"{config.dataset.featurization}_{config.mixture_model.mix_encoder.type}"

        config.root_dir = os.path.abspath(f"../results/cv_split/{experiment_name}/{config.dataset.name}/{config.dataset.property.lower().replace(' ', '_')}")
        
        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