#!/usr/bin/env python3
"""Create the six manuscript figures from the supplementary data files."""
from __future__ import annotations

import csv
import math
from pathlib import Path
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from common import DATA, FIGURES

mpl.rcParams.update({
    'font.family': 'serif',
    'mathtext.fontset': 'stix',
    'font.size': 10.5,
    'axes.labelsize': 11.5,
    'xtick.labelsize': 9.5,
    'ytick.labelsize': 9.5,
    'legend.fontsize': 9.3,
    'axes.linewidth': 0.8,
    'xtick.major.width': 0.8,
    'ytick.major.width': 0.8,
    'xtick.major.size': 4,
    'ytick.major.size': 4,
    'savefig.facecolor': 'white',
    'figure.facecolor': 'white',
    'axes.facecolor': 'white',
    'pdf.fonttype': 42,
    'ps.fonttype': 42,
})


def save(fig: plt.Figure, number: int, stem: str) -> None:
    FIGURES.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGURES / f'fig{number}.pdf', bbox_inches='tight')
    fig.savefig(FIGURES / f'{stem}.png', dpi=320, bbox_inches='tight')
    plt.close(fig)


def clean_axes(ax: plt.Axes, grid: bool = True) -> None:
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.set_axisbelow(True)
    if grid:
        ax.grid(axis='y', color='0.88', linewidth=0.6)


def read_tsv(path: Path) -> list[dict[str, str]]:
    with path.open() as fh:
        return list(csv.DictReader(fh, delimiter='\t'))


def figure_1() -> None:
    x = np.linspace(0, 4, 1601)
    y = 2 ** (x - np.floor(x))
    jumps = np.isclose(x % 1.0, 0.0, atol=1e-12) & (x > 0)
    y[jumps] = np.nan
    fig, ax = plt.subplots(figsize=(6.9, 3.55), constrained_layout=True)
    clean_axes(ax)
    ax.plot(x, y, color='black', linewidth=2.0)
    for k in range(1, 5):
        ax.axvline(k, color='0.85', linewidth=0.7)
    ax.set_xlim(0, 4)
    ax.set_ylim(0.98, 2.02)
    ax.set_xticks([0, 1, 2, 3, 4])
    ax.set_yticks([1.0, 1.5, 2.0])
    ax.set_xlabel(r'$x=\log_2 B$')
    ax.set_ylabel(r'$B_N a_N/p(N)$')
    save(fig, 1, 'dyadic_sawtooth')


def figure_2() -> None:
    c = np.linspace(-6, 6, 1201)
    y = 1.0 + np.exp(-2.0 * np.exp(-c))
    fig, ax = plt.subplots(figsize=(6.9, 3.55), constrained_layout=True)
    clean_axes(ax)
    ax.plot(c, y, color='black', linewidth=2.0)
    ax.set_xlim(-6, 6)
    ax.set_ylim(0.98, 2.02)
    ax.set_xticks([-6, -4, -2, 0, 2, 4, 6])
    ax.set_yticks([1.0, 1.5, 2.0])
    ax.set_xlabel(r'transition coordinate $c$')
    ax.set_ylabel(r'$K a_N/p(N)$')
    save(fig, 2, 'poisson_transition')


def hierarchy_data() -> tuple[np.ndarray, ...]:
    rows = read_tsv(DATA / 'figures' / 'fig_03_04_hierarchy_and_error_sources_N2000_N20000.tsv')
    x = np.array([float(r['log2_L_N_star']) for r in rows])
    e1 = np.array([float(r['one_level_absolute_relative_error_percent']) for r in rows])
    e2 = np.array([float(r['two_level_absolute_relative_error_percent']) for r in rows])
    beta = np.array([float(r['low_block_scale_beta_percent']) for r in rows])
    upper = np.array([float(r['upper_tail_scale_percent']) for r in rows])
    return x, e1, e2, beta, upper


def prep_log_error(ax: plt.Axes, x: np.ndarray) -> None:
    clean_axes(ax)
    for line in np.arange(7.5, 9.2, 1.0):
        ax.axvline(line, color='0.88', linewidth=0.7, linestyle=(0, (2, 2)), zorder=0)
    ax.set_yscale('log')
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(1e-9, 3e1)
    ax.set_xlabel(r'$\log_2 L_N^*$')
    ax.set_ylabel('absolute relative error (%)')


def figure_3() -> None:
    x, e1, e2, _, _ = hierarchy_data()
    fig, ax = plt.subplots(figsize=(6.9, 4.4), constrained_layout=True)
    prep_log_error(ax, x)
    ax.plot(x, e1, color='#0072B2', linewidth=1.8, linestyle=(0, (5, 2.5)),
            label=r'one level ($r=1$)')
    ax.plot(x, e2, color='#D55E00', linewidth=1.9, label=r'two levels ($r=2$)')
    ax.legend(frameon=False, loc='lower left', handlelength=2.8)
    save(fig, 3, 'dyadic_hierarchy')


def figure_4() -> None:
    x, _, e2, beta, upper = hierarchy_data()
    fig, ax = plt.subplots(figsize=(6.9, 4.4), constrained_layout=True)
    prep_log_error(ax, x)
    ax.plot(x, e2, color='black', linewidth=1.55, label='observed two-level error')
    ax.plot(x, beta, color='#009E73', linewidth=1.55, linestyle=(0, (4, 2)),
            label=r'lower-block scale $100\beta_J(t)$')
    ax.plot(x, upper, color='#CC79A7', linewidth=1.55, linestyle=(0, (1.5, 2.2)),
            label=r'upper-tail scale $100e^{-4tJ}/t$')

    diff = np.log(beta) - np.log(upper)
    crossings = np.where(diff[:-1] * diff[1:] <= 0)[0]
    for k, i in enumerate(crossings):
        denominator = abs(diff[i]) + abs(diff[i + 1])
        weight = abs(diff[i]) / denominator if denominator else 0.5
        xc = x[i] + weight * (x[i + 1] - x[i])
        yc = np.exp(np.log(beta[i]) + weight * (np.log(beta[i + 1]) - np.log(beta[i])))
        ax.plot(xc, yc, marker='o', markersize=5.3, markerfacecolor='white',
                markeredgecolor='0.2', markeredgewidth=0.9, linestyle='None',
                label='diagnostic crossing' if k == 0 else None, zorder=6)
    ax.legend(frameon=False, loc='lower left', handlelength=2.8)
    save(fig, 4, 'two_level_error_sources')


def figure_5() -> None:
    rows = read_tsv(DATA / 'figures' / 'fig_05_fixed_size_correction_N2000_N20000.tsv')
    t = np.array([float(r['t_N']) for r in rows])
    x = np.log2(np.log(1.0 / t) / t)
    observed = np.array([float(r['observed_fixed_size_error_percent']) for r in rows])
    prediction = np.array([float(r['second_order_prediction_percent']) for r in rows])

    fig, ax = plt.subplots(figsize=(6.9, 4.25), constrained_layout=True)
    clean_axes(ax)
    for line in np.arange(7.5, 9.2, 1.0):
        ax.axvline(line, color='0.88', linewidth=0.7, linestyle=(0, (2, 2)), zorder=0)
    ax.axhline(0, color='0.75', linewidth=0.8, zorder=0)
    ax.plot(x, observed, color='black', linewidth=1.8, label='observed fixed-size error')
    ax.plot(x, prediction, color='#0072B2', linewidth=1.65, linestyle=(0, (5, 2.5)),
            label='second-order diagnostic')
    ax.set_xlim(x.min(), x.max())
    margin = 0.25
    ax.set_ylim(min(observed.min(), prediction.min()) - margin,
                max(observed.max(), prediction.max()) + margin)
    ax.set_xlabel(r'$\log_2 L_N^*$')
    ax.set_ylabel('signed relative error (%)')
    ax.legend(frameon=False, loc='upper right', handlelength=2.8)
    save(fig, 5, 'fixed_size_correction')


def figure_6() -> None:
    rows = read_tsv(DATA / 'figures' / 'fig_06_exact_density_N100_N20000.tsv')
    x = np.array([float(r['log2_B_N']) for r in rows])
    exact = np.array([float(r['normalized_exact_density']) for r in rows])
    prediction = np.array([float(r['normalized_first_order_density']) for r in rows])
    phase = np.array([float(r['limiting_dyadic_phase']) for r in rows])

    fig, ax = plt.subplots(figsize=(6.9, 4.1), constrained_layout=True)
    clean_axes(ax)
    for line in np.arange(math.ceil(x.min()), math.floor(x.max()) + 1):
        ax.axvline(line, color='0.90', linewidth=0.7, zorder=0)
    ax.plot(x, phase, color='#D55E00', linewidth=1.25, linestyle=(0, (1.2, 2.0)),
            label=r'limiting phase $2^{\{\log_2 B_N\}}$')
    ax.plot(x, prediction, color='#0072B2', linewidth=1.5, linestyle=(0, (5, 2.5)),
            label='first-order prediction')
    ax.plot(x, exact, color='black', linewidth=2.0, label='exact normalized density')
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(0.88, 2.05)
    ax.set_xlabel(r'$\log_2 B_N$')
    ax.set_ylabel(r'$B_N a_N/p(N)$')
    ax.legend(frameon=False, loc='upper left', handlelength=2.8)
    save(fig, 6, 'exact_density')


def main() -> None:
    figure_1()
    figure_2()
    figure_3()
    figure_4()
    figure_5()
    figure_6()
    print(f'Figures written to {FIGURES}')


if __name__ == '__main__':
    main()
