from hyper_tunning import load_data
import os
from ray.tune.schedulers import ASHAScheduler
from ray.tune import CLIReporter
from ray import tune
from functools import partial
from hyper_tunning import train_melu
import numpy as np


def main(num_samples, max_num_epochs=20, gpus_per_trial=2):
    data_dir = os.path.abspath("/media/external_10TB/10TB/maheri/define_task_melu_data")
    load_data(data_dir)
    config = {
        # "l1": tune.sample_from(lambda _: 2 ** np.random.randint(2, 9)),
        # "l2": tune.sample_from(lambda _: 2 ** np.random.randint(2, 9)),
        # "lr": tune.loguniform(1e-4, 1e-1),
        # "batch_size": tune.choice([2, 4, 8, 16])
        "transformer": tune.choice(['kronoker']),
        "meta_algo": tune.choice(['gbml']),
        "first_order": tune.choice([False]),
        "adapt_transform": tune.choice([True, False]),
        # "local_lr":tune.choice([5e-6,5e-4,5e-3]),
        # "lr":tune.choice([5e-5,5e-4]),
        "local_lr": tune.loguniform(5e-6, 5e-3),
        "lr": tune.loguniform(5e-5, 5e-3),
        "batch_size": tune.choice([16, 32, 64]),
        "inner": tune.choice([7, 5, 4, 3, 1]),
        "test_state": tune.choice(["user_and_item_cold_state"]),

        "embedding_dim": tune.choice([16, 32, 64]),
        "first_fc_hidden_dim": tune.choice([32, 64, 128]),
        "second_fc_hidden_dim": tune.choice([32, 64]),

        'cluster_h1_dim': tune.choice([256, 128, 64]),
        'cluster_h2_dim': tune.choice([128, 64, 32]),
        'cluster_final_dim': tune.choice([64, 32]),
        'cluster_dropout_rate': tune.choice([0, 0.01, 0.1]),
        'cluster_k': tune.choice([3, 5, 7, 9, 11]),
        'temperature': tune.choice([0.1, 0.5, 1.0, 2.0, 10.0]),
        'trainer_dropout_rate': tune.choice([0, 0.01, 0.1]),
    }

    scheduler = ASHAScheduler(
        metric="loss",
        mode="min",
        max_t=30,
        grace_period=10,
        reduction_factor=2)
    reporter = CLIReporter(
        # parameter_columns=["l1", "l2", "lr", "batch_size"],
        metric_columns=["loss", "ndcg1", "ndcg3", "training_iteration"])
    result = tune.run(
        partial(train_melu, data_dir=data_dir),
        resources_per_trial={"cpu": 4, "gpu": gpus_per_trial},
        config=config,
        num_samples=num_samples,
        scheduler=scheduler,
        progress_reporter=reporter,
        log_to_file=True,
        # resume=True,
        local_dir="./hyper_tunning_all_cold",
        name="melu_all_cold_clustered",

    )

    best_trial = result.get_best_trial("loss", "min", "last")
    print("Best trial config: {}".format(best_trial.config))
    print("Best trial final validation loss: {}".format(
        best_trial.last_result["loss"]))
    print("Best trial final validation ndcg1: {}".format(
        best_trial.last_result["ndcg1"]))
    print("Best trial final validation ndcg3: {}".format(
        best_trial.last_result["ndcg3"]))

    #
    print("=======================================================")
    print(result.results_df)
    print("=======================================================\n")

    # best_trained_model = Net(best_trial.config["l1"], best_trial.config["l2"])
    # device = "cpu"
    # if torch.cuda.is_available():
    #     device = "cuda:0"
    #     if gpus_per_trial > 1:
    #         best_trained_model = nn.DataParallel(best_trained_model)
    # best_trained_model.to(device)
    #
    # best_checkpoint_dir = best_trial.checkpoint.value
    # model_state, optimizer_state = torch.load(os.path.join(
    #     best_checkpoint_dir, "checkpoint"))
    # best_trained_model.load_state_dict(model_state)
    #
    # test_acc = test_accuracy(best_trained_model, device)
    # print("Best trial test set accuracy: {}".format(test_acc))


if __name__ == "__main__":
    # You can change the number of GPUs per trial here:
    main(num_samples=150, max_num_epochs=25, gpus_per_trial=1)