from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd
from scipy import sparse

Matrix = np.ndarray | sparse.spmatrix


@dataclass(frozen=True)
class CustomS3:
    s3class: str
    data: dict[str, Any]


@dataclass(frozen=True)
class CustomS4:
    s4class: str
    data: dict[str, Any]


@dataclass(frozen=True)
class Assay:
    name: str
    features: np.ndarray
    data: Matrix
    key: str
    counts: Matrix | None = None
    scale_data: np.ndarray | None = None
    scaled_features: np.ndarray | None = None
    meta_features: pd.DataFrame | None = None
    variable_features: np.ndarray | None = None
    misc: dict[str, Any] | None = None


@dataclass(frozen=True)
class Reduction:
    name: str
    active_assay: list[str]
    key: str
    global_: bool
    cell_embeddings: np.ndarray
    feature_loadings: np.ndarray | None = None
    feature_loadings_projected: np.ndarray | None = None
    misc: dict[str, Any] | None = None
    jackstraw: dict[str, Any] | None = None


@dataclass(frozen=True)
class Graph:
    name: str
    matrix: sparse.spmatrix
    assay_used: str | None = None


@dataclass(frozen=True)
class SeuratObject:
    project: str
    active_assay: str
    version: str
    cell_names: np.ndarray
    meta_data: pd.DataFrame
    assays: dict[str, Assay]
    reductions: dict[str, Reduction]
    graphs: dict[str, Graph]
    images: dict[str, Any]
    misc: dict[str, Any]
    tools: dict[str, Any]
    commands: dict[str, Any]

    def summary(self) -> str:
        """Generate a summary of the SeuratObject.

        Returns:
            str: A formatted summary string containing key information about the object.
        """
        lines = []
        lines.append(f"SeuratObject: {self.project}")
        lines.append(f"Version: {self.version}")
        lines.append(f"Active assay: {self.active_assay}")
        lines.append(f"Number of cells: {len(self.cell_names)}")

        # Assay information
        lines.append(f"\nAssays ({len(self.assays)}):")
        for assay_name, assay in self.assays.items():
            n_features = assay.features.shape[0] if hasattr(assay.features, "shape") else len(assay.features)
            n_cells = (
                assay.data.shape[1]
                if hasattr(assay.data, "shape") and len(assay.data.shape) > 1
                else assay.data.shape[0]
            )
            active_marker = " (active)" if assay_name == self.active_assay else ""
            lines.append(f"  - {assay_name}: {n_features} features × {n_cells} cells{active_marker}")

        # Dimensional reductions
        if self.reductions:
            lines.append(f"\nDimensional reductions ({len(self.reductions)}):")
            for red_name, reduction in self.reductions.items():
                n_dims = reduction.cell_embeddings.shape[1] if len(reduction.cell_embeddings.shape) > 1 else 1
                lines.append(f"  - {red_name}: {n_dims} dimensions")

        # Graphs
        if self.graphs:
            lines.append(f"\nGraphs ({len(self.graphs)}):")
            for graph_name in self.graphs.keys():
                lines.append(f"  - {graph_name}")

        # Metadata
        if self.meta_data is not None and not self.meta_data.empty:
            lines.append(f"\nMetadata columns ({len(self.meta_data.columns)}):")
            col_list = ", ".join(self.meta_data.columns[:10])
            if len(self.meta_data.columns) > 10:
                col_list += f", ... ({len(self.meta_data.columns) - 10} more)"
            lines.append(f"  {col_list}")

        # Images
        if self.images:
            lines.append(f"\nImages: {len(self.images)}")

        return "\n".join(lines)

    def __str__(self) -> str:
        """Return a string representation of the SeuratObject."""
        return self.summary()

    def __repr__(self) -> str:
        """Return a detailed representation of the SeuratObject."""
        return self.summary()
