import numpy as np

# Simple example of what embeddings capture
# These vectors are simplified — real embeddings are 768-3072 dimensions

# "The cat sat on the mat" → [0.2, 0.8, 0.1, 0.9, ...]
# "A feline rested on the rug" → [0.21, 0.79, 0.12, 0.88, ...]
# "Stock prices rose today" → [0.9, 0.1, 0.8, 0.2, ...]

# The first two vectors are close (similar meaning)
# The third is far away (different topic)

# Cosine Similarity — the core metric for RAG
def cosine_similarity(a, b):
    """
    Measures the angle between two vectors.
    Returns values from -1 (opposite) to 1 (identical).
    
    cos(θ) = (a · b) / (||a|| × ||b||)
    """
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    return dot_product / (norm_a * norm_b)

# Example
vec_ai = np.array([0.2, 0.8, 0.1, 0.9])   # "artificial intelligence"
vec_ml = np.array([0.25, 0.75, 0.15, 0.85]) # "machine learning"
vec_finance = np.array([0.9, 0.1, 0.8, 0.2]) # "stock market"

print(cosine_similarity(vec_ai, vec_ml))      # ~0.98 (very similar)
print(cosine_similarity(vec_ai, vec_finance)) # ~0.35 (different)

# For most production systems:
# 1. Start with OpenAI text-embedding-3-small (cheap, good quality)
# 2. Use text-embedding-3-large if accuracy matters more than cost
# 3. Use local models (sentence-transformers) if you need data privacy

# Cost comparison (as of 2024):
# text-embedding-3-small: $0.02 / 1M tokens
# text-embedding-3-large: $0.13 / 1M tokens
# Local models: $0 (but you pay for GPU/compute)

def fixed_size_chunk(text, chunk_size=500, overlap=50):
    """
    Split text into fixed-size chunks with overlap.
    Overlap prevents losing context at chunk boundaries.
    """
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap  # Overlap creates sliding window
    return chunks

# Example
doc = "RAG combines retrieval and generation. " * 100  # Long document
chunks = fixed_size_chunk(doc, chunk_size=500, overlap=50)
print(f"Created {len(chunks)} chunks")

def semantic_chunk(text, embedding_model, threshold=0.5):
    """
    Split text at semantic boundaries — where meaning changes.
    Uses embeddings to detect topic shifts.
    """
    sentences = split_into_sentences(text)
    chunks = []
    current_chunk = [sentences[0]]
    
    for i in range(1, len(sentences)):
        # Compare current sentence with previous
        sim = cosine_similarity(
            embedding_model.embed(sentences[i]),
            embedding_model.embed(sentences[i-1])
        )
        
        if sim < threshold:
            # Big drop in similarity = topic change
            chunks.append(' '.join(current_chunk))
            current_chunk = [sentences[i]]
        else:
            current_chunk.append(sentences[i])
    
    chunks.append(' '.join(current_chunk))
    return chunks

def recursive_split(text, separators=None, chunk_size=1000):
    """
    Try to split on natural boundaries first (paragraphs, sentences),
    fall back to character-level splitting.
    """
    if separators is None:
        separators = ['\n\n', '\n', '. ', ' ', '']
    
    if len(text) <= chunk_size:
        return [text]
    
    separator = separators[0]
    remaining_separators = separators[1:]
    
    splits = text.split(separator)
    chunks = []
    current = ''
    
    for split in splits:
        if len(current) + len(split) <= chunk_size:
            current += separator + split if current else split
        else:
            if current:
                chunks.append(current)
            if len(split) > chunk_size:
                # Recursive call with smaller separators
                sub_chunks = recursive_split(split, remaining_separators, chunk_size)
                chunks.extend(sub_chunks)
            else:
                current = split
    
    if current:
        chunks.append(current)
    
    return chunks

def document_structure_chunk(markdown_text, chunk_size=1000):
    """
    Respect document structure: headers, sections, code blocks.
    Ideal for technical documentation.
    """
    chunks = []
    current_section = ''
    current_header = ''
    
    for line in markdown_text.split('\n'):
        if line.startswith('#'):
            # New section — save previous
            if current_section:
                chunks.append({
                    'header': current_header,
                    'content': current_section.strip()
                })
            current_header = line
            current_section = ''
        else:
            current_section += line + '\n'
            
            if len(current_section) > chunk_size:
                # Section too long — split it
                chunks.append({
                    'header': current_header,
                    'content': current_section.strip()
                })
                current_section = ''
    
    return chunks

# Decision tree:
# 1. Just learning/prototyping → ChromaDB (simple, local)
# 2. Production, don't want ops → Pinecone (managed)
# 3. Production, need filters → Qdrant or Weaviate
# 4. Already use PostgreSQL → pgvector
# 5. Billions of vectors → Milvus or Pinecone

def basic_retrieval(query, vector_store, k=5):
    """
    Simple similarity search — the baseline RAG approach.
    """
    query_embedding = embed(query)
    results = vector_store.search(query_embedding, k=k)
    return results

# This works but has limitations:
# - Only semantic search (misses exact keyword matches)
# - No understanding of query intent
# - No result filtering or re-ranking

def hybrid_search(query, vector_store, bm25_index, k=5, alpha=0.5):
    """
    Combine semantic search (embeddings) with keyword search (BM25).
    alpha: weight for semantic search (1-alpha for keyword search)
    """
    # Semantic search (understands meaning)
    semantic_results = vector_store.similarity_search(query, k=k*2)
    
    # Keyword search (exact matches, acronyms, proper nouns)
    keyword_results = bm25_index.search(query, k=k*2)
    
    # Reciprocal Rank Fusion (RRF) — merge results
    combined_scores = {}
    for rank, doc in enumerate(semantic_results):
        combined_scores[doc.id] = combined_scores.get(doc.id, 0) + 1 / (60 + rank)
    
    for rank, doc in enumerate(keyword_results):
        combined_scores[doc.id] = combined_scores.get(doc.id, 0) + 1 / (60 + rank)
    
    # Sort by combined score
    sorted_results = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_results[:k]

def rerank(query, documents, top_k=5):
    """
    Cross-encoder re-ranking: scores each (query, document) pair.
    More accurate than bi-encoder embeddings but slower.
    """
    from sentence_transformers import CrossEncoder
    
    model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
    
    # Score each query-document pair
    pairs = [(query, doc.text) for doc in documents]
    scores = model.predict(pairs)
    
    # Sort by score
    ranked = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, score in ranked[:top_k]]

def query_expansion(query, llm, num_queries=3):
    """
    Generate multiple query variations to improve recall.
    """
    prompt = f"""
    Given the user question: "{query}"
    
    Generate {num_queries} alternative ways to ask this question
    that might retrieve different relevant documents.
    Return only the queries, one per line.
    """
    
    expanded = llm.generate(prompt)
    return [query] + expanded.split('\n')

# Example:
# Original: "How does backpropagation work?"
# Expanded:
#   1. "How does backpropagation work?"
#   2. "Backpropagation algorithm neural networks"
#   3. "Gradient descent weight updates deep learning"
#   4. "Chain rule derivatives neural network training"

def build_rag_prompt(query, retrieved_chunks):
    """
    Construct the prompt with retrieved context.
    """
    context = "\n\n---\n\n".join([
        f"[Source: {chunk.metadata['source']}]\n{chunk.text}"
        for chunk in retrieved_chunks
    ])
    
    return f"""You are a helpful assistant. Answer the question based ONLY on the 
following context. If the context doesn't contain the answer, say "I don't have 
enough information to answer this question."

Context:
{context}

Question: {query}

Answer:"""

def generate_with_citations(query, chunks, llm):
    """
    Generate answers with inline citations.
    """
    prompt = f"""Answer the question using the provided context.
Always cite your sources using [Source N] notation.

Context chunks:
[1] {chunks[0].text}
[2] {chunks[1].text}
[3] {chunks[2].text}

Question: {query}

Answer with citations:"""
    
    return llm.generate(prompt)

# Output: "RAG works by retrieving relevant documents first [1], 
# then generating an answer grounded in those documents [2][3]."

def hyde_retrieval(query, llm, vector_store, k=5):
    """
    Step 1: Generate a hypothetical answer (even if it's wrong)
    Step 2: Embed the hypothetical answer
    Step 3: Search for documents similar to the hypothetical answer
    
    Why? The hypothetical answer is closer in embedding space to the 
    actual answer documents than the raw question is.
    """
    # Step 1: Generate hypothetical answer
    hyde_prompt = f"""Write a detailed paragraph that would answer this question:
    
    {query}
    
    Write as if you're explaining the answer from a textbook."""
    
    hypothetical_doc = llm.generate(hyde_prompt)
    
    # Step 2-3: Embed and search
    hyde_embedding = embed(hypothetical_doc)
    results = vector_store.search(hyde_embedding, k=k)
    
    return results

def self_rag(query, llm, vector_store):
    """
    Self-RAG asks four questions at each step:
    1. [Retrieve] — Do I need to retrieve?
    2. [IsRel] — Is this retrieved document relevant?
    3. [IsSup] — Does my generation support this document?
    4. [IsUse] — Is this generation useful?
    """
    
    # Step 1: Should we retrieve?
    retrieve_prompt = f"""Question: {query}
    
    Is external retrieval needed to answer this question?
    Answer: [Retrieve] or [No Retrieve]"""
    
    decision = llm.generate(retrieve_prompt)
    
    if "[No Retrieve]" in decision:
        # Direct answer — no retrieval needed
        return llm.generate(f"Answer: {query}")
    
    # Step 2: Retrieve and filter
    documents = vector_store.search(query, k=5)
    
    relevant_docs = []
    for doc in documents:
        # Check relevance
        rel_prompt = f"""Question: {query}
        Document: {doc.text}
        Is this document relevant to the question?
        Answer: [IsRel] or [Not Relevant]"""
        
        if "[IsRel]" in llm.generate(rel_prompt):
            relevant_docs.append(doc)
    
    # Step 3: Generate with reflection
    generation = llm.generate(f"Answer {query} using: {relevant_docs}")
    
    # Step 4: Check if generation is supported
    support_prompt = f"""Generation: {generation}
    Documents: {relevant_docs}
    Is the generation supported by the documents?
    Answer: [IsFully Supported] or [Is Partially Supported] or [Is Not Supported]"""
    
    support = llm.generate(support_prompt)
    
    if "Not Supported" in support:
        # Retry with different approach
        return self_rag(query, llm, vector_store)
    
    return generation

def corrective_rag(query, llm, vector_store):
    """
    CRAG evaluates retrieval quality:
    - Correct → Use retrieved documents
    - Ambiguous → Refine search with knowledge strips
    - Incorrect → Fall back to web search
    """
    
    documents = vector_store.search(query, k=5)
    
    # Evaluate retrieval quality
    eval_prompt = f"""Question: {query}
    Retrieved documents: {[d.text[:200] for d in documents]}
    
    Rate the retrieval quality:
    - [Correct] — Documents directly answer the question
    - [Ambiguous] — Documents are partially relevant
    - [Incorrect] — Documents are not relevant"""
    
    evaluation = llm.generate(eval_prompt)
    
    if "[Correct]" in evaluation:
        # Use documents as-is
        return generate_from_documents(query, documents, llm)
    
    elif "[Ambiguous]" in evaluation:
        # Extract useful knowledge and refine
        knowledge_strips = extract_knowledge_strips(documents)
        refined_query = refine_query(query, knowledge_strips, llm)
        return corrective_rag(refined_query, llm, vector_store)
    
    else:
        # Fall back to web search
        web_results = web_search(query)
        return generate_from_documents(query, web_results, llm)

def graph_rag(query, llm, knowledge_graph):
    """
    Graph RAG retrieves information through knowledge graph traversal.
    Good for questions requiring multiple hops of reasoning.
    """
    
    # Step 1: Extract entities from query
    entities = extract_entities(query, llm)
    # e.g., "How did Einstein's work on relativity influence GPS technology?"
    # → ["Einstein", "relativity", "GPS"]
    
    # Step 2: Traverse knowledge graph
    subgraph = knowledge_graph.get_subgraph(entities, depth=3)
    
    # Step 3: Find paths between entities
    paths = knowledge_graph.find_paths(
        source="Einstein",
        target="GPS",
        max_hops=4
    )
    
    # Step 4: Generate answer from graph context
    graph_context = format_paths_as_context(paths)
    return llm.generate(f"Answer: {query}\n\nKnowledge graph paths:\n{graph_context}")

def evaluate_rag_system(test_cases, rag_system):
    """
    Comprehensive RAG evaluation.
    """
    metrics = {
        'faithfulness': [],      # Is the answer grounded in context?
        'answer_relevancy': [],   # Does the answer address the question?
        'context_precision': [],  # Are the retrieved chunks relevant?
        'context_recall': [],     # Did we retrieve all relevant chunks?
    }
    
    for test in test_cases:
        result = rag_system.query(test['question'])
        
        # Faithfulness: LLM judges if answer is supported by context
        faithfulness = judge_faithfulness(
            answer=result['answer'],
            context=result['context']
        )
        
        # Answer relevancy: LLM judges if answer addresses question
        relevancy = judge_relevancy(
            question=test['question'],
            answer=result['answer']
        )
        
        # Context precision: What fraction of retrieved chunks are relevant
        precision = sum(1 for chunk in result['chunks'] 
                       if chunk_is_relevant(chunk, test['ground_truth'])) / len(result['chunks'])
        
        # Context recall: What fraction of ground truth is covered
        recall = sum(1 for gt in test['ground_truth']
                    if any(chunk_covers(chunk, gt) for chunk in result['chunks'])) / len(test['ground_truth'])
        
        metrics['faithfulness'].append(faithfulness)
        metrics['answer_relevancy'].append(relevancy)
        metrics['context_precision'].append(precision)
        metrics['context_recall'].append(recall)
    
    return {k: np.mean(v) for k, v in metrics.items()}

from ragas import evaluate
from ragas.metrics import (
    faithfulness,
    answer_relevancy,
    context_precision,
    context_recall,
)

# Prepare evaluation dataset
eval_dataset = {
    'question': [...],
    'answer': [...],           # RAG system's answer
    'contexts': [...],         # Retrieved contexts
    'ground_truth': [...]      # Expected answers
}

# Run evaluation
result = evaluate(
    dataset=eval_dataset,
    metrics=[
        faithfulness,        # Answer grounded in context?
        answer_relevancy,    # Answer addresses question?
        context_precision,   # Retrieved chunks relevant?
        context_recall,      # All relevant chunks found?
    ]
)

print(result)
# {'faithfulness': 0.87, 'answer_relevancy': 0.92, 
#  'context_precision': 0.78, 'context_recall': 0.83}

PRODUCTION_CHECKLIST = {
    "Indexing": [
        "Chunk size optimized for your use case",
        "Metadata stored with each chunk (source, date, section)",
        "Incremental indexing pipeline (not full rebuild)",
        "Document preprocessing (clean HTML, extract text, handle images)",
    ],
    "Retrieval": [
        "Hybrid search (semantic + keyword)",
        "Re-ranking step for precision",
        "Query expansion for recall",
        "Filtering by metadata (date, source, language)",
    ],
    "Generation": [
        "Prompt template with citation instructions",
        "Context window management (don't exceed LLM limits)",
        "Fallback for when no relevant context is found",
        "Answer length control",
    ],
    "Evaluation": [
        "Automated metrics (faithfulness, relevancy, precision, recall)",
        "Human evaluation on representative sample",
        "A/B testing framework for prompt/model changes",
        "Regression test suite",
    ],
    "Infrastructure": [
        "Vector database backup and recovery",
        "Embedding cache for repeated queries",
        "Rate limiting and cost monitoring",
        "Logging and observability",
    ]
}