# -*- coding: utf-8 -*-
"""
KG baseline model on the fine-tuning classification task, assuming the model embeddings are pre-trained.

Run with:
python -m src.stonkgs.models.kg_baseline_model
"""

import logging
import os
from collections import Counter
from typing import Dict, List, Optional

import click  # type: ignore
import mlflow
import numpy as np
import pandas as pd
import pytorch_lightning as pl
import torch
from sklearn.metrics import f1_score
from sklearn.model_selection import KFold, StratifiedShuffleSplit

from stonkgs.constants import (
    CELL_LINE_DIR,
    CORRECT_DIR,
    DISEASE_DIR,
    EMBEDDINGS_PATH,
    KG_BL_OUTPUT_DIR,
    LOCATION_DIR,
    MLFLOW_FINETUNING_TRACKING_URI,
    RANDOM_WALKS_PATH,
    RELATION_TYPE_DIR,
    SPECIES_DIR,
)

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
# Disable alembic info
logging.getLogger("alembic").setLevel(logging.WARNING)


class KGEClassificationModel(pl.LightningModule):
    """KGE baseline model."""

    def __init__(
        self,
        num_classes,
        class_weights,
        d_in: int = 768,
        lr: float = 1e-3,
    ):
        """Initialize the components of the KGE based classification model.

        :param num_classes: number of classes
        :param class_weights: class weights
        :param d_in: dimensions
        :param lr: learning rate

        The model consists of
        1) "Max-Pooling" (embedding-dimension-wise max)
        2) Dropout
        3) Linear layer (d_in x num_classes)
        4) Softmax
        (Not part of the model, but of the class: class_weights for the cross_entropy function)
        """
        super(KGEClassificationModel, self).__init__()

        # Model architecture
        self.pooling = torch.max
        self.dropout = torch.nn.Dropout(0.1)
        self.linear = torch.nn.Linear(d_in, num_classes)
        self.softmax = torch.nn.Softmax(dim=1)

        # Other class-specific parameters
        # Class weights for CE loss
        self.class_weights = torch.tensor(class_weights)
        # Learning rate
        self.lr = lr

        # Add an attribute to save the predictions later on (will be overridden anyways)
        self.predicted_labels = torch.empty(size=(1,))

        # Log the additional parameters
        self.log_dict({"num_classes": num_classes, "class_weights": class_weights, "lr": lr})

    def forward(self, x):
        """Perform forward pass consisting of pooling (dimension-wise max), and a linear layer followed by softmax.

        :param x: embedding sequences (random walk embeddings) for the given triples
        :return: class probabilities for the given triples
        """
        h_pooled = self.pooling(x, dim=1).values
        dropout_output = self.dropout(h_pooled)
        linear_output = self.linear(dropout_output)
        y_pred = self.softmax(linear_output)
        # Note that the forward pass returns class probabilities.
        return y_pred

    def configure_optimizers(self):
        """Configure optimizer and scheduler."""
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
        return optimizer

    def training_step(self, batch, batch_idx):
        """Perform one training step on one batch using CE loss."""
        train_inputs, train_labels = batch
        train_outputs = self.forward(train_inputs)
        loss_fct = torch.nn.CrossEntropyLoss(reduction="mean", weight=self.class_weights)
        loss = loss_fct(train_outputs, train_labels)

        # Log loss at each training step
        self.log("loss", loss.item(), on_step=True)

        return loss

    def test_step(self, batch, batch_nb):
        """Perform one test step on a given batch and return the weighted-averaged f1 score over all batches."""
        test_inputs, test_labels = batch
        test_class_probs = self.forward(test_inputs)
        # Class probabilities to class labels
        test_predictions = torch.argmax(test_class_probs, dim=1)
        # logger.info(f'Predicted labels: {test_predictions}')

        return {"labels": test_labels, "predictions": test_predictions}

    def test_epoch_end(self, outputs):
        """Return average and std weighted-averaged f1-score over all batches."""
        # Get the weighted-averaged f1-score
        all_labels = torch.cat([x["labels"] for x in outputs])
        all_predictions = torch.cat([x["predictions"] for x in outputs])
        test_f1 = f1_score(all_labels, all_predictions, average="weighted")

        # Log the final f1 score
        self.log("f1_score_weighted", test_f1)

        # Use this (weird) fix to save the predicted labels
        self.predicted_labels = all_predictions

        return {"test_f1": test_f1}


class INDRAEntityDataset(torch.utils.data.Dataset):
    """Custom dataset class for INDRA data."""

    def __init__(self, embedding_dict, random_walk_dict, sources, targets, labels, max_len=254):
        """Initialize INDRA Dataset based on random walk embeddings for 2 nodes in each triple."""
        self.max_length = max_len
        # Two entities (source, target) of each triple
        self.sources = sources
        self.targets = targets
        # Initialize dictionary of node name -> embedding vector
        self.embedding_dict = embedding_dict
        # Add the null vector to the embedding dict for "out-of-vocabulary" nodes
        self.embedding_dict[-1] = np.zeros(np.shape(next(iter(self.embedding_dict.values()))))
        # Initialize dictionary of node name -> random walk node names
        self.random_walk_dict = random_walk_dict
        # Get the embedding sequences for each triple
        self.embeddings = self.get_embeddings()
        # Assumes that the labels are numerically encoded
        self.labels = labels

    def __getitem__(self, idx):
        """Get embeddings and labels for given indices."""
        # Get embeddings (of random walk sequences of source + target) for given indices
        item = torch.tensor(self.embeddings[idx, :, :], dtype=torch.float)
        # Get labels for given indices
        labels = torch.tensor(self.labels[idx], dtype=torch.long)
        return item, labels

    def __len__(self):
        """Return the length of the dataset."""
        return len(self.labels)

    def get_embeddings(self):
        """Get the embedding sequences for each triple in the dataset (node emb from sources + targets random walks).

        :return: embedding sequences for each triple in the dataset
        """
        # Number of total examples in the dataset
        number_of_triples = len(self.sources)
        # Get the embedding dimension by accessing a random element
        embedding_dim = len(next(iter(self.embedding_dict.values())))

        # Initialize the embedding array of dimension n x random_walk_length x embedding_dim
        embeddings = np.empty((number_of_triples, self.max_length, embedding_dim))

        # 1. Iterate through all triples: Get random walks for sources and targets using random_walk_dict
        for idx, (source, target) in enumerate(zip(self.sources, self.targets)):
            # Get the random walk sequences
            random_walk_source = self.random_walk_dict[source]
            random_walk_target = self.random_walk_dict[target]

            # 2. Concatenate and the random walks
            # The total random walk has the length max_length. Therefore its split half into the random walk of source
            # and half target.
            random_walk = random_walk_source.tolist() + random_walk_target.tolist()  # noqa: N400
            # 3. Get embeddings for each node using embedding_dict stated by its index in each random walk
            embeds_random_walk = np.stack(
                [self.embedding_dict[node] for node in random_walk], axis=0
            )
            # The final embedding sequence for a given triple has the dimension max_length x embedding_dim
            embeddings[idx, :, :] = embeds_random_walk

        return embeddings


def prepare_df(embedding_path: str, sep: str = "\t") -> Dict[str, List[str]]:
    """Prepare dataframe to node->embeddings/random walks."""
    # Load embeddings
    df = pd.read_csv(
        embedding_path,
        sep=sep,
        header=None,
        index_col=0,
    )
    # Node id -> embeddings
    return {index: row.values for index, row in df.iterrows()}


def get_train_test_splits(
    data: pd.DataFrame,
    label_column_name: str = "class",
    random_seed: int = 42,
    n_splits: int = 5,
    max_dataset_size: int = 100000,
) -> List:
    """Return deterministic train/test indices for n_splits based on the fine-tuning dataset that is passed."""
    # Leave out the label in the dataset
    data_no_labels = data.drop(label_column_name, axis=1)
    labels = data[label_column_name]

    # Cut the dataset down to max_dataset_size (deterministically!) using StratifiedShuffleSplit if needed:
    # (this is not an actual train/test split, this is just for getting a dataset of size max_dataset_size in a
    # stratified and deterministic manner)
    if len(data) > max_dataset_size:
        splitter = StratifiedShuffleSplit(
            n_splits=1,
            train_size=max_dataset_size,
            random_state=random_seed,
        )
        for train_index, _ in splitter.split(data_no_labels, labels):
            data_no_labels = data_no_labels.iloc[train_index, :].reset_index(drop=True)
            labels = labels.iloc[train_index].reset_index(drop=True)

    # Generate the actual train/test splits here:
    # Implement non-stratified train/test splits with no validation split
    # It is shuffled deterministically (determined by random_seed)
    skf = KFold(n_splits=n_splits, random_state=random_seed, shuffle=True)

    # Return a list of dictionaries for train and test indices
    return [
        {"train_idx": train_idx, "test_idx": test_idx}
        for train_idx, test_idx in skf.split(data_no_labels, labels)
    ]


def run_kg_baseline_classification_cv(
    triples_path: str,
    embedding_path: str = EMBEDDINGS_PATH,
    random_walks_path: str = RANDOM_WALKS_PATH,
    logging_uri_mlflow: Optional[str] = MLFLOW_FINETUNING_TRACKING_URI,
    n_splits: int = 5,
    epochs: int = 100,
    train_batch_size: int = 8,
    test_batch_size: int = 64,
    lr: float = 1e-3,
    label_column_name: str = "class",
    log_steps: int = 500,
    task_name: str = "",
    max_dataset_size: int = 100000,
) -> Dict[str, float]:
    """Run KG baseline classification."""
    # Step 1. load the tsv file with the annotation types you want to test and make the splits
    triples_df = pd.read_csv(
        triples_path,
        sep="\t",
        usecols=[
            "source",
            "target",
            label_column_name,
        ],
    )

    # Prepare embeddings and random walks
    embeddings_dict = prepare_df(embedding_path)
    random_walks_dict = prepare_df(random_walks_path)

    # Filter out any triples that contain a node that is not in the embeddings_dict
    original_length = len(triples_df)
    triples_df = triples_df[
        triples_df["source"].isin(embeddings_dict.keys())
        & triples_df["target"].isin(embeddings_dict.keys())
    ].reset_index(drop=True)
    new_length = len(triples_df)
    logger.info(
        f"{original_length - new_length} out of {original_length} triples are left out because they contain "
        f"nodes which are not present in the pre-training data"
    )

    # Numerically encode labels
    unique_tags = set(label for label in triples_df[label_column_name])
    tag2id = {label: number for number, label in enumerate(unique_tags)}
    id2tag = {value: key for key, value in tag2id.items()}

    # Get labels
    labels = pd.Series([int(tag2id[label]) for label in triples_df[label_column_name]])

    # Get the train/test split indices
    train_test_splits = get_train_test_splits(
        triples_df,
        n_splits=n_splits,
        label_column_name=label_column_name,
        max_dataset_size=max_dataset_size,
    )

    # Initialize f1-scores
    f1_scores = []

    # Initialize INDRA for KG baseline dataset
    kg_embeds = INDRAEntityDataset(
        embeddings_dict,
        random_walks_dict,
        triples_df["source"],
        triples_df["target"],
        labels,
    )

    mlflow.set_tracking_uri(logging_uri_mlflow)
    mlflow.set_experiment("KG Baseline for STonKGs")

    mlflow.pytorch.autolog()

    # Initialize a dataframe for all the predicted labels
    result_df = pd.DataFrame()

    # Train and test the model in a cv setting
    for idx, indices in enumerate(train_test_splits):
        # Sample elements randomly from a given list of ids, no replacement
        train_subsampler = torch.utils.data.SubsetRandomSampler(indices["train_idx"])
        test_subsampler = torch.utils.data.SubsetRandomSampler(indices["test_idx"])
        # CE class weights for the model based on training data class distribution,
        # based on the class counts (Inverse Number of Samples, INS)
        weights = [
            1
            / len(
                [
                    i
                    for i in triples_df.iloc[indices["train_idx"], :][label_column_name]
                    # note that we only employ train idx
                    if i == id2tag[id_num]
                ]
            )
            for id_num in range(len(unique_tags))
        ]

        # Define data loaders for training and testing data in this fold
        trainloader = torch.utils.data.DataLoader(
            kg_embeds,
            batch_size=train_batch_size,
            sampler=train_subsampler,
        )
        testloader = torch.utils.data.DataLoader(
            kg_embeds,
            batch_size=test_batch_size,
            sampler=test_subsampler,
        )

        model = KGEClassificationModel(
            num_classes=len(triples_df[label_column_name].unique()),
            class_weights=weights,
            lr=lr,
        )

        # Initialize pytorch lighting Trainer for the KG baseline model
        trainer = pl.Trainer(
            max_epochs=epochs,
            log_every_n_steps=log_steps,
        )

        # Train and predict in a separate run for each split
        with mlflow.start_run():
            # Fit on training split
            trainer.fit(model, train_dataloader=trainloader)
            # Predict on test split
            test_results = trainer.test(model, test_dataloaders=testloader)

            # Save the predicted + true labels
            partial_result_df = pd.DataFrame(
                {
                    "split": idx,
                    "index": indices["test_idx"].tolist(),
                    "predicted_label": model.predicted_labels.tolist(),
                    "true_label": labels[indices["test_idx"]].tolist(),
                },
            )
            result_df = result_df.append(
                partial_result_df,
                ignore_index=True,
            )

            # Log some details about the datasets used in training and testing
            mlflow.log_param("label dict", str(tag2id))
            mlflow.log_param("training dataset size", str(len(trainloader.dataset)))
            mlflow.log_param("training class dist", str(Counter(trainloader.dataset.labels)))
            mlflow.log_param("test dataset size", str(len(testloader.dataset)))
            mlflow.log_param("test class dist", str(Counter(testloader.dataset.labels)))

        # Append f1 score per split based on the weighted average
        f1_scores.append(test_results[0]["test_f1"])

    # Map the labels in the result df back to their original names
    result_df = result_df.replace({"predicted_label": id2tag, "true_label": id2tag})
    # Save the result_df
    result_df.to_csv(
        os.path.join(KG_BL_OUTPUT_DIR, "predicted_labels_kg_" + task_name + "df.tsv"),
        index=False,
        sep="\t",
    )

    # Save the last model
    trainer.save_checkpoint(os.path.join(KG_BL_OUTPUT_DIR, "kg_baseline.ckpt"))

    # Log the mean and std f1 score from the cross validation procedure to mlflow
    with mlflow.start_run():
        mlflow.log_metric("f1_score_mean", np.mean(f1_scores))
        mlflow.log_metric("f1_score_std", np.std(f1_scores))

        # Log the task name as well
        mlflow.log_param("task name", task_name)

        # Also log how many triples were left out
        mlflow.log_param("original no. of triples", original_length)
        mlflow.log_param("no. of left out triples", original_length - new_length)

    # Log mean and std f1-scores from the cross validation procedure (average and std across all splits) to the
    # standard logger
    logger.info(f"Mean f1-score: {np.mean(f1_scores)}")
    logger.info(f"Std f1-score: {np.std(f1_scores)}")

    # Return the final f1 score mean and the std for all CV folds
    return {"f1_score_mean": np.mean(f1_scores), "f1_score_std": float(np.std(f1_scores))}


@click.command()
@click.option("-e", "--epochs", default=5, help="Number of epochs", type=int)
@click.option("--lr", default=1e-4, help="Learning rate", type=float)
@click.option(
    "--logging_dir",
    default=MLFLOW_FINETUNING_TRACKING_URI,
    help="Mlflow logging/tracking URI",
    type=str,
)
@click.option("--log_steps", default=500, help="Number of steps between each log", type=int)
@click.option("--batch_size", default=16, help="Batch size", type=int)
@click.option(
    "--max_dataset_size",
    default=100000,
    help="Maximum dataset size of the fine-tuning datasets",
    type=int,
)
def run_all_fine_tuning_tasks(
    epochs: int = 5,
    log_steps: int = 500,
    lr: float = 1e-3,
    logging_dir: Optional[str] = MLFLOW_FINETUNING_TRACKING_URI,
    batch_size: int = 16,
    max_dataset_size: int = 100000,  # effectively removes the max dataset size restriction
):
    """Run all fine-tuning tasks at once."""
    # Specify all directories and file names
    directories = [
        CELL_LINE_DIR,
        CORRECT_DIR,
        CORRECT_DIR,
        DISEASE_DIR,
        LOCATION_DIR,
        SPECIES_DIR,
        RELATION_TYPE_DIR,
        RELATION_TYPE_DIR,
    ]
    file_names = [
        "cell_line_no_duplicates.tsv",
        "correct_incorrect_binary_no_duplicates.tsv",
        "correct_incorrect_multiclass_no_duplicates.tsv",
        "disease_no_duplicates.tsv",
        "location_no_duplicates.tsv",
        "species_no_duplicates.tsv",
        "relation_type_no_duplicates.tsv",
        "relation_type_no_duplicates.tsv",
    ]
    task_names = [
        "cell_line",
        "correct_binary",
        "correct_multiclass",
        "disease",
        "location",
        "species",
        "interaction",
        "polarity",
    ]
    # Specify the column names of the target variable
    column_names = ["class"] * 6 + ["interaction"] + ["polarity"]

    for directory, file, column_name, task_name in zip(
        directories,
        file_names,
        column_names,
        task_names,
    ):
        # Run each of the six classification tasks
        run_kg_baseline_classification_cv(
            triples_path=os.path.join(directory, file),
            label_column_name=column_name,
            logging_uri_mlflow=logging_dir,
            epochs=epochs,
            lr=lr,
            log_steps=log_steps,
            train_batch_size=batch_size,
            task_name=task_name,
            max_dataset_size=max_dataset_size,
        )
        logger.info(f"Finished the {task_name} task")


if __name__ == "__main__":
    # Run all fine-tuning classification tasks at once
    run_all_fine_tuning_tasks()
