"""Hirzebruch L-polynomials -- numberdb.org/T336.

This generator stores the homogeneous Hirzebruch L-polynomials

    L_n(p1, p2, ...) = [degree n] prod_i sqrt(x_i) / tanh(sqrt(x_i)),

where pj is the j-th elementary symmetric polynomial in the Pontryagin roots
x_i. The table starts at n = 1 and stops at n = 6, because L_7 has a nonzero
p7 term and the database searches polynomials in at most six variables.

Run it with SageMath:

    $ sage -pip install numberdb          # once
    $ sage -python generate.py            # check the table against this code
    $ sage -python generate.py --publish  # fill the draft, with NUMBERDB_API_KEY set

The values are exact rational polynomials. The integrity checks compare the
first three nonconstant components with standard printed formulas, verify the
Pontryagin-root definition in seven roots, and check that substituting the
Pontryagin classes of T CP^(2m) gives signature 1 for 1 <= m <= 6.
"""

import os
import sys
from itertools import combinations

import numberdb.sage as numberdb
from sage.arith.misc import binomial
from sage.rings.integer_ring import ZZ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.rings.rational_field import QQ


TABLE = os.environ.get("NUMBERDB_TABLE", "T336")
MAX_DEGREE = 6
CHECK_DEGREE = 7

OUTPUT_RING = PolynomialRing(
    QQ, ["p%s" % i for i in range(1, MAX_DEGREE + 1)]
)
OUTPUT_P = OUTPUT_RING.gens()

CHECK_RING = PolynomialRing(
    QQ, ["p%s" % i for i in range(1, CHECK_DEGREE + 1)]
)
CHECK_P = CHECK_RING.gens()

ROOT_RING = PolynomialRing(
    QQ, ["x%s" % i for i in range(1, CHECK_DEGREE + 1)]
)
ROOTS = ROOT_RING.gens()

H_RING = PolynomialRing(QQ, "h")
H = H_RING.gen()


def _key_from_stdin():
    if os.environ.get("NUMBERDB_KEY_FROM_STDIN") != "1":
        return
    token = sys.stdin.read().strip()
    if "=" in token and token.split("=", 1)[0].isupper():
        token = token.split("=", 1)[1].strip().strip("'\"")
    if token:
        os.environ["NUMBERDB_API_KEY"] = token
        numberdb.configure(api_key=token)


def factorial(n):
    value = QQ(1)
    for k in range(2, int(n) + 1):
        value *= QQ(k)
    return value


def bernoulli_plus(n):
    """Bernoulli number with B_1 = +1/2. Only even values are used."""
    values = [QQ(0) for _ in range(int(n) + 1)]
    for m in range(int(n) + 1):
        values[m] = QQ(1) / QQ(m + 1)
        for j in range(m, 0, -1):
            values[j - 1] = QQ(j) * (values[j - 1] - values[j])
    return values[0]


def series_multiply(left, right, degree):
    product = [QQ(0) for _ in range(degree + 1)]
    for i, a in enumerate(left):
        if a == 0:
            continue
        for j, b in enumerate(right):
            if i + j > degree:
                break
            product[i + j] += QQ(a) * QQ(b)
    return product


def l_series_coefficients(degree):
    """Coefficients of sqrt(x) / tanh(sqrt(x)) through x^degree."""
    return [
        (QQ(2) ** (2 * k)) * bernoulli_plus(2 * k) / factorial(2 * k)
        for k in range(int(degree) + 1)
    ]


def log_series(series, degree):
    """Formal log of a series whose constant coefficient is 1."""
    if series[0] != 1:
        raise ValueError("log_series needs constant coefficient 1")
    u = list(series)
    u[0] -= QQ(1)
    power = [QQ(1)] + [QQ(0) for _ in range(degree)]
    out = [QQ(0) for _ in range(degree + 1)]
    for exponent in range(1, degree + 1):
        power = series_multiply(power, u, degree)
        sign = QQ(1) if exponent % 2 else -QQ(1)
        scale = sign / QQ(exponent)
        for k in range(degree + 1):
            out[k] += scale * power[k]
    return out


def weighted_degree(exponents):
    return sum((index + 1) * exponent for index, exponent in enumerate(exponents))


def ordinary_degree(exponents):
    return sum(exponents)


def monomial_from_exponents(ring, variables, exponents):
    monomial = ring.one()
    for variable, exponent in zip(variables, exponents):
        if exponent:
            monomial *= variable ** exponent
    return monomial


def truncate_weighted(polynomial, degree):
    ring = polynomial.parent()
    variables = ring.gens()
    total = ring.zero()
    for exponents, coefficient in polynomial.dict().items():
        if weighted_degree(exponents) <= degree:
            total += coefficient * monomial_from_exponents(
                ring, variables, exponents
            )
    return total


def truncate_ordinary(polynomial, degree):
    ring = polynomial.parent()
    variables = ring.gens()
    total = ring.zero()
    for exponents, coefficient in polynomial.dict().items():
        if ordinary_degree(exponents) <= degree:
            total += coefficient * monomial_from_exponents(
                ring, variables, exponents
            )
    return total


def homogeneous_weighted(polynomial, degree):
    ring = polynomial.parent()
    variables = ring.gens()
    total = ring.zero()
    for exponents, coefficient in polynomial.dict().items():
        if weighted_degree(exponents) == degree:
            total += coefficient * monomial_from_exponents(
                ring, variables, exponents
            )
    return total


def power_sums(ring, variables, degree):
    """Power sums in terms of elementary symmetric polynomials."""
    sums = [ring.zero()]
    for m in range(1, int(degree) + 1):
        total = ring.zero()
        for j in range(1, m):
            sign = QQ(1) if j % 2 else -QQ(1)
            total += sign * variables[j - 1] * sums[m - j]
        sign = QQ(1) if (m + 1) % 2 == 0 else -QQ(1)
        total += sign * QQ(m) * variables[m - 1]
        sums.append(truncate_weighted(total, degree))
    return sums


def total_l_polynomial():
    log_q = log_series(l_series_coefficients(CHECK_DEGREE), CHECK_DEGREE)
    p = power_sums(CHECK_RING, CHECK_P, CHECK_DEGREE)
    exponent = CHECK_RING.zero()
    for degree in range(1, CHECK_DEGREE + 1):
        exponent += log_q[degree] * p[degree]
    exponent = truncate_weighted(exponent, CHECK_DEGREE)

    total = CHECK_RING.one()
    power = CHECK_RING.one()
    for k in range(1, CHECK_DEGREE + 1):
        power = truncate_weighted(power * exponent, CHECK_DEGREE)
        total = truncate_weighted(total + power / factorial(k), CHECK_DEGREE)
    return total


TOTAL_L = total_l_polynomial()


def to_output_ring(polynomial):
    total = OUTPUT_RING.zero()
    for exponents, coefficient in polynomial.dict().items():
        if any(exponents[MAX_DEGREE:]):
            raise AssertionError("term beyond p6 survived: %s" % (polynomial,))
        monomial = OUTPUT_RING.one()
        for variable, exponent in zip(OUTPUT_P, exponents[:MAX_DEGREE]):
            if exponent:
                monomial *= variable ** exponent
        total += coefficient * monomial
    return total


def l_polynomial(n):
    return to_output_ring(homogeneous_weighted(TOTAL_L, int(n)))


def root_q(root):
    coeffs = l_series_coefficients(CHECK_DEGREE)
    return sum(coeffs[k] * root ** k for k in range(CHECK_DEGREE + 1))


def elementary_roots(k):
    if k == 0:
        return ROOT_RING.one()
    total = ROOT_RING.zero()
    for indexes in combinations(range(CHECK_DEGREE), int(k)):
        monomial = ROOT_RING.one()
        for index in indexes:
            monomial *= ROOTS[index]
        total += monomial
    return total


ROOT_ELEMENTARY = [elementary_roots(k) for k in range(1, CHECK_DEGREE + 1)]


def root_product_l():
    total = ROOT_RING.one()
    for root in ROOTS:
        total = truncate_ordinary(total * root_q(root), CHECK_DEGREE)
    return total


def substitute_check_to_roots(polynomial):
    return truncate_ordinary(ROOT_RING(polynomial(*ROOT_ELEMENTARY)), CHECK_DEGREE)


def complex_projective_space_substitution(polynomial, m):
    """Substitute p_j(T CP^(2m)) = binomial(2m + 1, j) h^j."""
    images = [
        binomial(2 * int(m) + 1, j) * H ** j
        for j in range(1, MAX_DEGREE + 1)
    ]
    return H_RING(polynomial(*images))


def iter_parameters():
    for n in range(1, MAX_DEGREE + 1):
        yield {"n": ZZ(n)}


def check_identities():
    expected = {
        1: OUTPUT_P[0] / QQ(3),
        2: (7 * OUTPUT_P[1] - OUTPUT_P[0] ** 2) / QQ(45),
        3: (
            62 * OUTPUT_P[2]
            - 13 * OUTPUT_P[0] * OUTPUT_P[1]
            + 2 * OUTPUT_P[0] ** 3
        ) / QQ(945),
    }
    for n, value in expected.items():
        found = l_polynomial(n)
        if found != value:
            raise AssertionError("L_%s printed formula failed: %s != %s"
                                 % (n, found, value))

    l7 = homogeneous_weighted(TOTAL_L, 7)
    p7_coefficient = l7.monomial_coefficient(CHECK_P[6])
    if p7_coefficient == 0:
        raise AssertionError("L_7 unexpectedly has no p7 term")

    via_roots = root_product_l()
    via_components = substitute_check_to_roots(TOTAL_L)
    if via_components != via_roots:
        raise AssertionError("root-product check failed")

    for m in range(1, MAX_DEGREE + 1):
        value = complex_projective_space_substitution(l_polynomial(m), m)
        if value[m] != 1:
            raise AssertionError(
                "signature of CP^(2*%s) failed: coefficient is %s"
                % (m, value[m])
            )

    lengths = [
        (len(str(l_polynomial(n))), n)
        for n in range(1, MAX_DEGREE + 1)
    ]
    longest = max(lengths, key=lambda item: item[0])
    print("integrity checks passed for %d Hirzebruch L-polynomials"
          % MAX_DEGREE)
    print("matched the printed formulas for L_1 through L_3")
    print("root-product and projective-space signature checks passed")
    print("L_7 has p7 coefficient %s" % p7_coefficient)
    print("longest polynomial has %d characters at n=%s"
          % (longest[0], longest[1]))


class HirzebruchLPolynomials(numberdb.Generator):

    table = TABLE
    parameters = ("n",)
    type = "Q[]"
    rigour = "exact"

    def enumerate(self):
        yield from iter_parameters()

    def value(self, params, digits):
        return l_polynomial(ZZ(params["n"]))


def fill_draft_once(generator, message):
    """Fill a fresh prose draft without the client's empty upsert probe."""
    from numberdb._generate import (
        _check_precision,
        _check_rigour,
        _producer,
        _run_name,
        _source_files,
    )
    from numberdb._write import Entries, attach, submit_entries, to_text

    table = generator.table
    run = _run_name(generator)
    entries = Entries(*generator.parameters)

    for params in generator.enumerate():
        params = dict(params)
        wanted = generator.digits_for(params)
        entry = generator._entry(params, wanted)
        value = entry["number"]
        identity = ",".join(str(params[name]) for name in generator.parameters)
        _check_rigour(generator, table, identity, value)

        written = to_text(value, wanted, generator.format)
        _check_precision(table, identity, written, wanted, lowering=False)

        record = dict(entry)
        record.pop("digits", None)
        entries.add(**params, **record, digits=wanted)

    answer = submit_entries(
        table,
        entries,
        message=message,
        produced_by=_producer(generator, os.environ.get("NUMBERDB_ASSISTED_BY", "")),
        upsert=False,
        run=run,
        rigour=generator.rigour,
    )

    for name, body in sorted(_source_files(generator).items()):
        attach(table, name, body, run=run, message=message, rigour=generator.rigour)

    return answer


def stored_values():
    table = numberdb.table(TABLE)
    numbers = table.get("Numbers") or {}
    values = {}
    if isinstance(numbers, dict):
        for key, value in numbers.items():
            if isinstance(value, dict):
                value = value.get("number")
            values[ZZ(key)] = OUTPUT_RING(str(value))
    elif isinstance(numbers, list):
        for row in numbers:
            params = row.get("params", row)
            value = row.get("number")
            values[ZZ(params["n"])] = OUTPUT_RING(str(value))
    return values


def check_stored_values():
    values = stored_values()
    if len(values) != MAX_DEGREE:
        raise AssertionError("stored table has %s values, expected %s"
                             % (len(values), MAX_DEGREE))
    for n in range(1, MAX_DEGREE + 1):
        expected = l_polynomial(n)
        if values.get(ZZ(n)) != expected:
            raise AssertionError("stored L_%s failed: %s != %s"
                                 % (n, values.get(ZZ(n)), expected))
    print("stored values match the independently checked generator")


def main():
    _key_from_stdin()
    check_identities()
    generator = HirzebruchLPolynomials()
    if os.environ.get("NUMBERDB_PUBLISH") == "1" or "--publish" in sys.argv:
        print(fill_draft_once(
            generator,
            message="fill Hirzebruch L-polynomials from exact characteristic series",
        ))
    else:
        report = generator.verify(sample=None)
        print(report)
        if not report.ok:
            sys.exit(1)
        check_stored_values()


if __name__ == "__main__":
    main()