"""
Enhanced Charts and Visualizations for Reports
Beautiful, interactive charts for all report types with modern styling
"""

import io
import base64
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from django.utils import timezone
import logging

logger = logging.getLogger(__name__)

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

# Color schemes for different chart types
COLORS = {
    'primary': '#3B82F6',
    'secondary': '#10B981', 
    'warning': '#F59E0B',
    'danger': '#EF4444',
    'info': '#06B6D4',
    'success': '#22C55E',
    'purple': '#8B5CF6',
    'pink': '#EC4899'
}

CHART_COLORS = [
    '#3B82F6', '#10B981', '#F59E0B', '#EF4444', 
    '#8B5CF6', '#EC4899', '#06B6D4', '#22C55E'
]

class EnhancedChartGenerator:
    """Generate beautiful, professional charts for all report types"""
    
    @staticmethod
    def create_portfolio_overview_chart(data):
        """Create comprehensive portfolio overview with multiple visualizations"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Loan Status Distribution', 'Monthly Disbursement Trend', 
                              'Collection Rate Trend', 'Risk Analysis'),
                specs=[[{"type": "pie"}, {"type": "scatter"}],
                       [{"type": "bar"}, {"type": "bar"}]]
            )
            
            # Pie chart for loan status
            if 'loan_status' in data:
                fig.add_trace(
                    go.Pie(
                        labels=list(data['loan_status'].keys()),
                        values=list(data['loan_status'].values()),
                        hole=0.4,
                        marker_colors=CHART_COLORS[:len(data['loan_status'])]
                    ),
                    row=1, col=1
                )
            
            # Monthly disbursement trend
            if 'monthly_disbursements' in data:
                months = data['monthly_disbursements'].get('months', [])
                amounts = data['monthly_disbursements'].get('amounts', [])
                fig.add_trace(
                    go.Scatter(
                        x=months, y=amounts,
                        mode='lines+markers',
                        name='Disbursements',
                        line=dict(color=COLORS['primary'], width=3),
                        marker=dict(size=8)
                    ),
                    row=1, col=2
                )
            
            # Collection rate trend
            if 'collection_rates' in data:
                months = data['collection_rates'].get('months', [])
                rates = data['collection_rates'].get('rates', [])
                fig.add_trace(
                    go.Bar(
                        x=months, y=rates,
                        name='Collection Rate %',
                        marker_color=COLORS['success']
                    ),
                    row=2, col=1
                )
            
            # Risk analysis
            if 'risk_categories' in data:
                categories = list(data['risk_categories'].keys())
                values = list(data['risk_categories'].values())
                fig.add_trace(
                    go.Bar(
                        x=categories, y=values,
                        name='Risk Distribution',
                        marker_color=[COLORS['success'], COLORS['warning'], COLORS['danger']]
                    ),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Portfolio Overview Dashboard",
                title_x=0.5,
                height=800,
                showlegend=True,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="portfolio_chart")
            
        except Exception as e:
            logger.error(f"Portfolio chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_loan_performance_chart(data):
        """Create detailed loan performance visualization"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Loan Volume Trend', 'Average Loan Size', 
                              'Completion Rate', 'Default Rate Analysis'),
                specs=[[{"secondary_y": True}, {"type": "bar"}],
                       [{"type": "scatter"}, {"type": "bar"}]]
            )
            
            if 'monthly_data' in data:
                months = data['monthly_data'].get('months', [])
                loan_counts = data['monthly_data'].get('loan_counts', [])
                amounts = data['monthly_data'].get('amounts', [])
                
                # Loan volume (bar) and amount (line)
                fig.add_trace(
                    go.Bar(x=months, y=loan_counts, name='Loan Count', 
                          marker_color=COLORS['primary']),
                    row=1, col=1
                )
                fig.add_trace(
                    go.Scatter(x=months, y=amounts, mode='lines+markers',
                             name='Total Amount', line=dict(color=COLORS['secondary'], width=3)),
                    row=1, col=1, secondary_y=True
                )
                
                # Average loan size
                avg_sizes = [amt/count if count > 0 else 0 for amt, count in zip(amounts, loan_counts)]
                fig.add_trace(
                    go.Bar(x=months, y=avg_sizes, name='Avg Loan Size',
                          marker_color=COLORS['info']),
                    row=1, col=2
                )
            
            if 'performance_metrics' in data:
                months = data['performance_metrics'].get('months', [])
                completion_rates = data['performance_metrics'].get('completion_rates', [])
                default_rates = data['performance_metrics'].get('default_rates', [])
                
                # Completion rate trend
                fig.add_trace(
                    go.Scatter(x=months, y=completion_rates, mode='lines+markers',
                             name='Completion Rate %', 
                             line=dict(color=COLORS['success'], width=3)),
                    row=2, col=1
                )
                
                # Default rate analysis
                fig.add_trace(
                    go.Bar(x=months, y=default_rates, name='Default Rate %',
                          marker_color=COLORS['danger']),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Loan Performance Analytics",
                title_x=0.5,
                height=800,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="performance_chart")
            
        except Exception as e:
            logger.error(f"Performance chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_client_analysis_chart(data):
        """Create client behavior and demographics analysis"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Client Demographics', 'Loan Size Distribution', 
                              'Repayment Behavior', 'Client Risk Profile'),
                specs=[[{"type": "pie"}, {"type": "histogram"}],
                       [{"type": "box"}, {"type": "bar"}]]
            )
            
            # Demographics pie chart
            if 'demographics' in data:
                fig.add_trace(
                    go.Pie(
                        labels=list(data['demographics'].keys()),
                        values=list(data['demographics'].values()),
                        hole=0.3,
                        marker_colors=CHART_COLORS
                    ),
                    row=1, col=1
                )
            
            # Loan size distribution
            if 'loan_sizes' in data:
                fig.add_trace(
                    go.Histogram(
                        x=data['loan_sizes'],
                        nbinsx=20,
                        name='Loan Size Distribution',
                        marker_color=COLORS['primary']
                    ),
                    row=1, col=2
                )
            
            # Repayment behavior box plot
            if 'repayment_days' in data:
                fig.add_trace(
                    go.Box(
                        y=data['repayment_days'],
                        name='Days to Repay',
                        marker_color=COLORS['success']
                    ),
                    row=2, col=1
                )
            
            # Risk profile
            if 'risk_profile' in data:
                categories = list(data['risk_profile'].keys())
                values = list(data['risk_profile'].values())
                colors = [COLORS['success'], COLORS['warning'], COLORS['danger']]
                fig.add_trace(
                    go.Bar(
                        x=categories, y=values,
                        name='Risk Distribution',
                        marker_color=colors[:len(categories)]
                    ),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Client Analysis Dashboard",
                title_x=0.5,
                height=800,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="client_chart")
            
        except Exception as e:
            logger.error(f"Client chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_financial_summary_chart(data):
        """Create comprehensive financial summary visualization"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Revenue Streams', 'Monthly Cash Flow', 
                              'Expense Breakdown', 'Profitability Trend'),
                specs=[[{"type": "pie"}, {"type": "bar"}],
                       [{"type": "pie"}, {"type": "scatter"}]]
            )
            
            # Revenue streams
            if 'revenue_streams' in data:
                fig.add_trace(
                    go.Pie(
                        labels=list(data['revenue_streams'].keys()),
                        values=list(data['revenue_streams'].values()),
                        hole=0.4,
                        marker_colors=CHART_COLORS
                    ),
                    row=1, col=1
                )
            
            # Monthly cash flow
            if 'cash_flow' in data:
                months = data['cash_flow'].get('months', [])
                inflow = data['cash_flow'].get('inflow', [])
                outflow = data['cash_flow'].get('outflow', [])
                
                fig.add_trace(
                    go.Bar(x=months, y=inflow, name='Cash Inflow',
                          marker_color=COLORS['success']),
                    row=1, col=2
                )
                fig.add_trace(
                    go.Bar(x=months, y=[-x for x in outflow], name='Cash Outflow',
                          marker_color=COLORS['danger']),
                    row=1, col=2
                )
            
            # Expense breakdown
            if 'expenses' in data:
                fig.add_trace(
                    go.Pie(
                        labels=list(data['expenses'].keys()),
                        values=list(data['expenses'].values()),
                        hole=0.3,
                        marker_colors=CHART_COLORS[2:]
                    ),
                    row=2, col=1
                )
            
            # Profitability trend
            if 'profitability' in data:
                months = data['profitability'].get('months', [])
                profit = data['profitability'].get('profit', [])
                
                fig.add_trace(
                    go.Scatter(
                        x=months, y=profit,
                        mode='lines+markers',
                        name='Net Profit',
                        line=dict(color=COLORS['primary'], width=3),
                        fill='tonexty'
                    ),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Financial Summary Dashboard",
                title_x=0.5,
                height=800,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="financial_chart")
            
        except Exception as e:
            logger.error(f"Financial chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_risk_analysis_chart(data):
        """Create comprehensive risk analysis visualization"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Portfolio at Risk', 'Delinquency Aging', 
                              'Risk Score Distribution', 'Recovery Rate Trend'),
                specs=[[{"type": "bar"}, {"type": "bar"}],
                       [{"type": "histogram"}, {"type": "scatter"}]]
            )
            
            # Portfolio at Risk
            if 'par_data' in data:
                categories = list(data['par_data'].keys())
                values = list(data['par_data'].values())
                colors = [COLORS['success'], COLORS['warning'], COLORS['danger']]
                
                fig.add_trace(
                    go.Bar(
                        x=categories, y=values,
                        name='Portfolio at Risk',
                        marker_color=colors[:len(categories)]
                    ),
                    row=1, col=1
                )
            
            # Delinquency aging
            if 'aging_data' in data:
                age_buckets = list(data['aging_data'].keys())
                amounts = list(data['aging_data'].values())
                
                fig.add_trace(
                    go.Bar(
                        x=age_buckets, y=amounts,
                        name='Delinquent Amount',
                        marker_color=COLORS['warning']
                    ),
                    row=1, col=2
                )
            
            # Risk score distribution
            if 'risk_scores' in data:
                fig.add_trace(
                    go.Histogram(
                        x=data['risk_scores'],
                        nbinsx=15,
                        name='Risk Score Distribution',
                        marker_color=COLORS['info']
                    ),
                    row=2, col=1
                )
            
            # Recovery rate trend
            if 'recovery_trend' in data:
                months = data['recovery_trend'].get('months', [])
                rates = data['recovery_trend'].get('rates', [])
                
                fig.add_trace(
                    go.Scatter(
                        x=months, y=rates,
                        mode='lines+markers',
                        name='Recovery Rate %',
                        line=dict(color=COLORS['success'], width=3)
                    ),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Risk Analysis Dashboard",
                title_x=0.5,
                height=800,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="risk_chart")
            
        except Exception as e:
            logger.error(f"Risk chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_repayment_analysis_chart(data):
        """Create repayment pattern analysis"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Payment Methods', 'Daily Collections', 
                              'Repayment Timing', 'Collection Efficiency'),
                specs=[[{"type": "pie"}, {"type": "bar"}],
                       [{"type": "histogram"}, {"type": "scatter"}]]
            )
            
            # Payment methods distribution
            if 'payment_methods' in data:
                fig.add_trace(
                    go.Pie(
                        labels=list(data['payment_methods'].keys()),
                        values=list(data['payment_methods'].values()),
                        hole=0.4,
                        marker_colors=CHART_COLORS
                    ),
                    row=1, col=1
                )
            
            # Daily collections
            if 'daily_collections' in data:
                dates = data['daily_collections'].get('dates', [])
                amounts = data['daily_collections'].get('amounts', [])
                
                fig.add_trace(
                    go.Bar(
                        x=dates, y=amounts,
                        name='Daily Collections',
                        marker_color=COLORS['success']
                    ),
                    row=1, col=2
                )
            
            # Repayment timing (days from due date)
            if 'repayment_timing' in data:
                fig.add_trace(
                    go.Histogram(
                        x=data['repayment_timing'],
                        nbinsx=20,
                        name='Days from Due Date',
                        marker_color=COLORS['primary']
                    ),
                    row=2, col=1
                )
            
            # Collection efficiency trend
            if 'efficiency_trend' in data:
                months = data['efficiency_trend'].get('months', [])
                efficiency = data['efficiency_trend'].get('efficiency', [])
                
                fig.add_trace(
                    go.Scatter(
                        x=months, y=efficiency,
                        mode='lines+markers',
                        name='Collection Efficiency %',
                        line=dict(color=COLORS['info'], width=3)
                    ),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Repayment Analysis Dashboard",
                title_x=0.5,
                height=800,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="repayment_chart")
            
        except Exception as e:
            logger.error(f"Repayment chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_branch_comparison_chart(data):
        """Create branch performance comparison"""
        try:
            fig = make_subplots(
                rows=2, cols=2,
                subplot_titles=('Branch Performance', 'Loan Volume by Branch', 
                              'Collection Rates', 'Risk Levels'),
                specs=[[{"type": "bar"}, {"type": "bar"}],
                       [{"type": "bar"}, {"type": "bar"}]]
            )
            
            if 'branches' in data:
                branches = data['branches']
                branch_names = list(branches.keys())
                
                # Branch performance (total disbursed)
                performance = [branches[b].get('total_disbursed', 0) for b in branch_names]
                fig.add_trace(
                    go.Bar(
                        x=branch_names, y=performance,
                        name='Total Disbursed',
                        marker_color=COLORS['primary']
                    ),
                    row=1, col=1
                )
                
                # Loan volume
                volumes = [branches[b].get('loan_count', 0) for b in branch_names]
                fig.add_trace(
                    go.Bar(
                        x=branch_names, y=volumes,
                        name='Loan Count',
                        marker_color=COLORS['secondary']
                    ),
                    row=1, col=2
                )
                
                # Collection rates
                collection_rates = [branches[b].get('collection_rate', 0) for b in branch_names]
                fig.add_trace(
                    go.Bar(
                        x=branch_names, y=collection_rates,
                        name='Collection Rate %',
                        marker_color=COLORS['success']
                    ),
                    row=2, col=1
                )
                
                # Risk levels (PAR)
                risk_levels = [branches[b].get('par_30', 0) for b in branch_names]
                fig.add_trace(
                    go.Bar(
                        x=branch_names, y=risk_levels,
                        name='PAR 30 %',
                        marker_color=COLORS['warning']
                    ),
                    row=2, col=2
                )
            
            fig.update_layout(
                title_text="Branch Performance Comparison",
                title_x=0.5,
                height=800,
                template="plotly_white"
            )
            
            return fig.to_html(include_plotlyjs='cdn', div_id="branch_chart")
            
        except Exception as e:
            logger.error(f"Branch chart generation failed: {e}")
            return "<div class='alert alert-warning'>Chart generation failed</div>"
    
    @staticmethod
    def create_matplotlib_chart(data, chart_type='bar'):
        """Create matplotlib chart and return as base64 string"""
        try:
            plt.figure(figsize=(12, 8))
            
            if chart_type == 'bar' and 'categories' in data and 'values' in data:
                bars = plt.bar(data['categories'], data['values'], 
                              color=CHART_COLORS[:len(data['categories'])])
                plt.title(data.get('title', 'Chart'), fontsize=16, fontweight='bold')
                plt.xlabel(data.get('xlabel', ''), fontsize=12)
                plt.ylabel(data.get('ylabel', ''), fontsize=12)
                
                # Add value labels on bars
                for bar in bars:
                    height = bar.get_height()
                    plt.text(bar.get_x() + bar.get_width()/2., height,
                            f'{height:,.0f}', ha='center', va='bottom')
            
            elif chart_type == 'line' and 'x' in data and 'y' in data:
                plt.plot(data['x'], data['y'], marker='o', linewidth=3, markersize=8)
                plt.title(data.get('title', 'Chart'), fontsize=16, fontweight='bold')
                plt.xlabel(data.get('xlabel', ''), fontsize=12)
                plt.ylabel(data.get('ylabel', ''), fontsize=12)
                plt.grid(True, alpha=0.3)
            
            elif chart_type == 'pie' and 'labels' in data and 'values' in data:
                plt.pie(data['values'], labels=data['labels'], autopct='%1.1f%%',
                       colors=CHART_COLORS[:len(data['labels'])], startangle=90)
                plt.title(data.get('title', 'Chart'), fontsize=16, fontweight='bold')
            
            plt.tight_layout()
            
            # Convert to base64
            buffer = io.BytesIO()
            plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
            buffer.seek(0)
            image_base64 = base64.b64encode(buffer.getvalue()).decode()
            plt.close()
            
            return f"data:image/png;base64,{image_base64}"
            
        except Exception as e:
            logger.error(f"Matplotlib chart generation failed: {e}")
            plt.close()
            return None