#!/usr/bin/env python3
"""
Simple Deep Search Agent
========================

This module implements a minimal deep learning based search agent.
It uses a TF‑IDF vectorizer to transform documents and queries into
feature vectors and a single linear layer (logistic regression)
to predict relevance scores.  The agent can be trained on a small
synthetic dataset and used to retrieve the top‑k most relevant
documents for a given query.

Author: Artur Kuzakhmetov
Date: 30.06.2026
"""

import sys
from typing import List, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split


class SearchAgent(nn.Module):
    """
    A simple search agent based on a linear classifier.
    """

    def __init__(self, documents: List[str], device: torch.device = None):
        """
        Parameters
        ----------
        documents : List[str]
            List of document texts.
        device : torch.device, optional
            Device to run the model on. Defaults to CUDA if available.
        """
        super().__init__()
        self.documents = documents
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Fit TF‑IDF vectorizer on documents
        self.vectorizer = TfidfVectorizer()
        self.doc_vectors = self.vectorizer.fit_transform(self.documents).toarray()
        self.doc_vectors = torch.tensor(self.doc_vectors, dtype=torch.float32, device=self.device)

        # Linear layer: input_dim -> 1 (relevance score)
        self.linear = nn.Linear(self.doc_vectors.shape[1], 1).to(self.device)

    def forward(self, query_vec: torch.Tensor) -> torch.Tensor:
        """
        Forward pass: compute relevance scores for all documents given a query vector.

        Parameters
        ----------
        query_vec : torch.Tensor
            Tensor of shape (1, feature_dim).

        Returns
        -------
        torch.Tensor
            Tensor of shape (num_documents,) with relevance scores.
        """
        # Compute dot product between query and each document vector
        scores = torch.matmul(self.doc_vectors, query_vec.t()).squeeze(1)
        return scores

    def train_agent(
        self,
        queries: List[str],
        labels: List[List[int]],
        epochs: int = 10,
        lr: float = 0.01,
        batch_size: int = 4,
        verbose: bool = True,
    ) -> None:
        """
        Train the agent on query‑document relevance pairs.

        Parameters
        ----------
        queries : List[str]
            List of query texts.
        labels : List[List[int]]
            List of relevance labels for each query. Each inner list contains
            indices of relevant documents (0‑based).
        epochs : int, default 10
            Number of training epochs.
        lr : float, default 0.01
            Learning rate.
        batch_size : int, default 4
            Batch size.
        verbose : bool, default True
            Whether to print training progress.
        """
        # Vectorize queries
        query_vectors = self.vectorizer.transform(queries).toarray()
        query_vectors = torch.tensor(query_vectors, dtype=torch.float32, device=self.device)

        # Prepare training data
        # For each query, create a target vector of relevance scores (1 for relevant, 0 otherwise)
        targets = []
        for rel_indices in labels:
            target = torch.zeros(self.doc_vectors.shape[0], device=self.device)
            target[rel_indices] = 1.0
            targets.append(target)
        targets = torch.stack(targets)  # shape: (num_queries, num_documents)

        # Loss and optimizer
        criterion = nn.BCEWithLogitsLoss()
        optimizer = optim.Adam(self.parameters(), lr=lr)

        dataset = torch.utils.data.TensorDataset(query_vectors, targets)
        loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

        self.train()
        for epoch in range(1, epochs + 1):
            epoch_loss = 0.0
            for batch_q, batch_t in loader:
                optimizer.zero_grad()
                outputs = self.forward(batch_q)  # shape: (batch_size, num_documents)
                loss = criterion(outputs, batch_t)
                loss.backward()
                optimizer.step()
                epoch_loss += loss.item() * batch_q.size(0)
            epoch_loss /= len(dataset)
            if verbose:
                print(f"Epoch {epoch}/{epochs} - Loss: {epoch_loss:.4f}")

    def evaluate(
        self,
        queries: List[str],
        labels: List[List[int]],
        threshold: float = 0.5,
    ) -> Tuple[float, float]:
        """
        Evaluate the agent on a test set.

        Parameters
        ----------
        queries : List[str]
            List of query texts.
        labels : List[List[int]]
            List of relevance labels for each query.
        threshold : float, default 0.5
            Threshold to convert scores to binary predictions.

        Returns
        -------
        Tuple[float, float]
            (precision, recall)
        """
        self.eval()
        with torch.no_grad():
            query_vectors = self.vectorizer.transform(queries).toarray()
            query_vectors = torch.tensor(query_vectors, dtype=torch.float32, device=self.device)
            outputs = self.forward(query_vectors)  # shape: (num_queries, num_documents)
            preds = (outputs > threshold).int()

        # Compute precision and recall
        total_relevant = 0
        total_predicted = 0
        total_correct = 0

        for i, rel_indices in enumerate(labels):
            pred_indices = preds[i].nonzero(as_tuple=True)[0].cpu().numpy().tolist()
            total_relevant += len(rel_indices)
            total_predicted += len(pred_indices)
            total_correct += len(set(pred_indices) & set(rel_indices))

        precision = total_correct / total_predicted if total_predicted > 0 else 0.0
        recall = total_correct / total_relevant if total_relevant > 0 else 0.0
        return precision, recall

    def search(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
        """
        Retrieve top‑k documents for a given query.

        Parameters
        ----------
        query : str
            Query text.
        top_k : int, default 3
            Number of documents to return.

        Returns
        -------
        List[Tuple[str, float]]
            List of (document_text, score) tuples sorted by descending score.
        """
        self.eval()
        with torch.no_grad():
            q_vec = self.vectorizer.transform([query]).toarray()
            q_vec = torch.tensor(q_vec, dtype=torch.float32, device=self.device)
            scores = self.forward(q_vec).cpu().numpy().flatten()
        top_indices = np.argsort(scores)[::-1][:top_k]
        return [(self.documents[i], float(scores[i])) for i in top_indices]


def demo():
    """
    Demo usage of the SearchAgent with a tiny synthetic dataset.
    """
    # Sample documents
    docs = [
        "Deep learning models can learn complex patterns from data.",
        "Search engines index web pages to provide relevant results.",
        "Natural language processing enables computers to understand text.",
        "Python is a popular programming language for data science.",
        "Machine learning algorithms improve over time with more data.",
        "Artificial intelligence encompasses machine learning and deep learning.",
        "Information retrieval is a key component of search systems.",
        "Neural networks consist of layers of interconnected nodes.",
        "Data preprocessing is essential before training models.",
        "Evaluation metrics help assess model performance.",
    ]

    # Sample queries and relevance labels (indices of relevant docs)
    queries = [
        "What is deep learning?",
        "How do search engines work?",
        "Explain natural language processing.",
        "Why use Python for data science?",
        "What is machine learning?",
    ]

    relevance = [
        [0, 5],          # Relevant docs for query 0
        [1, 6],          # Relevant docs for query 1
        [2, 7],          # Relevant docs for query 2
        [3, 9],          # Relevant docs for query 3
        [4, 5],          # Relevant docs for query 4
    ]

    # Split into train/test
    train_q, test_q, train_rel, test_rel = train_test_split(
        queries, relevance, test_size=0.4, random_state=42
    )

    agent = SearchAgent(docs)
    print("Training agent...")
    agent.train_agent(train_q, train_rel, epochs=20, lr=0.01, verbose=True)

    print("\nEvaluating agent...")
    precision, recall = agent.evaluate(test_q, test_rel)
    print(f"Precision: {precision:.2f}, Recall: {recall:.2f}")

    # Search example
    query = "Tell me about deep learning and AI."
    print(f"\nSearching for: '{query}'")
    results = agent.search(query, top_k=5)
    for doc, score in results:
        print(f"Score: {score:.4f} | Doc: {doc}")


if __name__ == "__main__":
    demo()