#!/usr/bin/env python3
"""Recompute the second-order diagnostic used in Figure 5 and Table 4."""
from __future__ import annotations

import csv
import math
from common import (
    DATA, load_compact_counts, partition_numbers, saddle_parameter,
    saddle_variance_and_k3, boltzmann_derivatives,
)

MILESTONES = [10000, 12000, 19584, 20000]


def quantities(n: int, counts: dict[int, int], partitions: list[int]) -> dict[str, float | int]:
    t = saddle_parameter(n)
    variance, k3 = saddle_variance_and_k3(t)
    R, Rp, Rpp = boltzmann_derivatives(t)
    exact = counts[n] / partitions[n]
    observed = 100.0 * (2.0 * R / exact - 1.0)
    curvature = 100.0 * Rpp / (2.0 * variance * R)
    skew = 100.0 * k3 * Rp / (2.0 * variance * variance * R)
    L = math.log(1.0 / t) / t
    J = 2 ** math.floor(math.log2(L) + 0.5)
    Lambda = math.exp(-t * J) / t
    c = -math.log(Lambda)
    return {
        'n': n, 't': t, 'J': J, 'Lambda': Lambda, 'c': c,
        'observed': observed, 'curvature': curvature, 'skew': skew,
        'prediction': curvature + skew,
    }


def main() -> None:
    counts = load_compact_counts()
    partitions = partition_numbers(20000)
    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_05_fixed_size_correction_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', 'J_N_star', 'Lambda_N_star',
            'transition_coordinate_c', 'observed_fixed_size_error_percent',
            'curvature_term_percent', 'skew_term_percent',
            'second_order_prediction_percent'
        ])
        for n in range(2000, 20001, 50):
            q = quantities(n, counts, partitions)
            w.writerow([
                n, f"{q['t']:.17g}", q['J'], f"{q['Lambda']:.17g}",
                f"{q['c']:.17g}", f"{q['observed']:.15g}",
                f"{q['curvature']:.15g}", f"{q['skew']:.15g}",
                f"{q['prediction']:.15g}",
            ])

    table_path = tables_dir / 'table_04_second_order_fixed_size.tsv'
    with table_path.open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 't_N', 'observed_fixed_size_error_percent',
            'curvature_term_percent', 'skew_term_percent',
            'second_order_prediction_percent'
        ])
        for n in MILESTONES:
            q = quantities(n, counts, partitions)
            w.writerow([
                n, f"{q['t']:.17g}", f"{q['observed']:.15g}",
                f"{q['curvature']:.15g}", f"{q['skew']:.15g}",
                f"{q['prediction']:.15g}",
            ])

    print(grid_path)
    print(table_path)


if __name__ == '__main__':
    main()
