from __future__ import annotations

from collections import Counter

import numpy as np
from umap import UMAP

from services.artifacts import ArtifactStore, _split_lecturers
from services.pipeline import detect_language, embed_query, find_closest_topic_with_similarity


def _build_lecturer_viz_data(
    store: ArtifactStore,
) -> tuple[list[str], np.ndarray, list[int | None], list[int], list[str | None], object | None]:
    """Shared helper: build per-lecturer mean embeddings, UMAP coords, cluster labels.

    Returns (names, coords, clusters, counts, langs, reducer_model).
    Applies _split_lecturers so multi-author rows are expanded correctly.
    Lecturer embeddings are computed per-language to avoid mixing embedding spaces.
    """
    metadata = store.metadata
    lecturer_col = store.lecturer_col
    topic_col = store.topic_col
    lang_col = store.lang_col

    if lecturer_col not in metadata.columns:
        return [], np.zeros((0, 2), dtype=np.float32), [], [], [], None

    lecturer_to_indices: dict[str, list[int]] = {}
    for idx, cell in enumerate(metadata[lecturer_col].astype(str).tolist()):
        for name in _split_lecturers(cell):
            lecturer_to_indices.setdefault(name, []).append(idx)

    if not lecturer_to_indices:
        return [], np.zeros((0, 2), dtype=np.float32), [], [], [], None

    lang_vals = metadata[lang_col] if (lang_col and lang_col in metadata.columns) else None

    names = list(lecturer_to_indices.keys())
    emb_list: list[np.ndarray] = []
    clusters: list[int | None] = []
    counts: list[int] = []
    langs: list[str | None] = []

    topic_vals = metadata[topic_col] if (topic_col and topic_col in metadata.columns) else None

    for name in names:
        idxs = np.asarray(lecturer_to_indices[name], dtype=np.int64)
        counts.append(int(idxs.size))

        # Determine dominant language for this lecturer
        dom_lang: str | None = None
        if lang_vals is not None:
            lv = [str(x).strip().lower() for x in lang_vals.iloc[idxs].tolist() if str(x).strip().lower() in {"en", "id"}]
            if lv:
                dom_lang = Counter(lv).most_common(1)[0][0]

        # Average embeddings only within the dominant language space
        if dom_lang is not None and lang_vals is not None:
            same_lang_mask = lang_vals.iloc[idxs].astype(str).str.strip().str.lower() == dom_lang
            lang_idxs = idxs[same_lang_mask.to_numpy()]
            if lang_idxs.size > 0:
                vec = store.doc_embeddings[lang_idxs].astype(np.float32, copy=False).mean(axis=0)
            else:
                vec = store.doc_embeddings[idxs].astype(np.float32, copy=False).mean(axis=0)
        else:
            vec = store.doc_embeddings[idxs].astype(np.float32, copy=False).mean(axis=0)

        emb_list.append(vec)
        langs.append(dom_lang)

        if topic_vals is not None:
            raw = [
                x for x in topic_vals.iloc[idxs].tolist()
                if x is not None and str(x).strip() != "" and str(x).lower() not in {"nan", "none"}
            ]
            if raw:
                try:
                    ints = [int(float(x)) for x in raw if int(float(x)) >= 0]
                    if ints:
                        clusters.append(int(Counter(ints).most_common(1)[0][0]))
                    else:
                        clusters.append(None)
                except Exception:
                    clusters.append(None)
            else:
                clusters.append(None)
        else:
            clusters.append(None)

    X = np.stack(emb_list, axis=0).astype(np.float32, copy=False)
    if X.shape[0] == 1:
        coords = np.zeros((1, 2), dtype=np.float32)
        reducer = None
    else:
        n_neighbors = min(15, max(2, X.shape[0] - 1))
        try:
            reducer = UMAP(
                n_components=2,
                n_neighbors=n_neighbors,
                min_dist=0.08,
                metric="cosine",
                random_state=42,
                transform_seed=42,
            )
            coords = reducer.fit_transform(X).astype(np.float32, copy=False)
        except Exception:
            # Fallback to a simple 2D projection so the endpoint never hard-crashes
            reducer = None
            # Normalize rows
            Xn = X - X.mean(axis=0, keepdims=True)
            norms = np.linalg.norm(Xn, axis=1, keepdims=True)
            norms = np.where(norms == 0, 1, norms)
            Xn = Xn / norms
            # Truncate or pad to 2 dims
            dim = Xn.shape[1]
            if dim >= 2:
                coords = Xn[:, :2].astype(np.float32, copy=False)
            else:
                coords = np.pad(Xn, ((0, 0), (0, 2 - dim)), mode="constant").astype(np.float32, copy=False)

    return names, coords, clusters, counts, langs, reducer


def build_publication_points_2d(
    store: ArtifactStore, max_points: int = 500, title: str | None = None
) -> tuple[list[dict], dict | None]:
    """Return UMAP-reduced 2D positions for individual publications.

    If title is provided, embed and project it as a query_point.
    """
    metadata = store.metadata
    lecturer_col = store.lecturer_col
    topic_col = store.topic_col
    lang_col = store.lang_col

    title_col: str | None = None
    for candidate in ("judul_publikasi", "title"):
        if candidate in metadata.columns:
            title_col = candidate
            break

    n = min(len(metadata), max_points)
    if n == 0:
        return [], None

    X = store.doc_embeddings[:n].astype(np.float32, copy=False)
    reducer_model: object | None = None
    if X.shape[0] == 1:
        coords = np.zeros((1, 2), dtype=np.float32)
    else:
        n_neighbors = min(15, max(2, X.shape[0] - 1))
        try:
            reducer_model = UMAP(
                n_components=2,
                n_neighbors=n_neighbors,
                min_dist=0.08,
                metric="cosine",
                random_state=42,
                transform_seed=42,
            )
            coords = reducer_model.fit_transform(X).astype(np.float32, copy=False)
        except Exception:
            # Fallback to a simple 2D projection so the endpoint never hard-crashes
            reducer_model = None
            Xn = X - X.mean(axis=0, keepdims=True)
            norms = np.linalg.norm(Xn, axis=1, keepdims=True)
            norms = np.where(norms == 0, 1, norms)
            Xn = Xn / norms
            dim = Xn.shape[1]
            if dim >= 2:
                coords = Xn[:, :2].astype(np.float32, copy=False)
            else:
                coords = np.pad(Xn, ((0, 0), (0, 2 - dim)), mode="constant").astype(np.float32, copy=False)

    topic_vals = metadata[topic_col] if (topic_col and topic_col in metadata.columns) else None
    lec_vals = metadata[lecturer_col] if (lecturer_col and lecturer_col in metadata.columns) else None
    title_vals = metadata[title_col] if title_col else None
    lang_vals = metadata[lang_col] if (lang_col and lang_col in metadata.columns) else None

    points: list[dict] = []
    for i in range(n):
        judul = str(title_vals.iloc[i]) if title_vals is not None else ""
        if judul.lower() in {"nan", "none"}:
            judul = ""

        dosen = str(lec_vals.iloc[i]) if lec_vals is not None else ""
        if dosen.lower() in {"nan", "none"}:
            dosen = ""

        cluster: int | None = None
        if topic_vals is not None:
            raw = topic_vals.iloc[i]
            if raw is not None and str(raw).strip() and str(raw).lower() not in {"nan", "none"}:
                try:
                    val = int(float(raw))
                    cluster = val if val >= 0 else None
                except (TypeError, ValueError):
                    cluster = None

        lang_str: str | None = None
        if lang_vals is not None:
            lv = str(lang_vals.iloc[i]).strip().lower()
            if lv in {"en", "id"}:
                lang_str = lv

        points.append({
            "pub_id": int(i),
            "judul": judul[:120],
            "dosen": dosen,
            "cluster": cluster,
            "lang": lang_str,
            "x": float(coords[i, 0]),
            "y": float(coords[i, 1]),
        })

    q = (title or "").strip()
    if not q or reducer_model is None:
        return points, None

    lang = detect_language(q)
    q_emb = embed_query(store=store, text=q, lang=lang)
    topic_id, topic_sim = find_closest_topic_with_similarity(store, q_emb, lang=lang)
    q_xy = reducer_model.transform(q_emb.reshape(1, -1)).astype(np.float32, copy=False)

    query_point = {
        "title": q,
        "x": float(q_xy[0, 0]),
        "y": float(q_xy[0, 1]),
        "closest_topic_id": (None if topic_id is None else int(topic_id)),
        "topic_similarity": (None if topic_sim is None else float(topic_sim)),
        "detected_lang": lang,
    }

    return points, query_point


def build_lecturer_points_2d(store: ArtifactStore) -> list[dict]:
    names, coords, clusters, counts, langs, _ = _build_lecturer_viz_data(store)
    return [
        {
            "nama_dosen": name,
            "x": float(coords[i, 0]),
            "y": float(coords[i, 1]),
            "cluster": (None if clusters[i] is None else int(clusters[i])),
            "n_pubs": int(counts[i]),
            "lang": langs[i],
        }
        for i, name in enumerate(names)
    ]


def build_lecturer_points_2d_with_query(
    store: ArtifactStore, title: str | None
) -> tuple[list[dict], dict | None]:
    names, coords, clusters, counts, langs, reducer = _build_lecturer_viz_data(store)

    if not names:
        return [], None

    points = [
        {
            "nama_dosen": name,
            "x": float(coords[i, 0]),
            "y": float(coords[i, 1]),
            "cluster": (None if clusters[i] is None else int(clusters[i])),
            "n_pubs": int(counts[i]),
            "lang": langs[i],
        }
        for i, name in enumerate(names)
    ]

    q = (title or "").strip()
    if not q:
        return points, None

    lang = detect_language(q)
    q_emb = embed_query(store=store, text=q, lang=lang)
    topic_id, topic_sim = find_closest_topic_with_similarity(store, q_emb, lang=lang)

    if reducer is None:
        q_xy = np.zeros((1, 2), dtype=np.float32)
    else:
        q_xy = reducer.transform(q_emb.reshape(1, -1)).astype(np.float32, copy=False)

    query_point = {
        "title": q,
        "x": float(q_xy[0, 0]),
        "y": float(q_xy[0, 1]),
        "closest_topic_id": (None if topic_id is None else int(topic_id)),
        "topic_similarity": (None if topic_sim is None else float(topic_sim)),
        "detected_lang": lang,
    }

    return points, query_point