from __future__ import annotations

from pathlib import Path
from typing import Any

import torch

from .io import (
    build_metadata_index,
    load_metadata_records,
    load_vectors,
    resolve_image_paths,
)
from .plotting import create_3d_figure, save_interactive_html
from .types import EmbeddingDataset, PCAResult


class CitallExplorer:
    """High-level API to build interactive 3D PCA visualizations."""

    def __init__(self, dataset: EmbeddingDataset):
        self.dataset = dataset
        self.pca_result: PCAResult | None = None

    @classmethod
    def from_files(
        cls,
        vectors_pt: str | Path,
        metadata_path: str | Path | None = None,
        metadata_key: str | None = None,
        images_dir: str | Path | None = None,
        vector_key: str | None = None,
        id_key: str | None = None,
        strict_metadata: bool = False,
    ) -> "CitallExplorer":
        vectors, ids = load_vectors(vectors_pt, vector_key=vector_key, id_key=id_key)

        metadata_by_id: dict[str, dict[str, Any]] = {}
        if metadata_path is not None:
            if metadata_key is None:
                raise ValueError("metadata_key is required when metadata_path is provided.")
            records = load_metadata_records(metadata_path)
            metadata_by_id = build_metadata_index(records, metadata_key=metadata_key)

            if strict_metadata:
                missing = [x for x in ids if x not in metadata_by_id]
                if missing:
                    preview = ", ".join(missing[:5])
                    raise ValueError(
                        "Some embedding ids are missing metadata entries. "
                        f"Missing count={len(missing)}, first examples={preview}"
                    )

        image_paths_by_id = resolve_image_paths(ids, images_dir=images_dir)

        dataset = EmbeddingDataset(
            vectors=vectors,
            ids=ids,
            metadata_by_id=metadata_by_id,
            image_paths_by_id=image_paths_by_id,
        )
        return cls(dataset)

    def compute_pca(self, n_components: int = 3) -> PCAResult:
        if n_components != 3:
            raise ValueError("This method currently supports n_components=3 only.")

        x = self.dataset.vectors
        if not torch.is_floating_point(x):
            x = x.float()
        x = x.to(dtype=torch.float32)

        if x.ndim != 2:
            raise ValueError(f"Expected vectors shape (N, D), got {tuple(x.shape)}")
        if x.shape[0] < 3:
            raise ValueError("At least 3 samples are required for a meaningful 3D PCA.")

        x_centered = x - x.mean(dim=0, keepdim=True)
        min_dim = min(x_centered.shape)
        q = min(max(n_components + 2, 5), min_dim)

        try:
            _, s, v = torch.pca_lowrank(x_centered, q=q, center=False)
            components = v[:, :n_components]
            coords = x_centered @ components
            explained = (s ** 2) / max(x_centered.shape[0] - 1, 1)
            ratio = (explained[:n_components] / explained.sum()).tolist()
        except RuntimeError:
            # Fallback for edge cases where low-rank PCA fails.
            u, s, _ = torch.linalg.svd(x_centered, full_matrices=False)
            coords = u[:, :n_components] * s[:n_components]
            explained = (s ** 2) / max(x_centered.shape[0] - 1, 1)
            ratio = (explained[:n_components] / explained.sum()).tolist()

        self.pca_result = PCAResult(coords_3d=coords, explained_variance_ratio=ratio)
        return self.pca_result

    def create_figure(
        self,
        title: str = "citall: PCA 3D Explorer",
        hover_fields: list[str] | None = None,
        point_size: int = 5,
        color_by: str | None = None,
        legend_title: str | None = None,
    ):
        if self.pca_result is None:
            self.compute_pca(n_components=3)

        assert self.pca_result is not None
        return create_3d_figure(
            coords_3d=self.pca_result.coords_3d,
            ids=self.dataset.ids,
            metadata_by_id=self.dataset.metadata_by_id,
            image_paths_by_id=self.dataset.image_paths_by_id,
            explained_variance_ratio=self.pca_result.explained_variance_ratio,
            title=title,
            hover_fields=hover_fields,
            point_size=point_size,
            color_by=color_by,
            legend_title=legend_title,
        )

    def save_html(
        self,
        output_html: str | Path,
        title: str = "citall: PCA 3D Explorer",
        hover_fields: list[str] | None = None,
        point_size: int = 5,
        color_by: str | None = None,
        legend_title: str | None = None,
    ) -> Path:
        fig, payloads = self.create_figure(
            title=title,
            hover_fields=hover_fields,
            point_size=point_size,
            color_by=color_by,
            legend_title=legend_title,
        )
        output_html = Path(output_html)
        save_interactive_html(fig=fig, payloads=payloads, output_html=output_html)
        return output_html

    def summary(self) -> dict[str, Any]:
        ratio = self.pca_result.explained_variance_ratio if self.pca_result else None
        return {
            "num_points": len(self.dataset.ids),
            "vector_shape": tuple(self.dataset.vectors.shape),
            "metadata_rows_matched": sum(
                1 for item_id in self.dataset.ids if item_id in self.dataset.metadata_by_id
            ),
            "images_matched": sum(
                1 for item_id in self.dataset.ids if item_id in self.dataset.image_paths_by_id
            ),
            "pca_explained_variance_ratio": ratio,
        }


def generate_pca_html(
    vectors_pt: str | Path,
    output_html: str | Path,
    metadata_path: str | Path | None = None,
    metadata_key: str | None = None,
    images_dir: str | Path | None = None,
    vector_key: str | None = None,
    id_key: str | None = None,
    strict_metadata: bool = False,
    title: str = "citall: PCA 3D Explorer",
    hover_fields: list[str] | None = None,
    point_size: int = 5,
    color_by: str | None = None,
    legend_title: str | None = None,
) -> Path:
    explorer = CitallExplorer.from_files(
        vectors_pt=vectors_pt,
        metadata_path=metadata_path,
        metadata_key=metadata_key,
        images_dir=images_dir,
        vector_key=vector_key,
        id_key=id_key,
        strict_metadata=strict_metadata,
    )
    explorer.compute_pca(n_components=3)
    return explorer.save_html(
        output_html=output_html,
        title=title,
        hover_fields=hover_fields,
        point_size=point_size,
        color_by=color_by,
        legend_title=legend_title,
    )
