"""Text-driven retrieval helpers for NeuroVLM.

These utilities accept natural language queries and return the most
relevant publications or NeuroWiki entries along with formatted context
blocks suitable for prompting an LLM.
"""

from __future__ import annotations

import re
from typing import List, Tuple, Literal, Union

import numpy as np
import pandas as pd
import torch

from neurovlm.retrieval_resources import (
    _load_pubmed_dataframe,
    _load_latent_text,
    _load_latent_wiki,
    _load_latent_cogatlas,
    _load_latent_cogatlas_disorder,
    _load_latent_cogatlas_task,
    _load_neuro_wiki,
    _load_cogatlas_dataset,
    _load_cogatlas_disorder_dataset,
    _load_cogatlas_task_dataset,
    _load_specter,
    _proj_head_text_infonce
)

__all__ = ["search_papers_from_text", "search_wiki_from_text", "search_cogatlas_from_text", "generate_llm_response_from_text"]


@torch.no_grad()
def search_papers_from_text(
    query: str,
    top_k: int = 5,
    show_titles: bool = False,
) -> Tuple[str, List[str]]:
    """Return a context block of top papers and their titles for a text query."""
    if not isinstance(query, str):
        raise TypeError("query must be a string for text-based retrieval")

    df = _load_pubmed_dataframe()
    latent_text, latent_pmids = _load_latent_text()

    specter = _load_specter()
    proj_head = _proj_head_text_infonce()

    encoded_query = specter(query)[0].detach().to("cpu")
    proj_query = proj_head(encoded_query)
    proj_query = proj_query / proj_query.norm()

    proj_text = proj_head(latent_text)
    proj_text = proj_text / proj_text.norm(dim=1)[:, None]
    cos_sim = proj_text @ proj_query

    inds = torch.argsort(cos_sim, descending=True)
    inds_top = inds[:top_k].tolist()
    pmids_top = [latent_pmids[i] for i in inds_top]

    if "pmid" in df.columns:
        pmid_lookup = df.drop_duplicates("pmid").set_index("pmid", drop=False)
        selected_rows = []
        for pmid in pmids_top:
            try:
                row = pmid_lookup.loc[pmid]
                if isinstance(row, pd.DataFrame):
                    row = row.iloc[0]
                selected_rows.append(row)
            except KeyError:
                continue
        rows = pd.DataFrame(selected_rows) if selected_rows else df.iloc[inds_top]
    else:
        rows = df.iloc[inds_top]

    pieces = []
    title_col = "name" if "name" in rows.columns else ("title" if "title" in rows.columns else None)
    titles: List[str]
    if title_col is not None:
        titles = rows[title_col].astype(str).tolist()
    else:
        titles = ["Untitled"] * len(rows)

    for idx, (_, row) in enumerate(rows.iterrows(), start=1):
        title = (
            str(row["name"]) if "name" in rows.columns else str(row.get("title", "Untitled"))
        )
        desc = str(row.get("description", "")).replace("\n", " ")
        desc = re.sub(r"\s+", " ", desc).strip()
        pieces.append(f"[{idx}] {title}\n{desc}\n")

    if show_titles:
        print("Top matches:")
        for idx, title in enumerate(titles, start=1):
            print(f"{idx}. {title}")

    papers_context = "\n".join(pieces)
    return papers_context, titles


@torch.no_grad()
def search_wiki_from_text(
    query: str,
    top_k: int = 2,
    show_titles: bool = False,
) -> Tuple[str, List[str]]:
    """Return a context block of top NeuroWiki entries for a text query."""
    if not isinstance(query, str):
        raise TypeError("query must be a string for text-based retrieval")

    df = _load_neuro_wiki()
    latent_wiki, latent_ids = _load_latent_wiki()
    if not (df['id'] == latent_ids).all():
        raise ValueError("Mismatch between DataFrame 'id' column and latent_ids: ensure they are aligned.")
    specter = _load_specter()
    proj_head = _proj_head_text_infonce()

    encoded_query = specter(query)[0].detach().to("cpu")
    encoded_query = encoded_query / encoded_query.norm()
    proj_query = proj_head(encoded_query)
    proj_query = proj_query / proj_query.norm()

    wiki_embed = latent_wiki / latent_wiki.norm(dim=1)[:, None]
    proj_wiki = proj_head(wiki_embed)
    proj_wiki = proj_wiki / proj_wiki.norm(dim=1)[:, None]
    cos_sim = proj_wiki @ proj_query

    inds = torch.argsort(cos_sim, descending=True)
    inds_top = inds[:top_k].tolist()
    ids_top = [latent_ids[i] for i in inds_top]
    missing_columns = [col for col in ("id", "title", "summary") if col not in df.columns]
    if missing_columns:
        raise ValueError(
            f"NeuroWiki DataFrame is missing required columns: {', '.join(missing_columns)}"
        )

    id_lookup = df.drop_duplicates("id").set_index("id", drop=False)
    selected_rows = []
    for entry_id in ids_top:
        try:
            row = id_lookup.loc[entry_id]
            if isinstance(row, pd.DataFrame):
                row = row.iloc[0]
            selected_rows.append(row)
        except KeyError:
            continue

    rows = pd.DataFrame(selected_rows) if selected_rows else df.iloc[inds_top]

    pieces = []
    titles: List[str] = []
    for idx, (_, row) in enumerate(rows.iterrows(), start=1):
        title = str(row["title"]) if pd.notna(row["title"]) else "Untitled"
        titles.append(title)
        summary = str(row["summary"]) if pd.notna(row["summary"]) else ""
        summary = re.sub(r"\s+", " ", summary).strip()
        pieces.append(f"[{idx}] {title}\n{summary}\n")

    if show_titles:
        print("Top matches:")
        for idx, title in enumerate(titles, start=1):
            print(f"{idx}. {title}")

    wiki_context = "\n".join(pieces)
    return wiki_context, titles


@torch.no_grad()
def search_cogatlas_from_text(
    query: Union[str, torch.Tensor],
    top_k: int = 5,
    show_titles: bool = False,
    category: Literal["cogatlas", "cogatlas_task", "cogatlas_disorder"] = "cogatlas",
) -> Tuple[str, List[str], np.ndarray]:
    """Return a context block of top Cognitive Atlas terms for a text query or embedding.

    Parameters
    ----------
    query : str or torch.Tensor
        Either a text string or a pre-computed tensor embedding (768-dim from SPECTER).
    top_k : int
        Number of top results to return.
    show_titles : bool
        Whether to print the top matches.
    category : str
        Which category to search: "cogatlas" (concepts), "cogatlas_task", or "cogatlas_disorder".

    Returns
    -------
    Tuple[str, List[str], np.ndarray]
        Context string, list of terms, and cosine similarity scores.
    """
    if not isinstance(query, (str, torch.Tensor)):
        raise TypeError("query must be a string or torch.Tensor for text-based retrieval")

    # Load appropriate dataset and embeddings based on category
    if category == "cogatlas":
        df = _load_cogatlas_dataset()
        latent_cogatlas, latent_terms = _load_latent_cogatlas()
    elif category == "cogatlas_task":
        df = _load_cogatlas_task_dataset(filtered=True)
        latent_cogatlas, latent_terms = _load_latent_cogatlas_task()
    elif category == "cogatlas_disorder":
        df = _load_cogatlas_disorder_dataset()
        latent_cogatlas, latent_terms = _load_latent_cogatlas_disorder()
    else:
        raise ValueError(f"Unknown category: {category}")

    proj_head = _proj_head_text_infonce()

    # Handle string vs tensor query
    if isinstance(query, str):
        specter = _load_specter()
        encoded_query = specter(query)[0].detach().to("cpu")
        encoded_query = encoded_query / encoded_query.norm()
        proj_query = proj_head(encoded_query)
        proj_query = proj_query / proj_query.norm()
    else:
        # Query is already a tensor embedding
        encoded_query = query / query.norm()
        proj_query = proj_head(encoded_query)
        proj_query = proj_query / proj_query.norm()

    cogatlas_embed = latent_cogatlas / latent_cogatlas.norm(dim=1)[:, None]
    proj_cogatlas = proj_head(cogatlas_embed)
    proj_cogatlas = proj_cogatlas / proj_cogatlas.norm(dim=1)[:, None]
    cos_sim = proj_cogatlas @ proj_query

    inds = torch.argsort(cos_sim, descending=True)
    inds_top = inds[:top_k].tolist()
    cos_sim_top = cos_sim[inds_top].detach().cpu().numpy()
    terms_top = [latent_terms[i] for i in inds_top]

    missing_columns = [col for col in ("term", "definition") if col not in df.columns]
    if missing_columns:
        raise ValueError(
            f"CogAtlas DataFrame is missing required columns: {', '.join(missing_columns)}"
        )

    # Ensure term column is lowercase for matching
    df["term"] = df["term"].str.lower()
    term_lookup = df.drop_duplicates("term").set_index("term", drop=False)
    selected_rows = []
    for term in terms_top:
        try:
            row = term_lookup.loc[term]
            if isinstance(row, pd.DataFrame):
                row = row.iloc[0]
            selected_rows.append(row)
        except KeyError:
            print(f"Term '{term}' not found in CogAtlas DataFrame.")
            continue

    rows = pd.DataFrame(selected_rows) if selected_rows else df.iloc[inds_top]

    pieces = []
    terms: List[str] = []
    for idx, (_, row) in enumerate(rows.iterrows(), start=1):
        term = str(row["term"]) if pd.notna(row["term"]) else "Untitled"
        terms.append(term)
        definition = str(row["definition"]) if pd.notna(row["definition"]) else ""
        definition = re.sub(r"\s+", " ", definition).strip()
        pieces.append(f"[{idx}] {term}\n{definition}\n")

    if show_titles:
        print("Top matches:")
        for idx, term in enumerate(terms, start=1):
            print(f"{idx}. {term}")

    cogatlas_context = "\n".join(pieces)
    return cogatlas_context, terms, cos_sim_top


@torch.no_grad()
def generate_llm_response_from_text(
    query: str,
    top_k_papers: int = 5,
    top_k_cogatlas_concepts: int = 5,
    top_k_cogatlas_disorders: int = 5,
    top_k_cogatlas_tasks: int = 5,
    backend: Literal["ollama", "huggingface"] = "ollama",
    model_name: str | None = None,
):
    """
    Generate an LLM response for a given natural language query.

    For text input, focuses on papers and CogAtlas terms (concepts, disorders, tasks).
    NeuroWiki entries are NOT included for text-based queries.

    Parameters
    ----------
    query : str
        A natural language query to be encoded.
    top_k_papers : int
        Number of top papers to retrieve.
    top_k_cogatlas_concepts : int
        Number of top CogAtlas concept terms to retrieve.
    top_k_cogatlas_disorders : int
        Number of top CogAtlas disorder terms to retrieve.
    top_k_cogatlas_tasks : int
        Number of top CogAtlas task terms to retrieve.
    backend : {"ollama", "huggingface"}, optional
        Which LLM backend to use. Default: "ollama" (faster, requires Ollama installed).
    model_name : str, optional
        Model name. If None, uses backend defaults.

    Returns
    -------
    str
        The generated LLM response.

    Examples
    --------
    >>> # Use Ollama (default, fast)
    >>> output = generate_llm_response_from_text("default mode network")

    >>> # Use HuggingFace with small model
    >>> output = generate_llm_response_from_text(
    ...     "default mode network",
    ...     backend="huggingface",
    ...     model_name="Qwen/Qwen2.5-0.5B-Instruct"
    ... )
    """
    from neurovlm.llm_summary import generate_response

    # Retrieve top k papers
    papers_context, paper_titles = search_papers_from_text(
        query, top_k=top_k_papers, show_titles=False
    )

    # Retrieve top k for each CogAtlas category
    cogatlas_concepts_context, cogatlas_concepts_terms, _ = search_cogatlas_from_text(
        query, top_k=top_k_cogatlas_concepts, show_titles=False, category="cogatlas"
    )

    cogatlas_disorders_context, cogatlas_disorders_terms, _ = search_cogatlas_from_text(
        query, top_k=top_k_cogatlas_disorders, show_titles=False, category="cogatlas_disorder"
    )

    cogatlas_tasks_context, cogatlas_tasks_terms, _ = search_cogatlas_from_text(
        query, top_k=top_k_cogatlas_tasks, show_titles=False, category="cogatlas_task"
    )

    # Combine all cogatlas contexts
    cogatlas_combined_context = "\n".join([
        "Concepts:\n" + cogatlas_concepts_context,
        "Disorders:\n" + cogatlas_disorders_context,
        "Tasks:\n" + cogatlas_tasks_context,
    ])

    return generate_response(
        query=query,
        papers_context=papers_context,
        wiki_context=None,  # No wiki for text input
        cogatlas_context=cogatlas_combined_context,
        user_prompt=query,
        backend=backend,
        model_name=model_name,
    )
