import yaml
import argparse
from pathlib import Path
import numpy as np
import torch
import random
import pytorch_lightning as pl
from torch_geometric.data import DataLoader
from pytorch_lightning.loggers import WandbLogger
from stereograph.data import get_features_df, add_prediction_to_dataframe, get_graphs_from_file
from stereograph.load_model import load_model


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(
        "--checkpoint", "-ckpt", type=Path, help="Path to a checkpoint file to load the best checkpoint for inference "
    )
    parser.add_argument("--gamma_test", type=Path, help="Path to the test file")
    parser.add_argument("--proton_test", type=Path, help="Path to the test file")

    
  

    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, checkpoint_path, debug=False, reload_graphs=False):
    output_dir = Path(config["output_dir"], config["exp_name"])
    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,
        )
        logger.experiment.config.update(config)
    else:
        logger = None


 
    task_type = config["network"]["task_type"]
    gamma_test_file=args.gamma_test
    
    #load test files
    gamma_test_graphs=get_graphs_from_file(gamma_test_file, config)
    
    #Create the graph dataset
    if task_type == "classification": #we use both protons and gammas for cassification, otherwise onlt gamma files are used 
         proton_test_file=args.proton_test
         proton_test_graphs=get_graphs_from_file(proton_test_file, config)
         if not args.proton_test:
            raise ValueError("For classification task, proton_file must be provided.")
         test_dataset = gamma_test_graphs + proton_test_graphs
    else:
        test_dataset = gamma_test_graphs 


    # Initialize model
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Model checkpoint not found at {checkpoint_path}")
    model = load_model(config, checkpoint_path)

    test_loader = DataLoader(test_dataset, batch_size=config["network"]["batch_size"], shuffle=False)

    trainer = pl.Trainer(
        logger=logger,
        accelerator="auto",
    )

    predictions = trainer.predict(model, test_loader)

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

    output_filepath = {
        "classification": output_dir / "dl2_class_test.h5",
        "energy": output_dir / "dl2_energy_test.h5",
        "multitask": output_dir / "dl2_multitask_test.h5",
        "direction": output_dir / "dl2_direction_test.h5",
    }[config["network"]["task_type"]]

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


if __name__ == "__main__":
    args = parse_arguments()
    config = yaml.safe_load(args.config_file.open())
    output_dir = Path(config["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)

    if not args.checkpoint.exists():
        raise FileNotFoundError(f"Checkpoint file {args.checkpoint} not found")

    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."
        )
    else:
        main(config, args.checkpoint, debug=args.debug)
