#!/usr/bin/env python3
"""
Test the n-gram freshness component of novelty sensor.
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from salience_os_seed.core.sensors.novelty import NoveltySensor
from salience_os_seed.core.sensors.base import MedianMADNormalizer

print("="*100)
print("N-GRAM FRESHNESS COMPONENT TEST")
print("="*100)
print()

# ==============================================================================
# Test freshness calculation
# ==============================================================================
print("Testing freshness calculation...")
print()

normalizer = MedianMADNormalizer()
sensor = NoveltySensor(normalizer, ngram=4)

# Test with tokens
test_cases = [
    ["Hello", "MONIKA", "this", "is", "new"],
    ["Hello", "MONIKA", "this", "is", "new"],  # Repeat
    ["Different", "words", "should", "have", "high"],
    ["Hello", "MONIKA", "this", "is", "new"],  # Repeat again
    ["Totally", "unique", "sentence", "here", "xyz"],
]

print("Testing n-gram freshness:")
print(f"  N-gram size: {sensor._ngram}")
print()

for i, tokens in enumerate(test_cases):
    # Call internal freshness method
    freshness = sensor._update_freshness(tokens)
    count = sensor._ngram_counter.get("::".join(tokens[-sensor._ngram:]), 0)
    
    print(f"  Case {i+1}: {' '.join(tokens)}")
    print(f"    N-gram: {' '.join(tokens[-sensor._ngram:])}")
    print(f"    Freshness: {freshness:.4f}")
    print(f"    Count after update: {count}")
    print()

print("N-gram counter state:")
print(f"  Total unique n-grams: {len(sensor._ngram_counter)}")
print(f"  Top 10 n-grams:")
for phrase, count in sorted(sensor._ngram_counter.items(), key=lambda x: -x[1])[:10]:
    print(f"    '{phrase}': {count:.2f}")
print()

# ==============================================================================
# Test with context.tokens from actual state
# ==============================================================================
print("="*100)
print("Testing with actual state context.tokens...")
print()

from salience_os_seed.conversation.session import ConversationSession, ConversationConfig

session = ConversationSession(ConversationConfig())

# Create fresh sensor
sensor2 = NoveltySensor(normalizer, ngram=4)

test_texts = [
    "Hello MONIKA can you help me",
    "What is your name today",
    "Hello MONIKA can you help me",  # Exact repeat
    "Never seen this before xyz quantum",
]

for text in test_texts:
    state = session._build_state(text, speaker="user")
    state["prediction"]["surprisal"] = 5.0  # Add surprisal above baseline
    
    tokens = state["context"]["tokens"]
    print(f"Text: '{text}'")
    print(f"Tokens: {tokens}")
    print(f"N-gram: {tokens[-sensor2._ngram:] if len(tokens) >= sensor2._ngram else 'TOO SHORT'}")
    
    # Measure
    novelty = sensor2._measure(state, {}, {})
    
    # Get freshness separately
    if len(tokens) >= sensor2._ngram:
        phrase = "::".join(tokens[-sensor2._ngram:])
        count = sensor2._ngram_counter.get(phrase, 0)
        freshness = 1.0 / (1.0 + count)
        print(f"Freshness: {freshness:.4f} (count={count})")
    else:
        print(f"Freshness: 0.0 (too few tokens)")
    
    print(f"Novelty: {novelty:.4f}")
    print()

# ==============================================================================
# Test edge cases
# ==============================================================================
print("="*100)
print("EDGE CASE TESTING")
print("="*100)
print()

print("Edge Case 1: Too few tokens")
sensor3 = NoveltySensor(normalizer, ngram=4)
short_tokens = ["Hi", "there"]
freshness = sensor3._update_freshness(short_tokens)
print(f"  Tokens: {short_tokens}")
print(f"  Freshness: {freshness} (expected 0.0 for < 4 tokens)")
print()

print("Edge Case 2: Empty tokens")
empty_tokens = []
freshness = sensor3._update_freshness(empty_tokens)
print(f"  Tokens: {empty_tokens}")
print(f"  Freshness: {freshness} (expected 0.0 for empty)")
print()

print("Edge Case 3: Exact 4 tokens")
exact_tokens = ["one", "two", "three", "four"]
freshness = sensor3._update_freshness(exact_tokens)
print(f"  Tokens: {exact_tokens}")
print(f"  Freshness: {freshness} (expected 1.0 for first occurrence)")
print()

# ==============================================================================
# Test complete novelty calculation
# ==============================================================================
print("="*100)
print("COMPLETE NOVELTY CALCULATION")
print("="*100)
print()

import math

sensor4 = NoveltySensor(normalizer, baseline_surprisal=3.5, ngram=4)

test_state = {
    "prediction": {"surprisal": 5.0},  # Above baseline
    "context": {"tokens": ["Hello", "world", "this", "is", "new"]},
}

print("Manual calculation:")
print(f"  Surprisal: {test_state['prediction']['surprisal']}")
print(f"  Baseline: {sensor4._baseline_surprisal}")

surprisal = 5.0
surprisal_delta = max(surprisal - 3.5, 0.0)
print(f"  Surprisal delta: {surprisal_delta}")

delta_scale = math.tanh(surprisal_delta)
print(f"  Delta scale (tanh): {delta_scale}")

# Calculate freshness
tokens = test_state["context"]["tokens"]
if len(tokens) >= 4:
    phrase = "::".join(tokens[-4:])
    count = sensor4._ngram_counter.get(phrase, 0)
    freshness = 1.0 / (1.0 + count)
    print(f"  N-gram: '{phrase}'")
    print(f"  Count: {count}")
    print(f"  Freshness: {freshness}")
else:
    freshness = 0.0
    print(f"  Freshness: {freshness} (too few tokens)")

novelty_raw = math.sqrt(delta_scale * freshness)
print(f"  Novelty raw: {novelty_raw}")
print()

actual_novelty = sensor4._measure(test_state, {}, {})
print(f"Actual sensor measurement: {actual_novelty}")
print()

print("="*100)
print("SUMMARY")
print("="*100)
print()

print("Freshness component requires:")
print("  1. At least 4 tokens (default n-gram size)")
print("  2. Decreases with repeated n-grams: 1/(1+count)")
print()

print("Novelty formula:")
print("  novelty = sqrt(tanh(max(surprisal - 3.5, 0)) * freshness)")
print()

print("For novelty > 0, ALL conditions must be met:")
print("  ✓ surprisal > 3.5")
print("  ✓ len(tokens) >= ngram_size (4)")
print("  ✓ freshness > 0 (n-gram not over-repeated)")
print()

print("If ANY condition fails, novelty = 0.0")
print()

print("="*100)