visualization.py

import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots

def plot_tariff_price_comparison(tariff_data, price_data):
    """
    Create a visualization comparing tariff rates and price changes.
    """
    # Process data for visualization
    # Aggregate by category
    tariff_by_category = tariff_data.groupby('Category')['Tariff_Rate'].mean().reset_index()
    
    # Calculate price change for each category
    price_changes = []
    
    for category in tariff_by_category['Category']:
        category_price_data = price_data[price_data['Category'] == category].sort_values('Year')
        
        if len(category_price_data) > 0:
            first_price = category_price_data['Price_Index'].iloc[0]
            last_price = category_price_data['Price_Index'].iloc[-1]
            
            price_change = ((last_price / first_price) - 1) * 100
            price_changes.append({
                'Category': category,
                'Price_Change': price_change
            })
    
    price_change_df = pd.DataFrame(price_changes)
    
    # Merge tariff and price change data
    comparison_data = pd.merge(
        tariff_by_category,
        price_change_df,
        on='Category',
        how='inner'
    )
    
    # Sort by tariff rate to make the visualization clearer
    comparison_data = comparison_data.sort_values('Tariff_Rate')
    
    # Create figure with secondary y-axis
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    
    # Add traces
    fig.add_trace(
        go.Bar(
            x=comparison_data['Category'],
            y=comparison_data['Tariff_Rate'],
            name="Average Tariff Rate (%)",
            marker_color='rgb(55, 83, 109)',
            opacity=0.7
        ),
        secondary_y=False,
    )
    
    fig.add_trace(
        go.Scatter(
            x=comparison_data['Category'],
            y=comparison_data['Price_Change'],
            name="Price Change (%)",
            mode='markers+lines',
            marker=dict(
                color='rgb(235, 52, 52)',
                size=10,
                line=dict(width=2)
            ),
            line=dict(width=3)
        ),
        secondary_y=True,
    )
    
    # Add figure title
    fig.update_layout(
        title_text="Tariff Rates vs. Price Changes by Category",
        hovermode="x unified",
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        height=600
    )
    
    # Set x-axis title
    fig.update_xaxes(
        title_text="Product Category",
        tickangle=45
    )
    
    # Set y-axes titles
    fig.update_yaxes(
        title_text="Average Tariff Rate (%)",
        secondary_y=False
    )
    fig.update_yaxes(
        title_text="Price Change (%)",
        secondary_y=True
    )
    
    return fig

def plot_historical_tariffs(historical_data):
    """
    Create a visualization of historical tariff rates over time.
    """
    # Calculate average tariff rates across all categories by year
    avg_tariff_by_year = historical_data.groupby('Year')['Tariff_Rate'].mean().reset_index()
    
    # Calculate average price index across all categories by year
    avg_price_by_year = historical_data.groupby('Year')['Price_Index'].mean().reset_index()
    
    # Create figure with secondary y-axis
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    
    # Add traces
    fig.add_trace(
        go.Scatter(
            x=avg_tariff_by_year['Year'],
            y=avg_tariff_by_year['Tariff_Rate'],
            name="Average Tariff Rate (%)",
            mode='lines+markers',
            line=dict(width=3, color='rgb(55, 83, 109)'),
            marker=dict(size=8)
        ),
        secondary_y=False,
    )
    
    fig.add_trace(
        go.Scatter(
            x=avg_price_by_year['Year'],
            y=avg_price_by_year['Price_Index'],
            name="Average Price Index",
            mode='lines+markers',
            line=dict(width=3, color='rgb(235, 52, 52)'),
            marker=dict(size=8)
        ),
        secondary_y=True,
    )
    
    # Add figure title
    fig.update_layout(
        title_text="Historical Tariff Rates and Price Indices Over Time",
        hovermode="x unified",
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        height=500
    )
    
    # Add annotations for key tariff events
    annotations = [
        dict(
            x=2018, y=avg_tariff_by_year[avg_tariff_by_year['Year'] == 2018]['Tariff_Rate'].values[0] + 1,
            xref="x", yref="y",
            text="Section 301 Tariffs<br>on Chinese Tech",
            showarrow=True,
            arrowhead=2,
            ax=0, ay=-40
        ),
        dict(
            x=2019, y=avg_tariff_by_year[avg_tariff_by_year['Year'] == 2019]['Tariff_Rate'].values[0] + 1,
            xref="x", yref="y",
            text="Expanded Tariffs<br>on Consumer Electronics",
            showarrow=True,
            arrowhead=2,
            ax=0, ay=-40
        ),
        dict(
            x=2022, y=avg_tariff_by_year[avg_tariff_by_year['Year'] == 2022]['Tariff_Rate'].values[0] + 1,
            xref="x", yref="y",
            text="Semiconductor<br>Focused Tariffs",
            showarrow=True,
            arrowhead=2,
            ax=0, ay=-40
        )
    ]
    
    fig.update_layout(annotations=annotations)
    
    # Set x-axis title
    fig.update_xaxes(
        title_text="Year",
        dtick=1  # Show every year
    )
    
    # Set y-axes titles
    fig.update_yaxes(
        title_text="Average Tariff Rate (%)",
        secondary_y=False
    )
    fig.update_yaxes(
        title_text="Average Price Index",
        secondary_y=True
    )
    
    return fig

def plot_category_impact(tariff_data, price_data, categories):
    """
    Create a visualization comparing the impact of tariffs across different categories.
    """
    if not categories:
        return go.Figure()
    
    # Filter data for selected categories
    filtered_tariff = tariff_data[tariff_data['Category'].isin(categories)]
    filtered_price = price_data[price_data['Category'].isin(categories)]
    
    # Calculate metrics for visualization
    impact_data = []
    
    for category in categories:
        cat_tariff = filtered_tariff[filtered_tariff['Category'] == category]
        cat_price = filtered_price[filtered_price['Category'] == category]
        
        if not cat_tariff.empty and not cat_price.empty:
            # Calculate average tariff
            avg_tariff = cat_tariff['Tariff_Rate'].mean()
            
            # Calculate tariff change
            tariff_by_year = cat_tariff.groupby('Year')['Tariff_Rate'].mean()
            first_tariff = tariff_by_year.iloc[0]
            last_tariff = tariff_by_year.iloc[-1]
            tariff_change = last_tariff - first_tariff
            
            # Calculate price change
            price_by_year = cat_price.groupby('Year')['Price_Index'].mean()
            first_price = price_by_year.iloc[0]
            last_price = price_by_year.iloc[-1]
            price_change = ((last_price / first_price) - 1) * 100
            
            # Calculate price sensitivity (price change per 1% tariff change)
            # Avoid division by zero
            if abs(tariff_change) > 0.001:
                price_sensitivity = price_change / tariff_change
            else:
                price_sensitivity = 0
            
            impact_data.append({
                'Category': category,
                'Avg_Tariff': avg_tariff,
                'Tariff_Change': tariff_change,
                'Price_Change': price_change,
                'Price_Sensitivity': price_sensitivity
            })
    
    impact_df = pd.DataFrame(impact_data)
    
    # Create bubble chart
    fig = px.scatter(
        impact_df,
        x='Tariff_Change',
        y='Price_Change',
        size='Avg_Tariff',
        color='Price_Sensitivity',
        hover_name='Category',
        text='Category',
        size_max=30,
        color_continuous_scale=px.colors.sequential.Viridis,
        labels={
            'Tariff_Change': 'Tariff Rate Change (percentage points)',
            'Price_Change': 'Price Change (%)',
            'Avg_Tariff': 'Average Tariff Rate (%)',
            'Price_Sensitivity': 'Price Sensitivity (% change per 1% tariff)'
        },
        title='Tariff Impact by Product Category'
    )
    
    # Update layout
    fig.update_layout(
        height=600,
        xaxis=dict(
            zeroline=True,
            zerolinewidth=1,
            zerolinecolor='gray',
            gridcolor='lightgray'
        ),
        yaxis=dict(
            zeroline=True,
            zerolinewidth=1,
            zerolinecolor='gray',
            gridcolor='lightgray'
        ),
        plot_bgcolor='white'
    )
    
    # Add a horizontal line at y=0
    fig.add_shape(
        type="line",
        xref="x",
        yref="y",
        x0=min(impact_df['Tariff_Change']) - 1,
        y0=0,
        x1=max(impact_df['Tariff_Change']) + 1,
        y1=0,
        line=dict(
            color="gray",
            width=1,
            dash="dash",
        )
    )
    
    # Add a vertical line at x=0
    fig.add_shape(
        type="line",
        xref="x",
        yref="y",
        x0=0,
        y0=min(impact_df['Price_Change']) - 5,
        x1=0,
        y1=max(impact_df['Price_Change']) + 5,
        line=dict(
            color="gray",
            width=1,
            dash="dash",
        )
    )
    
    return fig

def plot_price_sensitivity(tariff_data, price_data):
    """
    Create a visualization showing the relationship between tariff rates and prices.
    """
    # Process data for visualization
    # Group by year to get annual average tariff and price
    tariff_by_year = tariff_data.groupby('Year')['Tariff_Rate'].mean().reset_index()
    price_by_year = price_data.groupby('Year')['Price_Index'].mean().reset_index()
    
    # Merge the data
    sensitivity_data = pd.merge(
        tariff_by_year,
        price_by_year,
        on='Year',
        how='inner'
    )
    
    # Create scatter plot with trend line
    fig = px.scatter(
        sensitivity_data,
        x='Tariff_Rate',
        y='Price_Index',
        hover_name='Year',
        trendline='ols',
        trendline_color_override='red',
        labels={
            'Tariff_Rate': 'Tariff Rate (%)',
            'Price_Index': 'Price Index',
            'Year': 'Year'
        },
        title='Price Sensitivity to Tariff Rates'
    )
    
    # Update layout
    fig.update_layout(
        height=400,
        plot_bgcolor='white',
        xaxis=dict(gridcolor='lightgray'),
        yaxis=dict(gridcolor='lightgray')
    )
    
    # Add text annotations for each point (year)
    for i, row in sensitivity_data.iterrows():
        fig.add_annotation(
            x=row['Tariff_Rate'],
            y=row['Price_Index'],
            text=str(row['Year']),
            showarrow=False,
            xshift=10,
            yshift=10
        )
    
    return fig