#!/usr/bin/env python3
"""Build the master count table, Table 1, and Figure 6 data."""
from __future__ import annotations

import csv
import math
from decimal import Decimal, getcontext
from common import DATA, load_compact_counts, partition_numbers, saddle_parameter

MILESTONES = [1500, 6328, 10000, 12000, 19584, 20000]


def main() -> None:
    counts = load_compact_counts()
    nmax = max(counts)
    partitions = partition_numbers(nmax)
    getcontext().prec = 50

    counts_dir = DATA / 'counts'
    tables_dir = DATA / 'tables'
    figures_dir = DATA / 'figures'
    counts_dir.mkdir(parents=True, exist_ok=True)
    tables_dir.mkdir(parents=True, exist_ok=True)
    figures_dir.mkdir(parents=True, exist_ok=True)

    with (counts_dir / 'nim_losing_counts_N0_N20000.tsv').open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 'oeis_index_k', 'parity', 'losing_count_a_N',
            'partition_number_p_N', 'density_a_N_over_p_N'
        ])
        for n in range(nmax + 1):
            if n % 2 == 0:
                a = counts[n]
                k: int | str = n // 2
            else:
                a = 0
                k = ''
            density = Decimal(a) / Decimal(partitions[n])
            w.writerow([n, k, 'even' if n % 2 == 0 else 'odd', a,
                        partitions[n], format(density, '.24E')])

    with (tables_dir / 'table_01_first_order_milestones.tsv').open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 'oeis_index_k', 't_N', 'L_N_star', 'log2_L_N_star',
            'J_N_star', 'Lambda_N_star', 'exact_density', 'first_order_density',
            'signed_relative_error_percent', 'B_N', 'log2_B_N',
            'B_N_times_exact_density', 'B_N_times_first_order_density'
        ])
        for n in MILESTONES:
            t = saddle_parameter(n)
            L = math.log(1.0 / t) / t
            J = 2 ** math.floor(math.log2(L) + 0.5)
            Lambda = math.exp(-t * J) / t
            prediction = (1.0 + math.exp(-2.0 * Lambda)) / J
            exact = counts[n] / partitions[n]
            tau = math.pi / math.sqrt(6.0 * n)
            B = math.log(1.0 / tau) / tau
            w.writerow([n, n // 2, f'{t:.17g}', f'{L:.17g}', f'{math.log2(L):.17g}',
                        J, f'{Lambda:.17g}', f'{exact:.17g}', f'{prediction:.17g}',
                        f'{100.0 * (prediction / exact - 1.0):.15g}', f'{B:.17g}',
                        f'{math.log2(B):.17g}', f'{B * exact:.17g}',
                        f'{B * prediction:.17g}'])

    with (figures_dir / 'fig_06_exact_density_N100_N20000.tsv').open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 'B_N', 'log2_B_N', 'J_N', 'Lambda_N',
            'exact_density', 'first_order_density', 'normalized_exact_density',
            'normalized_first_order_density', 'limiting_dyadic_phase'
        ])
        for n in range(100, nmax + 1, 2):
            tau = math.pi / math.sqrt(6.0 * n)
            B = math.log(1.0 / tau) / tau
            J = 2 ** math.floor(math.log2(B) + 0.5)
            Lambda = math.exp(-tau * J) / tau
            exact = counts[n] / partitions[n]
            prediction = (1.0 + math.exp(-2.0 * Lambda)) / J
            phase = 2.0 ** (math.log2(B) - math.floor(math.log2(B)))
            w.writerow([n, f'{B:.17g}', f'{math.log2(B):.17g}', J,
                        f'{Lambda:.17g}', f'{exact:.17g}', f'{prediction:.17g}',
                        f'{B * exact:.17g}', f'{B * prediction:.17g}', f'{phase:.17g}'])

    print('Count data and first-order tables rebuilt.')


if __name__ == '__main__':
    main()
