"""Spending and Debt Charts — Hyperscaler Capex, Tech Debt, NVIDIA Revenue"""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from src.data.ai_infrastructure import (
    hyperscaler_capex_annual,
    hyperscaler_capex_quarterly,
    hyperscaler_capex_meta,
)
from src.utils.styling import get_theme, EXPORT_DPI, get_company_colors, GRAY_DARK
from src.utils.export import save_chart_tight


def plot_hyperscaler_capex() -> str:
    plt.rcParams.update(get_theme())
    fig, ax = plt.subplots(figsize=(14, 8))

    company_colors = get_company_colors()
    companies = ["Microsoft", "Alphabet", "Meta", "Amazon"]

    # Build combined timeline: annual 2020-2023 + quarterly 2024-Q1 through 2026-Q1
    # Periods: 2020, 2021, 2022, 2023, 2024-Q1, 2024-Q2, 2024-Q3, 2024-Q4,
    #           2025-Q1, 2025-Q2, 2026-Q1
    annual_years = [2020, 2021, 2022, 2023]
    quarterly_periods = [
        (2024, "Q1"), (2024, "Q2"), (2024, "Q3"), (2024, "Q4"),
        (2025, "Q1"), (2025, "Q2"),
        (2026, "Q1"),
    ]

    x_labels = [str(y) for y in annual_years] + [
        f"{y}-{q}" for y, q in quarterly_periods
    ]
    n_periods = len(x_labels)
    x_positions = list(range(n_periods))

    # Organize annual data by company (2020-2023)
    annual_data = {c: {} for c in companies}
    for entry in hyperscaler_capex_annual:
        if entry["year"] in annual_years:
            annual_data[entry["company"]][entry["year"]] = entry["capex_billions"]

    # Organize quarterly data by company
    quarterly_data = {c: {} for c in companies}
    for entry in hyperscaler_capex_quarterly:
        key = (entry["year"], entry["quarter"])
        quarterly_data[entry["company"]][key] = entry["capex_billions"]

    # Build per-company value arrays
    data = {}
    for c in companies:
        vals = []
        # Annual portion
        for y in annual_years:
            vals.append(annual_data[c].get(y, 0))
        # Quarterly portion
        for y, q in quarterly_periods:
            vals.append(quarterly_data[c].get((y, q), 0))
        data[c] = np.array(vals)

    # Stacked area
    y_offset = np.zeros(n_periods)
    for c in companies:
        values = data[c]
        ax.fill_between(x_positions, y_offset, y_offset + values,
                        alpha=0.8, color=company_colors.get(c, "#666"), label=c)
        ax.plot(x_positions, y_offset + values, color=company_colors.get(c, "#666"), linewidth=1)
        y_offset += values

    # Total annotations
    for i, (pos, label, total) in enumerate(zip(x_positions, x_labels, y_offset)):
        total_label = f"${total:.0f}B"
        # Mark 2026 as projected
        if label.startswith("2026"):
            total_label += "*"
        # Only annotate every period (avoid crowding)
        ax.text(pos, total + 4, total_label, ha="center", fontsize=9, fontweight="bold")

    # Dashed vertical line between last quarterly data (2026-Q1) and remaining 2026
    ax.axvline(x=10.5, color=GRAY_DARK, linestyle="--", alpha=0.4, linewidth=1)

    ax.set_xticks(x_positions)
    ax.set_xticklabels(x_labels, rotation=45, ha="right", fontsize=9)
    ax.set_title("Hyperscaler AI Infrastructure Capex — 2020 to 2026-Q1",
                 fontsize=16, fontweight="bold")
    ax.set_xlabel("Period", fontsize=12)
    ax.set_ylabel("Capex (Billions USD)", fontsize=12)
    ax.legend(loc="upper left", fontsize=10)
    ax.grid(True, alpha=0.3)
    ax.text(0.02, 0.97, "*2026 = guided/projected", transform=ax.transAxes,
            fontsize=9, style="italic", color="gray", va="top")

    # Granularity note
    ax.text(0.50, 0.03, "2020-2023: annual | 2024-2026-Q1: quarterly",
            transform=ax.transAxes, fontsize=9, style="italic",
            color="gray", ha="center", va="bottom")

    # AI-related capex share note
    ax.text(0.98, 0.03, "80-90% of 2025/2026 capex is AI-related",
            transform=ax.transAxes, fontsize=9, style="italic",
            color="gray", ha="right", va="bottom")

    path = save_chart_tight(fig, "05_hyperscaler_capex.png")
    plt.close(fig)
    return path


def main():
    path = plot_hyperscaler_capex()
    print(f"Chart saved: {path}")


if __name__ == "__main__":
    main()
