from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parent
DATA = ROOT / "data"
IMAGES = ROOT / "images"

COLORS = {
    "navy": "#17365D",
    "blue": "#2F75B5",
    "orange": "#ED7D31",
    "gray": "#667085",
    "grid": "#D9E2F3",
    "shade": "#DCE6F1",
}


def finish(ax):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis="y", color=COLORS["grid"], linewidth=0.8, alpha=0.7)
    ax.set_axisbelow(True)


comparison = pd.read_csv(DATA / "model_comparison.csv")

# Adjusted R-squared comparison
pivot = comparison.pivot(index="outcome", columns="model", values="adjusted_r_squared")
outcomes = ["Gini", "Labor share"]
ols = [pivot.loc[o, "OLS"] for o in outcomes]
lp = [pivot.loc[o, "LP median regression"] for o in outcomes]

fig, ax = plt.subplots(figsize=(9.2, 5.4))
x = np.arange(len(outcomes))
width = 0.31
bars1 = ax.bar(x - width / 2, ols, width, label="OLS", color=COLORS["gray"])
bars2 = ax.bar(x + width / 2, lp, width, label="Tax-period LP", color=COLORS["blue"])
fig.suptitle("Tax-period LP improves adjusted R²", y=0.98,
             fontsize=17, weight="bold", color=COLORS["navy"])
fig.text(0.5, 0.925,
         "The improvement remains after penalizing additional period coefficients.",
         ha="center", fontsize=10.5, color=COLORS["gray"])
ax.set_ylabel("Adjusted R²")
ax.set_xticks(x, outcomes)
ax.set_ylim(0, 1.02)
ax.legend(frameon=False, loc="upper left")
for bars in (bars1, bars2):
    for bar in bars:
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.025,
                f"{bar.get_height():.3f}", ha="center", va="bottom", fontsize=10)
for i, (a, b) in enumerate(zip(ols, lp)):
    ax.text(i, max(a, b) + 0.11, f"+{b-a:.3f}", ha="center",
            color=COLORS["blue"], weight="bold", fontsize=10)
finish(ax)
fig.tight_layout(rect=[0, 0, 1, 0.88])
fig.savefig(IMAGES / "adjusted_r_squared_comparison.png", dpi=180, bbox_inches="tight")
plt.close(fig)


# Mean absolute error comparison, indexed so different outcome units can share a chart.
mae_pivot = comparison.pivot(index="outcome", columns="model", values="mean_absolute_error")
ols_index = [100.0, 100.0]
lp_index = [100.0 * mae_pivot.loc[o, "LP median regression"] / mae_pivot.loc[o, "OLS"]
            for o in outcomes]

fig, ax = plt.subplots(figsize=(9.2, 5.4))
bars1 = ax.bar(x - width / 2, ols_index, width, label="OLS baseline", color=COLORS["gray"])
bars2 = ax.bar(x + width / 2, lp_index, width, label="Tax-period LP", color=COLORS["orange"])
fig.suptitle("Absolute error falls in both comparisons", y=0.98,
             fontsize=17, weight="bold", color=COLORS["navy"])
fig.text(0.5, 0.925,
         "Mean absolute error indexed to each outcome's OLS result (OLS = 100).",
         ha="center", fontsize=10.5, color=COLORS["gray"])
ax.set_ylabel("MAE index")
ax.set_xticks(x, outcomes)
ax.set_ylim(0, 115)
ax.legend(frameon=False, loc="upper right")
for i, value in enumerate(lp_index):
    ax.text(i + width / 2, value + 3, f"{value:.1f}", ha="center", fontsize=10)
    ax.text(i, 8, f"{100-value:.1f}% lower", ha="center",
            color=COLORS["orange"], weight="bold", fontsize=10)
finish(ax)
fig.tight_layout(rect=[0, 0, 1, 0.88])
fig.savefig(IMAGES / "mae_comparison.png", dpi=180, bbox_inches="tight")
plt.close(fig)


REGIMES = [
    ("1954-08-16", "1968-06-28"),
    ("1969-12-30", "1971-12-10"),
    ("1976-10-04", "1978-11-06"),
    ("1982-09-03", "1984-07-18"),
    ("1986-10-22", "1990-11-05"),
    ("1993-08-10", "2001-06-07"),
    ("2003-05-28", "2013-01-02"),
    ("2017-12-22", "2022-08-16"),
]


def fit_chart(csv_name, title, output_name, y_label, regimes):
    df = pd.read_csv(DATA / csv_name)
    df["date"] = pd.to_datetime(df["date"])
    fig, ax = plt.subplots(figsize=(10.6, 5.6))
    for idx, (start, end) in enumerate(regimes, start=1):
        start_dt = pd.Timestamp(start)
        end_dt = pd.Timestamp(end)
        ax.axvspan(start_dt, end_dt, color=COLORS["shade"], alpha=0.45, linewidth=0)
        midpoint = start_dt + (end_dt - start_dt) / 2
        ax.text(midpoint, 0.96, str(idx), transform=ax.get_xaxis_transform(),
                ha="center", va="top", fontsize=8, weight="bold", color=COLORS["navy"])
    ax.plot(df["date"], df["actual"], color=COLORS["navy"], linewidth=2.0, label="Actual")
    ax.plot(df["date"], df["fitted"], color=COLORS["orange"], linewidth=2.0,
            label="Tax-period LP fit")
    fig.suptitle(title, y=0.98, fontsize=17, weight="bold", color=COLORS["navy"])
    fig.text(0.5, 0.925, "Shaded areas mark modeled tax-policy regimes.",
             ha="center", fontsize=10.5, color=COLORS["gray"])
    ax.set_ylabel(y_label)
    ax.set_xlabel("Year")
    ax.legend(frameon=False, ncol=2, loc="lower right")
    finish(ax)
    fig.tight_layout(rect=[0, 0, 1, 0.88])
    fig.savefig(IMAGES / output_name, dpi=180, bbox_inches="tight")
    plt.close(fig)


fit_chart("gini_lp_fit.csv", "Gini: actual versus tax-period LP fit",
          "gini_actual_vs_lp_fit.png", "Gini index", REGIMES)
fit_chart("labor_share_lp_fit.csv", "Labor share: actual versus tax-period LP fit",
          "labor_share_actual_vs_lp_fit.png", "Labor share", REGIMES[-4:])