# %%
"""
AI Agent Testing and Development Notebook
==========================================

This notebook helps you test, debug, and improve your AI agent.
Use this for development before submitting to the evaluation system.
"""

# %%
import sys
import os
import pandas as pd
from langchain_core.messages import HumanMessage
import matplotlib.pyplot as plt
import time

# Add current directory to path for imports
sys.path.append('.')

from agent import build_graph, extract_final_answer
from config import config, check_config

# %%
# Configuration Check
print("🔧 Configuration Status")
print("=" * 50)
is_valid = check_config()
print("=" * 50)

if not is_valid:
    print("⚠️ Some features may not work without proper configuration.")
    print("See README.md for setup instructions.")

# %%
# Initialize Agent
print("🤖 Initializing Agent...")

try:
    agent_graph = build_graph(
        provider=config.agent.llm_provider,
        use_reasoning=config.agent.use_reasoning
    )
    print("✅ Agent initialized successfully!")
except Exception as e:
    print(f"❌ Agent initialization failed: {e}")
    print("Check your API keys and configuration.")

# %%
# Test Questions Database
TEST_QUESTIONS = [
    # Mathematical Questions
    {
        "category": "Math",
        "difficulty": "Easy",
        "question": "What is 15 * 23?",
        "expected_answer": "345"
    },
    {
        "category": "Math", 
        "difficulty": "Medium",
        "question": "If I save $485 per month for 3.5 years, how much will I have saved?",
        "expected_answer": "20370"
    },
    
    # Knowledge Questions
    {
        "category": "Knowledge",
        "difficulty": "Easy", 
        "question": "What is the capital of France?",
        "expected_answer": "Paris"
    },
    {
        "category": "Knowledge",
        "difficulty": "Medium",
        "question": "Who wrote 'Pride and Prejudice'?",
        "expected_answer": "Jane Austen"
    },
    
    # Reasoning Questions
    {
        "category": "Reasoning",
        "difficulty": "Medium",
        "question": "If all cats are animals and some animals are pets, can we conclude that some cats are pets?",
        "expected_answer": "No"  # This requires careful logical reasoning
    },
    
    # Current Events (requires web search)
    {
        "category": "Current",
        "difficulty": "Hard",
        "question": "What is the current weather in New York City?",
        "expected_answer": None  # Can't predict this
    }
]

# %%
# Single Question Testing Function
def test_single_question(question_text, expected_answer=None, show_details=True):
    """Test a single question with detailed output."""
    
    print(f"❓ Question: {question_text}")
    print("-" * 50)
    
    try:
        # Time the response
        start_time = time.time()
        
        # Get agent response
        result = agent_graph.invoke({"messages": [HumanMessage(content=question_text)]})
        response_time = time.time() - start_time
        
        # Extract answer
        raw_response = result['messages'][-1].content
        final_answer = extract_final_answer(raw_response)
        
        # Display results
        if show_details:
            print(f"🤖 Raw Response: {raw_response}")
            print(f"✅ Final Answer: {final_answer}")
            print(f"⏱️ Response Time: {response_time:.2f}s")
            
            if expected_answer:
                is_correct = expected_answer.lower() in final_answer.lower()
                print(f"🎯 Expected: {expected_answer}")
                print(f"📊 Correct: {'✅ Yes' if is_correct else '❌ No'}")
        
        return {
            "question": question_text,
            "answer": final_answer,
            "expected": expected_answer,
            "response_time": response_time,
            "raw_response": raw_response,
            "success": True
        }
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return {
            "question": question_text,
            "answer": f"Error: {e}",
            "expected": expected_answer,
            "response_time": None,
            "raw_response": None,
            "success": False
        }

# %%
# Test Single Question (Interactive)
# Change this question to test specific cases
test_question = "What is 25 * 4?"

print("🧪 Single Question Test")
print("=" * 50)
result = test_single_question(test_question, expected_answer="100")

# %%
# Batch Testing Function
def run_batch_tests(questions, verbose=False):
    """Run multiple test questions and analyze results."""
    
    results = []
    
    print(f"🧪 Running {len(questions)} test questions...")
    print("=" * 50)
    
    for i, test in enumerate(questions, 1):
        print(f"\n[{i}/{len(questions)}] Testing: {test['category']} - {test['difficulty']}")
        
        result = test_single_question(
            test['question'], 
            test['expected_answer'],
            show_details=verbose
        )
        
        # Add metadata
        result.update({
            "category": test['category'],
            "difficulty": test['difficulty']
        })
        
        results.append(result)
        
        if not verbose:
            status = "✅" if result['success'] else "❌"
            print(f"   {status} {result['answer'][:50]}...")
    
    return results

# %%
# Run Batch Tests
print("🎯 Batch Testing")
print("=" * 50)

batch_results = run_batch_tests(TEST_QUESTIONS, verbose=False)

# %%
# Results Analysis
def analyze_results(results):
    """Analyze test results and generate insights."""
    
    df = pd.DataFrame(results)
    
    # Basic Statistics
    total_tests = len(results)
    successful_tests = df['success'].sum()
    success_rate = (successful_tests / total_tests) * 100 if total_tests > 0 else 0
    
    # Response Time Analysis
    valid_times = df[df['response_time'].notna()]['response_time']
    avg_response_time = valid_times.mean() if len(valid_times) > 0 else 0
    
    print("📊 Test Results Analysis")
    print("=" * 50)
    print(f"Total Tests: {total_tests}")
    print(f"Successful: {successful_tests}")
    print(f"Success Rate: {success_rate:.1f}%")
    print(f"Average Response Time: {avg_response_time:.2f}s")
    
    # Category Performance
    print("\n📈 Performance by Category:")
    category_stats = df.groupby('category').agg({
        'success': ['count', 'sum'],
        'response_time': 'mean'
    }).round(2)
    
    for category in df['category'].unique():
        cat_data = df[df['category'] == category]
        total = len(cat_data)
        success = cat_data['success'].sum()
        rate = (success / total) * 100 if total > 0 else 0
        avg_time = cat_data['response_time'].mean()
        
        print(f"  {category}: {success}/{total} ({rate:.1f}%) - {avg_time:.2f}s avg")
    
    # Failed Tests
    failed_tests = df[~df['success']]
    if len(failed_tests) > 0:
        print(f"\n❌ Failed Tests ({len(failed_tests)}):")
        for _, test in failed_tests.iterrows():
            print(f"  - {test['question'][:50]}...")
            print(f"    Error: {test['answer']}")
    
    return df

# %%
# Analyze Results
analysis_df = analyze_results(batch_results)

# %%
# Visualization
def plot_results(df):
    """Create visualizations of test results."""
    
    if len(df) == 0:
        print("No data to plot")
        return
    
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 8))
    
    # Success Rate by Category
    category_success = df.groupby('category')['success'].agg(['count', 'sum'])
    category_success['rate'] = (category_success['sum'] / category_success['count']) * 100
    
    ax1.bar(category_success.index, category_success['rate'], color='skyblue')
    ax1.set_title('Success Rate by Category (%)')
    ax1.set_ylabel('Success Rate (%)')
    ax1.tick_params(axis='x', rotation=45)
    
    # Response Time Distribution
    valid_times = df[df['response_time'].notna()]['response_time']
    if len(valid_times) > 0:
        ax2.hist(valid_times, bins=10, color='lightgreen', alpha=0.7)
        ax2.set_title('Response Time Distribution')
        ax2.set_xlabel('Response Time (seconds)')
        ax2.set_ylabel('Frequency')
    
    # Success vs Difficulty
    difficulty_success = df.groupby('difficulty')['success'].agg(['count', 'sum'])
    difficulty_success['rate'] = (difficulty_success['sum'] / difficulty_success['count']) * 100
    
    ax3.bar(difficulty_success.index, difficulty_success['rate'], color='orange')
    ax3.set_title('Success Rate by Difficulty (%)')
    ax3.set_ylabel('Success Rate (%)')
    
    # Response Time by Category
    if len(valid_times) > 0:
        df_valid = df[df['response_time'].notna()]
        category_times = df_valid.groupby('category')['response_time'].mean()
        
        ax4.bar(category_times.index, category_times.values, color='coral')
        ax4.set_title('Average Response Time by Category')
        ax4.set_ylabel('Response Time (seconds)')
        ax4.tick_params(axis='x', rotation=45)
    
    plt.tight_layout()
    plt.show()

# %%
# Generate Plots
if len(analysis_df) > 0:
    plot_results(analysis_df)
else:
    print("No successful tests to visualize")

# %%
# Custom Test Section
print("🎯 Custom Testing Section")
print("=" * 50)
print("Use this section to test specific scenarios:")

# Add your custom test questions here
custom_questions = [
    "What is the square root of 144?",
    "Name three planets in our solar system.",
    "If I have 10 apples and give away 3, how many do I have left?"
]

for question in custom_questions:
    print(f"\n🧪 Testing: {question}")
    result = test_single_question(question, show_details=False)
    print(f"Answer: {result['answer']}")

# %%
# Performance Optimization Tips
print("""
🚀 Performance Optimization Tips:

1. **Tool Selection Strategy**:
   - Use calculator for math questions
   - Use web search for current events
   - Use knowledge base for historical facts

2. **Response Time Optimization**:
   - Choose faster LLM providers (Groq > Google > HuggingFace)
   - Reduce context window size
   - Implement caching for repeated questions

3. **Accuracy Improvements**:
   - Better system prompts
   - Few-shot examples in prompts
   - Error handling and retry logic

4. **Debugging Strategies**:
   - Test individual tools separately
   - Check API key configurations
   - Monitor rate limits and quotas

5. **Agent Architecture**:
   - Implement question classification
   - Use different strategies per question type
   - Add confidence scoring
""")

# %%
# Export Results for Analysis
if len(batch_results) > 0:
    # Save results to CSV
    results_df = pd.DataFrame(batch_results)
    results_df.to_csv('test_results.csv', index=False)
    print("📁 Results saved to 'test_results.csv'")
    
    # Generate summary report
    summary = {
        "total_tests": len(batch_results),
        "successful_tests": sum(r['success'] for r in batch_results),
        "success_rate": (sum(r['success'] for r in batch_results) / len(batch_results)) * 100,
        "avg_response_time": sum(r['response_time'] for r in batch_results if r['response_time']) / len([r for r in batch_results if r['response_time']]),
        "categories_tested": list(set(r['category'] for r in batch_results)),
        "config_used": {
            "provider": config.agent.llm_provider,
            "use_reasoning": config.agent.use_reasoning,
            "enable_web_search": config.agent.enable_web_search
        }
    }
    
    import json
    with open('test_summary.json', 'w') as f:
        json.dump(summary, f, indent=2)
    
    print("📊 Summary saved to 'test_summary.json'")

print("\n🎉 Testing Complete!")
print("Use these results to improve your agent before final submission.")

# %%