visualization.py

"""
Visualization utilities for the Investor Decision Support System
Provides functions for creating interactive charts and plots
"""

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 typing import Dict, List, Tuple, Optional
import matplotlib.pyplot as plt
import seaborn as sns


class FinancialVisualizer:
    """Class for creating financial visualizations"""
    
    def __init__(self, theme: str = "plotly_white"):
        """
        Initialize visualizer
        
        Args:
            theme: Plotly theme to use
        """
        self.theme = theme
        self.colors = px.colors.qualitative.Set3
        
    def plot_efficient_frontier(self, efficient_frontier_data: pd.DataFrame, 
                              optimal_portfolio: Dict = None) -> go.Figure:
        """
        Create efficient frontier plot
        
        Args:
            efficient_frontier_data: DataFrame with portfolio data
            optimal_portfolio: Dictionary with optimal portfolio data
            
        Returns:
            Plotly figure
        """
        fig = go.Figure()
        
        # Plot random portfolios
        fig.add_trace(go.Scatter(
            x=efficient_frontier_data['volatility'],
            y=efficient_frontier_data['return'],
            mode='markers',
            marker=dict(
                size=8,
                color=efficient_frontier_data['sharpe_ratio'],
                colorscale='Viridis',
                showscale=True,
                colorbar=dict(title="Sharpe Ratio")
            ),
            name='Random Portfolios',
            hovertemplate='<b>Volatility:</b> %{x:.3f}<br>' +
                         '<b>Return:</b> %{y:.3f}<br>' +
                         '<b>Sharpe Ratio:</b> %{marker.color:.3f}<extra></extra>'
        ))
        
        # Plot optimal portfolio if provided
        if optimal_portfolio and 'metrics' in optimal_portfolio:
            metrics = optimal_portfolio['metrics']
            fig.add_trace(go.Scatter(
                x=[metrics['volatility']],
                y=[metrics['expected_return']],
                mode='markers',
                marker=dict(
                    size=15,
                    color='red',
                    symbol='star'
                ),
                name='Optimal Portfolio',
                hovertemplate='<b>Optimal Portfolio</b><br>' +
                             '<b>Volatility:</b> %{x:.3f}<br>' +
                             '<b>Return:</b> %{y:.3f}<br>' +
                             '<b>Sharpe Ratio:</b> ' + f"{metrics['sharpe_ratio']:.3f}<extra></extra>"
            ))
        
        fig.update_layout(
            title='Efficient Frontier',
            xaxis_title='Volatility (Risk)',
            yaxis_title='Expected Return',
            template=self.theme,
            height=500
        )
        
        return fig
    
    def plot_portfolio_weights(self, weights: Dict[str, float]) -> go.Figure:
        """
        Create pie chart for portfolio weights
        
        Args:
            weights: Dictionary with stock symbols and weights
            
        Returns:
            Plotly figure
        """
        symbols = list(weights.keys())
        weight_values = list(weights.values())
        
        fig = go.Figure(data=[go.Pie(
            labels=symbols,
            values=weight_values,
            hovertemplate='<b>%{label}</b><br>Weight: %{percent}<extra></extra>'
        )])
        
        fig.update_layout(
            title='Portfolio Allocation',
            template=self.theme,
            height=400
        )
        
        return fig
    
    def plot_capm_analysis(self, capm_results: pd.DataFrame, sml_data: pd.DataFrame) -> go.Figure:
        """
        Create CAPM Security Market Line plot
        
        Args:
            capm_results: DataFrame with CAPM analysis results
            sml_data: DataFrame with SML data points
            
        Returns:
            Plotly figure
        """
        fig = go.Figure()
        
        # Plot Security Market Line
        fig.add_trace(go.Scatter(
            x=sml_data['beta'],
            y=sml_data['expected_return'],
            mode='lines',
            name='Security Market Line',
            line=dict(color='blue', dash='dash')
        ))
        
        # Plot individual stocks
        for symbol in capm_results.index:
            fig.add_trace(go.Scatter(
                x=[capm_results.loc[symbol, 'beta']],
                y=[capm_results.loc[symbol, 'actual_return']],
                mode='markers+text',
                text=[symbol],
                textposition='top center',
                name=symbol,
                hovertemplate=f'<b>{symbol}</b><br>' +
                             '<b>Beta:</b> %{x:.3f}<br>' +
                             '<b>Actual Return:</b> %{y:.3f}<br>' +
                             '<b>Expected Return:</b> ' + f"{capm_results.loc[symbol, 'expected_return']:.3f}<extra></extra>"
            ))
        
        fig.update_layout(
            title='CAPM Analysis - Security Market Line',
            xaxis_title='Beta (Systematic Risk)',
            yaxis_title='Expected Return',
            template=self.theme,
            height=500
        )
        
        return fig
    
    def plot_stock_forecasts(self, symbol: str, historical_prices: pd.Series, 
                           forecast_data: pd.DataFrame) -> go.Figure:
        """
        Create stock price forecast plot
        
        Args:
            symbol: Stock symbol
            historical_prices: Historical price data
            forecast_data: Forecast data with confidence intervals
            
        Returns:
            Plotly figure
        """
        fig = go.Figure()
        
        # Plot historical prices
        fig.add_trace(go.Scatter(
            x=historical_prices.index,
            y=historical_prices.values,
            mode='lines',
            name='Historical Prices',
            line=dict(color='blue')
        ))
        
        # Plot forecast
        fig.add_trace(go.Scatter(
            x=forecast_data['date'],
            y=forecast_data['forecast'],
            mode='lines',
            name='Forecast',
            line=dict(color='red', dash='dash')
        ))
        
        # Plot confidence intervals if available
        if 'lower_bound' in forecast_data.columns and 'upper_bound' in forecast_data.columns:
            fig.add_trace(go.Scatter(
                x=forecast_data['date'],
                y=forecast_data['upper_bound'],
                mode='lines',
                line=dict(width=0),
                showlegend=False,
                hoverinfo='skip'
            ))
            
            fig.add_trace(go.Scatter(
                x=forecast_data['date'],
                y=forecast_data['lower_bound'],
                mode='lines',
                line=dict(width=0),
                fill='tonexty',
                fillcolor='rgba(255,0,0,0.2)',
                name='Confidence Interval',
                hoverinfo='skip'
            ))
        
        fig.update_layout(
            title=f'{symbol} Price Forecast',
            xaxis_title='Date',
            yaxis_title='Price',
            template=self.theme,
            height=500
        )
        
        return fig
    
    def plot_correlation_heatmap(self, correlation_matrix: pd.DataFrame) -> go.Figure:
        """
        Create correlation heatmap
        
        Args:
            correlation_matrix: Correlation matrix DataFrame
            
        Returns:
            Plotly figure
        """
        fig = go.Figure(data=go.Heatmap(
            z=correlation_matrix.values,
            x=correlation_matrix.columns,
            y=correlation_matrix.index,
            colorscale='RdBu',
            zmid=0,
            text=np.round(correlation_matrix.values, 2),
            texttemplate='%{text}',
            textfont={"size": 10},
            hovertemplate='<b>%{y} vs %{x}</b><br>Correlation: %{z:.3f}<extra></extra>'
        ))
        
        fig.update_layout(
            title='Stock Correlation Matrix',
            template=self.theme,
            height=500
        )
        
        return fig
    
    def plot_risk_return_scatter(self, individual_metrics: pd.DataFrame) -> go.Figure:
        """
        Create risk-return scatter plot for individual stocks
        
        Args:
            individual_metrics: DataFrame with individual stock metrics
            
        Returns:
            Plotly figure
        """
        fig = go.Figure()
        
        fig.add_trace(go.Scatter(
            x=individual_metrics['Volatility'],
            y=individual_metrics['Expected Return'],
            mode='markers+text',
            text=individual_metrics.index,
            textposition='top center',
            marker=dict(
                size=10,
                color=individual_metrics['Sharpe Ratio'],
                colorscale='Viridis',
                showscale=True,
                colorbar=dict(title="Sharpe Ratio")
            ),
            hovertemplate='<b>%{text}</b><br>' +
                         '<b>Volatility:</b> %{x:.3f}<br>' +
                         '<b>Expected Return:</b> %{y:.3f}<br>' +
                         '<b>Sharpe Ratio:</b> %{marker.color:.3f}<extra></extra>'
        ))
        
        fig.update_layout(
            title='Risk-Return Analysis',
            xaxis_title='Volatility (Risk)',
            yaxis_title='Expected Return',
            template=self.theme,
            height=500
        )
        
        return fig
    
    def plot_portfolio_performance(self, portfolio_metrics: Dict, 
                                 individual_metrics: pd.DataFrame) -> go.Figure:
        """
        Create portfolio performance comparison plot
        
        Args:
            portfolio_metrics: Dictionary with portfolio metrics
            individual_metrics: DataFrame with individual stock metrics
            
        Returns:
            Plotly figure
        """
        fig = go.Figure()
        
        # Plot individual stocks
        fig.add_trace(go.Scatter(
            x=individual_metrics['Volatility'],
            y=individual_metrics['Expected Return'],
            mode='markers',
            marker=dict(size=8, color='lightblue'),
            name='Individual Stocks',
            hovertemplate='<b>%{text}</b><br>' +
                         '<b>Volatility:</b> %{x:.3f}<br>' +
                         '<b>Expected Return:</b> %{y:.3f}<extra></extra>'
        ))
        
        # Plot portfolio
        fig.add_trace(go.Scatter(
            x=[portfolio_metrics['volatility']],
            y=[portfolio_metrics['expected_return']],
            mode='markers',
            marker=dict(size=15, color='red', symbol='star'),
            name='Portfolio',
            hovertemplate='<b>Portfolio</b><br>' +
                         '<b>Volatility:</b> %{x:.3f}<br>' +
                         '<b>Expected Return:</b> %{y:.3f}<br>' +
                         '<b>Sharpe Ratio:</b> ' + f"{portfolio_metrics['sharpe_ratio']:.3f}<extra></extra>"
        ))
        
        fig.update_layout(
            title='Portfolio vs Individual Stocks',
            xaxis_title='Volatility (Risk)',
            yaxis_title='Expected Return',
            template=self.theme,
            height=500
        )
        
        return fig
    
    def create_kpi_cards(self, metrics: Dict) -> List[Dict]:
        """
        Create KPI card data for Streamlit
        
        Args:
            metrics: Dictionary with financial metrics
            
        Returns:
            List of dictionaries with KPI data
        """
        kpi_cards = []
        
        # Define KPI mappings
        kpi_mappings = {
            'expected_return': {'label': 'Expected Return', 'format': '{:.2%}', 'color': 'green'},
            'volatility': {'label': 'Volatility', 'format': '{:.2%}', 'color': 'orange'},
            'sharpe_ratio': {'label': 'Sharpe Ratio', 'format': '{:.3f}', 'color': 'blue'},
            'beta': {'label': 'Beta', 'format': '{:.3f}', 'color': 'purple'},
            'alpha': {'label': 'Alpha', 'format': '{:.2%}', 'color': 'red'}
        }
        
        for key, value in metrics.items():
            if key in kpi_mappings:
                mapping = kpi_mappings[key]
                formatted_value = mapping['format'].format(value)
                kpi_cards.append({
                    'label': mapping['label'],
                    'value': formatted_value,
                    'color': mapping['color']
                })
        
        return kpi_cards