"""Exact reconstruction helpers from T683; provenance in source-provenance.json."""
import sys
import json
from pathlib import Path
from functools import lru_cache
import numberdb.sage as numberdb
from sage.libs.pari import pari
from sage.rings.rational_field import QQ
from sage.rings.integer_ring import ZZ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.rings.qqbar import QQbar, _init_qqbar
from sage.rings.complex_mpfr import ComplexField
from sage.rings.complex_interval_field import ComplexIntervalField

# Sage's named import does not run this initializer, which exact complex
# conjugation needs. Initialize it explicitly without importing sage.all.
_init_qqbar()

BASE = Path(__file__).resolve().parent
ROSTER = json.loads((BASE / 'roster.json').read_text())
BY_FORM = {e['label']: (r, e) for r in ROSTER for e in r['embeddings']}
WORKING_DIGITS = (160, 220)
P = PolynomialRing(QQ, 'x')


def rational_decimal(text):
    from fractions import Fraction
    x = Fraction(text)
    return QQ(x.numerator) / QQ(x.denominator)


def coeff_poly(value):
    return P(str(pari.lift(value)).replace('y', 'x'))


@lru_cache(maxsize=None)
def algebra(label):
    row = next(r for r in ROSTER if r['label'] == label)
    F = P(row['field_poly'])
    cs = [P([QQ(c) for c in v]) for v in row['an_power_basis']]
    roots = [z for z, mult in F.roots(QQbar) if mult == 1]
    assert len(roots) == row['degree']
    embedded = {}
    for e in row['embeddings']:
        box = [[rational_decimal(c) for c in edge] for edge in e['root_rectangle']]
        found = [z for z in roots if box[0][0] < z.real() < box[0][1]
                 and box[1][0] < z.imag() < box[1][1]]
        assert len(found) == 1, e['label']
        z = found[0]
        is_real = all(c(z).imag() == 0 for c in cs)
        embedded[e['label']] = (z, is_real)
    zs = [z for z, real in embedded.values()]
    assert all(zs[i] != zs[j] for i in range(len(zs)) for j in range(i))
    return F, cs, embedded


def reconstruct(row):
    N, k = row['level'], row['weight']
    spec = [N, k, row['character_discriminant']]
    mf = pari.mfinit(spec, 0)
    F, cs, embedded = algebra(row['label'])
    h = P([QQ(c) for c in row['pari_generator_in_catalogue_field']])
    choices = []
    for f in pari.mfeigenbasis(mf):
        fp = pari.mfparams(f)[3]
        if str(fp) != row['pari_field_poly']:
            continue
        cf = list(pari.mfcoefs(f, row['sturm_bound']))
        if all(coeff_poly(a)(h) % F == c for a, c in zip(cf, cs)):
            choices.append(f)
    assert len(choices) == 1, ('exact reconstruction', row['label'])
    f = choices[0]
    if row['fricke_eigenvalue'] is not None:
        A = pari.mfatkininit(mf, N, precision=128)
        assert A[2] == 1
        af = pari.mfatkin(A, f)
        ac = list(pari.mfcoefs(af, row['sturm_bound']))
        fc = list(pari.mfcoefs(f, row['sturm_bound']))
        assert all(a == row['fricke_eigenvalue'] * b for a, b in zip(ac, fc))
    return mf, f, h, embedded