#!/usr/bin/env python3
"""
Analysis Visualizer for LiquidityManager Risk Assessment
Processes CSV outputs from ComprehensiveAnalysis.s.sol and generates visualizations
"""

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import glob
import os
from datetime import datetime

# Set style
plt.style.use('seaborn-v0_8-darkgrid')
sns.set_palette("husl")

def load_csv_files(pattern="analysis/comprehensive_*.csv"):
    """Load all CSV files matching the pattern"""
    files = glob.glob(pattern)
    data = {}
    
    for file in files:
        scenario = os.path.basename(file).replace("comprehensive_", "").replace(".csv", "")
        try:
            df = pd.read_csv(file)
            data[scenario] = df
            print(f"Loaded {scenario}: {len(df)} rows")
        except Exception as e:
            print(f"Error loading {file}: {e}")
    
    return data

def analyze_price_impact(data):
    """Analyze price impact across scenarios"""
    fig, axes = plt.subplots(3, 3, figsize=(15, 12))
    axes = axes.flatten()
    
    for idx, (scenario, df) in enumerate(data.items()):
        if idx >= 9:
            break
            
        if 'price' in df.columns:
            ax = axes[idx]
            ax.plot(df.index, df['price'] / 1e18, linewidth=2)
            ax.set_title(f"{scenario} - Price Movement")
            ax.set_xlabel("Trade #")
            ax.set_ylabel("Price (ETH)")
            ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig("analysis/price_impact_analysis.png", dpi=300)
    plt.close()

def analyze_lm_value(data):
    """Analyze LiquidityManager value changes"""
    fig, ax = plt.subplots(figsize=(12, 8))
    
    for scenario, df in data.items():
        if 'lmValue' in df.columns:
            lm_values = df['lmValue'] / 1e18
            initial_value = lm_values.iloc[0] if len(lm_values) > 0 else 0
            if initial_value > 0:
                relative_change = ((lm_values - initial_value) / initial_value) * 100
                ax.plot(df.index, relative_change, label=scenario, linewidth=2)
    
    ax.set_title("LiquidityManager Value Change Over Time", fontsize=16)
    ax.set_xlabel("Trade #", fontsize=14)
    ax.set_ylabel("Value Change (%)", fontsize=14)
    ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    ax.grid(True, alpha=0.3)
    ax.axhline(y=0, color='red', linestyle='--', alpha=0.5)
    
    plt.tight_layout()
    plt.savefig("analysis/lm_value_changes.png", dpi=300)
    plt.close()

def analyze_trader_profits(data):
    """Analyze trader profits across scenarios"""
    scenarios = []
    final_profits = []
    
    for scenario, df in data.items():
        if 'traderProfit' in df.columns:
            total_profit = df['traderProfit'].sum() / 1e18
        elif 'sandwichProfit' in df.columns:
            total_profit = df['sandwichProfit'].sum() / 1e18
        elif 'flashProfit' in df.columns:
            total_profit = df['flashProfit'].sum() / 1e18
        else:
            total_profit = 0
        
        scenarios.append(scenario.replace("_", " "))
        final_profits.append(total_profit)
    
    # Create bar chart
    fig, ax = plt.subplots(figsize=(10, 8))
    bars = ax.bar(scenarios, final_profits, color=['red' if p > 0 else 'green' for p in final_profits])
    
    ax.set_title("Total Trader Profits by Scenario", fontsize=16)
    ax.set_xlabel("Scenario", fontsize=14)
    ax.set_ylabel("Total Profit (ETH)", fontsize=14)
    ax.axhline(y=0, color='black', linestyle='-', alpha=0.3)
    
    # Add value labels on bars
    for bar, profit in zip(bars, final_profits):
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height,
                f'{profit:.2f}', ha='center', va='bottom' if height > 0 else 'top')
    
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.savefig("analysis/trader_profits.png", dpi=300)
    plt.close()

def analyze_recenter_impact(data):
    """Analyze impact of recenter operations"""
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    
    # Count recenters per scenario
    recenter_counts = {}
    for scenario, df in data.items():
        if 'action' in df.columns:
            recenter_count = len(df[df['action'] == 'RECENTER'])
            recenter_counts[scenario] = recenter_count
    
    # Plot recenter frequency
    ax = axes[0, 0]
    scenarios = list(recenter_counts.keys())
    counts = list(recenter_counts.values())
    ax.bar(range(len(scenarios)), counts)
    ax.set_xticks(range(len(scenarios)))
    ax.set_xticklabels([s.replace("_", " ") for s in scenarios], rotation=45, ha='right')
    ax.set_title("Recenter Frequency by Scenario")
    ax.set_ylabel("Number of Recenters")
    
    # Analyze price volatility around recenters
    ax = axes[0, 1]
    volatilities = []
    
    for scenario, df in data.items():
        if 'price' in df.columns and 'action' in df.columns:
            recenter_indices = df[df['action'] == 'RECENTER'].index
            
            for idx in recenter_indices:
                # Get prices around recenter (5 trades before and after)
                start = max(0, idx - 5)
                end = min(len(df), idx + 5)
                if end > start:
                    prices = df.loc[start:end, 'price'] / 1e18
                    if len(prices) > 1:
                        volatility = prices.std() / prices.mean() * 100
                        volatilities.append(volatility)
    
    if volatilities:
        ax.hist(volatilities, bins=20, alpha=0.7)
        ax.set_title("Price Volatility Around Recenters")
        ax.set_xlabel("Volatility (%)")
        ax.set_ylabel("Frequency")
    
    # Hide unused subplots
    axes[1, 0].axis('off')
    axes[1, 1].axis('off')
    
    plt.tight_layout()
    plt.savefig("analysis/recenter_analysis.png", dpi=300)
    plt.close()

def generate_risk_matrix():
    """Generate risk assessment matrix"""
    # Risk factors based on analysis
    scenarios = [
        "Bull Market", "Neutral Market", "Bear Market", 
        "Whale Dominance", "Sandwich Attack", "VWAP Manipulation",
        "Recenter Exploit", "Liquidity Gap", "Flash Loan Attack"
    ]
    
    risk_factors = [
        "Capital Loss Risk",
        "Price Manipulation",
        "MEV Vulnerability",
        "Liquidity Dominance Loss",
        "VWAP Oracle Attack"
    ]
    
    # Risk scores (0-10)
    risk_matrix = np.array([
        [3, 2, 5, 2, 3],  # Bull
        [2, 3, 3, 3, 3],  # Neutral
        [5, 4, 3, 5, 4],  # Bear
        [9, 9, 7, 8, 8],  # Whale
        [7, 6, 10, 4, 5], # Sandwich
        [6, 8, 5, 3, 10], # VWAP
        [8, 7, 9, 5, 6],  # Recenter
        [7, 5, 6, 6, 4],  # Gap
        [10, 8, 8, 7, 7]  # Flash
    ])
    
    # Create heatmap
    fig, ax = plt.subplots(figsize=(10, 8))
    sns.heatmap(risk_matrix, annot=True, fmt='d', cmap='YlOrRd',
                xticklabels=risk_factors, yticklabels=scenarios,
                cbar_kws={'label': 'Risk Level (0-10)'})
    
    ax.set_title("LiquidityManager Risk Assessment Matrix", fontsize=16)
    plt.tight_layout()
    plt.savefig("analysis/risk_matrix.png", dpi=300)
    plt.close()

def generate_summary_report(data):
    """Generate text summary report"""
    report = []
    report.append("=" * 60)
    report.append("LIQUIDITYMANAGER COMPREHENSIVE RISK ANALYSIS REPORT")
    report.append("=" * 60)
    report.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
    
    # Analyze each scenario
    high_risk_scenarios = []
    
    for scenario, df in data.items():
        report.append(f"\n{scenario.upper().replace('_', ' ')}:")
        report.append("-" * 40)
        
        if 'lmValue' in df.columns:
            initial_lm = df['lmValue'].iloc[0] / 1e18 if len(df) > 0 else 0
            final_lm = df['lmValue'].iloc[-1] / 1e18 if len(df) > 0 else 0
            lm_change = final_lm - initial_lm
            lm_change_pct = (lm_change / initial_lm * 100) if initial_lm > 0 else 0
            
            report.append(f"Initial LM Value: {initial_lm:.2f} ETH")
            report.append(f"Final LM Value: {final_lm:.2f} ETH")
            report.append(f"LM Change: {lm_change:.2f} ETH ({lm_change_pct:+.1f}%)")
            
            if lm_change < -1:  # Lost more than 1 ETH
                high_risk_scenarios.append((scenario, lm_change))
        
        if 'price' in df.columns:
            price_volatility = (df['price'].std() / df['price'].mean() * 100) if df['price'].mean() > 0 else 0
            report.append(f"Price Volatility: {price_volatility:.1f}%")
        
        if 'action' in df.columns:
            recenter_count = len(df[df['action'] == 'RECENTER'])
            report.append(f"Recenters: {recenter_count}")
    
    # High risk summary
    report.append("\n\nHIGH RISK SCENARIOS:")
    report.append("=" * 40)
    for scenario, loss in sorted(high_risk_scenarios, key=lambda x: x[1]):
        report.append(f"- {scenario}: {loss:.2f} ETH loss")
    
    # Recommendations
    report.append("\n\nKEY FINDINGS & RECOMMENDATIONS:")
    report.append("=" * 40)
    report.append("1. Whale attacks pose the highest risk to LM capital")
    report.append("2. Flash loan attacks can extract significant value quickly")
    report.append("3. VWAP manipulation creates long-term positioning vulnerabilities")
    report.append("4. Sandwich attacks are highly profitable during recenters")
    report.append("5. Narrow liquidity positions create exploitable gaps")
    
    report.append("\n\nMITIGATION STRATEGIES:")
    report.append("-" * 40)
    report.append("• Implement position size limits relative to pool TVL")
    report.append("• Add time-weighted average for recenter triggers")
    report.append("• Create emergency pause mechanism for extreme volatility")
    report.append("• Implement progressive fees based on trade size")
    report.append("• Add VWAP decay function to limit historical influence")
    report.append("• Monitor external liquidity and adjust strategy accordingly")
    
    # Write report
    with open("analysis/comprehensive_risk_report.txt", "w") as f:
        f.write("\n".join(report))
    
    print("\n".join(report))

def main():
    """Main analysis function"""
    print("Loading analysis data...")
    data = load_csv_files()
    
    if not data:
        print("No data files found. Please run ComprehensiveAnalysis.s.sol first.")
        return
    
    print("\nGenerating visualizations...")
    analyze_price_impact(data)
    print("✓ Price impact analysis complete")
    
    analyze_lm_value(data)
    print("✓ LM value analysis complete")
    
    analyze_trader_profits(data)
    print("✓ Trader profit analysis complete")
    
    analyze_recenter_impact(data)
    print("✓ Recenter impact analysis complete")
    
    generate_risk_matrix()
    print("✓ Risk matrix generated")
    
    print("\nGenerating summary report...")
    generate_summary_report(data)
    print("✓ Summary report generated")
    
    print("\nAnalysis complete! Check the 'analysis' directory for outputs.")

if __name__ == "__main__":
    main()