import argparse
from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim

import ray
from ray import tune
from ray.tune import Checkpoint as RayCheckpoint, CLIReporter
from ray.tune.schedulers import ASHAScheduler

from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Loss, Accuracy
from ignite.handlers import Checkpoint, DiskSaver

from utils import Net, get_data_loaders, get_test_loader, load_data


def train_with_ignite(config, data_dir=None, num_epochs=10, num_workers=8):
    device = config.get("device", "cpu")
    net = Net(config["l1"], config["l2"])
    net = net.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=config["lr"], momentum=0.9)

    train_loader, val_loader = get_data_loaders(config["batch_size"], data_dir, num_workers)

    trainer = create_supervised_trainer(net, optimizer, criterion, device=device)

    metrics = {"accuracy": Accuracy(), "loss": Loss(criterion)}
    evaluator = create_supervised_evaluator(net, metrics=metrics, device=device)

    checkpoint_dir = tune.get_context().get_trial_dir()
    save_handler = DiskSaver(checkpoint_dir, create_dir=True)
    to_save = {"net": net, "optimizer": optimizer, "trainer": trainer}
    checkpoint_handler = Checkpoint(to_save, save_handler=save_handler, filename_pattern="checkpoint.pt")
    trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler)

    @trainer.on(Events.EPOCH_COMPLETED)
    def log_and_report(engine):
        epoch = engine.state.epoch

        evaluator.run(val_loader)
        val_loss = evaluator.state.metrics["loss"]
        val_accuracy = evaluator.state.metrics["accuracy"]

        print(f"Epoch {epoch}: loss={val_loss:.4f}, accuracy={val_accuracy:.4f}")

        checkpoint = RayCheckpoint.from_directory(checkpoint_dir)
        tune.report(
            {"loss": val_loss, "accuracy": val_accuracy},
            checkpoint=checkpoint,
        )

    trainer.run(train_loader, num_epochs)


def test_accuracy(model, device, data_dir):
    test_loader = get_test_loader(4, data_dir)
    model = model.to(device)

    evaluator = create_supervised_evaluator(model, metrics={"accuracy": Accuracy()}, device=device)
    evaluator.run(test_loader)

    return evaluator.state.metrics["accuracy"]


def tune_cifar(num_trials=10, num_epochs=10, gpus_per_trial=1, cpus_per_trial=4, data_dir="./data"):
    data_dir = Path(data_dir).resolve()

    config = {
        "l1": tune.choice([2**i for i in range(2, 9)]),
        "l2": tune.choice([2**i for i in range(2, 9)]),
        "lr": tune.loguniform(1e-4, 1e-1),
        "batch_size": tune.choice([256, 512, 1024]),
        "device": "cuda" if gpus_per_trial > 0 else "cpu",
    }

    scheduler = ASHAScheduler(
        max_t=num_epochs,
        grace_period=1,
        reduction_factor=2,
    )

    print("Downloading CIFAR10 dataset...")

    load_data(data_dir)

    reporter = CLIReporter(
        metric_columns={"loss": "loss", "accuracy": "accuracy"},
        parameter_columns={"l1": "l1", "l2": "l2", "lr": "lr", "batch_size": "batch_size"},
    )

    tuner = tune.Tuner(
        tune.with_resources(
            tune.with_parameters(
                train_with_ignite,
                data_dir=data_dir,
                num_epochs=num_epochs,
            ),
            resources={"cpu": cpus_per_trial, "gpu": gpus_per_trial},
        ),
        tune_config=tune.TuneConfig(
            metric="loss",
            mode="min",
            scheduler=scheduler,
            num_samples=num_trials,
        ),
        param_space=config,
        run_config=ray.tune.RunConfig(
            storage_path=None,
            name="cifar10_ray_tune",
            progress_reporter=reporter,
        ),
    )

    results = tuner.fit()

    best_result = results.get_best_result("accuracy", "max", filter_nan_and_inf=False)
    print(f"Best trial config: {best_result.config}")
    print(f"Best trial final validation loss: {best_result.metrics['loss']}")
    print(f"Best trial final validation accuracy: {best_result.metrics['accuracy']}")

    best_trained_model = Net(best_result.config["l1"], best_result.config["l2"])
    device = best_result.config["device"]
    best_trained_model = best_trained_model.to(device)

    best_checkpoint = best_result.checkpoint
    with best_checkpoint.as_directory() as checkpoint_dir:
        checkpoint_path = Path(checkpoint_dir) / "checkpoint.pt"
        best_checkpoint_data = torch.load(checkpoint_path, map_location=device)
        best_trained_model.load_state_dict(best_checkpoint_data["net"])

        test_acc = test_accuracy(best_trained_model, device, data_dir)
        print(f"Best trial test set accuracy: {test_acc}")

    return results


def main():
    parser = argparse.ArgumentParser(description="Ignite + Ray Tune CIFAR10 Example")
    parser.add_argument("--num_trials", type=int, default=10, help="Number of trials for Ray Tune")
    parser.add_argument("--num_epochs", type=int, default=10, help="Number of epochs per trial")
    parser.add_argument("--gpus_per_trial", type=float, default=1, help="GPUs per trial")
    parser.add_argument("--cpus_per_trial", type=int, default=4, help="CPUs per trial")
    parser.add_argument(
        "--data_dir",
        type=str,
        default="./data",
        help="Path to CIFAR10 dataset",
    )
    args = parser.parse_args()

    ray.init(include_dashboard=False)
    tune_cifar(
        num_trials=args.num_trials,
        num_epochs=args.num_epochs,
        gpus_per_trial=args.gpus_per_trial,
        cpus_per_trial=args.cpus_per_trial,
        data_dir=args.data_dir,
    )


if __name__ == "__main__":
    main()
