#!/usr/bin/env python3
"""
Test script to verify color visualization is working correctly
"""

import numpy as np
import re
from app import create_visualization

def test_color_visualization():
    """Test that the visualization creates colored HTML"""
    print("🎨 Testing Color Visualization")
    print("=" * 50)

    # Test with sample data
    tokens = ["The", "quick", "brown", "fox", "jumps", "over", "lazy", "dog"]
    perplexities = np.array([1.2, 2.5, 8.3, 3.1, 15.7, 2.0, 12.4, 1.8])

    print(f"📝 Tokens: {tokens}")
    print(f"📊 Perplexities: {perplexities}")

    # Generate visualization
    html = create_visualization(tokens, perplexities)

    # Check that HTML was generated
    assert len(html) > 100, "HTML output too short"
    print("✅ HTML generated successfully")

    # Check for color information in HTML
    color_pattern = r'rgba?\(\d+,\s*\d+,\s*\d+(?:,\s*[\d.]+)?\)'
    colors_found = re.findall(color_pattern, html)

    print(f"🎨 Colors found in HTML: {len(colors_found)}")
    for i, color in enumerate(colors_found[:5]):  # Show first 5
        print(f"   Color {i+1}: {color}")

    assert len(colors_found) > 0, "No colors found in HTML output"
    print("✅ Color information found in HTML")

    # Check for span elements with style attributes
    span_pattern = r'<span style="[^"]*background-color[^"]*"[^>]*>'
    spans_found = re.findall(span_pattern, html)

    print(f"🏷️  Styled spans found: {len(spans_found)}")
    assert len(spans_found) >= len(tokens) - 2, "Not enough styled spans found"  # Allow for some filtering
    print("✅ Styled spans with background colors found")

    # Check for tooltip information
    assert 'title="Perplexity:' in html, "No tooltip information found"
    print("✅ Tooltip information found")

    # Verify different colors for different perplexity ranges
    # Extract RGB values
    rgb_values = []
    for color in colors_found:
        # Extract numbers from rgba(r,g,b,a) or rgb(r,g,b)
        numbers = re.findall(r'\d+', color)
        if len(numbers) >= 3:
            rgb_values.append((int(numbers[0]), int(numbers[1]), int(numbers[2])))

    if len(rgb_values) >= 2:
        # Check that we have different colors (not all the same)
        unique_colors = set(rgb_values)
        print(f"🌈 Unique colors found: {len(unique_colors)}")
        assert len(unique_colors) > 1, "All tokens have the same color"
        print("✅ Multiple different colors found")

        # Check color range makes sense
        red_values = [r for r, g, b in rgb_values]
        green_values = [g for r, g, b in rgb_values]

        print(f"🔴 Red range: {min(red_values)} - {max(red_values)}")
        print(f"🟢 Green range: {min(green_values)} - {max(green_values)}")

        # Should have variation in color channels
        assert max(red_values) - min(red_values) > 20, "Not enough red variation"
        print("✅ Sufficient color variation found")

    return html

def test_edge_cases():
    """Test edge cases for color visualization"""
    print("\n🧪 Testing Edge Cases")
    print("=" * 50)

    # Test with very high perplexities
    tokens = ["unusual", "words", "here"]
    high_perplexities = np.array([100.0, 200.0, 50.0])

    html = create_visualization(tokens, high_perplexities)
    assert len(html) > 50, "HTML too short for high perplexities"
    print("✅ High perplexity values handled")

    # Test with very low perplexities
    low_perplexities = np.array([0.1, 0.2, 0.15])
    html = create_visualization(tokens, low_perplexities)
    assert len(html) > 50, "HTML too short for low perplexities"
    print("✅ Low perplexity values handled")

    # Test with single token
    single_token = ["word"]
    single_perplexity = np.array([5.0])
    html = create_visualization(single_token, single_perplexity)
    assert len(html) > 50, "HTML too short for single token"
    print("✅ Single token handled")

    # Test with empty input
    empty_html = create_visualization([], np.array([]))
    assert "No tokens" in empty_html, "Empty case not handled properly"
    print("✅ Empty input handled")

def test_color_gradient():
    """Test that color gradient works as expected"""
    print("\n🌈 Testing Color Gradient")
    print("=" * 50)

    # Create tokens with ascending perplexities
    tokens = [f"token_{i}" for i in range(10)]
    perplexities = np.array([i * 2.0 + 1.0 for i in range(10)])  # 1, 3, 5, 7, 9, 11, 13, 15, 17, 19

    html = create_visualization(tokens, perplexities)

    # Extract all RGB values in order
    color_pattern = r'rgba?\((\d+),\s*(\d+),\s*(\d+)(?:,\s*[\d.]+)?\)'
    colors_found = re.findall(color_pattern, html)

    if len(colors_found) >= 5:
        # Convert to numeric values
        rgb_values = [(int(r), int(g), int(b)) for r, g, b in colors_found]

        # Check that low perplexity tokens are more green
        low_perp_color = rgb_values[0]  # First token (lowest perplexity)
        high_perp_color = rgb_values[-1]  # Last token (highest perplexity)

        print(f"🟢 Low perplexity color (perp={perplexities[0]:.1f}): RGB{low_perp_color}")
        print(f"🔴 High perplexity color (perp={perplexities[-1]:.1f}): RGB{high_perp_color}")

        # Low perplexity should be more green (higher green value)
        # High perplexity should be more red (higher red value)
        if low_perp_color[1] > high_perp_color[1]:  # Green component
            print("✅ Low perplexity tokens are greener")
        else:
            print("⚠️  Expected low perplexity to be greener")

        if high_perp_color[0] > low_perp_color[0]:  # Red component
            print("✅ High perplexity tokens are redder")
        else:
            print("⚠️  Expected high perplexity to be redder")

def main():
    """Run all color visualization tests"""
    print("🎨 Color Visualization Test Suite")
    print("=" * 60)

    try:
        # Test basic functionality
        html = test_color_visualization()

        # Test edge cases
        test_edge_cases()

        # Test color gradient
        test_color_gradient()

        print("\n" + "=" * 60)
        print("🎉 All color visualization tests passed!")
        print("🌈 The tokens should now appear with colored backgrounds!")
        print("=" * 60)

        # Save a sample HTML file for manual inspection
        with open("sample_visualization.html", "w") as f:
            f.write(f"""
<!DOCTYPE html>
<html>
<head>
    <title>Sample Perplexity Visualization</title>
</head>
<body>
    <h1>Sample Perplexity Visualization</h1>
    <p>This is what the colored visualization should look like:</p>
    {html}
</body>
</html>
""")
        print("💾 Sample visualization saved to 'sample_visualization.html'")
        print("   Open this file in your browser to see the colors!")

        return True

    except Exception as e:
        print(f"\n❌ Color visualization test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)