from __future__ import annotations

import numpy as np
import pandas as pd
from sklearn.manifold import TSNE

from .console import console
from .config import ProjectPaths
from .io import load_pickle, write_csv


def generate_atlas(
    paths: ProjectPaths | None = None,
    *,
    perplexity: int = 30,
    random_seed: int = 42,
) -> dict:
    paths = paths or ProjectPaths()
    console.info("Loading model and dataset to build atlas coordinates")
    model = load_pickle(paths.resolve_model_input())
    dataset = load_pickle(paths.resolve_dataset_input())

    item_embeddings = model.item_embeddings
    _, _, item_id_map, _ = dataset.mapping()
    internal_to_original_anime_id = {internal_id: anime_id for anime_id, internal_id in item_id_map.items()}
    console.info(f"Atlas input embeddings: {len(item_embeddings)} items")

    if len(item_embeddings) <= 1:
        coordinates = np.zeros((len(item_embeddings), 2))
    else:
        safe_perplexity = min(perplexity, max(1, len(item_embeddings) - 1))
        if safe_perplexity >= len(item_embeddings):
            safe_perplexity = len(item_embeddings) - 1
        console.info(f"Running t-SNE with perplexity={safe_perplexity}")
        tsne = TSNE(
            n_components=2,
            random_state=random_seed,
            perplexity=safe_perplexity,
            max_iter=1000,
        )
        coordinates = tsne.fit_transform(item_embeddings)

    rows = []
    for internal_id, (x_coord, y_coord) in enumerate(coordinates):
        anime_id = internal_to_original_anime_id.get(internal_id)
        if anime_id is None:
            continue
        rows.append(
            {
                "anime_id": int(anime_id),
                "x_coord": float(x_coord),
                "y_coord": float(y_coord),
            }
        )

    atlas_df = pd.DataFrame(rows, columns=["anime_id", "x_coord", "y_coord"])
    if not atlas_df.empty:
        atlas_df = atlas_df.sort_values("anime_id")
    write_csv(atlas_df, paths.atlas, paths.legacy_atlas, quoting=0)
    console.success(f"Wrote atlas coordinates for {len(atlas_df)} anime to {paths.atlas}")
    return {
        "atlas_path": str(paths.atlas),
        "anime_count": int(len(atlas_df)),
    }