"""Produce the publication figures (APA style, no titles; titles are added in
the manuscript).

Figure 1: forest plot of country-specific coefficients (ESCS, girl, immigrant).
Figure 2: school mean digital self-efficacy by school mean ESCS, by system.

Input : data/processed/analysis_dataset.csv, output/results/res_percountry.json
Output: output/figures/Figure1_gradients.png, Figure2_school_escs.png (300 dpi)
"""
import json
from pathlib import Path

import matplotlib
import numpy as np
import pandas as pd

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
FIGURES = ROOT / "output" / "figures"
FIGURES.mkdir(parents=True, exist_ok=True)

plt.rcParams.update({
    "font.family": "serif",
    "font.serif": ["Times New Roman", "Liberation Serif", "DejaVu Serif"],
    "font.size": 11,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.linewidth": 0.8,
})

ORDER = ["Argentina", "Brazil", "Chile", "Dominican Rep.", "Panama", "Uruguay"]


def figure1(pc: dict) -> None:
    fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.6), sharey=True)
    panels = [("ESCS_c", "ESCS (within school)"), ("girl", "Girl"),
              ("immig", "Immigrant background")]
    for ax, (key, xlab) in zip(axes, panels):
        ys = np.arange(len(ORDER))[::-1]
        for y, c in zip(ys, ORDER):
            b, se = pc[c]["params"][key], pc[c]["bse"][key]
            ax.errorbar(b, y, xerr=1.96 * se, fmt="o", color="black",
                        ecolor="black", ms=4.5, capsize=2.5, lw=1)
        ax.axvline(0, color="gray", lw=0.8, ls="--")
        ax.set_yticks(ys)
        ax.set_yticklabels(ORDER)
        ax.set_xlabel(xlab)
    fig.text(0.5, -0.03, "Coefficient on digital self-efficacy (95% CI)",
             ha="center", fontsize=11)
    plt.tight_layout()
    plt.savefig(FIGURES / "Figure1_gradients.png", dpi=300,
                bbox_inches="tight")
    plt.close()


def figure2(m: pd.DataFrame) -> None:
    s = (m.groupby(["country", "CNTSCHID"])
          .agg(escs=("ESCS", "mean"), eff=("ICTEFFIC", "mean"),
               n=("CNTSTUID", "size"))
          .reset_index())
    fig, axes = plt.subplots(2, 3, figsize=(10.5, 6.4),
                             sharex=True, sharey=True)
    for ax, c in zip(axes.flat, ORDER):
        g = s[s.country == c]
        ax.scatter(g.escs, g.eff, s=np.sqrt(g.n) * 2.2,
                   facecolors="none", edgecolors="gray", lw=0.7)
        b, a = np.polyfit(g.escs, g.eff, 1, w=g.n)
        xs = np.linspace(g.escs.min(), g.escs.max(), 10)
        ax.plot(xs, a + b * xs, color="black", lw=1.4)
        ax.text(0.04, 0.93, c, transform=ax.transAxes, fontsize=11, va="top")
        ax.axhline(0, color="gray", lw=0.5, ls=":")
    for ax in axes[1]:
        ax.set_xlabel("School mean ESCS")
    for ax in axes[:, 0]:
        ax.set_ylabel("School mean digital self-efficacy")
    plt.tight_layout()
    plt.savefig(FIGURES / "Figure2_school_escs.png", dpi=300,
                bbox_inches="tight")
    plt.close()


def main() -> None:
    m = pd.read_csv(ROOT / "data" / "processed" / "analysis_dataset.csv")
    pc = json.load(open(ROOT / "output" / "results" / "res_percountry.json"))
    figure1(pc)
    figure2(m)
    print("figures written to", FIGURES)


if __name__ == "__main__":
    main()
