"""
Learned Retrieval Model (Research-Grade).

This module implements a pairwise neural ranking strategy.
It learns to rank candidates by optimizing a neural network (RankerNet)
that scores feature vectors (dense score, symbolic score, trust, etc.).

Objective: Pairwise Margin Ranking Loss
minimize sum max(0, -y * (score_pos - score_neg) + margin)
where y=1 implies pos should be ranked higher than neg.
"""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from dataclasses import dataclass
from typing import List, Any, Optional

from hieromem_core.retrieval.dual import DualRetriever, DualRetrievalResult, DualRetrievalConfig
from hieromem_core.evaluation.framework import EvaluationExample


@dataclass
class LearnedConfig(DualRetrievalConfig):
    learning_rate: float = 0.001
    epochs: int = 50
    margin: float = 1.0
    batch_size: int = 32
    hidden_dim: int = 64


class RankerNet(nn.Module):
    """
    Simple MLP for scoring candidate feature vectors.
    Input: [Dense, Symbolic, Trust, Entities, Bias] -> 5 dim
    Output: scalars score
    """
    def __init__(self, input_dim=5, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1)
        )

    def forward(self, x):
        return self.net(x)


class LearnedRetriever(DualRetriever):
    """
    Retriever that learns to rank candidates using a neural network.
    """
    def __init__(self, memory, reasoner=None, model_path: Optional[str] = None):
        super().__init__(memory, reasoner)
        self.feature_names = ["Dense", "Symbolic", "Trust", "Entities", "Bias"]
        # Input dim = 5 roughly corresponds to the features we extract
        self.model = RankerNet(input_dim=len(self.feature_names))
        self.optimizer = None  # Lazy init in train

        self.heuristic_mode = True
        self.heuristic_weights = np.array([1.0, 0.5, 0.2, 0.1, 0.0])  # [Dense, Symbolic, Trust, Entities, Bias]

        if model_path:
            self.load(model_path)
            self.heuristic_mode = False

    def extract_features(self, query_vector: np.ndarray, candidates: List[Any],
                         dense_scores: dict, symbolic_scores: dict) -> np.ndarray:
        """
        Extract feature matrix for candidates.
        Shape: (num_candidates, num_features)
        """
        features = []
        for ep in candidates:
            # 1. Dense Score (Cosine)
            d_score = dense_scores.get(ep.id, 0.0)

            # 2. Symbolic Score
            s_score = symbolic_scores.get(ep.id, 0.0)

            # 3. Graph Trust (Normalized)
            raw_trust = float(ep.meta.get("trust", 1.0) if ep.meta else 1.0)
            trust = np.log1p(raw_trust)

            # 4. Entity Overlap (Proxy: Node count)
            node_count = ep.G.number_of_nodes() if hasattr(ep, 'G') and ep.G else 0
            ent_score = np.log1p(node_count)

            # 5. Bias
            bias = 1.0

            features.append([d_score, s_score, trust, ent_score, bias])

        return np.array(features, dtype=np.float32)

    def predict_scores(self, feature_matrix: np.ndarray) -> np.ndarray:
        """
        Predict relevance scores using the neural net or heuristics.
        """
        if feature_matrix.size == 0:
            return np.array([])

        if self.heuristic_mode:
            return feature_matrix @ self.heuristic_weights

        self.model.eval()
        with torch.no_grad():
            inp = torch.tensor(feature_matrix, dtype=torch.float32)
            scores = self.model(inp).squeeze(-1).numpy()
        return scores

    def train(self, examples: List[EvaluationExample], config: LearnedConfig = None):
        """
        Train the ranker using Pairwise Margin Ranking Loss.
        We form pairs (positive, negative) from the examples.
        """
        if not examples:
            return

        cfg = config or LearnedConfig()
        self.model.train()
        self.optimizer = optim.Adam(self.model.parameters(), lr=cfg.learning_rate)
        criterion = nn.MarginRankingLoss(margin=cfg.margin)

        print(f"Training LearnedRetriever (Neural Ranker) on {len(examples)} examples...")

        # Pre-process data: generate candidates & features once (or per epoch if dynamic)
        # For efficiency/simplicity here, we do it once.
        training_pairs = []  # List of (feat_pos, feat_neg)

        for ex in examples:
            # 1. Get candidates
            q_vec = (
                np.array(ex.query_vector) if ex.query_vector is not None
                else np.random.rand(self.memory.encoder.dimension)
            )
            candidates = self.memory.resonant_retrieve(q_vec, k=20)

            # Ensure we have the positives
            cand_map = {ep.id: ep for ep in candidates}
            positives = []

            # Check existing candidates
            for ep in candidates:
                if ep.id in ex.relevant_ids:
                    positives.append(ep)

            # If positives missing from retrieval, add them manually
            for true_id in ex.relevant_ids:
                if true_id not in cand_map:
                    ep = self.memory._get_episode_by_id(true_id)
                    if ep:
                        positives.append(ep)
                        candidates.append(ep)
                        cand_map[ep.id] = ep

            if not positives:
                continue

            # Compute features for all candidates
            # We need d_scores and s_scores
            c_vectors = np.stack([ep.e for ep in candidates])
            # safe cosine
            norm_q = np.linalg.norm(q_vec) + 1e-9
            norm_c = np.linalg.norm(c_vectors, axis=1) + 1e-9
            d_scores = (c_vectors @ q_vec) / (norm_c * norm_q)
            d_map = {ep.id: s for ep, s in zip(candidates, d_scores)}
            s_map = {ep.id: 0.0 for ep in candidates}  # Placeholder for symbolic

            feats = self.extract_features(q_vec, candidates, d_map, s_map)
            feat_map = {ep.id: f for ep, f in zip(candidates, feats)}

            # Create pairs: (Positive, Negative)
            # A negative is any candidate NOT in relevant_ids
            negatives = [ep for ep in candidates if ep.id not in ex.relevant_ids]

            if not negatives:
                continue

            for pos_ep in positives:
                for neg_ep in negatives:
                    training_pairs.append((feat_map[pos_ep.id], feat_map[neg_ep.id]))

        if not training_pairs:
            print("No valid training pairs generated.")
            return

        # Training Loop
        X_pos = torch.tensor(np.array([p[0] for p in training_pairs]), dtype=torch.float32)
        X_neg = torch.tensor(np.array([p[1] for p in training_pairs]), dtype=torch.float32)
        # Target is 1s because X_pos should rank higher than X_neg
        targets = torch.ones(X_pos.shape[0])

        dataset_size = X_pos.shape[0]
        batch_size = cfg.batch_size

        for epoch in range(cfg.epochs):
            perm = torch.randperm(dataset_size)
            epoch_loss = 0.0

            for i in range(0, dataset_size, batch_size):
                idx = perm[i:i + batch_size]
                batch_pos = X_pos[idx]
                batch_neg = X_neg[idx]
                batch_y = targets[idx]

                self.optimizer.zero_grad()

                score_pos = self.model(batch_pos).squeeze(-1)
                score_neg = self.model(batch_neg).squeeze(-1)

                loss = criterion(score_pos, score_neg, batch_y)
                loss.backward()
                self.optimizer.step()

                epoch_loss += loss.item() * len(idx)

            if epoch % 10 == 0:
                print(f"Epoch {epoch}: Avg Pairwise Loss {epoch_loss/dataset_size:.4f}")

        print("Training Complete.")
        self.heuristic_mode = False

    def retrieve(self, query_vector, **kwargs) -> DualRetrievalResult:
        """
        Rank candidates using the trained neural model.
        """
        # 1. Base Retrieval (Dense + Symbolic candidates)
        base_result: DualRetrievalResult = super().retrieve(query_vector, **kwargs)

        if not base_result.episodes:
            return base_result

        candidates = base_result.episodes
        # 2. Extract Features
        features = self.extract_features(query_vector, candidates,
                                         base_result.dense_scores,
                                         base_result.symbolic_scores)

        # 3. Predict Neural Scores
        new_scores = self.predict_scores(features)

        # 4. Re-rank
        sorted_indices = np.argsort(-new_scores)
        sorted_eps = [candidates[i] for i in sorted_indices]
        sorted_weights = new_scores[sorted_indices]

        # Update result
        base_result.episodes = sorted_eps
        base_result.weights = sorted_weights

        return base_result

    def save(self, path: str):
        torch.save(self.model.state_dict(), path)

    def load(self, path: str):
        self.model.load_state_dict(torch.load(path, map_location='cpu'))  # nosec B614
        self.heuristic_mode = False
