# Copyright 2024 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Run hyperparameter search and store results with a command-line script."""

import numpy as np
import sys
import os
import time
import argparse
import logging

logging.getLogger().setLevel(logging.INFO)
from importlib import import_module
import pandas as pd
from pathlib import Path
import matplotlib.pyplot as plt
from sklearn.model_selection import GridSearchCV
from qcml.utils.hyperparam_search import (
    read_data,
    construct_hyperparameter_grid,
)
from qcml.bench.hyperparameter_settings import hyper_parameter_settings

np.random.seed(42)

logging.info("cpu count:" + str(os.cpu_count()))


if __name__ == "__main__":
    # Create an argument parser
    parser = argparse.ArgumentParser(
        description="Run experiments with hyperparameter search.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--classifier-name",
        help="Classifier to run",
    )

    parser.add_argument(
        "--dataset-path",
        help="Path to the dataset",
    )

    parser.add_argument(
        "--results-path", default=".", help="Path to store the experiment results"
    )

    parser.add_argument(
        "--clean",
        help="True or False. Remove previous results if it exists",
        dest="clean",
        default=False,
        type=bool,
    )

    parser.add_argument(
        "--hyperparameter-scoring",
        type=list,
        nargs="+",
        default=["accuracy", "roc_auc"],
        help="Scoring for hyperparameter search.",
    )

    parser.add_argument(
        "--hyperparameter-refit",
        type=str,
        default="accuracy",
        help="Refit scoring for hyperparameter search.",
    )

    parser.add_argument(
        "--plot-loss",
        help="True or False. Plot loss history for single fit",
        dest="plot_loss",
        default=False,
        type=bool,
    )

    parser.add_argument(
        "--n-jobs", type=int, default=-1, help="Number of parallel threads to run"
    )

    # Parse the arguments along with any extra arguments that might be model specific
    args, unknown_args = parser.parse_known_args()

    if any(arg is None for arg in [args.classifier_name, args.dataset_path]):
        msg = "\n================================================================================"
        msg += "\nA classifier from qml.benchmarks.model and dataset path are required. E.g., \n \n"
        msg += "python run_hyperparameter_search \ \n--classifier DataReuploadingClassifier \ \n--dataset-path train.csv\n"
        msg += "\nCheck all arguments for the script with \n"
        msg += "python run_hyperparameter_search --help\n"
        msg += "================================================================================"
        raise ValueError(msg)

    # Add model specific arguments to override the default hyperparameter grid
    hyperparam_grid = construct_hyperparameter_grid(
        hyper_parameter_settings, args.classifier_name
    )
    for hyperparam in hyperparam_grid:
        hp_type = type(hyperparam_grid[hyperparam][0])
        parser.add_argument(
            f"--{hyperparam}",
            type=hp_type,
            nargs="+",
            default=hyperparam_grid[hyperparam],
            help=f"{hyperparam} grid values for {args.classifier_name}",
        )

    args = parser.parse_args(unknown_args, namespace=args)

    for hyperparam in hyperparam_grid:
        override = getattr(args, hyperparam)
        if override is not None:
            hyperparam_grid[hyperparam] = override
    logging.info(
        "Running hyperparameter search experiment with the following settings\n"
    )
    logging.info(args.classifier_name)
    logging.info(args.dataset_path)
    logging.info(" ".join(args.hyperparameter_scoring))
    logging.info(args.hyperparameter_refit)
    logging.info(
        "Hyperparam grid:"
        + " ".join(
            [
                (str(key) + str(":") + str(hyperparam_grid[key]))
                for key in hyperparam_grid.keys()
            ]
        )
    )

    experiment_path = args.results_path
    results_path = os.path.join(experiment_path, "results")

    if not os.path.exists(results_path):
        os.makedirs(results_path)

    ###################################################################
    # Get the classifier, dataset and search methods from the arguments
    ###################################################################
    Classifier = getattr(import_module("qml_benchmarks.models"), args.classifier_name)
    classifier_name = Classifier.__name__

    # Run the experiments save the results
    train_dataset_filename = os.path.join(args.dataset_path)
    X, y = read_data(train_dataset_filename)

    dataset_path_obj = Path(args.dataset_path)
    results_filename_stem = " ".join(
        [Classifier.__name__ + "_" + dataset_path_obj.stem + "_GridSearchCV"]
    )

    # If we have already run this experiment then continue
    if os.path.isfile(os.path.join(results_path, results_filename_stem + ".csv")):
        if args.clean is False:
            msg = "\n================================================================================="
            msg += "\nResults exist in " + os.path.join(
                results_path, results_filename_stem + ".csv"
            )
            msg += "\nSpecify --clean True to override results or new --results-path"
            msg += "\n================================================================================="
            logging.warning(msg)
            sys.exit(msg)
        else:
            logging.warning(
                "Cleaning existing results for ",
                os.path.join(results_path, results_filename_stem + ".csv"),
            )

    ###########################################################################
    # Single fit to check everything works
    ###########################################################################
    classifier = Classifier()
    a = time.time()
    classifier.fit(X, y)
    b = time.time()
    acc_train = classifier.score(X, y)
    logging.info(
        " ".join(
            [
                classifier_name,
                "Dataset path",
                args.dataset_path,
                "Train acc:",
                str(acc_train),
                "Time single run",
                str(b - a),
            ]
        )
    )
    if hasattr(classifier, "loss_history_"):
        if args.plot_loss:
            plt.plot(classifier.loss_history_)
            plt.xlabel("Iterations")
            plt.ylabel("Loss")
            plt.show()

    if hasattr(classifier, "n_qubits_"):
        logging.info(" ".join(["Num qubits", f"{classifier.n_qubits_}"]))

    ###########################################################################
    # Hyperparameter search
    ###########################################################################
    gs = GridSearchCV(
        estimator=classifier,
        param_grid=hyperparam_grid,
        scoring=args.hyperparameter_scoring,
        refit=args.hyperparameter_refit,
        verbose=3,
        n_jobs=-1,
    ).fit(X, y)
    logging.info("Best hyperparams")
    logging.info(gs.best_params_)

    df = pd.DataFrame.from_dict(gs.cv_results_)
    df.to_csv(os.path.join(results_path, results_filename_stem + ".csv"))

    best_df = pd.DataFrame(
        list(gs.best_params_.items()), columns=["hyperparameter", "best_value"]
    )

    # Save best hyperparameters to a CSV file
    best_df.to_csv(
        os.path.join(results_path, results_filename_stem + "-best-hyperparameters.csv"),
        index=False,
    )
