"""
    ML models for classification
"""
from typing import Dict, List, Tuple, Optional, Union
from datasets import Dataset
import evaluate
import numpy as np
import pandas as pd

def data_idx_split(idx: List[int],
                   ratio: float = 0.2,
                   valratio: float = 0.5,
                   random_state: int = 42) -> Dict[str, np.ndarray]:
    """
    Split indices into training, validation, and test sets.

    This function splits a set of indices into training, validation, and test 
    sets using specified ratios. First, it divides the indices into a training 
    set and a temporary set (for validation and testing). Then it further splits 
    the temporary set into validation and test sets.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    idx : Union[pd.Index, np.ndarray, List[int]]
        The indices to split (e.g., DataFrame indices).
    ratio : float, optional (default=0.2)
        The proportion of the total indices to allocate to the combined 
        validation and test sets.
    valratio : float, optional (default=0.5)
        The proportion of the temporary (validation + test) set to allocate to 
        the test set.
    random_state : int, optional (default=42)
        The random seed to ensure reproducibility of the splits.

    ---------------------------------------------------------------------------
    OUTPUT:
    indices : Dict[str, np.ndarray]
        A dictionary with keys 'train', 'valid', and 'test', each containing the 
        corresponding index arrays.

    ---------------------------------------------------------------------------
    Copyright (C) - 2025 - Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Step 1: Split into train and temporary (valid+test)
    from sklearn.model_selection import train_test_split

    train_indices, temp_indices = train_test_split(
        idx, test_size=ratio, random_state=random_state
    )

    if valratio > 0:
        # Step 2: Split temporary data into validation and test sets
        valid_indices, test_indices = train_test_split(
            temp_indices, test_size=valratio, random_state=random_state
            )
    else:
        valid_indices = temp_indices
        test_indices = None

    indices = {
        'train': train_indices,
        'valid': valid_indices,
        'test': test_indices
    }

    return indices


def list2hfdict(text_list: List[str],
              labels: Optional[List] = None) -> Dataset:
    """
    Convert a list of text entries into a Hugging Face Dataset format.

    This function takes a list of input texts and optionally a list of target labels, 
    and converts them into a Hugging Face Dataset object.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text_list : List[str]
        A list of strings representing the input texts.
    labels : Optional[List]
        An optional list of labels corresponding to each text entry.

    ---------------------------------------------------------------------------
    OUTPUT:
    output : Dataset
        A Hugging Face Dataset object containing 'text' and optionally 'label' fields.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if labels is not None:
        if len(text_list) != len(labels):
            raise ValueError(
                "Length of text_list and labels must be the same.")
        output = Dataset.from_dict({'text': text_list, 'label': labels})
    else:
        output = Dataset.from_dict({'text': text_list})

    return output


def train_val_test_split(df: pd.DataFrame,
                         covariate: str,
                         target: str,
                         indices: Dict[str, np.ndarray],
                         ratio: float = 0.3,
                         valratio: float = 0.5) -> Tuple[Dataset, Dataset, Dataset]:
    """
    Split a DataFrame into train, validation, and test datasets.

    This function splits the provided DataFrame into train, validation, and test 
    sets based on indices and specified ratios. It extracts the covariate and 
    target columns for each set and converts them into Hugging Face `Dataset` 
    objects.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    df : pd.DataFrame
        The DataFrame containing the data to split.
    covariate : str
        The name of the column to be used as the input feature (e.g., text).
    target : str
        The name of the column to be used as the label.
    indices : Dict[str, np.ndarray]
        A dictionary with keys 'train', 'valid', and 'test' containing index 
        lists for each split.
    ratio : float, optional (default=0.3)
        The proportion of the dataset to allocate to the test and validation 
        sets combined.
    valratio : float, optional (default=0.5)
        The proportion of the non-training set to allocate to the validation 
        set (the rest is test data).

    ---------------------------------------------------------------------------
    OUTPUT:
    train_data : Dataset
        A Hugging Face Dataset object containing the training data.
    valid_data : Dataset
        A Hugging Face Dataset object containing the validation data.
    test_data : Dataset
        A Hugging Face Dataset object containing the test data.

    ---------------------------------------------------------------------------
    Copyright (C) - 2025 - Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Get the corresponding data from the indices
    train_texts = df.loc[indices['train'], covariate].values.tolist()
    train_labels = df.loc[indices['train'], target].values.tolist()

    valid_texts = df.loc[indices['valid'], covariate].values.tolist()
    valid_labels = df.loc[indices['valid'], target].values.tolist()

    test_texts = df.loc[indices['test'], covariate].values.tolist()
    test_labels = df.loc[indices['test'], target].values.tolist()

    # Convert into Hugging Face dataset format
    train_data = Dataset.from_dict(
        {'text': train_texts, 'label': train_labels})
    valid_data = Dataset.from_dict(
        {'text': valid_texts, 'label': valid_labels})
    test_data = Dataset.from_dict({'text': test_texts, 'label': test_labels})

    return train_data, valid_data, test_data


def logreg_fit(X: np.ndarray,
               y: np.ndarray,
               train_idx: np.ndarray):
    """
    Fit a Logistic Regression model on a subset of data.

    This function trains a Logistic Regression model using a subset of the 
    input data, where the subset is defined by the indices provided in 
    `train_idx`. It returns the trained model.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    X : np.ndarray, shape (n_samples, n_features)
        Feature matrix for training the model.
    y : np.ndarray, shape (n_samples,)
        Target labels corresponding to the feature matrix X.
    train_idx : np.ndarray, shape (n_samples_subset,)
         Indices used to restrict the input data (X and y) for training.

    ---------------------------------------------------------------------------
    OUTPUT:
    model : LogisticRegression
         A trained Logistic Regression model fitted on the restricted data.

    ---------------------------------------------------------------------------
    Copyright (C) - 2025 - Samuel Davenport
    ---------------------------------------------------------------------------
    """
    X_restricted = X[train_idx]
    y_restricted = y[train_idx]

    from sklearn.linear_model import LogisticRegression

    # Train Logistic Regression model
    model = LogisticRegression()
    model.fit(X_restricted, y_restricted)

    return model


def logreg_metrics(X: np.ndarray,
                   y: np.ndarray,
                   model,
                   test_indices: np.ndarray,
                   doprint: int = 0) -> float:
    """
    Compute and optionally display the accuracy of a Logistic Regression model.

    This function evaluates the performance of a given Logistic Regression model 
    on a test subset of the data, specified by `test_indices`. It computes the 
    accuracy score and, if requested, prints both the accuracy and a detailed 
    classification report.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    X : np.ndarray, shape (n_samples, n_features)
        Feature matrix containing the input data.
    y : np.ndarray, shape (n_samples,)
        Target labels corresponding to the feature matrix X.
    model : LogisticRegression
        A trained Logistic Regression model that implements the `predict` method.
    test_indices : np.ndarray, shape (n_samples_subset,)
        Indices specifying which samples to use as the test set.
    doprint : int, optional (default=0)
        If set to 1, prints the accuracy and the classification report.

    ---------------------------------------------------------------------------
    OUTPUT:
    accuracy : float
        The accuracy of the model on the test subset.

    ---------------------------------------------------------------------------
    Copyright (C) - 2025 - Samuel Davenport
    ---------------------------------------------------------------------------
    """
    X_test = X[test_indices]
    y_test = y[test_indices]
    y_pred = model.predict(X_test)

    from sklearn.metrics import accuracy_score

    accuracy = accuracy_score(y_test, y_pred)

    # Print accuracy and classification report if doprint is 1
    if doprint == 1:
        print("Accuracy:", accuracy)
        from sklearn.metrics import classification_report
        print(classification_report(y_test, y_pred))

    return accuracy

def accuracy_df(vec2, vec1):
    """
    Calculate detailed classification metrics for each unique label.

    Parameters:
    -----------
    vec2 : array-like  
        Predicted/comparison labels
    vec1 : array-like
        True/reference labels

    Returns:
    --------
    out : pandas.DataFrame
        DataFrame with columns 'Label', 'Accuracy', 'Precision', 'Recall', 'F1_Score', 
        'Sample_Count', 'Correct_Predictions', 'Predicted_Count'
    """
    if type(vec1) == np.ndarray:
        tdf = pd.DataFrame({
            "vec1": vec1,
            "vec2": vec2
            })
        vec1 = tdf['vec1']
        vec2 = tdf['vec2']
    
    labels = vec1.unique().tolist()
    label_accuracy = []
    label_precision = []
    label_recall = []
    label_f1 = []
    sample_counts = []
    correct_counts = []
    predicted_counts = []

    for label in labels:
        # True positives: both true and predicted are this label
        true_positives = ((vec1 == label) & (vec2 == label)).sum()

        # False positives: predicted as this label but true label is different
        false_positives = ((vec1 != label) & (vec2 == label)).sum()

        # False negatives: true label is this but predicted as different
        false_negatives = ((vec1 == label) & (vec2 != label)).sum()

        # Counts
        # Total true instances of this label
        sample_count = (vec1 == label).sum()
        # Total predicted instances of this label
        predicted_count = (vec2 == label).sum()

        # Accuracy (for this label vs all others)
        if sample_count > 0:
            accuracy = true_positives / sample_count
        else:
            accuracy = 0.0

        # Precision: TP / (TP + FP)
        if predicted_count > 0:
            precision = true_positives / predicted_count
        else:
            precision = 0.0

        # Recall: TP / (TP + FN)
        if sample_count > 0:
            recall = true_positives / sample_count
        else:
            recall = 0.0

        # F1-Score: 2 * (Precision * Recall) / (Precision + Recall)
        if precision + recall > 0:
            f1 = 2 * (precision * recall) / (precision + recall)
        else:
            f1 = 0.0

        # Append results
        label_accuracy.append(accuracy)
        label_precision.append(precision)
        label_recall.append(recall)
        label_f1.append(f1)
        sample_counts.append(sample_count)
        correct_counts.append(true_positives)
        predicted_counts.append(predicted_count)

    # Calculate overall metrics
    overall_accuracy = (vec1 == vec2).sum() / len(vec1)
    
    # Overall precision, recall, F1 (macro-averaged)
    overall_precision = sum(label_precision) / len(label_precision) if label_precision else 0.0
    overall_recall = sum(label_recall) / len(label_recall) if label_recall else 0.0
    overall_f1 = 2 * (overall_precision * overall_recall) / (overall_precision + overall_recall) if (overall_precision + overall_recall) > 0 else 0.0
    
    overall_sample_count = len(vec1)
    overall_correct_count = (vec1 == vec2).sum()
    overall_predicted_count = len(vec2)

    # Add overall row
    labels.append('Overall')
    label_accuracy.append(overall_accuracy)
    label_precision.append(overall_precision)
    label_recall.append(overall_recall)
    label_f1.append(overall_f1)
    sample_counts.append(overall_sample_count)
    correct_counts.append(overall_correct_count)
    predicted_counts.append(overall_predicted_count)

    out = pd.DataFrame({
        'Label': labels,
        'Accuracy': label_accuracy,
        'Precision': label_precision,
        'Recall': label_recall,
        'F1_Score': label_f1,
        'Sample_Count': sample_counts,
        'Correct_Predictions': correct_counts,
        'Predicted_Count': predicted_counts
    })

    return out

def within_n_accuracy(y_true, y_pred, n):
    """
    Computes the within-2 accuracy for an ordinal classification task.

    Parameters:
    ----------
    y_true : array-like
        Ground truth labels.
    y_pred : array-like
        Predicted labels.

    Returns:
    -------
    float
        Within-2 accuracy score.
    """
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    
    correct = np.abs(y_true - y_pred) <= n
    return np.mean(correct)

def df2hfdict(df: pd.DataFrame,
            covariate: str,
            target: Optional[str] = None,
            indices: Optional[Union[List[int], pd.Index]] = None) -> Dataset:
    """
    Convert a pandas DataFrame into a Hugging Face DatasetDict format.

    This function extracts specified covariate and target columns from a DataFrame, optionally filtered by indices, 
    and converts the result into a Hugging Face Dataset object.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    df : pd.DataFrame
        The input DataFrame containing the covariate and target columns.
    covariate : str
        The name of the column in the DataFrame representing the input texts.
    target : Optional[str]
        The name of the column in the DataFrame representing the target labels.
    indices : Optional[Union[List[int], pd.Index]]
        A list of indices specifying which rows of the DataFrame to include. 
        If None, all rows are used.

    ---------------------------------------------------------------------------
    OUTPUT:
    output : Dataset
        A Hugging Face Dataset object containing 'text' and 'label' fields.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if indices is None:
        indices = df.index

    # Get the corresponding data from the indices
    texts = df.loc[indices, covariate].tolist()

    if target is not None:
        labels = df.loc[indices, target].tolist()
        # Convert into Hugging Face dataset format
        output = Dataset.from_dict(
            {'text': texts, 'label': labels})
    else:
        output = Dataset.from_dict(
            {'text': texts})

    return output

def compute_metrics(eval_pred):
    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return {
        "accuracy": accuracy_score(labels, predictions),
        "precision": precision_score(labels, predictions, average="macro"),
        "recall": recall_score(labels, predictions, average="macro"),
        "f1": f1_score(labels, predictions, average="macro"),
    }

def compute_all_metrics(predictions, labels):
    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

    return {
        "accuracy": accuracy_score(labels, predictions),
        "precision": precision_score(labels, predictions, average="macro", zero_division=0),
        "recall": recall_score(labels, predictions, average="macro", zero_division=0),
        "f1": f1_score(labels, predictions, average="macro", zero_division=0),
    }

def accuracy(predictions, labels):
    matches = np.sum(np.array(predictions) == np.array(labels))
    out = matches/len(predictions)
    return out

def compute_metrics_dep(eval_pred: Tuple[np.ndarray, np.ndarray]) -> Dict[str, float]:
    """
    Compute evaluation metrics for model predictions.

    This function calculates accuracy, precision, recall, and F1-score for the 
    provided predictions and labels. The metrics are computed using a macro-average 
    to account for class imbalance across multiple classes.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    eval_pred : Tuple[np.ndarray, np.ndarray]
        A tuple containing two elements:
        - predictions (np.ndarray): The raw prediction scores or logits output by the model.
        - labels (np.ndarray): The true labels corresponding to the predictions.

    ---------------------------------------------------------------------------
    OUTPUT:
    metrics : Dict[str, float]
        A dictionary containing the computed metrics:
        - 'accuracy' (float): The overall accuracy of the predictions.
        - 'precision' (float): The macro-averaged precision score.
        - 'recall' (float): The macro-averaged recall score.
        - 'f1' (float): The macro-averaged F1 score.

    ---------------------------------------------------------------------------
    Copyright (C) - 2025 - Samuel Davenport
    ---------------------------------------------------------------------------
    """
    accuracy_metric = evaluate.load("accuracy")
    precision_metric = evaluate.load("precision")
    recall_metric = evaluate.load("recall")
    f1_metric = evaluate.load("f1")

    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    precision = precision_metric.compute(
        predictions=predictions, references=labels, average="macro"
    )["precision"]
    recall = recall_metric.compute(
        predictions=predictions, references=labels, average="macro"
    )["recall"]
    f1 = f1_metric.compute(predictions=predictions, references=labels, average="macro")[
        "f1"
    ]
    accuracy = accuracy_metric.compute(predictions=predictions, references=labels)[
        "accuracy"
    ]
    return {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1}

def compute_metrics_multilabel(eval_pred, threshold=0.5):
    """
    Compute metrics for multi-label classification tasks.
    
    This function calculates various performance metrics for multi-label classification
    by applying a threshold to the model's logits to convert them to binary predictions.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    eval_pred : tuple
        A tuple containing:
        - predictions (np.ndarray): Raw logits from the model, shape (n_samples, n_labels)
        - labels (np.ndarray): True binary labels, shape (n_samples, n_labels)
    
    threshold : float, optional (default=0.5)
        Threshold for converting logits to binary predictions after sigmoid.
    
    ---------------------------------------------------------------------------
    OUTPUT:
    metrics : dict
        A dictionary containing computed metrics:
        - 'accuracy': Exact match ratio (all labels must match)
        - 'f1_micro': Micro-averaged F1 score
        - 'f1_macro': Macro-averaged F1 score
        - 'f1_weighted': Weighted F1 score
        - 'precision_micro': Micro-averaged precision
        - 'precision_macro': Macro-averaged precision
        - 'recall_micro': Micro-averaged recall
        - 'recall_macro': Macro-averaged recall
        - 'hamming_accuracy': Fraction of correct labels (label-wise accuracy)
    
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    predictions, labels = eval_pred
    
    # Apply sigmoid to convert logits to probabilities
    sigmoid = lambda x: 1 / (1 + np.exp(-x))
    probs = sigmoid(predictions)
    
    # Apply threshold to get binary predictions
    preds = (probs > threshold).astype(int)
    
    # Exact match accuracy (all labels must be correct)
    exact_match = accuracy_score(labels, preds)
    
    # Hamming accuracy (label-wise accuracy)
    hamming_acc = np.mean(labels == preds)
    
    # F1 scores
    f1_micro = f1_score(labels, preds, average='micro', zero_division=0)
    f1_macro = f1_score(labels, preds, average='macro', zero_division=0)
    f1_weighted = f1_score(labels, preds, average='weighted', zero_division=0)
    
    # Precision scores
    precision_micro = precision_score(labels, preds, average='micro', zero_division=0)
    precision_macro = precision_score(labels, preds, average='macro', zero_division=0)
    
    # Recall scores
    recall_micro = recall_score(labels, preds, average='micro', zero_division=0)
    recall_macro = recall_score(labels, preds, average='macro', zero_division=0)
    
    return {
        'accuracy': exact_match,
        'hamming_accuracy': hamming_acc,
        'f1_micro': f1_micro,
        'f1_macro': f1_macro,
        'f1_weighted': f1_weighted,
        'precision_micro': precision_micro,
        'precision_macro': precision_macro,
        'recall_micro': recall_micro,
        'recall_macro': recall_macro,
    }


def compute_multilabel_metrics(labels_1, labels_2):
    """
    Compute bidirectional F1 scores and Jaccard similarity for multi-label classification.
    
    Args:
        labels_1: List of label sets for first annotator
        labels_2: List of label sets for second annotator
        
    Returns:
        dict: Dictionary containing all computed metrics
    """
    from sklearn.metrics import f1_score, precision_recall_fscore_support, jaccard_score
    from sklearn.preprocessing import MultiLabelBinarizer
    
    # Convert to binary format
    mlb = MultiLabelBinarizer()
    all_labels = labels_1 + labels_2
    mlb.fit(all_labels)
    
    binary_1 = mlb.transform(labels_1)
    binary_2 = mlb.transform(labels_2)
    
    # Direction 1: labels_1 as ground truth, labels_2 as predictions
    f1_micro_1 = f1_score(binary_1, binary_2, average='micro', zero_division=0)
    f1_macro_1 = f1_score(binary_1, binary_2, average='macro', zero_division=0)
    f1_samples_1 = f1_score(binary_1, binary_2, average='samples', zero_division=0)
    
    precision_1, recall_1, f1_1, support_1 = precision_recall_fscore_support(
        binary_1, binary_2, average=None, zero_division=0
    )
    
    # Direction 2: labels_2 as ground truth, labels_1 as predictions
    f1_micro_2 = f1_score(binary_2, binary_1, average='micro', zero_division=0)
    f1_macro_2 = f1_score(binary_2, binary_1, average='macro', zero_division=0)
    f1_samples_2 = f1_score(binary_2, binary_1, average='samples', zero_division=0)
    
    precision_2, recall_2, f1_2, support_2 = precision_recall_fscore_support(
        binary_2, binary_1, average=None, zero_division=0
    )
    
    # Symmetric averages
    f1_micro_sym = (f1_micro_1 + f1_micro_2) / 2
    f1_macro_sym = (f1_macro_1 + f1_macro_2) / 2
    f1_samples_sym = (f1_samples_1 + f1_samples_2) / 2
    precision_sym = (precision_1 + precision_2) / 2
    recall_sym = (recall_1 + recall_2) / 2
    f1_sym = (f1_1 + f1_2) / 2
    
    # Jaccard similarity (symmetric by nature)
    jaccard_micro = jaccard_score(binary_1, binary_2, average='micro', zero_division=0)
    jaccard_macro = jaccard_score(binary_1, binary_2, average='macro', zero_division=0)
    jaccard_samples = jaccard_score(binary_1, binary_2, average='samples', zero_division=0)
    jaccard_per_label = jaccard_score(binary_1, binary_2, average=None, zero_division=0)
    
    return {
        'mlb': mlb,
        'binary_1': binary_1,
        'binary_2': binary_2,
        'f1_micro_1': f1_micro_1,
        'f1_macro_1': f1_macro_1,
        'f1_samples_1': f1_samples_1,
        'precision_1': precision_1,
        'recall_1': recall_1,
        'f1_1': f1_1,
        'support_1': support_1,
        'f1_micro_2': f1_micro_2,
        'f1_macro_2': f1_macro_2,
        'f1_samples_2': f1_samples_2,
        'precision_2': precision_2,
        'recall_2': recall_2,
        'f1_2': f1_2,
        'support_2': support_2,
        'f1_micro_sym': f1_micro_sym,
        'f1_macro_sym': f1_macro_sym,
        'f1_samples_sym': f1_samples_sym,
        'precision_sym': precision_sym,
        'recall_sym': recall_sym,
        'f1_sym': f1_sym,
        'jaccard_micro': jaccard_micro,
        'jaccard_macro': jaccard_macro,
        'jaccard_samples': jaccard_samples,
        'jaccard_per_label': jaccard_per_label
    }