"""
Cluster set container.

This module provides a container for managing groups of tracks
that move together (formations, convoys, etc.).
"""

from __future__ import annotations

from typing import (
    Dict,
    Iterable,
    Iterator,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Union,
)

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

from pytcl.clustering.dbscan import dbscan
from pytcl.clustering.kmeans import kmeans
from pytcl.containers.track_list import TrackList
from pytcl.trackers.multi_target import Track


class TrackCluster(NamedTuple):
    """
    A cluster of related tracks.

    Attributes
    ----------
    id : int
        Unique cluster identifier.
    track_ids : Tuple[int, ...]
        Immutable tuple of track IDs in this cluster.
    centroid : ndarray
        Cluster center position.
    covariance : ndarray
        Cluster spread covariance matrix.
    time : float
        Time at which cluster was computed.
    """

    id: int
    track_ids: Tuple[int, ...]
    centroid: NDArray[np.float64]
    covariance: NDArray[np.float64]
    time: float


class ClusterStats(NamedTuple):
    """
    Statistics for a cluster.

    Attributes
    ----------
    n_tracks : int
        Number of tracks in the cluster.
    mean_separation : float
        Average distance between tracks and centroid.
    max_separation : float
        Maximum distance from any track to centroid.
    velocity_coherence : float
        Measure of how aligned velocities are (0-1).
        1.0 means perfectly aligned, 0.0 means random directions.
    """

    n_tracks: int
    mean_separation: float
    max_separation: float
    velocity_coherence: float


def compute_cluster_centroid(
    tracks: Iterable[Track],
    state_indices: Tuple[int, int] = (0, 2),
) -> NDArray[np.float64]:
    """
    Compute the centroid of a group of tracks.

    Parameters
    ----------
    tracks : Iterable[Track]
        Tracks to compute centroid for.
    state_indices : tuple of int, optional
        Indices of x and y in state vector (default: (0, 2)).

    Returns
    -------
    centroid : ndarray
        Centroid position [x, y].
    """
    track_list = list(tracks)
    if len(track_list) == 0:
        return np.array([0.0, 0.0])

    ix, iy = state_indices
    positions = np.array([[t.state[ix], t.state[iy]] for t in track_list])
    return np.mean(positions, axis=0)


def compute_cluster_covariance(
    tracks: Iterable[Track],
    state_indices: Tuple[int, int] = (0, 2),
) -> NDArray[np.float64]:
    """
    Compute the covariance of track positions in a cluster.

    Parameters
    ----------
    tracks : Iterable[Track]
        Tracks to compute covariance for.
    state_indices : tuple of int, optional
        Indices of x and y in state vector (default: (0, 2)).

    Returns
    -------
    covariance : ndarray
        Position covariance matrix (2x2).
    """
    track_list = list(tracks)
    if len(track_list) < 2:
        return np.eye(2)

    ix, iy = state_indices
    positions = np.array([[t.state[ix], t.state[iy]] for t in track_list])
    return np.cov(positions.T)


def cluster_tracks_dbscan(
    tracks: TrackList,
    eps: float,
    min_samples: int = 2,
    state_indices: Tuple[int, int] = (0, 2),
) -> "ClusterSet":
    """
    Cluster tracks using DBSCAN algorithm.

    Parameters
    ----------
    tracks : TrackList
        Tracks to cluster.
    eps : float
        Maximum distance between two tracks to be considered neighbors.
    min_samples : int, optional
        Minimum number of tracks to form a cluster (default: 2).
    state_indices : tuple of int, optional
        Indices of x and y in state vector (default: (0, 2)).

    Returns
    -------
    ClusterSet
        Set of track clusters.
    """
    if len(tracks) == 0:
        return ClusterSet()

    # Extract positions
    positions = tracks.positions(indices=state_indices)

    # Run DBSCAN
    result = dbscan(positions, eps=eps, min_samples=min_samples)

    # Build clusters from labels
    clusters = []
    track_list = list(tracks)
    time = track_list[0].time if track_list else 0.0

    for cluster_id in range(result.n_clusters):
        mask = result.labels == cluster_id
        cluster_tracks = [track_list[i] for i in range(len(track_list)) if mask[i]]
        track_ids = tuple(t.id for t in cluster_tracks)

        centroid = compute_cluster_centroid(cluster_tracks, state_indices)
        covariance = compute_cluster_covariance(cluster_tracks, state_indices)

        clusters.append(
            TrackCluster(
                id=cluster_id,
                track_ids=track_ids,
                centroid=centroid,
                covariance=covariance,
                time=time,
            )
        )

    return ClusterSet(clusters)


def cluster_tracks_kmeans(
    tracks: TrackList,
    n_clusters: int,
    state_indices: Tuple[int, int] = (0, 2),
    rng: Optional[np.random.Generator] = None,
) -> "ClusterSet":
    """
    Cluster tracks using K-means algorithm.

    Parameters
    ----------
    tracks : TrackList
        Tracks to cluster.
    n_clusters : int
        Number of clusters to form.
    state_indices : tuple of int, optional
        Indices of x and y in state vector (default: (0, 2)).
    rng : numpy.random.Generator, optional
        Random number generator for reproducibility.

    Returns
    -------
    ClusterSet
        Set of track clusters.
    """
    if len(tracks) == 0:
        return ClusterSet()

    if n_clusters > len(tracks):
        n_clusters = len(tracks)

    # Extract positions
    positions = tracks.positions(indices=state_indices)

    # Run K-means
    result = kmeans(positions, n_clusters=n_clusters, rng=rng)

    # Build clusters from labels
    clusters = []
    track_list = list(tracks)
    time = track_list[0].time if track_list else 0.0

    for cluster_id in range(n_clusters):
        mask = result.labels == cluster_id
        cluster_tracks = [track_list[i] for i in range(len(track_list)) if mask[i]]

        if len(cluster_tracks) == 0:
            continue

        track_ids = tuple(t.id for t in cluster_tracks)
        centroid = compute_cluster_centroid(cluster_tracks, state_indices)
        covariance = compute_cluster_covariance(cluster_tracks, state_indices)

        clusters.append(
            TrackCluster(
                id=cluster_id,
                track_ids=track_ids,
                centroid=centroid,
                covariance=covariance,
                time=time,
            )
        )

    return ClusterSet(clusters)


class ClusterSet:
    """
    Collection of track clusters.

    Provides:
    - Cluster creation from tracks
    - Cluster queries
    - Cluster merging/splitting

    Parameters
    ----------
    clusters : Iterable[TrackCluster], optional
        Initial clusters to add.

    Examples
    --------
    >>> import numpy as np
    >>> from pytcl.trackers.multi_target import Track, TrackStatus
    >>> from pytcl.containers.track_list import TrackList
    >>> # Create some tracks in two groups
    >>> t1 = Track(id=0, state=np.array([0, 0, 0, 0]),
    ...            covariance=np.eye(4), status=TrackStatus.CONFIRMED,
    ...            hits=5, misses=0, time=1.0)
    >>> t2 = Track(id=1, state=np.array([1, 0, 1, 0]),
    ...            covariance=np.eye(4), status=TrackStatus.CONFIRMED,
    ...            hits=5, misses=0, time=1.0)
    >>> t3 = Track(id=2, state=np.array([10, 0, 10, 0]),
    ...            covariance=np.eye(4), status=TrackStatus.CONFIRMED,
    ...            hits=5, misses=0, time=1.0)
    >>> tracks = TrackList([t1, t2, t3])
    >>> # Cluster using DBSCAN
    >>> clusters = cluster_tracks_dbscan(tracks, eps=5.0, min_samples=2)
    """

    def __init__(self, clusters: Optional[Iterable[TrackCluster]] = None) -> None:
        """Initialize cluster set."""
        if clusters is None:
            self._clusters: List[TrackCluster] = []
        else:
            self._clusters = list(clusters)

        # Build lookups
        self._id_to_idx: Dict[int, int] = {
            c.id: i for i, c in enumerate(self._clusters)
        }
        self._track_to_cluster: Dict[int, int] = {}
        for cluster in self._clusters:
            for track_id in cluster.track_ids:
                self._track_to_cluster[track_id] = cluster.id

    @classmethod
    def from_tracks(
        cls,
        tracks: TrackList,
        method: str = "dbscan",
        **kwargs,
    ) -> ClusterSet:
        """
        Create a ClusterSet by clustering tracks.

        Parameters
        ----------
        tracks : TrackList
            Tracks to cluster.
        method : str
            Clustering method: 'dbscan' or 'kmeans'.
        **kwargs
            Additional arguments passed to the clustering function.
            For DBSCAN: eps, min_samples, state_indices
            For K-means: n_clusters, state_indices, rng

        Returns
        -------
        ClusterSet
            New ClusterSet containing the computed clusters.

        Examples
        --------
        >>> clusters = ClusterSet.from_tracks(tracks, method='dbscan', eps=2.0)
        >>> clusters = ClusterSet.from_tracks(tracks, method='kmeans', n_clusters=3)
        """
        if method == "dbscan":
            return cluster_tracks_dbscan(tracks, **kwargs)
        elif method == "kmeans":
            return cluster_tracks_kmeans(tracks, **kwargs)
        else:
            raise ValueError(f"Unknown clustering method: {method}")

    def __len__(self) -> int:
        """Return number of clusters."""
        return len(self._clusters)

    def __iter__(self) -> Iterator[TrackCluster]:
        """Iterate over clusters."""
        return iter(self._clusters)

    def __getitem__(self, idx: Union[int, slice]) -> Union[TrackCluster, "ClusterSet"]:
        """
        Get cluster by index or slice.

        Parameters
        ----------
        idx : int or slice
            Index or slice to retrieve.

        Returns
        -------
        TrackCluster or ClusterSet
            Single cluster if int, ClusterSet if slice.
        """
        if isinstance(idx, int):
            return self._clusters[idx]
        else:
            return ClusterSet(self._clusters[idx])

    def __contains__(self, cluster_id: int) -> bool:
        """Check if cluster ID exists in set."""
        return cluster_id in self._id_to_idx

    def __repr__(self) -> str:
        """String representation."""
        return f"ClusterSet(n_clusters={len(self)})"

    def get_cluster(self, cluster_id: int) -> Optional[TrackCluster]:
        """
        Get cluster by ID.

        Parameters
        ----------
        cluster_id : int
            Cluster ID to find.

        Returns
        -------
        TrackCluster or None
            The cluster if found, None otherwise.
        """
        idx = self._id_to_idx.get(cluster_id)
        if idx is not None:
            return self._clusters[idx]
        return None

    def get_cluster_for_track(self, track_id: int) -> Optional[TrackCluster]:
        """
        Get the cluster containing a specific track.

        Parameters
        ----------
        track_id : int
            Track ID to find.

        Returns
        -------
        TrackCluster or None
            The cluster containing the track, or None if not found.
        """
        cluster_id = self._track_to_cluster.get(track_id)
        if cluster_id is not None:
            return self.get_cluster(cluster_id)
        return None

    def clusters_in_region(
        self, center: ArrayLike, radius: float
    ) -> List[TrackCluster]:
        """
        Get clusters with centroids within a spatial region.

        Parameters
        ----------
        center : array_like
            Center point [x, y].
        radius : float
            Radius of the region.

        Returns
        -------
        list of TrackCluster
            Clusters within the region.
        """
        center = np.asarray(center, dtype=np.float64)
        result = []

        for cluster in self._clusters:
            dist = np.linalg.norm(cluster.centroid - center)
            if dist <= radius:
                result.append(cluster)

        return result

    @property
    def cluster_ids(self) -> List[int]:
        """Get list of all cluster IDs."""
        return [c.id for c in self._clusters]

    @property
    def n_tracks_total(self) -> int:
        """Get total number of tracks across all clusters."""
        return sum(len(c.track_ids) for c in self._clusters)

    def cluster_stats(
        self,
        cluster_id: int,
        tracks: Optional[TrackList] = None,
        state_indices: Tuple[int, int] = (0, 2),
        velocity_indices: Tuple[int, int] = (1, 3),
    ) -> Optional[ClusterStats]:
        """
        Compute statistics for a cluster.

        Parameters
        ----------
        cluster_id : int
            ID of the cluster.
        tracks : TrackList, optional
            TrackList containing the tracks. Required for velocity coherence.
        state_indices : tuple of int, optional
            Indices of x and y in state vector (default: (0, 2)).
        velocity_indices : tuple of int, optional
            Indices of vx and vy in state vector (default: (1, 3)).

        Returns
        -------
        ClusterStats or None
            Statistics for the cluster, or None if cluster not found.
        """
        cluster = self.get_cluster(cluster_id)
        if cluster is None:
            return None

        n_tracks = len(cluster.track_ids)
        if n_tracks == 0:
            return ClusterStats(
                n_tracks=0,
                mean_separation=0.0,
                max_separation=0.0,
                velocity_coherence=0.0,
            )

        # Compute separations if we have track data
        mean_separation = 0.0
        max_separation = 0.0
        velocity_coherence = 0.0

        if tracks is not None:
            ix, iy = state_indices
            positions = []
            velocities = []

            for tid in cluster.track_ids:
                track = tracks.get_by_id(tid)
                if track is not None:
                    positions.append([track.state[ix], track.state[iy]])
                    ivx, ivy = velocity_indices
                    if len(track.state) > max(ivx, ivy):
                        velocities.append([track.state[ivx], track.state[ivy]])

            if len(positions) > 0:
                positions = np.array(positions)
                centroid = cluster.centroid

                # Compute separations
                separations = np.sqrt(np.sum((positions - centroid) ** 2, axis=1))
                mean_separation = float(np.mean(separations))
                max_separation = float(np.max(separations))

            # Compute velocity coherence
            if len(velocities) > 1:
                velocities = np.array(velocities)
                # Normalize velocities
                norms = np.linalg.norm(velocities, axis=1, keepdims=True)
                norms = np.maximum(norms, 1e-10)
                unit_velocities = velocities / norms

                # Mean velocity direction
                mean_vel = np.mean(unit_velocities, axis=0)
                mean_vel_norm = np.linalg.norm(mean_vel)

                # Coherence is magnitude of mean unit velocity
                velocity_coherence = float(mean_vel_norm)

        return ClusterStats(
            n_tracks=n_tracks,
            mean_separation=mean_separation,
            max_separation=max_separation,
            velocity_coherence=velocity_coherence,
        )

    def all_stats(
        self,
        tracks: Optional[TrackList] = None,
        state_indices: Tuple[int, int] = (0, 2),
        velocity_indices: Tuple[int, int] = (1, 3),
    ) -> Dict[int, ClusterStats]:
        """
        Compute statistics for all clusters.

        Parameters
        ----------
        tracks : TrackList, optional
            TrackList containing the tracks.
        state_indices : tuple of int, optional
            Indices of x and y in state vector (default: (0, 2)).
        velocity_indices : tuple of int, optional
            Indices of vx and vy in state vector (default: (1, 3)).

        Returns
        -------
        dict
            Mapping from cluster ID to ClusterStats.
        """
        result = {}
        for cluster in self._clusters:
            stats = self.cluster_stats(
                cluster.id, tracks, state_indices, velocity_indices
            )
            if stats is not None:
                result[cluster.id] = stats
        return result

    def add_cluster(self, cluster: TrackCluster) -> ClusterSet:
        """
        Add a cluster and return a new ClusterSet.

        Parameters
        ----------
        cluster : TrackCluster
            Cluster to add.

        Returns
        -------
        ClusterSet
            New ClusterSet with the cluster added.
        """
        return ClusterSet(self._clusters + [cluster])

    def remove_cluster(self, cluster_id: int) -> ClusterSet:
        """
        Remove a cluster by ID and return a new ClusterSet.

        Parameters
        ----------
        cluster_id : int
            ID of cluster to remove.

        Returns
        -------
        ClusterSet
            New ClusterSet without the specified cluster.
        """
        return ClusterSet([c for c in self._clusters if c.id != cluster_id])

    def merge_clusters(
        self,
        id1: int,
        id2: int,
        new_id: Optional[int] = None,
    ) -> ClusterSet:
        """
        Merge two clusters into one.

        Parameters
        ----------
        id1 : int
            ID of first cluster.
        id2 : int
            ID of second cluster.
        new_id : int, optional
            ID for the merged cluster. Defaults to id1.

        Returns
        -------
        ClusterSet
            New ClusterSet with merged cluster.

        Raises
        ------
        ValueError
            If either cluster ID is not found.
        """
        c1 = self.get_cluster(id1)
        c2 = self.get_cluster(id2)

        if c1 is None:
            raise ValueError(f"Cluster {id1} not found")
        if c2 is None:
            raise ValueError(f"Cluster {id2} not found")

        if new_id is None:
            new_id = id1

        # Merge track IDs
        merged_track_ids = tuple(set(c1.track_ids) | set(c2.track_ids))

        # Compute new centroid (weighted average)
        n1, n2 = len(c1.track_ids), len(c2.track_ids)
        new_centroid = (n1 * c1.centroid + n2 * c2.centroid) / (n1 + n2)

        # Combine covariances (simple average for now)
        new_covariance = (c1.covariance + c2.covariance) / 2

        merged = TrackCluster(
            id=new_id,
            track_ids=merged_track_ids,
            centroid=new_centroid,
            covariance=new_covariance,
            time=max(c1.time, c2.time),
        )

        # Build new cluster list
        new_clusters = [c for c in self._clusters if c.id not in (id1, id2)]
        new_clusters.append(merged)

        return ClusterSet(new_clusters)

    def split_cluster(
        self,
        cluster_id: int,
        track_ids_1: Iterable[int],
        track_ids_2: Iterable[int],
        new_id_1: Optional[int] = None,
        new_id_2: Optional[int] = None,
        tracks: Optional[TrackList] = None,
        state_indices: Tuple[int, int] = (0, 2),
    ) -> ClusterSet:
        """
        Split a cluster into two.

        Parameters
        ----------
        cluster_id : int
            ID of cluster to split.
        track_ids_1 : Iterable[int]
            Track IDs for first new cluster.
        track_ids_2 : Iterable[int]
            Track IDs for second new cluster.
        new_id_1 : int, optional
            ID for first new cluster. Defaults to cluster_id.
        new_id_2 : int, optional
            ID for second new cluster. Defaults to max(cluster_ids) + 1.
        tracks : TrackList, optional
            TrackList for computing centroids. If None, uses existing centroid.
        state_indices : tuple of int, optional
            Indices of x and y in state vector (default: (0, 2)).

        Returns
        -------
        ClusterSet
            New ClusterSet with split clusters.

        Raises
        ------
        ValueError
            If cluster ID is not found.
        """
        original = self.get_cluster(cluster_id)
        if original is None:
            raise ValueError(f"Cluster {cluster_id} not found")

        ids_1 = tuple(track_ids_1)
        ids_2 = tuple(track_ids_2)

        if new_id_1 is None:
            new_id_1 = cluster_id
        if new_id_2 is None:
            new_id_2 = max(self.cluster_ids) + 1 if self.cluster_ids else 0

        # Compute centroids
        if tracks is not None:
            tracks_1 = [tracks.get_by_id(tid) for tid in ids_1]
            tracks_1 = [t for t in tracks_1 if t is not None]
            tracks_2 = [tracks.get_by_id(tid) for tid in ids_2]
            tracks_2 = [t for t in tracks_2 if t is not None]

            centroid_1 = compute_cluster_centroid(tracks_1, state_indices)
            centroid_2 = compute_cluster_centroid(tracks_2, state_indices)
            cov_1 = compute_cluster_covariance(tracks_1, state_indices)
            cov_2 = compute_cluster_covariance(tracks_2, state_indices)
        else:
            # Use original centroid for both
            centroid_1 = original.centroid.copy()
            centroid_2 = original.centroid.copy()
            cov_1 = original.covariance.copy()
            cov_2 = original.covariance.copy()

        cluster_1 = TrackCluster(
            id=new_id_1,
            track_ids=ids_1,
            centroid=centroid_1,
            covariance=cov_1,
            time=original.time,
        )
        cluster_2 = TrackCluster(
            id=new_id_2,
            track_ids=ids_2,
            centroid=centroid_2,
            covariance=cov_2,
            time=original.time,
        )

        # Build new cluster list
        new_clusters = [c for c in self._clusters if c.id != cluster_id]
        new_clusters.extend([cluster_1, cluster_2])

        return ClusterSet(new_clusters)

    def copy(self) -> ClusterSet:
        """
        Create a copy of this ClusterSet.

        Returns
        -------
        ClusterSet
            A new ClusterSet with the same clusters.
        """
        return ClusterSet(self._clusters)


__all__ = [
    "TrackCluster",
    "ClusterSet",
    "ClusterStats",
    "cluster_tracks_dbscan",
    "cluster_tracks_kmeans",
    "compute_cluster_centroid",
]
