#!/usr/bin/env python3
"""
Step 3: Financial Metrics & Unit Economics
==========================================

Calculate and visualize key unit economic indicators:
- Gross Margin %
- Operating Margin %
- YoY Revenue Growth
- Rule of 40

Compare Toast's efficiency against legacy peers (VYX, PAR, FISV).
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from pathlib import Path
import warnings

warnings.filterwarnings('ignore')

# Set reproducibility
np.random.seed(42)

# Configure paths
SESSION_DIR = Path('/app/sandbox/session_20260203_085615_5f10d3d2357f')
RESULTS_DIR = SESSION_DIR / 'results'
FIGURES_DIR = SESSION_DIR / 'figures'

# Configure matplotlib
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.linewidth'] = 0.5
plt.style.use('seaborn-v0_8-whitegrid')

# Company colors for consistent visualization
COMPANY_COLORS = {
    'TOST': '#FF6B35',  # Orange - Toast (highlight)
    'VYX': '#4ECDC4',   # Teal - NCR Voyix
    'PAR': '#556B2F',   # Olive - PAR Technology
    'FISV': '#3498DB',  # Blue - Fiserv
}

COMPANY_NAMES = {
    'TOST': 'Toast Inc.',
    'VYX': 'NCR Voyix',
    'PAR': 'PAR Technology',
    'FISV': 'Fiserv Inc.'
}


def load_and_clean_data():
    """Load quarterly financial data and perform cleaning."""
    print("=" * 60)
    print("STEP 3: FINANCIAL METRICS & UNIT ECONOMICS")
    print("=" * 60)

    print("\n[1/6] Loading quarterly financial data...")
    df = pd.read_csv(RESULTS_DIR / 'financial_data_quarterly.csv')
    print(f"  Loaded {len(df)} rows, {len(df.columns)} columns")
    print(f"  Columns: {list(df.columns)}")

    # Convert Date to datetime
    df['Date'] = pd.to_datetime(df['Date'])

    # Filter out rows with missing key data
    df_clean = df.dropna(subset=['Total Revenue', 'Gross Profit', 'Operating Income'])
    print(f"  After cleaning: {len(df_clean)} rows with complete financial data")

    # Show tickers available
    tickers = df_clean['Ticker'].unique()
    print(f"  Tickers with data: {list(tickers)}")

    return df_clean.sort_values(['Ticker', 'Date'])


def calculate_margins(df):
    """Calculate Gross Margin % and Operating Margin %."""
    print("\n[2/6] Calculating margin metrics...")

    # Calculate margins
    df['Gross Margin %'] = (df['Gross Profit'] / df['Total Revenue']) * 100
    df['Operating Margin %'] = (df['Operating Income'] / df['Total Revenue']) * 100

    # Show summary
    for ticker in df['Ticker'].unique():
        ticker_data = df[df['Ticker'] == ticker]
        latest = ticker_data.iloc[-1] if len(ticker_data) > 0 else None
        if latest is not None:
            print(f"  {ticker}: Gross Margin = {latest['Gross Margin %']:.1f}%, "
                  f"Operating Margin = {latest['Operating Margin %']:.1f}%")

    return df


def calculate_yoy_revenue_growth(df):
    """Calculate Year-over-Year revenue growth (4-quarter lag)."""
    print("\n[3/6] Calculating YoY Revenue Growth...")

    df = df.copy()
    df['Revenue Growth YoY %'] = np.nan

    for ticker in df['Ticker'].unique():
        ticker_mask = df['Ticker'] == ticker
        ticker_data = df[ticker_mask].sort_values('Date')

        # Get revenue values
        revenues = ticker_data['Total Revenue'].values
        dates = ticker_data['Date'].values

        # Calculate 4-quarter lag growth
        growth_values = []
        for i, (current_rev, current_date) in enumerate(zip(revenues, dates)):
            if i >= 4:
                prior_rev = revenues[i - 4]
                if prior_rev > 0:
                    growth = ((current_rev - prior_rev) / prior_rev) * 100
                    growth_values.append(growth)
                else:
                    growth_values.append(np.nan)
            else:
                growth_values.append(np.nan)

        # Assign back to dataframe
        df.loc[ticker_mask, 'Revenue Growth YoY %'] = growth_values

        # Report available growth data
        valid_growth = [g for g in growth_values if not np.isnan(g)]
        if valid_growth:
            print(f"  {ticker}: Latest YoY Growth = {valid_growth[-1]:.1f}% "
                  f"({len(valid_growth)} quarters with YoY data)")
        else:
            print(f"  {ticker}: Insufficient history for YoY calculation")

    return df


def calculate_rule_of_40(df):
    """Calculate Rule of 40 = Revenue Growth (YoY) + Operating Margin."""
    print("\n[4/6] Calculating Rule of 40...")

    df['Rule of 40'] = df['Revenue Growth YoY %'] + df['Operating Margin %']

    for ticker in df['Ticker'].unique():
        ticker_data = df[df['Ticker'] == ticker]
        latest_rule = ticker_data['Rule of 40'].dropna()
        if len(latest_rule) > 0:
            latest_val = latest_rule.iloc[-1]
            status = "PASS" if latest_val >= 40 else "BELOW"
            print(f"  {ticker}: Rule of 40 = {latest_val:.1f} ({status})")
        else:
            print(f"  {ticker}: Rule of 40 = N/A (insufficient data)")

    return df


def plot_gross_margin_trends(df):
    """Create line chart of Gross Margin % over time."""
    print("\n[5/6] Generating visualizations...")
    print("  Creating gross_margin_trends.png...")

    fig, ax = plt.subplots(figsize=(10, 6))

    for ticker in sorted(df['Ticker'].unique()):
        ticker_data = df[df['Ticker'] == ticker].sort_values('Date')
        ax.plot(ticker_data['Date'], ticker_data['Gross Margin %'],
                marker='o', linewidth=2.5 if ticker == 'TOST' else 1.5,
                markersize=8 if ticker == 'TOST' else 5,
                label=f"{ticker} ({COMPANY_NAMES.get(ticker, ticker)})",
                color=COMPANY_COLORS.get(ticker, 'gray'),
                alpha=0.9 if ticker == 'TOST' else 0.7)

    ax.set_xlabel('Quarter', fontsize=11)
    ax.set_ylabel('Gross Margin (%)', fontsize=11)
    ax.set_title('Gross Margin Trends by Company', fontsize=14, fontweight='bold')
    ax.legend(loc='upper left', fontsize=9)
    ax.yaxis.set_major_formatter(mticker.PercentFormatter(decimals=0))
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(FIGURES_DIR / 'gross_margin_trends.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("    Saved: figures/gross_margin_trends.png")


def plot_operating_margin_trends(df):
    """Create line chart of Operating Margin % over time."""
    print("  Creating operating_margin_trends.png...")

    fig, ax = plt.subplots(figsize=(10, 6))

    for ticker in sorted(df['Ticker'].unique()):
        ticker_data = df[df['Ticker'] == ticker].sort_values('Date')
        ax.plot(ticker_data['Date'], ticker_data['Operating Margin %'],
                marker='o', linewidth=2.5 if ticker == 'TOST' else 1.5,
                markersize=8 if ticker == 'TOST' else 5,
                label=f"{ticker} ({COMPANY_NAMES.get(ticker, ticker)})",
                color=COMPANY_COLORS.get(ticker, 'gray'),
                alpha=0.9 if ticker == 'TOST' else 0.7)

    # Add zero line
    ax.axhline(y=0, color='black', linestyle='--', linewidth=0.8, alpha=0.5)

    ax.set_xlabel('Quarter', fontsize=11)
    ax.set_ylabel('Operating Margin (%)', fontsize=11)
    ax.set_title('Operating Margin Trends by Company', fontsize=14, fontweight='bold')
    ax.legend(loc='upper left', fontsize=9)
    ax.yaxis.set_major_formatter(mticker.PercentFormatter(decimals=0))
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(FIGURES_DIR / 'operating_margin_trends.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("    Saved: figures/operating_margin_trends.png")


def plot_rule_of_40_scatter(df):
    """Create Rule of 40 scatter plot for most recent quarter."""
    print("  Creating rule_of_40_scatter.png...")

    # Get most recent quarter data for each ticker
    latest_data = []
    for ticker in df['Ticker'].unique():
        ticker_data = df[df['Ticker'] == ticker].sort_values('Date')
        # Get latest row with Rule of 40 data
        valid_rows = ticker_data.dropna(subset=['Rule of 40'])
        if len(valid_rows) > 0:
            latest = valid_rows.iloc[-1]
            latest_data.append({
                'Ticker': ticker,
                'Revenue Growth YoY %': latest['Revenue Growth YoY %'],
                'Operating Margin %': latest['Operating Margin %'],
                'Rule of 40': latest['Rule of 40'],
                'Date': latest['Date']
            })

    if not latest_data:
        print("    WARNING: No companies have Rule of 40 data. Skipping scatter plot.")
        return

    latest_df = pd.DataFrame(latest_data)

    fig, ax = plt.subplots(figsize=(10, 8))

    # Plot each company
    for _, row in latest_df.iterrows():
        ticker = row['Ticker']
        ax.scatter(row['Revenue Growth YoY %'], row['Operating Margin %'],
                   s=300 if ticker == 'TOST' else 200,
                   c=COMPANY_COLORS.get(ticker, 'gray'),
                   edgecolors='black', linewidth=1.5,
                   alpha=0.9, zorder=5)
        ax.annotate(ticker,
                   (row['Revenue Growth YoY %'], row['Operating Margin %']),
                   xytext=(8, 8), textcoords='offset points',
                   fontsize=11, fontweight='bold')

    # Add Rule of 40 threshold line
    x_range = np.array([-30, 50])
    y_rule40 = 40 - x_range  # Rule of 40: growth + margin = 40
    ax.plot(x_range, y_rule40, 'r--', linewidth=2, label='Rule of 40 Threshold', alpha=0.7)
    ax.fill_between(x_range, y_rule40, 60, alpha=0.1, color='green', label='Above Rule of 40')
    ax.fill_between(x_range, y_rule40, -30, alpha=0.1, color='red', label='Below Rule of 40')

    ax.set_xlabel('Revenue Growth YoY (%)', fontsize=12)
    ax.set_ylabel('Operating Margin (%)', fontsize=12)
    ax.set_title('Rule of 40 Analysis: Revenue Growth vs Operating Margin\n(Most Recent Quarter)',
                 fontsize=14, fontweight='bold')
    ax.legend(loc='lower right', fontsize=9)
    ax.grid(True, alpha=0.3)
    ax.axhline(y=0, color='gray', linestyle='-', linewidth=0.5, alpha=0.5)
    ax.axvline(x=0, color='gray', linestyle='-', linewidth=0.5, alpha=0.5)

    # Set reasonable axis limits
    ax.set_xlim(-35, 55)
    ax.set_ylim(-25, 40)

    plt.tight_layout()
    plt.savefig(FIGURES_DIR / 'rule_of_40_scatter.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("    Saved: figures/rule_of_40_scatter.png")


def save_results(df):
    """Save computed metrics to CSV."""
    print("\n[6/6] Saving results...")

    # Select and order columns for output
    output_cols = [
        'Date', 'Ticker', 'Total Revenue', 'Gross Profit', 'Operating Income',
        'Gross Margin %', 'Operating Margin %', 'Revenue Growth YoY %', 'Rule of 40'
    ]

    output_df = df[output_cols].copy()
    output_path = RESULTS_DIR / 'financial_metrics_calculated.csv'
    output_df.to_csv(output_path, index=False)
    print(f"  Saved: {output_path}")
    print(f"  Records: {len(output_df)} rows")

    return output_df


def generate_summary_insights(df):
    """Generate key insights for README update."""
    print("\n" + "=" * 60)
    print("KEY INSIGHTS SUMMARY")
    print("=" * 60)

    insights = {}

    # 1. Who has the highest gross margins?
    latest_quarters = df.groupby('Ticker').last().reset_index()
    gross_margin_leader = latest_quarters.loc[latest_quarters['Gross Margin %'].idxmax()]
    insights['gross_margin_leader'] = {
        'ticker': gross_margin_leader['Ticker'],
        'value': gross_margin_leader['Gross Margin %']
    }
    print(f"\n1. HIGHEST GROSS MARGIN:")
    print(f"   {gross_margin_leader['Ticker']}: {gross_margin_leader['Gross Margin %']:.1f}%")

    # Rank all companies by gross margin
    gm_ranking = latest_quarters[['Ticker', 'Gross Margin %']].sort_values('Gross Margin %', ascending=False)
    insights['gross_margin_ranking'] = gm_ranking.to_dict('records')
    ranking_str = ', '.join([f"{r['Ticker']}({r['Gross Margin %']:.1f}%)" for r in gm_ranking.to_dict('records')])
    print(f"   Ranking: {ranking_str}")

    # 2. Is Toast showing margin expansion (operating leverage)?
    toast_data = df[df['Ticker'] == 'TOST'].sort_values('Date')
    if len(toast_data) >= 2:
        toast_first = toast_data.iloc[0]['Operating Margin %']
        toast_latest = toast_data.iloc[-1]['Operating Margin %']
        margin_change = toast_latest - toast_first
        insights['toast_margin_expansion'] = {
            'first_quarter': toast_first,
            'latest_quarter': toast_latest,
            'change': margin_change,
            'expanding': margin_change > 0
        }
        trend = "EXPANDING" if margin_change > 0 else "CONTRACTING"
        print(f"\n2. TOAST OPERATING LEVERAGE:")
        print(f"   Q1: {toast_first:.1f}% -> Latest: {toast_latest:.1f}% ({'+' if margin_change > 0 else ''}{margin_change:.1f}pp)")
        print(f"   Status: {trend}")

    # 3. Rule of 40 comparison
    rule40_data = latest_quarters.dropna(subset=['Rule of 40'])
    if len(rule40_data) > 0:
        insights['rule_of_40'] = {}
        print(f"\n3. RULE OF 40 COMPARISON:")
        for _, row in rule40_data.sort_values('Rule of 40', ascending=False).iterrows():
            ticker = row['Ticker']
            rule40 = row['Rule of 40']
            status = "PASS" if rule40 >= 40 else "BELOW"
            insights['rule_of_40'][ticker] = {'value': rule40, 'status': status}
            print(f"   {ticker}: {rule40:.1f} ({status})")

        # Compare Toast vs legacy vendors
        toast_rule40 = insights['rule_of_40'].get('TOST', {}).get('value', np.nan)
        vyx_rule40 = insights['rule_of_40'].get('VYX', {}).get('value', np.nan)
        par_rule40 = insights['rule_of_40'].get('PAR', {}).get('value', np.nan)

        if not np.isnan(toast_rule40):
            print(f"\n   TOAST vs LEGACY VENDORS:")
            if not np.isnan(vyx_rule40):
                diff_vyx = toast_rule40 - vyx_rule40
                print(f"   TOST vs VYX: {'+' if diff_vyx > 0 else ''}{diff_vyx:.1f}pp advantage")
            if not np.isnan(par_rule40):
                diff_par = toast_rule40 - par_rule40
                print(f"   TOST vs PAR: {'+' if diff_par > 0 else ''}{diff_par:.1f}pp advantage")

    return insights


def main():
    """Main execution function."""
    # Load and clean data
    df = load_and_clean_data()

    # Calculate metrics
    df = calculate_margins(df)
    df = calculate_yoy_revenue_growth(df)
    df = calculate_rule_of_40(df)

    # Generate visualizations
    plot_gross_margin_trends(df)
    plot_operating_margin_trends(df)
    plot_rule_of_40_scatter(df)

    # Save results
    output_df = save_results(df)

    # Generate insights
    insights = generate_summary_insights(df)

    print("\n" + "=" * 60)
    print("STEP 3 COMPLETED SUCCESSFULLY")
    print("=" * 60)
    print("\nOutputs Generated:")
    print("  - results/financial_metrics_calculated.csv")
    print("  - figures/gross_margin_trends.png")
    print("  - figures/operating_margin_trends.png")
    print("  - figures/rule_of_40_scatter.png")

    return output_df, insights


if __name__ == '__main__':
    main()