"""
This experiment runs the hyperbolic tsne code changing several parameters:
- dataset: [LUKK, MYELOID8000, PLANARIA, MNIST, C_ELEGANS, WORDNET]
- Value of Theta in the approximation: [0.0, 0.1, ..., 1.0]
The code only computes the runs that do not have a folder.
"""

###########
# IMPORTS #
###########

from pathlib import Path

import csv
import json
import traceback

import numpy as np
from matplotlib import pyplot as plt
from scipy.sparse import issparse, save_npz

from hyperbolicTSNE import Datasets, load_data, initialization, SequentialOptimizer, HyperbolicTSNE
from hyperbolicTSNE.util import find_last_embedding
from hyperbolicTSNE.visualization import plot_poincare

#################################
# GENERAL EXPERIMENT PARAMETERS #
#################################

BASE_DIR = "../results/timings_per_theta"  # dir where results will be saved
DATASETS_DIR = "../datasets"  # directory to read the data from

# Constants
SEED = 42  # seed to initialize random processes
PERP = 30  # perplexity value to be used throughout the experiments
KNN_METHOD = "hnswlib"  # use hnswlib for determining nearest neighbors in high-dimensional space; note that this is
# an approximation, switch to "sklearn" for an exact method
VANILLA = False  # if vanilla is set to true, regular gradient descent without any modifications is performed; for
# vanilla set to false, the optimization makes use of momentum and gains
EXAG = 12  # the factor by which the attractive forces are amplified during early exaggeration
hd_params = {"perplexity": PERP}

# Variables
datasets = [
    Datasets.LUKK,
    Datasets.MYELOID8000,
    Datasets.PLANARIA,
    Datasets.MNIST,
    Datasets.C_ELEGANS,
    Datasets.WORDNET
]
thetas = [n / 20 for n in range(20, -1, -1)]  # The different theta values to be used in the acceleration experiment

###################
# EXPERIMENT LOOP #
###################

overview_created = False
for dataset in datasets:  # Iterate over the data sets

    rng = np.random.default_rng(seed=SEED)  # random number generator

    dataX, dataLabels, D, V = load_data(
        dataset,
        data_home=DATASETS_DIR,
        to_return="X_labels_D_V",  # Return the high-dimensional data, its labels, the NN-graph, the probabilty matrix
        hd_params=hd_params,
        knn_method=KNN_METHOD
    )

    n_samples = dataX.shape[0]

    X_embedded = initialization(  # create an initial embedding of the data into 2-dimensional space via PCA
        n_samples=n_samples,
        n_components=2,
        X=dataX,
        random_state=rng.integers(0, 1000000),
        method="pca"
    )

    for config_id, theta in enumerate(thetas):

        print(f"[theta_timings] Processing {dataset}, config_id ({config_id}) with Theta: {theta}")

        LR = (dataX.shape[0] * 1) / (EXAG * 1000)  # Compute the learning rate

        opt_params = SequentialOptimizer.sequence_poincare(
            learning_rate_ex=LR,  # specify learning rate for the early exaggeration
            learning_rate_main=LR,  # specify learning rate for the non-exaggerated gradient descent
            exaggeration=EXAG,
            vanilla=VANILLA,
            momentum_ex=0.5,  # momentum to be used during early exaggeration
            momentum=0.8,  # momentum to be used during non-exaggerated gradient descent
            exact=False,
            n_iter_check=10,  # Needed for early stopping criterion
            size_tol=0.999,  # Size of the embedding to be used as early stopping criterion
            angle=theta  # The theta value to be used in the acceleration
        )

        run_dir = Path(f"{BASE_DIR}/{dataset.name}/theta_{theta}/")

        if run_dir.exists():
            # Skip already computed embeddings
            print(f"[theta_timings] - Exists so not computing it: {run_dir}")
        else:
            run_dir.mkdir(parents=True, exist_ok=True)

            params = {
                "lr": LR,
                "perplexity": PERP,
                "seed": SEED,
                "sample_size": int(n_samples),
                "tsne_type": "accelerated",
                "splitting_strategy": "equal_length",
                "theta": theta
            }

            print(f"[theta_timings] - Starting configuration {config_id} with dataset {dataset.name}: {params}")

            opt_params["logging_dict"] = {
                "log_path": str(run_dir.joinpath("embeddings"))
            }

            # Save the high-dimensional neighborhood matrices for later use
            json.dump(params, open(run_dir.joinpath("params.json"), "w"))
            if issparse(D):
                save_npz(run_dir.joinpath("D.npz"), D)
            else:
                np.save(run_dir.joinpath("D.npy"), D)
            if issparse(V):
                save_npz(run_dir.joinpath("P.npz"), V)
            else:
                np.save(run_dir.joinpath("P.npy"), V)

            hdeo_hyper = HyperbolicTSNE(
                init=X_embedded,
                n_components=X_embedded.shape[1],
                metric="precomputed",
                verbose=2,
                opt_method=SequentialOptimizer,
                opt_params=opt_params
            )

            error_title = ""
            try:
                res_hdeo_hyper = hdeo_hyper.fit_transform((D, V))
            except ValueError:

                error_title = "_error"
                res_hdeo_hyper = find_last_embedding(opt_params["logging_dict"]["log_path"])
                traceback.print_exc(file=open(str(run_dir) + "traceback.txt", "w"))

                print("[theta_timings] - Run failed ...")

            else:  # we save the data if there were no errors

                print("[theta_timings] - Finished running, saving run data directory ...")

                # Save the final embedding coordinates
                np.save(run_dir.joinpath("final_embedding.npy"), res_hdeo_hyper)

                # Save a plot of the final embedding
                fig = plot_poincare(res_hdeo_hyper, labels=dataLabels)
                fig.savefig(run_dir.joinpath(f"final_embedding{error_title}.png"))
                plt.close(fig)

                np.save(run_dir.joinpath("logging_dict.npy"), opt_params["logging_dict"])

                # Write out timings csv
                timings = np.array(hdeo_hyper.optimizer.cf.results)
                with open(run_dir.joinpath("timings.csv"), "w", newline="") as timings_file:
                    timings_writer = csv.writer(timings_file)
                    timings_writer.writerow(["it_n", "time_type", "total_time"])

                    for n, row in enumerate(timings):
                        timings_writer.writerow([n, "tree_building", row[0]])
                        timings_writer.writerow([n, "tot_gradient", row[1]])
                        timings_writer.writerow([n, "neg_force", row[2]])
                        timings_writer.writerow([n, "pos_force", row[3]])

                    # Create or append to overview csv file after every run
                    with open(run_dir.joinpath(f"overview_part.csv"), "w", newline="") as overview_file:
                        overview_writer = csv.writer(overview_file)
                        overview_writer.writerow(["dataset", *params, "run", "run_directory", "error"])
                        overview_writer.writerow(
                            [dataset.name, *params.values(), str(run_dir).replace(str(BASE_DIR), "."),
                             error_title != ""])

                    print()