# visualization.py

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from typing import Dict, List, Union

class Visualizer:
    """Creates visualizations for concentration analysis"""
    
    def __init__(self):
        """Initialize the visualizer"""
        self.color_scheme = {
            'primary': '#1f77b4',
            'secondary': '#ff7f0e',
            'accent': '#2ca02c',
            'warning': '#d62728',
            'neutral': '#7f7f7f'
        }
    
    def create_concentration_chart(self, df: pd.DataFrame, top_n: int = 10) -> go.Figure:
        """Create bar chart showing top customers by revenue percentage"""
        top_customers = df.head(top_n).copy()
        
        fig = go.Figure()
        
        fig.add_trace(go.Bar(
            x=top_customers['customer'],
            y=top_customers['percentage'],
            text=top_customers['percentage'].apply(lambda x: f'{x:.1f}%'),
            textposition='outside',
            marker_color=self.color_scheme['primary'],
            hovertemplate='<b>%{x}</b><br>Revenue: $%{customdata:,.0f}<br>Percentage: %{y:.1f}%<extra></extra>',
            customdata=top_customers['revenue']
        ))
        
        fig.update_layout(
            title=f'Top {top_n} Customers by Revenue Percentage',
            xaxis_title='Customer',
            yaxis_title='Percentage of Total Revenue',
            template='plotly_white',
            showlegend=False,
            xaxis={'categoryorder': 'total descending'},
            yaxis={'range': [0, max(top_customers['percentage']) * 1.1]}
        )
        
        fig.update_xaxes(tickangle=45)
        
        return fig
    
    def create_pareto_chart(self, df: pd.DataFrame) -> go.Figure:
        """Create Pareto chart showing cumulative revenue distribution"""
        fig = make_subplots(specs=[[{"secondary_y": True}]])
        
        # Bar chart for individual percentages
        fig.add_trace(
            go.Bar(
                name='Revenue %',
                x=list(range(1, len(df) + 1)),
                y=df['percentage'],
                marker_color=self.color_scheme['primary'],
                opacity=0.7,
                hovertemplate='Customer #%{x}<br>Revenue %: %{y:.1f}%<extra></extra>'
            ),
            secondary_y=False
        )
        
        # Line chart for cumulative percentage
        fig.add_trace(
            go.Scatter(
                name='Cumulative %',
                x=list(range(1, len(df) + 1)),
                y=df['cumulative'],
                mode='lines+markers',
                line=dict(color=self.color_scheme['secondary'], width=3),
                marker=dict(size=8),
                hovertemplate='Customer #%{x}<br>Cumulative %: %{y:.1f}%<extra></extra>'
            ),
            secondary_y=True
        )
        
        # Add 80/20 reference line
        fig.add_hline(y=80, line_dash="dash", line_color="red", secondary_y=True,
                      annotation_text="80% Revenue", annotation_position="right")
        
        fig.update_layout(
            title='Customer Revenue Distribution (Pareto Analysis)',
            xaxis_title='Customer Rank',
            template='plotly_white',
            showlegend=True,
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
        )
        
        fig.update_yaxes(title_text="Individual Revenue %", secondary_y=False, range=[0, max(df['percentage']) * 1.1])
        fig.update_yaxes(title_text="Cumulative Revenue %", secondary_y=True, range=[0, 105])
        
        return fig
    
    def create_hhi_gauge(self, hhi: float) -> go.Figure:
        """Create gauge chart for HHI score"""
        fig = go.Figure(go.Indicator(
            mode="gauge+number",
            value=hhi,
            domain={'x': [0, 1], 'y': [0, 1]},
            title={'text': "HHI Score"},
            gauge={
                'axis': {'range': [0, 10000]},
                'bar': {'color': self.color_scheme['primary']},
                'steps': [
                    {'range': [0, 1500], 'color': '#2ecc71'},
                    {'range': [1500, 2500], 'color': '#f39c12'},
                    {'range': [2500, 10000], 'color': '#e74c3c'}
                ],
                'threshold': {
                    'line': {'color': "red", 'width': 4},
                    'thickness': 0.75,
                    'value': hhi
                }
            }
        ))
        
        fig.update_layout(
            height=400,
            template='plotly_white'
        )
        
        return fig
    
    def create_revenue_distribution_pie(self, df: pd.DataFrame, top_n: int = 10) -> go.Figure:
        """Create pie chart showing revenue distribution"""
        top_customers = df.head(top_n).copy()
        others = pd.DataFrame({
            'customer': ['Others'],
            'revenue': [df.iloc[top_n:]['revenue'].sum()],
            'percentage': [df.iloc[top_n:]['percentage'].sum()]
        })
        
        if not others['revenue'].iloc[0] == 0:
            chart_data = pd.concat([top_customers, others])
        else:
            chart_data = top_customers
        
        fig = go.Figure(data=[go.Pie(
            labels=chart_data['customer'],
            values=chart_data['revenue'],
            hole=.3,
            textinfo='label+percent',
            textposition='outside',
            hovertemplate='<b>%{label}</b><br>Revenue: $%{value:,.0f}<br>Percentage: %{percent}<extra></extra>'
        )])
        
        fig.update_layout(
            title=f'Revenue Distribution (Top {top_n} + Others)',
            template='plotly_white'
        )
        
        return fig
    
    def create_customer_segmentation(self, df: pd.DataFrame) -> go.Figure:
        """Create customer segmentation visualization"""
        segments = {'A': 0, 'B': 0, 'C': 0}
        cumulative_revenue = 0
        total_revenue = df['revenue'].sum()
        
        segment_data = []
        
        for idx, row in df.iterrows():
            cumulative_revenue += row['revenue']
            cumulative_percent = (cumulative_revenue / total_revenue) * 100
            
            if cumulative_percent <= 80:
                segment = 'A (Top 80%)'
                color = self.color_scheme['accent']
            elif cumulative_percent <= 95:
                segment = 'B (Next 15%)'
                color = self.color_scheme['primary']
            else:
                segment = 'C (Bottom 5%)'
                color = self.color_scheme['warning']
            
            segment_data.append({
                'customer': row['customer'],
                'revenue': row['revenue'],
                'segment': segment,
                'color': color
            })
        
        seg_df = pd.DataFrame(segment_data)
        
        fig = go.Figure()
        
        for segment in seg_df['segment'].unique():
            segment_customers = seg_df[seg_df['segment'] == segment]
            fig.add_trace(go.Bar(
                name=segment,
                x=segment_customers['customer'],
                y=segment_customers['revenue'],
                marker_color=segment_customers['color'].iloc[0],
                hovertemplate='<b>%{x}</b><br>Revenue: $%{y:,.0f}<br>Segment: ' + segment + '<extra></extra>'
            ))
        
        fig.update_layout(
            title='Customer Segmentation Analysis',
            xaxis_title='Customer',
            yaxis_title='Revenue ($)',
            barmode='stack',
            template='plotly_white',
            showlegend=True,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        fig.update_xaxes(tickangle=45, showticklabels=False)
        
        return fig
    
    def create_risk_assessment_chart(self, metrics: Dict) -> go.Figure:
        """Create risk assessment visualization"""
        # Define risk factors and their scores
        risk_factors = []
        
        # HHI Risk
        hhi = metrics['hhi']
        if hhi > 2500:
            hhi_risk = 100
        elif hhi > 1500:
            hhi_risk = 60
        else:
            hhi_risk = 30
        risk_factors.append(('HHI Concentration', hhi_risk))
        
        # Top Customer Risk
        top_customer = metrics['top_customer_percent']
        if top_customer > 25:
            customer_risk = 100
        elif top_customer > 15:
            customer_risk = 70
        elif top_customer > 10:
            customer_risk = 40
        else:
            customer_risk = 20
        risk_factors.append(('Top Customer Dependency', customer_risk))
        
        # Top 5 Customers Risk
        top_5 = metrics['top_5_percent']
        if top_5 > 60:
            top5_risk = 90
        elif top_5 > 40:
            top5_risk = 60
        else:
            top5_risk = 30
        risk_factors.append(('Top 5 Concentration', top5_risk))
        
        # Customer Count Risk
        count = metrics['customer_count']
        if count < 10:
            count_risk = 100
        elif count < 20:
            count_risk = 70
        elif count < 50:
            count_risk = 40
        else:
            count_risk = 20
        risk_factors.append(('Limited Customer Base', count_risk))
        
        # Gini Coefficient Risk
        gini = metrics['gini_coefficient']
        if gini > 0.7:
            gini_risk = 90
        elif gini > 0.5:
            gini_risk = 60
        else:
            gini_risk = 30
        risk_factors.append(('Revenue Inequality', gini_risk))
        
        # Create radar chart
        categories = [factor[0] for factor in risk_factors]
        values = [factor[1] for factor in risk_factors]
        
        fig = go.Figure()
        
        fig.add_trace(go.Scatterpolar(
            r=values,
            theta=categories,
            fill='toself',
            fillcolor='rgba(31, 119, 180, 0.2)',
            line=dict(color=self.color_scheme['primary']),
            hovertemplate='%{theta}<br>Risk Score: %{r}<extra></extra>'
        ))
        
        fig.update_layout(
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, 100],
                    ticksuffix='%',
                    tickfont=dict(size=10)
                ),
                angularaxis=dict(
                    tickfont=dict(size=12),
                    rotation=90,
                    direction='clockwise'
                )
            ),
            showlegend=False,
            title='Risk Assessment Profile',
            template='plotly_white'
        )
        
        return fig
    
    def create_dimension_analysis_chart(self, df: pd.DataFrame) -> go.Figure:
        """Create charts showing concentration analysis by dimension"""
        
        # Create main chart showing HHI by dimension
        fig = go.Figure()
        
        # Color based on risk level
        colors = []
        for risk_level in df['risk_level']:
            if risk_level == 'High':
                colors.append('#e74c3c')
            elif risk_level == 'Moderate':
                colors.append('#f39c12')
            else:
                colors.append('#2ecc71')
        
        # Add HHI bars
        fig.add_trace(go.Bar(
            x=df['dimension_value'],
            y=df['hhi'],
            marker_color=colors,
            text=df['hhi'].apply(lambda x: f"{x:.0f}"),
            textposition='outside',
            hovertemplate='<b>%{x}</b><br>HHI: %{y:.0f}<br>Risk Level: %{customdata}<extra></extra>',
            customdata=df['risk_level']
        ))
        
        # Add reference lines for risk thresholds
        fig.add_hline(y=1500, line_dash="dash", line_color="green",
                     annotation_text="Low Risk Threshold", annotation_position="right")
        fig.add_hline(y=2500, line_dash="dash", line_color="red",
                     annotation_text="High Risk Threshold", annotation_position="right")
        
        fig.update_layout(
            title='Customer Concentration (HHI) by Dimension',
            xaxis_title='Dimension Value',
            yaxis_title='HHI Score',
            template='plotly_white',
            showlegend=False
        )
        
        # Update x-axis for better readability
        fig.update_xaxes(tickangle=45)
        
        return fig
    
    def create_time_trend_chart(self, df: pd.DataFrame) -> go.Figure:
        """Create time trend analysis chart"""
        
        # Create subplot with two y-axes
        fig = make_subplots(specs=[[{"secondary_y": True}]])
        
        # Add HHI line
        fig.add_trace(
            go.Scatter(
                x=df['period'],
                y=df['hhi'],
                mode='lines+markers',
                name='HHI Score',
                line=dict(color=self.color_scheme['primary'], width=3),
                marker=dict(size=10),
                hovertemplate='<b>%{x}</b><br>HHI: %{y:.0f}<br>Risk: %{customdata}<extra></extra>',
                customdata=df['risk_level']
            ),
            secondary_y=False
        )
        
        # Add top customer percentage line
        fig.add_trace(
            go.Scatter(
                x=df['period'],
                y=df['top_customer_percent'],
                mode='lines+markers',
                name='Top Customer %',
                line=dict(color=self.color_scheme['secondary'], width=2, dash='dot'),
                marker=dict(size=8),
                hovertemplate='<b>%{x}</b><br>Top Customer: %{y:.1f}%<extra></extra>'
            ),
            secondary_y=True
        )
        
        # Add risk threshold reference lines
        fig.add_hline(y=1500, line_dash="dash", line_color="green", secondary_y=False,
                     annotation_text="Low Risk", annotation_position="right")
        fig.add_hline(y=2500, line_dash="dash", line_color="red", secondary_y=False,
                     annotation_text="High Risk", annotation_position="right")
        
        # Update layout
        fig.update_layout(
            title='Concentration Trends Over Time',
            xaxis_title='Time Period',
            template='plotly_white',
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="x unified"
        )
        
        # Set y-axis titles
        fig.update_yaxes(title_text="HHI Score", secondary_y=False)
        fig.update_yaxes(title_text="Top Customer %", secondary_y=True)
        
        return fig
    
    def create_erp_overview_dashboard(self, raw_data: pd.DataFrame, summary: Dict) -> go.Figure:
        """Create an overview dashboard of the ERP data"""
        
        # Check for necessary columns
        date_col = None
        for col in raw_data.columns:
            if 'date' in col.lower():
                date_col = col
                break
                
        # Create dashboard with multiple subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=(
                "Transactions by Month", 
                "Revenue by Channel", 
                "Revenue by Region", 
                "Top Products"
            ),
            specs=[
                [{"type": "bar"}, {"type": "pie"}],
                [{"type": "bar"}, {"type": "bar"}]
            ]
        )
        
        # 1. Transactions by Month (if date column available)
        if date_col:
            raw_data[date_col] = pd.to_datetime(raw_data[date_col])
            raw_data['month'] = raw_data[date_col].dt.strftime('%Y-%m')
            monthly_counts = raw_data.groupby('month').size().reset_index(name='count')
            monthly_counts = monthly_counts.sort_values('month')
            
            fig.add_trace(
                go.Bar(
                    x=monthly_counts['month'],
                    y=monthly_counts['count'],
                    marker_color=self.color_scheme['primary']
                ),
                row=1, col=1
            )
        
        # 2. Revenue by Channel (if channel column available)
        channel_col = None
        for col in raw_data.columns:
            if 'channel' in col.lower():
                channel_col = col
                break
                
        if channel_col:
            channel_revenue = raw_data.groupby(channel_col)['amount'].sum().reset_index()
            channel_revenue = channel_revenue.sort_values('amount', ascending=False)
            
            fig.add_trace(
                go.Pie(
                    labels=channel_revenue[channel_col],
                    values=channel_revenue['amount'],
                    hole=.3
                ),
                row=1, col=2
            )
        
        # 3. Revenue by Region (if region column available)
        region_col = None
        for col in raw_data.columns:
            if 'region' in col.lower():
                region_col = col
                break
                
        if region_col:
            region_revenue = raw_data.groupby(region_col)['amount'].sum().reset_index()
            region_revenue = region_revenue.sort_values('amount', ascending=False)
            
            fig.add_trace(
                go.Bar(
                    x=region_revenue[region_col],
                    y=region_revenue['amount'],
                    marker_color=self.color_scheme['secondary']
                ),
                row=2, col=1
            )
        
        # 4. Top Products (if product column available)
        product_col = None
        for col in raw_data.columns:
            if 'product' in col.lower():
                product_col = col
                break
                
        if product_col:
            product_revenue = raw_data.groupby(product_col)['amount'].sum().reset_index()
            product_revenue = product_revenue.sort_values('amount', ascending=False).head(10)
            
            fig.add_trace(
                go.Bar(
                    x=product_revenue['amount'],
                    y=product_revenue[product_col],
                    orientation='h',
                    marker_color=self.color_scheme['accent']
                ),
                row=2, col=2
            )
        
        # Update layout
        fig.update_layout(
            height=800,
            title_text="ERP Data Overview",
            showlegend=False,
            template='plotly_white'
        )
        
        return fig
    
    def create_product_dependency_chart(self, dependency_data: Dict) -> go.Figure:
        """Create chart showing product dependency of top customers"""
        # Convert dictionary to DataFrame
        df = pd.DataFrame.from_dict(dependency_data, orient='index').reset_index()
        df = df.rename(columns={'index': 'customer'})
        
        # Create figure
        fig = make_subplots(specs=[[{"secondary_y": True}]])
        
        # Add product HHI bars
        fig.add_trace(
            go.Bar(
                x=df['customer'],
                y=df['product_hhi'],
                name='Product HHI',
                marker_color=self.color_scheme['primary'],
                hovertemplate='<b>%{x}</b><br>Product HHI: %{y:.0f}<br>Risk Level: %{customdata}<extra></extra>',
                customdata=df['risk_level']
            ),
            secondary_y=False
        )
        
        # Add top product percentage line
        fig.add_trace(
            go.Scatter(
                x=df['customer'],
                y=df['top_product_percent'],
                mode='lines+markers',
                name='Top Product %',
                line=dict(color=self.color_scheme['secondary'], width=3),
                marker=dict(size=8),
                hovertemplate='<b>%{x}</b><br>Top Product: %{customdata}<br>Percentage: %{y:.1f}%<extra></extra>',
                customdata=df['top_product']
            ),
            secondary_y=True
        )
        
        # Add risk thresholds
        fig.add_hline(y=1500, line_dash="dash", line_color="green", secondary_y=False,
                     annotation_text="Low Risk", annotation_position="right")
        fig.add_hline(y=2500, line_dash="dash", line_color="red", secondary_y=False,
                     annotation_text="High Risk", annotation_position="right")
        
        # Update layout
        fig.update_layout(
            title='Product Dependency Analysis of Top Customers',
            xaxis_title='Customer',
            template='plotly_white',
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
        )
        
        fig.update_yaxes(title_text="Product HHI", secondary_y=False)
        fig.update_yaxes(title_text="Top Product %", secondary_y=True, range=[0, 105])
        
        fig.update_xaxes(tickangle=45)
        
        return fig
    
    def create_revenue_volatility_chart(self, volatility_data: Dict) -> go.Figure:
        """Create chart showing revenue volatility of top customers"""
        customers = list(volatility_data.keys())
        
        # Create a figure with subplots
        fig = go.Figure()
        
        # Add revenue trend lines for each customer
        for customer in customers:
            # Get period data
            period_data = volatility_data[customer]['period_data']
            
            # Sort by period
            period_data = sorted(period_data, key=lambda x: x['period'])
            
            # Extract data for plotting
            periods = [p['period'] for p in period_data]
            revenues = [p['amount'] for p in period_data]
            
            # Calculate coefficient of variation
            cv = volatility_data[customer]['coefficient_of_variation']
            
            # Add line
            fig.add_trace(go.Scatter(
                x=periods,
                y=revenues,
                mode='lines+markers',
                name=f"{customer} (CV: {cv:.1f}%)",
                hovertemplate='<b>%{x}</b><br>Revenue: $%{y:,.2f}<extra></extra>'
            ))
        
        # Update layout
        fig.update_layout(
            title='Revenue Volatility of Top Customers',
            xaxis_title='Time Period',
            yaxis_title='Revenue',
            template='plotly_white',
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="x unified"
        )
        
        # Format y-axis as currency
        fig.update_yaxes(tickprefix='$', tickformat=',')
        
        return fig
    
    def create_geographic_exposure_chart(self, geographic_data: Dict) -> go.Figure:
        """Create geographic exposure visualization"""
        # Extract customer region data
        customer_regions = geographic_data['customer_regions']
        
        # Convert to DataFrame
        df = pd.DataFrame.from_dict(customer_regions, orient='index').reset_index()
        df = df.rename(columns={'index': 'customer'})
        
        # Create figure
        fig = make_subplots(specs=[[{"secondary_y": True}]])
        
        # Add geographic HHI bars
        fig.add_trace(
            go.Bar(
                x=df['customer'],
                y=df['geographic_hhi'],
                name='Geographic HHI',
                marker_color=self.color_scheme['primary'],
                hovertemplate='<b>%{x}</b><br>Geographic HHI: %{y:.0f}<br>Risk Level: %{customdata}<extra></extra>',
                customdata=df['risk_level']
            ),
            secondary_y=False
        )
        
        # Add primary region percentage line
        fig.add_trace(
            go.Scatter(
                x=df['customer'],
                y=df['primary_region_percent'],
                mode='lines+markers',
                name='Primary Region %',
                line=dict(color=self.color_scheme['secondary'], width=3),
                marker=dict(size=8),
                hovertemplate='<b>%{x}</b><br>Primary Region: %{customdata}<br>Percentage: %{y:.1f}%<extra></extra>',
                customdata=df['primary_region']
            ),
            secondary_y=True
        )
        
        # Add risk thresholds
        fig.add_hline(y=1500, line_dash="dash", line_color="green", secondary_y=False,
                     annotation_text="Low Risk", annotation_position="right")
        fig.add_hline(y=2500, line_dash="dash", line_color="red", secondary_y=False,
                     annotation_text="High Risk", annotation_position="right")
        
        # Update layout
        fig.update_layout(
            title='Geographic Exposure of Top Customers',
            xaxis_title='Customer',
            template='plotly_white',
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
        )
        
        fig.update_yaxes(title_text="Geographic HHI", secondary_y=False)
        fig.update_yaxes(title_text="Primary Region %", secondary_y=True, range=[0, 105])
        
        fig.update_xaxes(tickangle=45)
        
        return fig
    
    def create_forecast_chart(self, forecast_data: Dict) -> go.Figure:
        """Create revenue forecast chart for top customers"""
        # Create figure
        fig = go.Figure()
        
        # Check if forecast data is empty
        if not forecast_data:
            # Return an empty figure with a message
            fig.add_annotation(
                text="No forecast data available",
                xref="paper", yref="paper",
                x=0.5, y=0.5,
                showarrow=False,
                font=dict(size=16)
            )
            return fig
        
        # Add data for each customer
        for customer, data in forecast_data.items():
            try:
                # Extract historical data
                historical_periods = [h['period'] for h in data['historical']]
                historical_revenues = [h['amount'] for h in data['historical']]
                
                # Extract forecast data
                forecast_periods = [f['period'] for f in data['forecast']]
                forecast_revenues = [f['amount'] for f in data['forecast']]
                
                # Add historical line
                fig.add_trace(go.Scatter(
                    x=historical_periods,
                    y=historical_revenues,
                    mode='lines+markers',
                    name=f"{customer} (Historical)",
                    line=dict(color=self.color_scheme['primary']),
                    hovertemplate='<b>%{x}</b><br>Revenue: $%{y:,.2f}<extra></extra>'
                ))
                
                # Add forecast line
                fig.add_trace(go.Scatter(
                    x=forecast_periods,
                    y=forecast_revenues,
                    mode='lines+markers',
                    name=f"{customer} (Forecast)",
                    line=dict(color=self.color_scheme['primary'], dash='dash'),
                    marker=dict(symbol='circle-open'),
                    hovertemplate='<b>%{x}</b><br>Forecast: $%{y:,.2f}<extra></extra>'
                ))
                
                # Add trend annotation without using all_periods/all_revenues
                trend = data['trend']
                trend_color = '#00cc44' if trend == 'Increasing' else '#ff4b4b' if trend == 'Decreasing' else '#7f7f7f'
                
                # Only add annotation if there's forecast data
                if forecast_periods and forecast_revenues:
                    # Add annotation for trend at the last forecast point
                    fig.add_annotation(
                        x=forecast_periods[-1],
                        y=forecast_revenues[-1],
                        text=trend,
                        showarrow=True,
                        arrowhead=1,
                        arrowsize=1,
                        arrowwidth=2,
                        arrowcolor=trend_color,
                        font=dict(color=trend_color),
                        xanchor='left',
                        yanchor='bottom'
                    )
            except Exception as e:
                print(f"Error plotting forecast for {customer}: {str(e)}")
                continue
        
        # Update layout
        fig.update_layout(
            title='Revenue Forecast for Top Customers',
            xaxis_title='Time Period',
            yaxis_title='Revenue',
            template='plotly_white',
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="x unified"
        )
        
        # Add a vertical line between historical and forecast, but only if we have data
        if len(forecast_data) > 0:
            # Get the first customer's data to find the boundary
            try:
                first_customer = list(forecast_data.keys())[0]
                if data['historical'] and len(data['historical']) > 0:
                    historical_end = data['historical'][-1]['period']
                    
                    fig.add_vline(
                        x=historical_end,
                        line_dash="dot",
                        line_color="black",
                        annotation_text="Forecast Start",
                        annotation_position="top"
                    )
            except Exception as e:
                print(f"Error adding boundary line: {str(e)}")
        
        # Format y-axis as currency
        fig.update_yaxes(tickprefix='$', tickformat=',')
        
        return fig