#!/usr/bin/env python3
"""Reconstruct exact zero-XOR partition counts from modular residue files.

Each residue file is named <prime>.txt and contains two columns:

    k residue

where k is the OEIS A048833 index and total size is N=2k.
"""
from __future__ import annotations

import argparse
from math import prod
from pathlib import Path


def partition_numbers(nmax: int) -> list[int]:
    p = [0] * (nmax + 1)
    p[0] = 1
    for n in range(1, nmax + 1):
        total = 0
        k = 1
        while True:
            g1 = k * (3 * k - 1) // 2
            if g1 > n:
                break
            sign = 1 if k % 2 else -1
            total += sign * p[n - g1]
            g2 = k * (3 * k + 1) // 2
            if g2 <= n:
                total += sign * p[n - g2]
            k += 1
        p[n] = total
    return p


def read_residues(directory: Path, kmax: int) -> tuple[list[int], list[list[int]]]:
    files = sorted(directory.glob('*.txt'), key=lambda path: int(path.stem))
    if not files:
        raise ValueError(f'No residue files found in {directory}')

    moduli: list[int] = []
    residues: list[list[int]] = []
    for path in files:
        modulus = int(path.stem)
        values: list[int] = []
        for line in path.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            k_s, value_s = line.split()
            k = int(k_s)
            if k != len(values):
                raise ValueError(f'Unexpected row index in {path}: {k}')
            values.append(int(value_s))
        if len(values) != kmax + 1:
            raise ValueError(f'{path} has {len(values)} rows; expected {kmax + 1}')
        moduli.append(modulus)
        residues.append(values)
    return moduli, residues


def reconstruct(moduli: list[int], residues: list[list[int]], kmax: int) -> tuple[list[int], int]:
    modulus_product = prod(moduli)
    coefficients: list[int] = []
    for modulus in moduli:
        partial = modulus_product // modulus
        coefficients.append(partial * pow(partial, -1, modulus))

    values: list[int] = []
    for k in range(kmax + 1):
        value = sum(
            residues[i][k] * coefficients[i]
            for i in range(len(moduli))
        ) % modulus_product
        values.append(value)
    return values, modulus_product


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument('--residue-dir', type=Path, required=True)
    parser.add_argument('--total-size-max', type=int, default=20000)
    parser.add_argument('--output-dir', type=Path, required=True)
    args = parser.parse_args()

    if args.total_size_max % 2:
        raise ValueError('total-size-max must be even')

    nmax = args.total_size_max
    kmax = nmax // 2
    moduli, residues = read_residues(args.residue_dir, kmax)
    values, modulus_product = reconstruct(moduli, residues, kmax)
    partitions = partition_numbers(nmax)

    if modulus_product <= partitions[nmax]:
        raise ValueError(
            f'CRT product is too small: {modulus_product.bit_length()} bits, '
            f'but p({nmax}) has {partitions[nmax].bit_length()} bits'
        )
    for k, value in enumerate(values):
        if not 0 <= value <= partitions[2 * k]:
            raise AssertionError(f'Count outside [0,p(2k)] at k={k}')

    args.output_dir.mkdir(parents=True, exist_ok=True)

    compact = args.output_dir / f'oeis_A048833_k0_k{kmax}.txt'
    with compact.open('w') as fh:
        fh.write(f'# A048833(k) for k = 0..{kmax}; total size N = 2k.\n')
        fh.write(
            f'# Computed by finite Walsh-character recurrences modulo {len(moduli)} '
            'primes and CRT.\n'
        )
        fh.write(
            f'# The {modulus_product.bit_length()}-bit CRT product exceeds p({nmax}), '
            f'which has {partitions[nmax].bit_length()} bits.\n'
        )
        for k, value in enumerate(values):
            fh.write(f'{k} {value}\n')

    table = args.output_dir / f'oeis_A048833_k0_k{kmax}.tsv'
    with table.open('w') as fh:
        fh.write('oeis_index_k\ttotal_size_N\tA048833_k\tpartition_number_p_2k\tdensity_A048833_k_over_p_2k\n')
        for k, value in enumerate(values):
            n = 2 * k
            fh.write(f'{k}\t{n}\t{value}\t{partitions[n]}\t{value / partitions[n]:.18e}\n')

    report = args.output_dir / 'crt_reconstruction_report.txt'
    report.write_text(
        'CRT reconstruction report\n'
        '=========================\n'
        f'Range: k=0,...,{kmax}, equivalently N=0,...,{nmax} over even N.\n'
        f'Prime moduli: {len(moduli)}.\n'
        f'CRT product: {modulus_product.bit_length()} bits.\n'
        f'p({nmax}): {partitions[nmax].bit_length()} bits.\n'
        'The CRT reconstruction is unique because 0 <= a_N <= p(N).\n'
    )

    print(compact)
    print(table)
    print(report)


if __name__ == '__main__':
    main()
