# controllers/clustering.py
"""
Clustering Controller
=====================

Controller for clustering and novelty detection operations on audio embeddings.

Orchestrates between CLI/API views and core clustering functions.
Handles UMAP/PCA reduction, HDBSCAN/k-means clustering, novelty detection,
and visualization generation.

Example:
    from bioamla.controllers.clustering import ClusteringController
    import numpy as np

    controller = ClusteringController()

    # Cluster embeddings
    embeddings = np.load("embeddings.npy")
    result = controller.cluster(embeddings, method="hdbscan")
    print(f"Found {result.data.n_clusters} clusters")

    # Reduce and visualize
    result = controller.reduce_for_visualization(embeddings)
    x, y = result.data["x"], result.data["y"]
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from .base import BaseController, ControllerResult


@dataclass
class ClusteringSummary:
    """Summary of clustering results."""

    n_clusters: int
    n_samples: int
    n_noise: int
    noise_percentage: float
    silhouette_score: float
    method: str
    labels: List[int] = field(default_factory=list)
    cluster_sizes: Dict[int, int] = field(default_factory=dict)


@dataclass
class NoveltyDetectionSummary:
    """Summary of novelty detection results."""

    n_samples: int
    n_novel: int
    n_known: int
    novel_percentage: float
    method: str
    threshold: Optional[float]
    novel_indices: List[int] = field(default_factory=list)


@dataclass
class ClusterAnalysis:
    """Detailed cluster analysis results."""

    n_clusters: int
    n_samples: int
    n_noise: int
    silhouette_score: float
    calinski_harabasz_score: float
    cluster_stats: Dict[int, Dict[str, Any]]


class ClusteringController(BaseController):
    """
    Controller for clustering and novelty detection operations.

    Provides high-level methods for:
    - Embedding clustering (HDBSCAN, k-means, DBSCAN)
    - Dimensionality reduction (UMAP, t-SNE, PCA)
    - Novelty/outlier detection
    - Cluster quality metrics
    - Cluster visualization
    - Cluster export
    """

    def __init__(self):
        """Initialize clustering controller."""
        super().__init__()
        self._clusterer = None
        self._reducer = None
        self._novelty_detector = None

    # =========================================================================
    # Clustering
    # =========================================================================

    def cluster(
        self,
        embeddings: np.ndarray,
        method: str = "hdbscan",
        n_clusters: Optional[int] = None,
        min_cluster_size: int = 5,
        min_samples: int = 3,
        **kwargs,
    ) -> ControllerResult[ClusteringSummary]:
        """
        Cluster embeddings.

        Args:
            embeddings: Input embeddings of shape (n_samples, n_features)
            method: Clustering method ("hdbscan", "kmeans", "dbscan", "agglomerative")
            n_clusters: Number of clusters (for k-means/agglomerative)
            min_cluster_size: Minimum cluster size (for HDBSCAN)
            min_samples: Minimum samples per cluster
            **kwargs: Additional arguments for the clusterer

        Returns:
            Result with clustering summary
        """
        # Start run tracking
        self._start_run(
            name=f"Clustering: {method}",
            action="cluster",
            parameters={
                "method": method,
                "n_clusters": n_clusters,
                "min_cluster_size": min_cluster_size,
                "min_samples": min_samples,
                "n_samples": len(embeddings),
                "n_features": embeddings.shape[1] if len(embeddings.shape) > 1 else 1,
            },
        )

        try:
            from bioamla.core.analysis.clustering import AudioClusterer, ClusteringConfig

            config = ClusteringConfig(
                method=method,
                n_clusters=n_clusters or 10,
                min_cluster_size=min_cluster_size,
                min_samples=min_samples,
            )

            clusterer = AudioClusterer(config=config, **kwargs)
            labels = clusterer.fit_predict(embeddings)

            # Compute metrics
            n_noise = int((labels == -1).sum())
            unique_labels = set(labels)
            n_clusters = len(unique_labels) - (1 if -1 in unique_labels else 0)

            # Cluster sizes
            cluster_sizes = {}
            for label in unique_labels:
                cluster_sizes[int(label)] = int((labels == label).sum())

            # Silhouette score
            silhouette = 0.0
            if n_clusters > 1:
                from sklearn.metrics import silhouette_score

                mask = labels >= 0
                if mask.sum() > n_clusters:
                    silhouette = silhouette_score(embeddings[mask], labels[mask])

            summary = ClusteringSummary(
                n_clusters=n_clusters,
                n_samples=len(embeddings),
                n_noise=n_noise,
                noise_percentage=n_noise / len(embeddings) * 100,
                silhouette_score=float(silhouette),
                method=method,
                labels=labels.tolist(),
                cluster_sizes=cluster_sizes,
            )

            self._clusterer = clusterer

            # Complete run with results
            self._complete_run(
                results={
                    "n_clusters": n_clusters,
                    "n_samples": len(embeddings),
                    "n_noise": n_noise,
                    "noise_percentage": n_noise / len(embeddings) * 100,
                    "silhouette_score": float(silhouette),
                },
            )

            return ControllerResult.ok(
                data=summary,
                message=f"Found {n_clusters} clusters with {method}",
                labels=labels,
                clusterer=clusterer,
            )
        except Exception as e:
            self._fail_run(str(e))
            return ControllerResult.fail(str(e))

    def find_optimal_k(
        self,
        embeddings: np.ndarray,
        k_range: Tuple[int, int] = (2, 20),
        method: str = "silhouette",
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Find optimal number of clusters for k-means.

        Args:
            embeddings: Input embeddings
            k_range: Range of k values to try (min, max)
            method: Method for evaluation ("silhouette", "elbow")

        Returns:
            Result with optimal k and scores
        """
        try:
            from bioamla.core.analysis.clustering import find_optimal_clusters

            optimal_k = find_optimal_clusters(
                embeddings,
                method=method,
                k_range=k_range,
            )

            return ControllerResult.ok(
                data={
                    "optimal_k": optimal_k,
                    "k_range": k_range,
                    "method": method,
                },
                message=f"Optimal k = {optimal_k} using {method} method",
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    def analyze_clusters(
        self,
        embeddings: np.ndarray,
        labels: np.ndarray,
        filepaths: Optional[List[str]] = None,
    ) -> ControllerResult[ClusterAnalysis]:
        """
        Perform detailed analysis of clustering results.

        Args:
            embeddings: Input embeddings
            labels: Cluster labels
            filepaths: Optional list of file paths for metadata

        Returns:
            Result with detailed cluster analysis
        """
        try:
            from bioamla.core.analysis.clustering import analyze_clusters

            metadata = None
            if filepaths:
                metadata = [{"filepath": fp} for fp in filepaths]

            analysis = analyze_clusters(embeddings, labels, metadata)

            result = ClusterAnalysis(
                n_clusters=analysis["n_clusters"],
                n_samples=analysis["n_samples"],
                n_noise=analysis["n_noise"],
                silhouette_score=analysis["silhouette_score"],
                calinski_harabasz_score=analysis["calinski_harabasz_score"],
                cluster_stats=analysis["cluster_stats"],
            )

            return ControllerResult.ok(
                data=result,
                message=f"Analyzed {result.n_clusters} clusters with silhouette={result.silhouette_score:.3f}",
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    # =========================================================================
    # Dimensionality Reduction
    # =========================================================================

    def reduce_dimensions(
        self,
        embeddings: np.ndarray,
        method: str = "umap",
        n_components: int = 2,
        output_path: Optional[str] = None,
        **kwargs,
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Reduce dimensionality of embeddings.

        Args:
            embeddings: Input embeddings of shape (n_samples, n_features)
            method: Reduction method ("umap", "tsne", "pca")
            n_components: Number of output dimensions
            output_path: Optional path to save reduced embeddings
            **kwargs: Additional arguments for the reducer

        Returns:
            Result with reduced embeddings
        """
        try:
            from bioamla.core.analysis.clustering import reduce_dimensions

            reduced = reduce_dimensions(
                embeddings,
                method=method,
                n_components=n_components,
                **kwargs,
            )

            if output_path:
                np.save(output_path, reduced)

            return ControllerResult.ok(
                data={
                    "original_shape": embeddings.shape,
                    "reduced_shape": reduced.shape,
                    "method": method,
                    "n_components": n_components,
                    "output_path": output_path,
                },
                message=f"Reduced from {embeddings.shape[-1]}D to {n_components}D using {method}",
                reduced_embeddings=reduced,
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    def reduce_for_visualization(
        self,
        embeddings: np.ndarray,
        method: str = "umap",
        labels: Optional[np.ndarray] = None,
        **kwargs,
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Reduce embeddings to 2D for visualization.

        Args:
            embeddings: Input embeddings
            method: Reduction method ("umap", "tsne", "pca")
            labels: Optional cluster labels for coloring
            **kwargs: Additional arguments for the reducer

        Returns:
            Result with 2D coordinates
        """
        result = self.reduce_dimensions(
            embeddings,
            method=method,
            n_components=2,
            **kwargs,
        )

        if result.success:
            coords = result.metadata["reduced_embeddings"]
            result.data["x"] = coords[:, 0].tolist()
            result.data["y"] = coords[:, 1].tolist()
            if labels is not None:
                result.data["labels"] = (
                    labels.tolist() if hasattr(labels, "tolist") else list(labels)
                )

        return result

    # =========================================================================
    # Novelty Detection
    # =========================================================================

    def detect_novelty(
        self,
        embeddings: np.ndarray,
        known_embeddings: Optional[np.ndarray] = None,
        known_labels: Optional[np.ndarray] = None,
        method: str = "distance",
        threshold: Optional[float] = None,
        contamination: float = 0.1,
    ) -> ControllerResult[NoveltyDetectionSummary]:
        """
        Detect novel/outlier samples in embeddings.

        Args:
            embeddings: Embeddings to check for novelty
            known_embeddings: Optional known embeddings to fit detector on
            known_labels: Optional labels for known embeddings
            method: Detection method ("distance", "isolation_forest", "lof")
            threshold: Novelty threshold (auto if None)
            contamination: Expected proportion of outliers

        Returns:
            Result with novelty detection summary
        """
        try:
            from bioamla.core.analysis.clustering import NoveltyDetector

            detector = NoveltyDetector(
                method=method,
                threshold=threshold,
                contamination=contamination,
            )

            # Fit on known embeddings or self
            if known_embeddings is not None:
                detector.fit(known_embeddings, known_labels)
            else:
                detector.fit(embeddings)

            # Predict novelty
            results = detector.predict(embeddings)

            novel_indices = [r.sample_idx for r in results if r.is_novel]
            n_novel = len(novel_indices)

            summary = NoveltyDetectionSummary(
                n_samples=len(embeddings),
                n_novel=n_novel,
                n_known=len(embeddings) - n_novel,
                novel_percentage=n_novel / len(embeddings) * 100,
                method=method,
                threshold=threshold,
                novel_indices=novel_indices,
            )

            # Create boolean mask
            is_novel = np.array([r.is_novel for r in results])
            novelty_scores = np.array([r.novelty_score for r in results])

            self._novelty_detector = detector

            return ControllerResult.ok(
                data=summary,
                message=f"Detected {n_novel} novel samples ({summary.novel_percentage:.1f}%)",
                is_novel=is_novel,
                novelty_scores=novelty_scores,
                detector=detector,
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    def get_most_novel(
        self,
        embeddings: np.ndarray,
        n_samples: int = 10,
        known_embeddings: Optional[np.ndarray] = None,
        method: str = "distance",
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Get the most novel samples.

        Args:
            embeddings: Embeddings to check
            n_samples: Number of novel samples to return
            known_embeddings: Optional known embeddings to compare against
            method: Detection method

        Returns:
            Result with indices and scores of most novel samples
        """
        try:
            from bioamla.core.analysis.clustering import NoveltyDetector

            detector = NoveltyDetector(method=method)

            if known_embeddings is not None:
                detector.fit(known_embeddings)
            else:
                detector.fit(embeddings)

            results = detector.predict(embeddings)

            # Sort by novelty score
            sorted_results = sorted(results, key=lambda r: r.novelty_score, reverse=True)
            top_n = sorted_results[:n_samples]

            indices = [r.sample_idx for r in top_n]
            scores = [r.novelty_score for r in top_n]

            return ControllerResult.ok(
                data={
                    "indices": indices,
                    "scores": scores,
                    "n_samples": len(indices),
                },
                message=f"Found {len(indices)} most novel samples",
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    # =========================================================================
    # Cluster Similarity
    # =========================================================================

    def compute_similarity(
        self,
        embeddings: np.ndarray,
        labels: np.ndarray,
        metric: str = "cosine",
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Compute pairwise similarity between clusters.

        Args:
            embeddings: Input embeddings
            labels: Cluster labels
            metric: Distance metric ("cosine", "euclidean")

        Returns:
            Result with similarity matrix
        """
        try:
            from bioamla.core.analysis.clustering import compute_cluster_similarity

            similarity = compute_cluster_similarity(embeddings, labels, metric=metric)

            return ControllerResult.ok(
                data={
                    "n_clusters": similarity.shape[0],
                    "metric": metric,
                },
                message=f"Computed {similarity.shape[0]}x{similarity.shape[0]} similarity matrix",
                similarity_matrix=similarity,
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    def sort_clusters_by_similarity(
        self,
        embeddings: np.ndarray,
        labels: np.ndarray,
        reference_cluster: Optional[int] = None,
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Sort clusters by similarity to each other.

        Args:
            embeddings: Input embeddings
            labels: Cluster labels
            reference_cluster: Starting cluster (uses largest if None)

        Returns:
            Result with sorted cluster order
        """
        try:
            from bioamla.core.analysis.clustering import sort_clusters_by_similarity

            sorted_labels = sort_clusters_by_similarity(
                embeddings, labels, reference_label=reference_cluster
            )

            return ControllerResult.ok(
                data={
                    "sorted_labels": sorted_labels,
                    "reference_cluster": reference_cluster or "largest",
                },
                message=f"Sorted {len(sorted_labels)} clusters by similarity",
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    # =========================================================================
    # Export
    # =========================================================================

    def export_clusters(
        self,
        labels: np.ndarray,
        filepaths: List[str],
        output_dir: str,
        copy_files: bool = False,
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Export clustering results to directory structure.

        Args:
            labels: Cluster labels
            filepaths: List of file paths
            output_dir: Output directory
            copy_files: Whether to copy files to cluster directories

        Returns:
            Result with export info
        """
        try:
            from bioamla.core.analysis.clustering import export_clusters

            output_path = export_clusters(labels, filepaths, output_dir, copy_files=copy_files)

            unique_labels = set(labels)
            n_clusters = len(unique_labels) - (1 if -1 in unique_labels else 0)

            return ControllerResult.ok(
                data={
                    "output_dir": output_path,
                    "n_clusters": n_clusters,
                    "n_files": len(filepaths),
                    "files_copied": copy_files,
                },
                message=f"Exported {n_clusters} clusters to {output_path}",
            )
        except Exception as e:
            return ControllerResult.fail(str(e))

    def export_to_csv(
        self,
        labels: np.ndarray,
        filepaths: List[str],
        output_path: str,
        embeddings: Optional[np.ndarray] = None,
        reduced_embeddings: Optional[np.ndarray] = None,
    ) -> ControllerResult[Dict[str, Any]]:
        """
        Export clustering results to CSV.

        Args:
            labels: Cluster labels
            filepaths: List of file paths
            output_path: Output CSV path
            embeddings: Optional embeddings to include
            reduced_embeddings: Optional 2D coordinates to include

        Returns:
            Result with export info
        """
        try:
            import csv

            from bioamla.core.files import TextFile

            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)

            with TextFile(output_path, mode="w", newline="") as f:
                fieldnames = ["filepath", "cluster"]

                if reduced_embeddings is not None and reduced_embeddings.shape[1] >= 2:
                    fieldnames.extend(["x", "y"])

                writer = csv.DictWriter(f.handle, fieldnames=fieldnames)
                writer.writeheader()

                for i, (fp, label) in enumerate(zip(filepaths, labels)):
                    row = {
                        "filepath": fp,
                        "cluster": int(label),
                    }
                    if reduced_embeddings is not None and reduced_embeddings.shape[1] >= 2:
                        row["x"] = float(reduced_embeddings[i, 0])
                        row["y"] = float(reduced_embeddings[i, 1])
                    writer.writerow(row)

            return ControllerResult.ok(
                data={
                    "output_path": str(output_path),
                    "n_rows": len(labels),
                },
                message=f"Exported {len(labels)} rows to {output_path}",
            )
        except Exception as e:
            return ControllerResult.fail(str(e))
