#!/usr/bin/env python3
"""Independent N=1500 check using a direct distinct-part XOR knapsack.

The helper C++ program computes d_k modulo four primes, where d_k counts
partitions of 2k into distinct parts with XOR zero.  This script reconstructs
the d_k by CRT, applies a_{2n} = sum_k d_k p(n-k), and compares the result with
the packaged A048833 data.
"""
from __future__ import annotations

import argparse
import subprocess
import tempfile
from math import prod
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
SOURCE = Path(__file__).with_name('distinct_xor_knapsack.cpp')
DEFAULT_PRIMES = [2147483629, 2137483609, 2127483619, 2117483629]


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 distinct_partition_numbers(nmax: int) -> list[int]:
    q = [0] * (nmax + 1)
    q[0] = 1
    for part in range(1, nmax + 1):
        for n in range(nmax, part - 1, -1):
            q[n] += q[n - part]
    return q


def load_compact(path: Path) -> list[int]:
    values: list[int] = []
    for raw in path.read_text().splitlines():
        line = raw.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 index {k} in {path}')
        values.append(int(value_s))
    return values


def load_residues(path: Path, kmax: int) -> list[int]:
    values = load_compact(path)
    if len(values) != kmax + 1:
        raise ValueError(f'{path} has {len(values)} values, expected {kmax + 1}')
    return values


def crt_columns(moduli: list[int], residue_columns: list[list[int]]) -> tuple[list[int], int]:
    modulus_product = prod(moduli)
    coefficients = []
    for modulus in moduli:
        partial = modulus_product // modulus
        coefficients.append(partial * pow(partial, -1, modulus))
    values = []
    for k in range(len(residue_columns[0])):
        values.append(sum(residue_columns[i][k] * coefficients[i]
                          for i in range(len(moduli))) % modulus_product)
    return values, modulus_product


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument('--total-size', type=int, default=1500)
    parser.add_argument('--threads', type=int, default=1,
                        help='reserved for interface consistency; the direct DP is serial')
    parser.add_argument('--keep-work', type=Path)
    args = parser.parse_args()
    nmax = args.total_size
    if nmax <= 0 or nmax % 2:
        raise ValueError('--total-size must be a positive even integer')
    n = nmax // 2

    context = tempfile.TemporaryDirectory(prefix='nim_knapsack_') if args.keep_work is None else None
    work = Path(context.name) if context else args.keep_work
    work.mkdir(parents=True, exist_ok=True)
    executable = work / 'distinct_xor_knapsack'
    subprocess.run([
        'g++', '-O3', '-march=native', '-std=c++17', str(SOURCE), '-o', str(executable)
    ], check=True)

    residue_columns: list[list[int]] = []
    for prime in DEFAULT_PRIMES:
        output = work / f'{prime}.txt'
        print(f'computing direct knapsack modulo {prime} ...', flush=True)
        subprocess.run([str(executable), str(nmax), str(prime), str(output)], check=True)
        residue_columns.append(load_residues(output, n))

    d, modulus_product = crt_columns(DEFAULT_PRIMES, residue_columns)
    q = distinct_partition_numbers(nmax)
    if modulus_product <= max(q):
        raise AssertionError('CRT product does not dominate the distinct-partition bound')
    for k, value in enumerate(d):
        if not 0 <= value <= q[2 * k]:
            raise AssertionError(f'd_{k} lies outside [0,q_d(2k)]')

    p = partition_numbers(n)
    a = sum(d[k] * p[n - k] for k in range(n + 1))
    reference_path = ROOT / 'data' / 'counts' / 'oeis_A048833_k0_k10000.txt'
    reference = load_compact(reference_path)[n]
    if a != reference:
        raise AssertionError(f'independent count mismatch: {a} != {reference}')

    p_full = partition_numbers(nmax)[nmax]
    print(f'a_{nmax} = {a}')
    print(f'p({nmax}) = {p_full}')
    print(f'a_{nmax}/p({nmax}) = {a / p_full:.15f}')
    print(f'Independent direct-knapsack check passed at N={nmax}.')

    if context is not None:
        context.cleanup()


if __name__ == '__main__':
    main()
