"""
Hierarchical (Agglomerative) Clustering.

Hierarchical clustering builds a tree of clusters by iteratively
merging the closest pairs. This is useful for track fusion and
understanding cluster relationships.

References
----------
.. [1] S. C. Johnson, "Hierarchical clustering schemes,"
       Psychometrika, 1967.
"""

from enum import Enum
from typing import List, Literal, NamedTuple, Optional

import numpy as np
from numpy.typing import ArrayLike, NDArray


class LinkageType(Enum):
    """Linkage methods for hierarchical clustering."""

    SINGLE = "single"  # Minimum distance
    COMPLETE = "complete"  # Maximum distance
    AVERAGE = "average"  # Average distance
    WARD = "ward"  # Ward's minimum variance


class DendrogramNode(NamedTuple):
    """A node in the dendrogram (merge tree).

    Attributes
    ----------
    left : int
        Index of left child (negative for original samples).
    right : int
        Index of right child (negative for original samples).
    distance : float
        Distance/dissimilarity at which merge occurred.
    count : int
        Number of original samples in this cluster.
    """

    left: int
    right: int
    distance: float
    count: int


class HierarchicalResult(NamedTuple):
    """Result of hierarchical clustering.

    Attributes
    ----------
    labels : ndarray
        Cluster labels for each point, shape (n_samples,).
    n_clusters : int
        Number of clusters.
    linkage_matrix : ndarray
        Linkage matrix of shape (n_samples-1, 4).
        Each row [i, j, dist, count] represents a merge.
    dendrogram : list of DendrogramNode
        List of merge operations.
    """

    labels: NDArray[np.intp]
    n_clusters: int
    linkage_matrix: NDArray[np.floating]
    dendrogram: List[DendrogramNode]


def compute_distance_matrix(
    X: NDArray[np.floating],
) -> NDArray[np.floating]:
    """
    Compute pairwise Euclidean distance matrix.

    Parameters
    ----------
    X : ndarray
        Data points, shape (n_samples, n_features).

    Returns
    -------
    distances : ndarray
        Distance matrix, shape (n_samples, n_samples).
    """
    n = X.shape[0]
    distances = np.zeros((n, n))

    for i in range(n):
        for j in range(i + 1, n):
            d = np.sqrt(np.sum((X[i] - X[j]) ** 2))
            distances[i, j] = d
            distances[j, i] = d

    return distances


def _single_linkage(
    dist_i: NDArray,
    dist_j: NDArray,
    size_i: int,
    size_j: int,
) -> NDArray:
    """Single linkage: minimum of distances."""
    return np.minimum(dist_i, dist_j)


def _complete_linkage(
    dist_i: NDArray,
    dist_j: NDArray,
    size_i: int,
    size_j: int,
) -> NDArray:
    """Complete linkage: maximum of distances."""
    return np.maximum(dist_i, dist_j)


def _average_linkage(
    dist_i: NDArray,
    dist_j: NDArray,
    size_i: int,
    size_j: int,
) -> NDArray:
    """Average linkage: weighted average of distances."""
    return (size_i * dist_i + size_j * dist_j) / (size_i + size_j)


def _ward_linkage(
    dist_i: NDArray,
    dist_j: NDArray,
    size_i: int,
    size_j: int,
    size_k: NDArray,
    dist_ij: float,
) -> NDArray:
    """Ward's linkage: minimum variance merge."""
    total = size_i + size_j + size_k
    return np.sqrt(
        (
            (size_i + size_k) * dist_i**2
            + (size_j + size_k) * dist_j**2
            - size_k * dist_ij**2
        )
        / total
    )


def agglomerative_clustering(
    X: ArrayLike,
    n_clusters: Optional[int] = None,
    distance_threshold: Optional[float] = None,
    linkage: Literal["single", "complete", "average", "ward"] = "ward",
) -> HierarchicalResult:
    """
    Agglomerative hierarchical clustering.

    Recursively merges pairs of clusters that minimize the linkage
    criterion until the desired number of clusters is reached.

    Parameters
    ----------
    X : array_like
        Data points, shape (n_samples, n_features).
    n_clusters : int, optional
        Number of clusters to find. If None, uses distance_threshold.
    distance_threshold : float, optional
        Distance threshold for cluster merging. Merging stops when
        the minimum linkage distance exceeds this. If None, uses n_clusters.
    linkage : {'single', 'complete', 'average', 'ward'}
        Linkage criterion. Default 'ward'.

    Returns
    -------
    result : HierarchicalResult
        Clustering result with labels and dendrogram.

    Examples
    --------
    >>> import numpy as np
    >>> rng = np.random.default_rng(42)
    >>> X = np.vstack([rng.normal(0, 0.5, (20, 2)),
    ...                rng.normal(3, 0.5, (20, 2))])
    >>> result = agglomerative_clustering(X, n_clusters=2)
    >>> result.n_clusters
    2

    Notes
    -----
    Linkage methods:
    - 'single': Distance between closest points (can create chains)
    - 'complete': Distance between farthest points (tends to create compact clusters)
    - 'average': Average distance between all pairs
    - 'ward': Minimizes within-cluster variance (requires Euclidean distance)
    """
    X = np.asarray(X, dtype=np.float64)
    n_samples = X.shape[0]

    if n_samples == 0:
        return HierarchicalResult(
            labels=np.array([], dtype=np.intp),
            n_clusters=0,
            linkage_matrix=np.zeros((0, 4)),
            dendrogram=[],
        )

    if n_samples == 1:
        return HierarchicalResult(
            labels=np.array([0], dtype=np.intp),
            n_clusters=1,
            linkage_matrix=np.zeros((0, 4)),
            dendrogram=[],
        )

    if n_clusters is None and distance_threshold is None:
        n_clusters = 2

    # Compute initial distance matrix
    distances = compute_distance_matrix(X)

    # Initialize clusters
    # cluster_id -> list of original sample indices
    clusters = {i: [i] for i in range(n_samples)}
    cluster_sizes = {i: 1 for i in range(n_samples)}

    # Active cluster IDs
    active = set(range(n_samples))

    # Linkage matrix and dendrogram
    linkage_matrix = []
    dendrogram = []

    # Current cluster distance matrix (upper triangular, inf on diagonal)
    cluster_dist = distances.copy()
    np.fill_diagonal(cluster_dist, np.inf)

    next_cluster_id = n_samples

    # Merge until stopping criterion
    while len(active) > 1:
        # Check stopping criteria
        if n_clusters is not None and len(active) <= n_clusters:
            break

        # Find minimum distance pair among active clusters
        min_dist = np.inf
        merge_i, merge_j = -1, -1

        active_list = list(active)
        for idx_a, i in enumerate(active_list):
            for j in active_list[idx_a + 1 :]:
                if cluster_dist[i, j] < min_dist:
                    min_dist = cluster_dist[i, j]
                    merge_i, merge_j = i, j

        if merge_i == -1 or min_dist == np.inf:
            break

        if distance_threshold is not None and min_dist > distance_threshold:
            break

        # Ensure merge_i < merge_j for consistency
        if merge_i > merge_j:
            merge_i, merge_j = merge_j, merge_i

        # Record merge
        size_new = cluster_sizes[merge_i] + cluster_sizes[merge_j]
        linkage_matrix.append([merge_i, merge_j, min_dist, size_new])
        dendrogram.append(
            DendrogramNode(
                left=merge_i,
                right=merge_j,
                distance=min_dist,
                count=size_new,
            )
        )

        # Create new cluster
        new_cluster_id = next_cluster_id
        next_cluster_id += 1

        clusters[new_cluster_id] = clusters[merge_i] + clusters[merge_j]
        cluster_sizes[new_cluster_id] = size_new

        # Update distance matrix for new cluster
        # Expand matrix if needed
        if new_cluster_id >= cluster_dist.shape[0]:
            new_size = new_cluster_id + 1
            new_dist = np.full((new_size, new_size), np.inf)
            old_size = cluster_dist.shape[0]
            new_dist[:old_size, :old_size] = cluster_dist
            cluster_dist = new_dist

        # Compute distances from new cluster to all other active clusters
        for k in active:
            if k == merge_i or k == merge_j:
                continue

            if linkage == "single":
                new_d = _single_linkage(
                    cluster_dist[merge_i, k],
                    cluster_dist[merge_j, k],
                    cluster_sizes[merge_i],
                    cluster_sizes[merge_j],
                )
            elif linkage == "complete":
                new_d = _complete_linkage(
                    cluster_dist[merge_i, k],
                    cluster_dist[merge_j, k],
                    cluster_sizes[merge_i],
                    cluster_sizes[merge_j],
                )
            elif linkage == "average":
                new_d = _average_linkage(
                    cluster_dist[merge_i, k],
                    cluster_dist[merge_j, k],
                    cluster_sizes[merge_i],
                    cluster_sizes[merge_j],
                )
            elif linkage == "ward":
                new_d = _ward_linkage(
                    cluster_dist[merge_i, k],
                    cluster_dist[merge_j, k],
                    cluster_sizes[merge_i],
                    cluster_sizes[merge_j],
                    np.array([cluster_sizes[k]]),
                    cluster_dist[merge_i, merge_j],
                )[0]
            else:
                raise ValueError(f"Unknown linkage: {linkage}")

            cluster_dist[new_cluster_id, k] = new_d
            cluster_dist[k, new_cluster_id] = new_d

        # Update active set
        active.remove(merge_i)
        active.remove(merge_j)
        active.add(new_cluster_id)

        # Clean up old clusters
        del clusters[merge_i]
        del clusters[merge_j]
        del cluster_sizes[merge_i]
        del cluster_sizes[merge_j]

    # Assign labels based on final clusters
    labels = np.zeros(n_samples, dtype=np.intp)
    for cluster_label, cluster_id in enumerate(active):
        for sample_idx in clusters[cluster_id]:
            labels[sample_idx] = cluster_label

    # Convert linkage matrix to array
    if linkage_matrix:
        linkage_array = np.array(linkage_matrix)
    else:
        linkage_array = np.zeros((0, 4))

    return HierarchicalResult(
        labels=labels,
        n_clusters=len(active),
        linkage_matrix=linkage_array,
        dendrogram=dendrogram,
    )


def cut_dendrogram(
    linkage_matrix: ArrayLike,
    n_samples: int,
    n_clusters: Optional[int] = None,
    distance_threshold: Optional[float] = None,
) -> NDArray[np.intp]:
    """
    Cut a dendrogram to obtain cluster labels.

    Parameters
    ----------
    linkage_matrix : array_like
        Linkage matrix from agglomerative_clustering, shape (n-1, 4).
    n_samples : int
        Number of original samples.
    n_clusters : int, optional
        Number of clusters to form.
    distance_threshold : float, optional
        Maximum linkage distance.

    Returns
    -------
    labels : ndarray
        Cluster labels, shape (n_samples,).
    """
    linkage_matrix = np.asarray(linkage_matrix)

    if len(linkage_matrix) == 0:
        return np.zeros(n_samples, dtype=np.intp)

    if n_clusters is None and distance_threshold is None:
        n_clusters = 2

    # Build cluster membership by replaying merges
    # Initially each sample is its own cluster
    cluster_members = {i: {i} for i in range(n_samples)}
    next_id = n_samples

    n_current_clusters = n_samples

    for row in linkage_matrix:
        left, right, dist, _ = int(row[0]), int(row[1]), row[2], int(row[3])

        # Check stopping criteria
        if n_clusters is not None and n_current_clusters <= n_clusters:
            break
        if distance_threshold is not None and dist > distance_threshold:
            break

        # Merge clusters
        new_members = cluster_members[left] | cluster_members[right]
        cluster_members[next_id] = new_members
        del cluster_members[left]
        del cluster_members[right]

        next_id += 1
        n_current_clusters -= 1

    # Assign labels
    labels = np.zeros(n_samples, dtype=np.intp)
    for label, (cluster_id, members) in enumerate(cluster_members.items()):
        for sample_idx in members:
            labels[sample_idx] = label

    return labels


def fcluster(
    linkage_matrix: ArrayLike,
    n_samples: int,
    t: float,
    criterion: Literal["distance", "maxclust"] = "distance",
) -> NDArray[np.intp]:
    """
    Form flat clusters from hierarchical clustering.

    Compatible interface with scipy.cluster.hierarchy.fcluster.

    Parameters
    ----------
    linkage_matrix : array_like
        Linkage matrix from agglomerative_clustering.
    n_samples : int
        Number of original samples.
    t : float
        Threshold for forming clusters.
    criterion : {'distance', 'maxclust'}
        'distance': t is maximum cophenetic distance
        'maxclust': t is maximum number of clusters

    Returns
    -------
    labels : ndarray
        Cluster labels (1-indexed for scipy compatibility).
    """
    if criterion == "distance":
        labels = cut_dendrogram(linkage_matrix, n_samples, distance_threshold=t)
    elif criterion == "maxclust":
        labels = cut_dendrogram(linkage_matrix, n_samples, n_clusters=int(t))
    else:
        raise ValueError(f"Unknown criterion: {criterion}")

    # Convert to 1-indexed for scipy compatibility
    return labels + 1


__all__ = [
    "LinkageType",
    "DendrogramNode",
    "HierarchicalResult",
    "compute_distance_matrix",
    "agglomerative_clustering",
    "cut_dendrogram",
    "fcluster",
]
