#!/usr/bin/env python
"""
Example script demonstrating the usage of the distilled Qodo-Embed model.

This script shows how to:
1. Load the distilled model
2. Generate embeddings for code snippets
3. Compare performance with the original model
4. Demonstrate practical use cases like code search
"""

import logging
import time

import numpy as np
from model2vec import StaticModel
from sentence_transformers import SentenceTransformer

# Constants
CODE_SNIPPET_TRUNCATE_LENGTH = 100

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Sample code snippets for demonstration
CODE_SAMPLES = [
	"""
def binary_search(arr, target):
    left, right = 0, len(arr) - 1
    while left <= right:
        mid = (left + right) // 2
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    return -1
""",
	"""
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def inorder_traversal(root):
    result = []
    def dfs(node):
        if node:
            dfs(node.left)
            result.append(node.val)
            dfs(node.right)
    dfs(root)
    return result
""",
	"""
def merge_sort(arr):
    if len(arr) <= 1:
        return arr

    mid = len(arr) // 2
    left = merge_sort(arr[:mid])
    right = merge_sort(arr[mid:])

    return merge(left, right)

def merge(left, right):
    result = []
    i = j = 0

    while i < len(left) and j < len(right):
        if left[i] <= right[j]:
            result.append(left[i])
            i += 1
        else:
            result.append(right[j])
            j += 1

    result.extend(left[i:])
    result.extend(right[j:])
    return result
""",
	"""
def fibonacci_dp(n):
    if n <= 1:
        return n

    dp = [0] * (n + 1)
    dp[1] = 1

    for i in range(2, n + 1):
        dp[i] = dp[i-1] + dp[i-2]

    return dp[n]
""",
	"""
class LRUCache:
    def __init__(self, capacity):
        self.capacity = capacity
        self.cache = {}
        self.order = []

    def get(self, key):
        if key in self.cache:
            self.order.remove(key)
            self.order.append(key)
            return self.cache[key]
        return -1

    def put(self, key, value):
        if key in self.cache:
            self.order.remove(key)
        elif len(self.cache) >= self.capacity:
            oldest = self.order.pop(0)
            del self.cache[oldest]

        self.cache[key] = value
        self.order.append(key)
""",
]

# Code database for search demonstration
CODE_DATABASE = [
	"def quicksort(arr): # Implementation of quicksort algorithm",
	"class LinkedList: # Singly linked list implementation",
	"def bfs_traversal(graph, start): # Breadth-first search",
	"def dijkstra_algorithm(graph, start): # Shortest path algorithm",
	"class HashTable: # Hash table with collision handling",
	"def dynamic_programming_solution(): # DP approach to problem solving",
	"def recursive_factorial(n): # Recursive factorial calculation",
	"class BinarySearchTree: # BST implementation with insert/delete",
	"def matrix_multiplication(A, B): # Efficient matrix multiplication",
	"def string_matching_kmp(text, pattern): # KMP string matching",
]


def load_models() -> tuple[SentenceTransformer | None, StaticModel]:
	"""Load both original and distilled models for comparison."""
	logger.info("Loading original Qodo-Embed model...")
	try:
		original_model = SentenceTransformer("Qodo/Qodo-Embed-1-1.5B")
	except (ValueError, OSError, ImportError) as e:
		logger.warning(f"Could not load original model: {e}")
		logger.info("Proceeding with distilled model only...")
		original_model = None

	logger.info("Loading distilled model...")
	distilled_model = StaticModel.from_pretrained(".")

	return original_model, distilled_model


def demonstrate_basic_usage(distilled_model: StaticModel) -> None:
	"""Demonstrate basic embedding generation."""
	separator = "=" * 50
	logger.info("\n%s", separator)
	logger.info("BASIC USAGE DEMONSTRATION")
	logger.info("%s", separator)

	# Generate embeddings for code samples
	logger.info("Generating embeddings for %d code samples...", len(CODE_SAMPLES))

	start_time = time.time()
	embeddings = distilled_model.encode(CODE_SAMPLES)
	end_time = time.time()

	logger.info("Generated embeddings in %.3f seconds", end_time - start_time)
	logger.info("Embedding shape: %s", embeddings.shape)
	logger.info("Embedding dtype: %s", embeddings.dtype)

	# Show some statistics
	logger.info("Embedding statistics:")
	logger.info("  Mean: %.6f", np.mean(embeddings))
	logger.info("  Std:  %.6f", np.std(embeddings))
	logger.info("  Min:  %.6f", np.min(embeddings))
	logger.info("  Max:  %.6f", np.max(embeddings))


def cosine_similarity(v1: np.ndarray, v2: np.ndarray) -> float:
	"""Compute cosine similarity between two vectors."""
	dot_product = np.dot(v1, v2)
	norm_v1 = np.linalg.norm(v1)
	norm_v2 = np.linalg.norm(v2)
	return dot_product / (norm_v1 * norm_v2)


def calculate_semantic_similarity(emb1: np.ndarray, emb2: np.ndarray, samples: list[str]) -> None:
	"""Calculate semantic similarity between embeddings of different dimensions.

	This function demonstrates how to compare embeddings even when they have
	different dimensions by looking at relative similarities.
	"""
	logger.info("Calculating semantic similarities...")

	# Calculate pairwise similarities within each embedding space
	n_samples = len(samples)

	logger.info("Sample similarity comparisons:")
	for i in range(min(3, n_samples)):  # Show first 3 comparisons
		for j in range(i + 1, min(i + 3, n_samples)):
			sim1 = cosine_similarity(emb1[i], emb1[j])
			sim2 = cosine_similarity(emb2[i], emb2[j])

			logger.info(f"  Samples {i + 1} vs {j + 1}:")
			logger.info(f"    Original model similarity: {sim1:.4f}")
			logger.info(f"    Distilled model similarity: {sim2:.4f}")
			logger.info(f"    Difference: {abs(sim1 - sim2):.4f}")


def run_speed_benchmark(
	original_model: SentenceTransformer | None, distilled_model: StaticModel, samples: list[str]
) -> None:
	"""Run speed benchmark comparing original and distilled models."""
	separator = "=" * 50
	logger.info("\n%s", separator)
	logger.info("SPEED BENCHMARK")
	logger.info("%s", separator)

	n_runs = 5
	logger.info("Running %d iterations for each model...", n_runs)

	if original_model is not None:
		# Benchmark original model
		logger.info("Benchmarking original model...")
		start_time = time.time()
		for _ in range(n_runs):
			_ = original_model.encode(samples)
		original_time = time.time() - start_time
		original_speed = (len(samples) * n_runs) / original_time
		logger.info("Original model: %.2f texts/second", original_speed)
	else:
		original_speed = None
		logger.info("Original model not available for benchmarking")

	# Benchmark distilled model
	logger.info("Benchmarking distilled model...")
	start_time = time.time()
	for _ in range(n_runs):
		_ = distilled_model.encode(samples)
	distilled_time = time.time() - start_time
	distilled_speed = (len(samples) * n_runs) / distilled_time
	logger.info("Distilled model: %.2f texts/second", distilled_speed)

	# Compare speeds
	if original_speed is not None:
		speedup = distilled_speed / original_speed
		logger.info("Speedup: %.2fx", speedup)
	else:
		logger.info("Cannot calculate speedup without original model")


def demonstrate_similarity(
	original_embeddings: np.ndarray, distilled_embeddings: np.ndarray, samples: list[str]
) -> None:
	"""Demonstrate similarity between original and distilled embeddings."""
	separator = "=" * 50
	logger.info("\n%s", separator)
	logger.info("EMBEDDING SIMILARITY")
	logger.info("%s", separator)

	logger.info("Original embeddings shape: %s", original_embeddings.shape)
	logger.info("Distilled embeddings shape: %s", distilled_embeddings.shape)

	# Since dimensions might be different, we compare semantic relationships
	if original_embeddings.shape[1] != distilled_embeddings.shape[1]:
		logger.info("Different embedding dimensions - comparing semantic relationships...")
		calculate_semantic_similarity(original_embeddings, distilled_embeddings, samples)
	else:
		# Direct comparison if dimensions match
		logger.info("Same embedding dimensions - direct comparison possible")
		similarities = [cosine_similarity(original_embeddings[i], distilled_embeddings[i]) for i in range(len(samples))]
		avg_similarity = np.mean(similarities)
		logger.info("Average cosine similarity: %.4f", avg_similarity)


def demonstrate_code_search(distilled_model: StaticModel, query: str, code_database: list[str]) -> None:
	"""Demonstrate code search functionality with the distilled model."""
	separator = "=" * 50
	logger.info("\n%s", separator)
	logger.info("CODE SEARCH DEMO")
	logger.info("%s", separator)

	logger.info("Query: %s", query)
	logger.info("Searching in database of %d code snippets...", len(code_database))

	# Encode query and database
	query_embedding = distilled_model.encode([query])
	db_embeddings = distilled_model.encode(code_database)

	# Calculate similarities
	similarities = []
	for i, db_emb in enumerate(db_embeddings):
		sim = cosine_similarity(query_embedding[0], db_emb)
		similarities.append((i, sim))

	# Sort by similarity
	similarities.sort(key=lambda x: x[1], reverse=True)

	logger.info("Top 3 matches:")
	for i, (idx, sim) in enumerate(similarities[:3]):
		code_snippet = code_database[idx]
		if len(code_snippet) > CODE_SNIPPET_TRUNCATE_LENGTH:
			code_snippet = code_snippet[:CODE_SNIPPET_TRUNCATE_LENGTH] + "..."
		logger.info(f"{i + 1}. Similarity: {sim:.4f}")
		logger.info(f"   Code: {code_snippet}")


def main() -> None:
	"""Main demonstration function."""
	logger.info("Starting Qodo-Embed Model2Vec Distilled Model Demo")

	# Load models
	original_model, distilled_model = load_models()

	# Basic usage demonstration
	demonstrate_basic_usage(distilled_model)

	# Speed benchmark
	if original_model is not None:
		run_speed_benchmark(original_model, distilled_model, CODE_SAMPLES[:3])  # Use subset for speed

		# Similarity demonstration
		logger.info("Generating embeddings for similarity comparison...")
		original_embeddings = original_model.encode(CODE_SAMPLES[:3])
		distilled_embeddings = distilled_model.encode(CODE_SAMPLES[:3])
		demonstrate_similarity(original_embeddings, distilled_embeddings, CODE_SAMPLES[:3])

	# Code search demonstration
	search_queries = [
		"binary search implementation",
		"graph traversal algorithm",
		"dynamic programming solution",
	]

	for query in search_queries:
		demonstrate_code_search(distilled_model, query, CODE_DATABASE)

	separator = "=" * 50
	logger.info("\n%s", separator)
	logger.info("Model Summary: Qodo-Embed-M-1-1.5B-M2V-Distilled")
	logger.info("%s", separator)
	logger.info("✓ Significantly smaller model size")
	logger.info("✓ Faster inference speed")
	logger.info("✓ Maintains semantic understanding")
	logger.info("✓ Compatible with existing workflows")
	logger.info("✓ No GPU required for inference")


if __name__ == "__main__":
	main()