"""Advanced retrieval strategies for HIEROMEM."""

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import List, Mapping, Optional, Sequence

import numpy as np

from hieromem_core.symbolic import GraphReasoner, aggregate_symbolic_scores

if False:  # pragma: no cover - type checking hint
    from hieromem_core.memory import Memory
    from hieromem_core.episode import Episode
    import networkx as nx


@dataclass(slots=True)
class DualRetrievalResult:
    episodes: List["Episode"]
    weights: np.ndarray
    dense_scores: Mapping[str, float]
    symbolic_scores: Mapping[str, float]
    combined_scores: Mapping[str, float]
    iterations: int


@dataclass(slots=True)
class DualRetrievalConfig:
    alpha: float = 1.0
    beta: float = 50.0
    temperature: float = 1.0
    max_iter: int = 6
    tol: float = 1e-4
    k: int = 3


@dataclass(slots=True)
class DualRetrievalInput:
    """Container describing a single query for batch dual retrieval."""

    vector: np.ndarray
    query_graph: "nx.DiGraph" | None = None
    namespace: Optional[str] = None
    k: Optional[int] = None
    alpha: Optional[float] = None
    beta: Optional[float] = None
    temperature: Optional[float] = None
    max_iterations: Optional[int] = None
    tolerance: Optional[float] = None


class DualRetriever:
    """Combine dense resonance retrieval with symbolic alignment."""

    def __init__(
            self,
            memory: "Memory",
            reasoner: GraphReasoner | None = None) -> None:
        self.memory = memory
        self.reasoner = reasoner or GraphReasoner()

    # ------------------------------------------------------------------ #
    def retrieve(
        self,
        query_vector: np.ndarray,
        *,
        query_graph: "nx.DiGraph" | None = None,
        config: DualRetrievalConfig | None = None,
        namespace: str | None = None,
    ) -> DualRetrievalResult:
        config = config or DualRetrievalConfig(
            k=min(len(self.memory.episodes), 3))

        candidate_cap = config.k * 2 if config.k > 0 else 1
        candidate_cap = max(candidate_cap, 10)
        candidate_cap = min(candidate_cap, 100, len(self.memory.episodes))
        candidate_cap = max(candidate_cap, 1)
        dense_candidates = self.memory.resonant_retrieve(
            query_vector,
            k=candidate_cap,
            namespace=namespace,
        )
        if config.k <= 0 and not dense_candidates:
            return DualRetrievalResult(
                episodes=[],
                weights=np.zeros(0),
                dense_scores={},
                symbolic_scores={},
                combined_scores={},
                iterations=0,
            )

        # ----------------------------------------------------------- #
        # HYBRID RECALL: Merge Dense + Symbolic Candidates
        # ----------------------------------------------------------- #
        candidates_map = {ep.id: ep for ep in dense_candidates}

        # If we have a query graph, fetch episodes containing those entities directly
        if query_graph and hasattr(self.memory, "entity_index"):
            query_entities = []
            for node in query_graph.nodes():
                query_entities.append(str(node))

            symbolic_ids = self.memory.entity_index.lookup_many(query_entities)
            for eid in symbolic_ids:
                if eid not in candidates_map:
                    ep = self.memory._get_episode_by_id(eid)
                    if ep:
                        candidates_map[eid] = ep

            # Graph expansion from query entities (multi-hop retrieval)
            # This finds episodes connected via graph relationships, not just those mentioning the entity
            if hasattr(self.memory, "graph_store") and self.memory.graph_store is not None:
                for query_entity in query_entities:
                    try:
                        relations = self.memory.graph_store.query_relations(
                            query_entity, depth=2, namespace=namespace or self.memory.namespace
                        )
                        for rel in relations[:20]:  # Limit to avoid explosion
                            for key in ['src', 'dst']:
                                connected_entity = rel.get(key)
                                if connected_entity:
                                    connected_eps = self.memory.entity_index.lookup_many([str(connected_entity)])
                                    for eid in connected_eps:
                                        if eid not in candidates_map:
                                            ep = self.memory._get_episode_by_id(eid)
                                            if ep:
                                                candidates_map[eid] = ep
                    except Exception as e:  # nosec B110 - intentional: skip if entity not in graph
                        if self.memory.debug:
                            print(f"[DualRetriever] Query entity expansion error for '{query_entity}': {e}")

        # NEW: Use GraphStore for multi-hop expansion if available
        # This enables finding episodes connected through graph relationships
        # NOTE: Deduplication is handled by candidates_map being a dict (keyed by episode ID)
        if query_graph and hasattr(self.memory, "graph_store") and self.memory.graph_store is not None:
            # For each dense candidate, expand via graph to find related episodes
            expansion_limit = min(config.k * 3, 20)  # Limit expansion to avoid explosion
            expanded_entity_ids: set[str] = set()  # De-duplicate entity lookups
            graph_expansion_count = 0  # Track successful expansions

            for seed_ep in list(candidates_map.values())[:5]:  # Only expand from top-5 dense candidates
                try:
                    # Improved seeding: if seed episode has a graph, use its entity nodes
                    # This works even if episode hasn't been indexed in graph store yet
                    if hasattr(seed_ep, 'G') and seed_ep.G is not None and seed_ep.G.number_of_nodes() > 0:
                        seed_entities = list(seed_ep.G.nodes())[:10]  # Limit to top 10 entities
                        # Use query_relations for each entity instead of search_related with episode_id
                        for entity in seed_entities:
                            if str(entity) in expanded_entity_ids:
                                continue
                            expanded_entity_ids.add(str(entity))

                            relations = self.memory.graph_store.query_relations(
                                str(entity),
                                depth=2,
                                namespace=namespace or self.memory.namespace
                            )
                            # Extract connected entities from relations
                            for rel in relations[:expansion_limit]:
                                for key in ['src', 'dst']:
                                    connected_entity = rel.get(key)
                                    if connected_entity and str(connected_entity) not in expanded_entity_ids:
                                        expanded_entity_ids.add(str(connected_entity))
                                        # Look up episodes containing this entity
                                        entity_eps = self.memory.entity_index.lookup_many([str(connected_entity)])
                                        for eid in entity_eps:
                                            # Deduplication: candidates_map ensures each episode is added once
                                            if eid not in candidates_map:
                                                ep = self.memory._get_episode_by_id(eid)
                                                if ep:
                                                    candidates_map[eid] = ep
                                                    graph_expansion_count += 1
                    else:
                        # Fallback: use search_related with episode ID
                        # NOTE: search_related internally finds entities from the episode's edges
                        # in the graph store, so episode ID works as a valid seed
                        related = self.memory.graph_store.search_related(
                            seed_ep.id,
                            depth=2,
                            limit=expansion_limit,
                            namespace=namespace or self.memory.namespace
                        )
                        for item in related:
                            for key in ['source', 'target']:
                                entity = item.get(key)
                                if entity and str(entity) not in expanded_entity_ids:
                                    expanded_entity_ids.add(str(entity))
                                    entity_eps = self.memory.entity_index.lookup_many([str(entity)])
                                    for eid in entity_eps:
                                        # Deduplication: candidates_map ensures each episode is added once
                                        if eid not in candidates_map:
                                            ep = self.memory._get_episode_by_id(eid)
                                            if ep:
                                                candidates_map[eid] = ep
                                                graph_expansion_count += 1
                except Exception as e:
                    # Log for debugging if memory.debug is set
                    if self.memory.debug:
                        print(f"[DualRetriever] Graph expansion error for seed '{seed_ep.id}': {e}")

        # Final candidate pool
        all_candidates = list(candidates_map.values())
        if not all_candidates:
            return DualRetrievalResult(
                episodes=[],
                weights=np.zeros(0),
                dense_scores={},
                symbolic_scores={},
                combined_scores={},
                iterations=0,
            )

        # ----------------------------------------------------------- #
        # SCORING
        # ----------------------------------------------------------- #
        dense_ids = [ep.id for ep in all_candidates]
        dense_vectors = np.stack([ep.e for ep in all_candidates])
        candidate_graphs = {
            ep.id: getattr(
                ep,
                "G",
                None) for ep in all_candidates}

        with ThreadPoolExecutor(max_workers=2) as executor:
            dense_future = executor.submit(
                _dense_softmax, dense_vectors, query_vector, max(
                    config.temperature, 1e-6))
            symbolic_future = executor.submit(
                self.reasoner.batch_score, query_graph, candidate_graphs
            )
            dense_weights = dense_future.result()
            symbolic_scores = symbolic_future.result()
        symbolic_weights = aggregate_symbolic_scores(
            symbolic_scores, temperature=config.temperature)

        combined = dense_weights.copy()
        prev = None
        iterations = 0
        for iteration in range(1, config.max_iter + 1):
            sym_vector = np.array([symbolic_weights.get(ep_id, 0.0)
                                  for ep_id in dense_ids], dtype=float)
            sym_vector = _normalise(sym_vector + 1e-9)
            combined = config.alpha * dense_weights + config.beta * sym_vector
            combined = _normalise(combined)
            if prev is not None and np.linalg.norm(
                    combined - prev) < config.tol:
                iterations = iteration
                break
            prev = combined.copy()
            iterations = iteration

        cache_key = (
            "dual",
            namespace or self.memory.namespace,
            tuple(dense_ids),
            hash(np.asarray(query_vector, dtype=float).tobytes()),
        )
        combined = self.memory.stability_monitor.stabilise_weights(
            combined, key=cache_key)
        top_indices = np.argsort(-combined)[: config.k]
        final_eps = [all_candidates[i] for i in top_indices]
        final_weights = combined[top_indices]
        final_symbolic = {
            dense_ids[i]: symbolic_weights.get(
                dense_ids[i],
                0.0) for i in top_indices}

        final_dense = {
            dense_ids[i]: float(
                dense_weights[i]) for i in top_indices}

        final_combined = {
            dense_ids[i]: float(
                combined[i]) for i in top_indices}

        return DualRetrievalResult(
            episodes=final_eps,
            weights=final_weights.astype(float),
            dense_scores=final_dense,
            symbolic_scores=final_symbolic,
            combined_scores=final_combined,
            iterations=iterations,
        )


def _dense_softmax(
        dense_vectors: np.ndarray,
        query_vector: np.ndarray,
        temperature: float) -> np.ndarray:
    sims = dense_vectors @ query_vector / (
        np.linalg.norm(dense_vectors, axis=1) * (np.linalg.norm(query_vector) + 1e-9)
    )
    temp = max(temperature, 1e-6)
    return _normalise(np.exp(sims / temp))


def _normalise(values: Sequence[float] | np.ndarray) -> np.ndarray:
    arr = np.asarray(values, dtype=float)
    if arr.size == 0:
        return arr
    total = float(np.sum(arr))
    if total == 0:
        return np.ones_like(arr) / len(arr)
    return arr / total
