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/new_data_dir3")

    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', 'metasgd']),
        "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([1, 3, 5, 7]),
        "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.001, 0.1, 0.5, 1.0, 2.0, 10.0]),
        'trainer_dropout_rate': tune.choice([0, 0.01, 0.1]),

        'use_cuda': tune.choice([True]),
        # item
        'num_rate': tune.choice([6]),
        'num_genre': tune.choice([25]),
        'num_director': tune.choice([2186]),
        'num_actor': tune.choice([8030]),
        # user
        'num_gender': tune.choice([2]),
        'num_age': tune.choice([7]),
        'num_occupation': tune.choice([21]),
        'num_zipcode': tune.choice([3402]),

        'num_epoch': tune.choice([30]),
    }

    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": 8, "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_cold2",
        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=30, gpus_per_trial=1)