policy_chart_generator.py

"""
Chart Generation for Policy Impact Simulator
Creates visual charts and graphs for policy analysis results
"""

import matplotlib.pyplot as plt
import numpy as np
from typing import Dict, Any, List
import os
import base64
import io
from datetime import datetime

class PolicyChartGenerator:
    """Generates charts and graphs for policy impact analysis"""
    
    def __init__(self):
        # Set matplotlib to use non-interactive backend
        plt.switch_backend('Agg')
        
    def generate_scenario_comparison_chart(self, variants: Dict[str, Any], title: str = "Policy Scenario Analysis") -> str:
        """Generate a bar chart comparing scenario variants"""
        
        # Extract scenario data
        scenario_names = []
        costs = []
        colors = []
        
        color_map = {
            'best_case': '#28a745',  # Green
            'base_case': '#007bff',  # Blue  
            'worst_case': '#dc3545'  # Red
        }
        
        for scenario_type, data in variants.items():
            if isinstance(data, dict) and 'total_cost' in data:
                scenario_names.append(scenario_type.replace('_', ' ').title())
                costs.append(data['total_cost'])
                colors.append(color_map.get(scenario_type, '#6c757d'))
        
        # Create chart
        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.bar(scenario_names, costs, color=colors, alpha=0.8)
        
        # Customize chart
        ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
        ax.set_ylabel('Total Cost (₹ Crores)', fontsize=12)
        ax.set_xlabel('Scenario', fontsize=12)
        
        # Add value labels on bars
        for bar, cost in zip(bars, costs):
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height + max(costs) * 0.01,
                   f'₹{cost:,.0f}', ha='center', va='bottom', fontweight='bold')
        
        # Format y-axis
        ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'₹{x:,.0f}'))
        
        # Add grid
        ax.grid(True, alpha=0.3, axis='y')
        ax.set_axisbelow(True)
        
        # Tight layout
        plt.tight_layout()
        
        # Convert to base64 string
        return self._fig_to_base64(fig)
    
    def generate_yearly_breakdown_chart(self, yearly_data: List[Dict], title: str = "Yearly Impact Breakdown") -> str:
        """Generate a line chart showing yearly breakdown"""
        
        years = []
        impacts = []
        beneficiaries = []
        
        for year_data in yearly_data:
            years.append(f"Year {year_data.get('year', 0)}")
            impacts.append(year_data.get('impact', 0))
            beneficiaries.append(year_data.get('affected_beneficiaries', 0))
        
        # Create dual-axis chart
        fig, ax1 = plt.subplots(figsize=(12, 6))
        
        # Plot impact on primary axis
        color1 = '#007bff'
        ax1.set_xlabel('Years', fontsize=12)
        ax1.set_ylabel('Impact (₹ Crores)', fontsize=12, color=color1)
        line1 = ax1.plot(years, impacts, color=color1, marker='o', linewidth=3, markersize=8, label='Financial Impact')
        ax1.tick_params(axis='y', labelcolor=color1)
        ax1.grid(True, alpha=0.3)
        
        # Create secondary axis for beneficiaries
        ax2 = ax1.twinx()
        color2 = '#28a745'
        ax2.set_ylabel('Beneficiaries', fontsize=12, color=color2)
        line2 = ax2.plot(years, beneficiaries, color=color2, marker='s', linewidth=3, markersize=8, linestyle='--', label='Affected Population')
        ax2.tick_params(axis='y', labelcolor=color2)
        
        # Add title
        ax1.set_title(title, fontsize=16, fontweight='bold', pad=20)
        
        # Add value labels
        for i, (year, impact, beneficiary) in enumerate(zip(years, impacts, beneficiaries)):
            ax1.annotate(f'₹{impact:.1f}', (i, impact), textcoords="offset points", xytext=(0,10), ha='center', fontweight='bold')
            ax2.annotate(f'{beneficiary:,}', (i, beneficiary), textcoords="offset points", xytext=(0,-15), ha='center', fontweight='bold', color=color2)
        
        # Add legend
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        plt.tight_layout()
        return self._fig_to_base64(fig)
    
    def generate_ascii_chart(self, variants: Dict[str, Any], width: int = 50) -> str:
        """Generate ASCII chart for text-based display"""
        
        if not variants:
            return "No data available for chart generation"
        
        # Extract costs
        costs = []
        names = []
        for scenario_type, data in variants.items():
            if isinstance(data, dict) and 'total_cost' in data:
                costs.append(data['total_cost'])
                names.append(scenario_type.replace('_', ' ').title())
        
        if not costs:
            return "No cost data available"
        
        max_cost = max(costs) if costs else 1
        
        chart = "📊 Policy Scenario Comparison (₹ Crores)\n"
        chart += "=" * (width + 20) + "\n"
        
        for name, cost in zip(names, costs):
            # Calculate bar length
            bar_length = int((cost / max_cost) * width) if max_cost > 0 else 0
            bar = "█" * bar_length
            
            # Add emoji based on scenario type
            emoji = "🟢" if "Best" in name else "🔴" if "Worst" in name else "🔵"
            
            chart += f"{emoji} {name:12} │{bar:<{width}} ₹{cost:,.0f}\n"
        
        chart += "=" * (width + 20) + "\n"
        return chart
    
    def generate_implementation_timeline_chart(self, timeline_data: Dict[str, Any]) -> str:
        """Generate ASCII timeline chart for implementation phases"""
        
        phases = [
            "📋 Planning Phase (Months 1-3)",
            "⚖️ Legal Review (Months 2-4)", 
            "💰 Budget Approval (Months 4-6)",
            "🔄 System Updates (Months 6-8)",
            "📢 Communication (Months 8-10)",
            "🚀 Implementation (Months 10-12)"
        ]
        
        timeline = "🗓️ Implementation Timeline\n"
        timeline += "=" * 60 + "\n"
        
        for i, phase in enumerate(phases, 1):
            progress_bar = "▓" * (i * 2) + "░" * ((6 - i) * 2)
            timeline += f"  {phase}\n"
            timeline += f"  [{progress_bar}] {i}/6 phases\n\n"
        
        complexity = timeline_data.get('complexity', 'Medium')
        timeline += f"🔧 Implementation Complexity: {complexity}\n"
        timeline += f"⏱️ Estimated Duration: 12 months\n"
        timeline += f"💡 Key Success Factors: Budget approval, stakeholder buy-in, system readiness\n"
        
        return timeline
    
    def _fig_to_base64(self, fig) -> str:
        """Convert matplotlib figure to base64 string"""
        
        buffer = io.BytesIO()
        fig.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
        buffer.seek(0)
        
        # Convert to base64
        img_base64 = base64.b64encode(buffer.read()).decode('utf-8')
        
        # Close figure to free memory
        plt.close(fig)
        
        return f"data:image/png;base64,{img_base64}"
    
    def generate_comprehensive_report_charts(self, analysis_data: Dict[str, Any]) -> Dict[str, str]:
        """Generate all charts for a comprehensive policy analysis report"""
        
        charts = {}
        
        # Scenario comparison chart
        if 'variants' in analysis_data:
            charts['scenario_comparison'] = self.generate_scenario_comparison_chart(
                analysis_data['variants'],
                f"Policy Impact: {analysis_data.get('parameter_name', 'Analysis')}"
            )
            
            # ASCII version for text display
            charts['scenario_ascii'] = self.generate_ascii_chart(analysis_data['variants'])
        
        # Yearly breakdown chart
        if 'scenario_projections' in analysis_data:
            charts['yearly_breakdown'] = self.generate_yearly_breakdown_chart(
                analysis_data['scenario_projections'],
                "Financial Impact Over Time"
            )
        
        # Implementation timeline
        if 'implementation' in analysis_data:
            charts['implementation_timeline'] = self.generate_implementation_timeline_chart(
                analysis_data['implementation']
            )
        
        return charts

# Standalone functions for easy integration
def generate_policy_charts(analysis_data: Dict[str, Any]) -> Dict[str, str]:
    """Generate charts for policy analysis data"""
    generator = PolicyChartGenerator()
    return generator.generate_comprehensive_report_charts(analysis_data)

def generate_ascii_scenario_chart(variants: Dict[str, Any]) -> str:
    """Generate ASCII chart for immediate text display"""
    generator = PolicyChartGenerator()
    return generator.generate_ascii_chart(variants)