#!/usr/bin/env python3
"""Recompute the data used in Figures 3--4 and Table 3."""
from __future__ import annotations

import csv
import math
from common import DATA, saddle_parameter, boltzmann_zero_xor, one_level, two_level, low_block_scale

MILESTONES = [10000, 12000, 19584, 20000]


def quantities(n: int) -> dict[str, float | int]:
    t = saddle_parameter(n)
    L = math.log(1.0 / t) / t
    J = 2 ** math.floor(math.log2(L) + 0.5)
    exact = boltzmann_zero_xor(t)
    e1 = 100.0 * abs(one_level(t, J) / exact - 1.0)
    e2 = 100.0 * abs(two_level(t, J) / exact - 1.0)
    beta = 100.0 * low_block_scale(t, J)
    upper = 100.0 * math.exp(-4.0 * t * J) / t
    return {
        'n': n, 't': t, 'L': L, 'J': J, 'e1': e1, 'e2': e2,
        'beta': beta, 'upper': upper,
    }


def main() -> None:
    figures_dir = DATA / 'figures'
    tables_dir = DATA / 'tables'
    figures_dir.mkdir(parents=True, exist_ok=True)
    tables_dir.mkdir(parents=True, exist_ok=True)

    grid_path = figures_dir / 'fig_03_04_hierarchy_and_error_sources_N2000_N20000.tsv'
    with grid_path.open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 't_N', 'log2_L_N_star', 'J_N_star',
            'one_level_absolute_relative_error_percent',
            'two_level_absolute_relative_error_percent',
            'low_block_scale_beta_percent', 'upper_tail_scale_percent'
        ])
        for n in range(2000, 20001, 50):
            q = quantities(n)
            w.writerow([
                n, f"{q['t']:.17g}", f"{math.log2(float(q['L'])):.17g}", q['J'],
                f"{q['e1']:.15g}", f"{q['e2']:.15g}",
                f"{q['beta']:.15g}", f"{q['upper']:.15g}",
            ])

    table_path = tables_dir / 'table_03_two_level_error_budget.tsv'
    with table_path.open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 't_N', 'J_N_star',
            'observed_two_level_absolute_relative_error_percent',
            'low_block_scale_beta_percent', 'upper_tail_scale_percent'
        ])
        for n in MILESTONES:
            q = quantities(n)
            w.writerow([
                n, f"{q['t']:.17g}", q['J'], f"{q['e2']:.15g}",
                f"{q['beta']:.15g}", f"{q['upper']:.15g}",
            ])

    print(grid_path)
    print(table_path)


if __name__ == '__main__':
    main()
