"""
risk/neighborhoods/domains
~~~~~~~~~~~~~~~~~~~~~~~~~~
"""

from contextlib import suppress
from itertools import product
from typing import Tuple, Union

import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import linkage, fcluster
from scipy.optimize import minimize_scalar
from sklearn.metrics import calinski_harabasz_score, davies_bouldin_score, silhouette_score
from tqdm import tqdm

from risk.annotations import get_weighted_description
from risk.constants import GROUP_LINKAGE_METHODS, GROUP_DISTANCE_METRICS
from risk.log import logger


class LinkageThresholdError(Exception):
    """Exception raised for errors in the linkage threshold optimization process."""

    pass


def define_domains(
    top_annotations: pd.DataFrame,
    significant_neighborhoods_significance: np.ndarray,
    linkage_criterion: str,
    linkage_method: str,
    linkage_metric: str,
    linkage_threshold: Union[str, float],
) -> pd.DataFrame:
    """Define domains and assign nodes to these domains based on their significance scores and clustering,
    handling errors by assigning unique domains when clustering fails.

    Args:
        top_annotations (pd.DataFrame): DataFrame of top annotations data for the network nodes.
        significant_neighborhoods_significance (np.ndarray): The binary significance matrix below alpha.
        linkage_criterion (str): The clustering criterion for defining groups. Use "distance" for distance-based
            clustering or "maxclust" for a fixed number of clusters. Use "off" to skip clustering.
        linkage_method (str): The linkage method for clustering. Use "auto" to try multiple methods.
        linkage_metric (str): The linkage metric for clustering. Use "auto" to try multiple metrics.
        linkage_threshold (str, float): The linkage threshold for clustering, or one of "silhouette",
            "calinski_harabasz", or "davies_bouldin" to optimize the threshold.

    Returns:
        pd.DataFrame: DataFrame with the primary domain for each node.

    Raises:
        ValueError: If an improper value is passed for linkage_threshold. Acceptable values are "silhouette",
            "calinski_harabasz", "davies_bouldin", or a float value.
    """
    try:
        if linkage_criterion == "off":
            raise ValueError("Clustering is turned off.")

        # Transpose the matrix to cluster annotations
        m = significant_neighborhoods_significance[:, top_annotations["significant_annotations"]].T
        # Safeguard the matrix by replacing NaN, Inf, and -Inf values
        m = _safeguard_matrix(m)
        # Optimize silhouette score across different linkage methods and distance metrics
        best_linkage, best_metric, best_threshold = (
            _optimize_linkage_threshold_across_methods_and_metrics(
                m, linkage_criterion, linkage_method, linkage_metric, linkage_threshold
            )
        )
        # Perform hierarchical clustering
        Z = linkage(m, method=best_linkage, metric=best_metric)
        logger.warning(
            f"Linkage criterion: '{linkage_criterion}'\nLinkage method: '{best_linkage}'\nLinkage metric: '{best_metric}'"
        )
        logger.debug(f"Optimal linkage threshold: {round(best_threshold, 3)}")
        # Calculate the optimal threshold for clustering
        max_d_optimal = np.max(Z[:, 2]) * best_threshold
        # Assign domains to the annotations matrix
        domains = fcluster(Z, max_d_optimal, criterion=linkage_criterion)
        top_annotations["domain"] = 0
        top_annotations.loc[top_annotations["significant_annotations"], "domain"] = domains
    except ValueError:
        # If a ValueError is encountered, handle it by assigning unique domains
        n_rows = len(top_annotations)
        if linkage_criterion == "off":
            logger.warning(
                f"Clustering is turned off. Skipping clustering and assigning {n_rows} unique domains."
            )
        else:
            logger.error(
                f"Error encountered. Skipping clustering and assigning {n_rows} unique domains."
            )
        top_annotations["domain"] = range(1, n_rows + 1)  # Assign unique domains
    except LinkageThresholdError as e:
        # If a LinkageThresholdError is encountered, raise a ValueError with the original exception
        raise ValueError(e) from e

    # Create DataFrames to store domain information
    node_to_significance = pd.DataFrame(
        data=significant_neighborhoods_significance,
        columns=[top_annotations.index.values, top_annotations["domain"]],
    )
    node_to_domain = node_to_significance.T.groupby(level="domain").sum().T

    # Find the maximum significance score for each node
    t_max = node_to_domain.loc[:, 1:].max(axis=1)
    t_idxmax = node_to_domain.loc[:, 1:].idxmax(axis=1)
    t_idxmax[t_max == 0] = 0

    # Assign all domains where the score is greater than 0
    node_to_domain["all_domains"] = node_to_domain.loc[:, 1:].apply(
        lambda row: list(row[row > 0].index), axis=1
    )
    # Assign primary domain
    node_to_domain["primary_domain"] = t_idxmax

    return node_to_domain


def trim_domains(
    domains: pd.DataFrame,
    top_annotations: pd.DataFrame,
    min_cluster_size: int = 5,
    max_cluster_size: int = 1000,
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Trim domains that do not meet size criteria and find outliers.

    Args:
        domains (pd.DataFrame): DataFrame of domain data for the network nodes.
        top_annotations (pd.DataFrame): DataFrame of top annotations data for the network nodes.
        min_cluster_size (int, optional): Minimum size of a cluster to be retained. Defaults to 5.
        max_cluster_size (int, optional): Maximum size of a cluster to be retained. Defaults to 1000.

    Returns:
        Tuple[pd.DataFrame, pd.DataFrame]:
            - Trimmed domains (pd.DataFrame)
            - A DataFrame with domain labels (pd.DataFrame)
    """
    # Identify domains to remove based on size criteria
    domain_counts = domains["primary_domain"].value_counts()
    to_remove = set(
        domain_counts[(domain_counts < min_cluster_size) | (domain_counts > max_cluster_size)].index
    )

    # Add invalid domain IDs
    invalid_domain_id = 888888
    invalid_domain_ids = {0, invalid_domain_id}
    # Mark domains to be removed
    top_annotations["domain"] = top_annotations["domain"].replace(to_remove, invalid_domain_id)
    domains.loc[domains["primary_domain"].isin(to_remove), ["primary_domain"]] = invalid_domain_id

    # Normalize "num significant neighborhoods" by percentile for each domain and scale to 0-10
    top_annotations["normalized_value"] = top_annotations.groupby("domain")[
        "significant_neighborhood_significance_sums"
    ].transform(lambda x: (x.rank(pct=True) * 10).apply(np.ceil).astype(int))
    # Modify the lambda function to pass both full_terms and significant_significance_score
    top_annotations["combined_terms"] = top_annotations.apply(
        lambda row: " ".join([str(row["full_terms"])] * row["normalized_value"]), axis=1
    )

    # Perform the groupby operation while retaining the other columns and adding the weighting with significance scores
    domain_labels = (
        top_annotations.groupby("domain")
        .agg(
            full_terms=("full_terms", lambda x: list(x)),
            significance_scores=("significant_significance_score", lambda x: list(x)),
        )
        .reset_index()
    )
    domain_labels["combined_terms"] = domain_labels.apply(
        lambda row: get_weighted_description(
            pd.Series(row["full_terms"]), pd.Series(row["significance_scores"])
        ),
        axis=1,
    )

    # Rename the columns as necessary
    trimmed_domains_matrix = domain_labels.rename(
        columns={
            "domain": "id",
            "combined_terms": "normalized_description",
            "full_terms": "full_descriptions",
            "significance_scores": "significance_scores",
        }
    ).set_index("id")

    # Remove invalid domains
    valid_domains = domains[~domains["primary_domain"].isin(invalid_domain_ids)]
    valid_trimmed_domains_matrix = trimmed_domains_matrix[
        ~trimmed_domains_matrix.index.isin(invalid_domain_ids)
    ]
    return valid_domains, valid_trimmed_domains_matrix


def _safeguard_matrix(matrix: np.ndarray) -> np.ndarray:
    """Safeguard the matrix by replacing NaN, Inf, and -Inf values.

    Args:
        matrix (np.ndarray): Data matrix.

    Returns:
        np.ndarray: Safeguarded data matrix.
    """
    # Replace NaN with column mean
    nan_replacement = np.nanmean(matrix, axis=0)
    matrix = np.where(np.isnan(matrix), nan_replacement, matrix)
    # Replace Inf/-Inf with maximum/minimum finite values
    finite_max = np.nanmax(matrix[np.isfinite(matrix)])
    finite_min = np.nanmin(matrix[np.isfinite(matrix)])
    matrix = np.where(np.isposinf(matrix), finite_max, matrix)
    matrix = np.where(np.isneginf(matrix), finite_min, matrix)
    # Ensure rows have non-zero variance (optional step)
    row_variance = np.var(matrix, axis=1)
    matrix = matrix[row_variance > 0]
    return matrix


def _optimize_linkage_threshold_across_methods_and_metrics(
    m: np.ndarray,
    linkage_criterion: str,
    linkage_method: str,
    linkage_metric: str,
    linkage_threshold: Union[str, float],
) -> Tuple[str, str, float]:
    """Optimize the linkage method, metric, and threshold for hierarchical clustering. If the threshold is
    a string, optimize the threshold using the specified metric; otherwise, use the provided threshold.

    Args:
        m (np.ndarray): Data matrix.
        linkage_criterion (str): Criterion for fcluster (typically "distance").
        linkage_method (str): Linkage method for clustering, or "auto" to try multiple methods.
        linkage_metric (str): Distance metric for clustering, or "auto" to try multiple metrics.
        linkage_threshold (str, float): Either a numeric threshold or one of the following keywords:
            "silhouette", "calinski_harabasz", or "davies_bouldin" to trigger optimization.

    Returns:
        Tuple[str, str, float]:
            - The chosen linkage method.
            - The chosen linkage metric.
            - The optimized threshold (a float).

    Raises:
        ValueError: If linkage_threshold is neither one of the supported keywords nor convertible to float.
    """
    # Supported linkage threshold metrics
    supported_linkage_thresholds = {"silhouette", "calinski_harabasz", "davies_bouldin"}

    # If linkage_threshold is a string:
    if isinstance(linkage_threshold, str):
        if linkage_threshold in supported_linkage_thresholds:
            opt_metric = linkage_threshold
        else:
            try:
                threshold_float = float(linkage_threshold)
            except (TypeError, ValueError):
                raise LinkageThresholdError(
                    f"linkage_threshold must be one of {', '.join(supported_linkage_thresholds)} or a float value."
                )
            return linkage_method, linkage_metric, threshold_float
    else:
        # If not a string, try to convert it to float.
        try:
            threshold_float = float(linkage_threshold)
        except (TypeError, ValueError):
            raise LinkageThresholdError(
                f"linkage_threshold must be one of {', '.join(supported_linkage_thresholds)} or a float value."
            )
        return linkage_method, linkage_metric, threshold_float

    # Otherwise, perform optimization using the specified metric (opt_metric).
    best_overall_method = None
    best_overall_metric = None
    best_overall_threshold = None
    best_overall_score = -np.inf

    # Use the provided lists if "auto" is specified.
    methods = GROUP_LINKAGE_METHODS if linkage_method == "auto" else [linkage_method]
    metrics = GROUP_DISTANCE_METRICS if linkage_metric == "auto" else [linkage_metric]
    total_combinations = len(methods) * len(metrics)

    for method, metric in tqdm(
        product(methods, metrics),
        desc="Evaluating optimal linkage method and metric",
        total=total_combinations,
        bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
    ):
        with suppress(ValueError):
            Z = linkage(m, method=method, metric=metric)
            threshold, score = _find_optimal_linkage_threshold(
                Z, m, metric, linkage_criterion, opt_metric=opt_metric
            )
            if score > best_overall_score:
                best_overall_score = score
                best_overall_threshold = threshold
                best_overall_method = method
                best_overall_metric = metric

    if best_overall_method is None or best_overall_metric is None or best_overall_threshold is None:
        raise ValueError("Optimization failed to determine an optimal threshold.")
    return best_overall_method, best_overall_metric, best_overall_threshold


def _find_optimal_linkage_threshold(
    Z: np.ndarray,
    m: np.ndarray,
    linkage_metric: str,
    linkage_criterion: str,
    opt_metric: str = "silhouette",
) -> Tuple[float, float]:
    """Find the optimal linkage threshold coefficient for hierarchical clustering. The function optimizes
    the threshold value using the specified metric (opt_metric).

    Args:
        Z (np.ndarray): Linkage matrix generated by a hierarchical clustering algorithm.
        m (np.ndarray): Data matrix used for clustering.
        linkage_metric (str): Metric used to calculate distances between data points
            (e.g., "euclidean" or "cosine").
        linkage_criterion (str): Criterion to pass to `fcluster`, typically "distance".
        opt_metric (str, optional): Metric to optimize clustering quality. Options are:
            "silhouette", "calinski_harabasz", or "davies_bouldin". Defaults to "silhouette".

    Returns:
        Tuple[float, float]:
            - best_threshold (float): The optimal linkage threshold coefficient.
            - best_metric_value (float): The value of the clustering quality metric achieved
              at the optimal threshold (higher for "silhouette" and "calinski_harabasz",
              lower for "davies_bouldin").

    Raises:
        ValueError: If the `opt_metric` argument is not one of the supported metrics.
    """
    # Get the maximum distance in the linkage matrix
    max_d = np.max(Z[:, 2])
    resolution = 1e-6

    def compute_objective(coefficient: float) -> float:
        """Compute the objective function for optimization."""
        threshold_val = coefficient * max_d
        clusters = fcluster(Z, threshold_val, criterion=linkage_criterion)
        unique_clusters = np.unique(clusters)
        if len(unique_clusters) <= 1 or len(unique_clusters) == m.shape[0]:
            return 1e6
        try:
            if opt_metric == "silhouette":
                score = silhouette_score(m, clusters, metric=linkage_metric)
                return -score  # We want to maximize the score.
            elif opt_metric == "calinski_harabasz":
                score = calinski_harabasz_score(m, clusters)
                return -score
            elif opt_metric == "davies_bouldin":
                score = davies_bouldin_score(m, clusters)
                return score
            else:
                raise ValueError(f"Unknown optimization metric: {opt_metric}.")
        except Exception:
            return 1e6

    # Optimize the threshold using the specified metric
    res = minimize_scalar(
        compute_objective, bounds=(0.0, 1.0), method="bounded", options={"xatol": resolution}
    )

    best_threshold = res.x
    best_obj = res.fun
    # For silhouette and calinski_harabasz, the objective was negative.
    best_metric_value = -best_obj if opt_metric in ["silhouette", "calinski_harabasz"] else best_obj

    return best_threshold, float(best_metric_value)
