#!/usr/bin/env python3
"""
Demo script to show how to use the specialized data analysis agents.
These agents provide targeted capabilities for different types of data analysis tasks.
"""

import logging
import os
import pandas as pd
import numpy as np
from dotenv import load_dotenv

from src.ai.agent import AIAgent

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

# Load environment variables
load_dotenv()

def create_sample_data():
    """Create sample datasets for demonstration purposes."""
    log.info("Creating sample datasets...")
    
    # Sample Forex data with time series
    dates = pd.date_range(start='2023-01-01', end='2023-12-31', freq='D')
    np.random.seed(42)  # For reproducibility
    
    # Create a realistic EUR/USD dataset with trends and seasonality
    base = 1.10
    trend = np.linspace(0, 0.15, len(dates))  # Upward trend
    
    # Weekly seasonality (mid-week has higher values)
    day_of_week = np.array([d.dayofweek for d in dates])
    weekday_effect = np.sin(day_of_week * np.pi / 7) * 0.01
    
    # Monthly seasonality 
    day_of_month = np.array([d.day for d in dates])
    month_effect = np.sin(day_of_month * np.pi / 30) * 0.02
    
    # Quarterly effect
    quarter_effect = np.array([0.03 if d.quarter == 4 else 0 for d in dates])
    
    # Random noise
    noise = np.random.normal(0, 0.005, len(dates))
    
    # Combine all effects
    eur_usd = base + trend + weekday_effect + month_effect + quarter_effect + noise
    
    # Create the Forex DataFrame
    forex_df = pd.DataFrame({
        'date': dates,
        'EUR_USD': eur_usd,
        'GBP_USD': eur_usd * 1.2 + np.random.normal(0, 0.01, len(dates)),
        'USD_JPY': eur_usd * 110 + np.random.normal(0, 0.5, len(dates)),
        'USD_CHF': 1 / (eur_usd + np.random.normal(0, 0.005, len(dates))),
        'trading_volume': np.random.gamma(10, 500, len(dates)).astype(int),
        'volatility': np.abs(np.random.normal(0, 0.002, len(dates)))
    })
    
    # Add some missing values for realism
    indices = np.random.choice(len(forex_df), 20, replace=False)
    forex_df.loc[indices, 'EUR_USD'] = np.nan
    
    # Economic indicators data
    indicators_df = pd.DataFrame({
        'date': pd.date_range(start='2023-01-01', end='2023-12-01', freq='MS'),
        'eu_inflation': np.random.normal(2.5, 0.3, 12),
        'us_inflation': np.random.normal(3.0, 0.4, 12),
        'eu_interest_rate': np.concatenate([np.repeat(3.0, 6), np.repeat(3.25, 6)]),
        'us_interest_rate': np.concatenate([np.repeat(4.5, 3), np.repeat(4.75, 4), np.repeat(5.0, 5)]),
        'eu_gdp_growth': np.random.normal(1.2, 0.3, 12),
        'us_gdp_growth': np.random.normal(2.0, 0.4, 12)
    })
    
    # Market sentiment data
    sentiment_df = pd.DataFrame({
        'date': dates,
        'eur_sentiment': np.cumsum(np.random.normal(0, 0.1, len(dates))),
        'usd_sentiment': np.cumsum(np.random.normal(0, 0.1, len(dates))),
        'market_fear_index': 20 + np.random.gamma(2, 5, len(dates)),
        'analyst_bullish_pct': 50 + np.cumsum(np.random.normal(0, 0.5, len(dates))) % 40
    })
    
    return {
        'forex': forex_df,
        'economic_indicators': indicators_df,
        'market_sentiment': sentiment_df
    }

def run_agent_demo():
    """Run the specialized agent demo."""
    log.info("Starting specialized agent demo...")
    
    # Create sample data
    sample_data = create_sample_data()
    
    # Initialize the AI agent
    ai_agent = AIAgent()
    
    # Store sample data in the agent's memory
    for name, df in sample_data.items():
        ai_agent.data_store[name] = df
        log.info(f"Stored sample data '{name}' in agent memory: {df.shape[0]} rows, {df.shape[1]} columns")
    
    # Example queries for each specialized agent
    example_queries = {
        # Statistical analysis agent
        "statistics": [
            "Provide a comprehensive statistical analysis of the 'forex' dataset focusing on EUR/USD",
            "What are the correlations between different currency pairs in our data?",
            "Analyze the distribution of volatility in the forex dataset and identify outliers"
        ],
        
        # Time series agent
        "time_series": [
            "Analyze the EUR/USD exchange rate as a time series. Look for seasonality and trends.",
            "Forecast EUR/USD for the next 30 days based on historical patterns with confidence intervals",
            "Compare the stationarity of EUR/USD vs GBP/USD and recommend which might be easier to forecast"
        ],
        
        # Data transformation agent
        "data_transform": [
            "Merge the forex and economic indicators datasets on date, then show the relationship between interest rates and exchange rates",
            "Calculate monthly averages of all currency pairs and store as 'monthly_forex'",
            "Create a new indicator that shows the ratio of EUR/USD to USD/JPY over time"
        ],
        
        # Forex Market Intelligence agent
        "forex_intelligence": [
            "What are the current factors affecting EUR/USD exchange rate movements?",
            "Provide information about recent ECB and Federal Reserve interest rate decisions and their impact on forex markets",
            "What are analyst predictions for EUR/USD in the coming quarter?",
            "How have geopolitical tensions impacted major currency pairs recently?"
        ],
        
        # Analytics team
        "team": [
            "Analyze our EUR/USD data to identify trends, forecast the next 30 days, and complement with current market analyst opinions",
            "How do economic indicators correlate with currency movements, and what are experts predicting for upcoming indicator releases?",
            "Create a comprehensive report on the GBP/USD pair including statistical analysis, forecasting, and current market sentiment"
        ]
    }
    
    # Run demo queries for each agent type
    for agent_type, queries in example_queries.items():
        print(f"\n{'=' * 80}")
        print(f"DEMONSTRATING {agent_type.upper()} AGENT")
        print(f"{'=' * 80}")
        
        # Run one query as an example for each agent type
        query = queries[0]  # Just use the first query for demonstration
        print(f"\nQUERY: {query}")
        
        response = ai_agent.process_query(query, agent_type=agent_type)
        
        # Print a shortened version of the response for demonstration
        max_display_length = 500
        if len(response) > max_display_length:
            shortened_response = response[:max_display_length] + "...\n[Response truncated for display]"
            print(f"\nRESPONSE (shortened):\n{shortened_response}")
        else:
            print(f"\nRESPONSE:\n{response}")
            
        print(f"\n{'-' * 80}")
        
    # Show main agent with all tools
    print(f"\n{'=' * 80}")
    print(f"DEMONSTRATING MAIN AGENT (ALL TOOLS)")
    print(f"{'=' * 80}")
    
    main_query = "What insights can you provide about the EUR/USD exchange rate based on our data and current market conditions?"
    print(f"\nQUERY: {main_query}")
    
    response = ai_agent.process_query(main_query)
    max_display_length = 500
    if len(response) > max_display_length:
        shortened_response = response[:max_display_length] + "...\n[Response truncated for display]"
        print(f"\nRESPONSE (shortened):\n{shortened_response}")
    else:
        print(f"\nRESPONSE:\n{response}")

if __name__ == "__main__":
    run_agent_demo()