import argparse
from pathlib import Path
import numpy as np
import torch
import pytorch_lightning as pl
import random
from torch_geometric.data import DataLoader
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.utilities import rank_zero_only
from stereograph.config import load_config
from stereograph.data import get_features_df, add_prediction_to_dataframe, get_graphs_from_file
from stereograph.load_model import load_model
from stereograph.config import get_output_dir



TASK_TYPES = {"energy", "classification", "direction"}


def parse_arguments():
    parser = argparse.ArgumentParser(description="Train the GCN model")
    parser.add_argument("--config_file", "-c", type=Path, help="Path to the configuration file")
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Run the model in debug mode. Does not log to wandb.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite the output directory if it exists",
    )
    parser.add_argument(
        "--grid_search",
        action="store_true",
        help="Run model with a GridSearch over the hyperparameters",
    )
    parser.add_argument(
        "--checkpoint", "-ckpt", type=Path, help="Path to a checkpoint file to resume training from"
        
    )
    parser.add_argument("--gamma_file", type=Path, help="gamma-diffuse.dl2.h5 or gamma-diffuse.graphs") 
    parser.add_argument("--proton_file",type=Path, help="proton.dl2.h5 or proton.graphs")
    args = parser.parse_args()
    return args


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
 

def main(config, debug=False, reload_graphs=False):
    output_dir = get_output_dir(config)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Initialize seed
    set_seed(config["network"]["seed"])

    if not debug:
        logger = WandbLogger(
            entity="gammalearn",
            project="stereograph",
            name=config["exp_name"],
            log_model=False,
            save_dir=output_dir,
        )
        if rank_zero_only.rank == 0:
            logger.experiment.config.update(config)

    else:
        logger = None

    task_type = config["network"]["task_type"]
    if task_type not in TASK_TYPES:
        raise ValueError(f"Task type {task_type} not recognized. Must be one of {TASK_TYPES}")

    
    # Load graphs 
    gamma_train_graphs = get_graphs_from_file(args.gamma_file, config)
    proton_train_graphs = []

    if task_type in ["classification", "multitask"]:
        if not args.proton_file:
            raise ValueError("For classification task, proton_file must be provided.")
        proton_train_graphs = get_graphs_from_file(args.proton_file, config)

    # DATASET
    graph_data_list = gamma_train_graphs + proton_train_graphs

    # Shuffle and split the dataset
    np.random.shuffle(graph_data_list)
    train_size = int(config["data"]["train_ratio"] * len(graph_data_list))
    val_size = int(config["data"]["val_ratio"] * len(graph_data_list))

    train_dataset = graph_data_list[:train_size]
    val_dataset = graph_data_list[train_size:train_size + val_size]
     
    # Initialize model
    model = load_model(config)

    # Define the data loaders
    train_loader = DataLoader(train_dataset, batch_size=config["network"]["batch_size"], shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=config["network"]["batch_size"], shuffle=False)

    # Set up the model checkpoint to save only the best model
    checkpoint_callback = ModelCheckpoint(
        monitor="val_loss",  # Metric to monitor
        dirpath=output_dir.as_posix(),  # Save the checkpoint in the output directory
        filename="best-checkpoint",  # Filename for the best checkpoint
        save_top_k=1,  # Only save the best model
        mode="min",  # Mode can be 'min' or 'max', depending on the metric
        verbose=True,
    )

    early_stop_callback = EarlyStopping(
        monitor="val_loss",
        patience=config["network"]["patience"],
        mode="min",
    )

    trainer = pl.Trainer(
        max_epochs=config["network"]["epochs"],
        callbacks=[checkpoint_callback, early_stop_callback],
        logger=logger,
        accelerator="auto",
        devices=1,
    )

    trainer.fit(model, train_loader, val_loader, ckpt_path=args.checkpoint if args.checkpoint else None)

    # Save the model
    model_path = output_dir / "model.pth"
    torch.save(model.state_dict(), model_path)

    predictions = trainer.predict(model, val_loader)

    df = get_features_df(val_loader, task_type)
    df = add_prediction_to_dataframe(df, predictions, task_type)

    output_filepath = {
        "classification": output_dir / "dl2_class_validation.h5",
        "energy": output_dir / "dl2_energy_validation.h5",
        "multitask": output_dir / "dl2_multitask_validation.h5",
        "direction": output_dir / "dl2_direction_validation.h5",
    }

    df.to_hdf(output_filepath[task_type], key="/dl2/event/subarray/Graph", mode="w")


def run_grid_search(config_dict, debug, wandb_outdir=None):
    import wandb

    sweep_config = {
        "method": "bayes",  # You can use 'grid', 'random', or 'bayes'
        "metric": {"name": "val_loss", "goal": "minimize"},
        "parameters": {
            "network.hidden_dim": {"values": [64, 128, 512]},
            "network.learning_rate": {"distribution": "uniform", "min": 1e-5, "max": 1e-2},
            "network.dropout_rate": {
                "values": [
                    0.0,
                    0.1,
                    0.2,
                    0.3,
                    0.5,
                ]
            },
            "network.batch_size": {"values": [32, 64, 128, 512]},
            "network.num_layers": {"values": [3, 4, 5, 6]},
        },
    }

    sweep_id = wandb.sweep(sweep_config, project="stereograph", entity="gammalearn")

    def sweep_train():
        with wandb.init(dir=wandb_outdir):
            config_dict["network"]["hidden_dim"] = wandb.config.get("network.hidden_dim")
            config_dict["network"]["learning_rate"] = wandb.config.get("network.learning_rate")
            config_dict["network"]["dropout_rate"] = wandb.config.get("network.dropout_rate")
            config_dict["network"]["batch_size"] = wandb.config.get("network.batch_size")
            config_dict["network"]["num_layers"] = wandb.config.get("network.num_layers")

            main(config_dict, debug=False)

    wandb.agent(sweep_id, function=sweep_train)


if __name__ == "__main__":
    args = parse_arguments()
    config = load_config(args.config_file)
    output_dir = get_output_dir(config)
    output_dir.mkdir(parents=True, exist_ok=True)
    if not args.overwrite and len(list(output_dir.iterdir())) > 0:
        raise FileExistsError(
            f"Output directory {output_dir} already exists. Use --overwrite to overwrite it."
        )

    if args.grid_search:
        run_grid_search(config, debug=args.debug,  wandb_outdir=output_dir.joinpath("wandb_runs"))
    else:
        main(config, debug=args.debug)
