"""Shared completed-zeta Loewner construction used by active certificates."""

from __future__ import annotations

from fractions import Fraction

import mpmath as mp
from flint import arb, arb_mat, arb_series, ctx


def as_arb(value: Fraction | int) -> arb:
    if isinstance(value, int):
        return arb(value)
    return arb(value.numerator) / value.denominator


def completed_log_derivative(rate: Fraction) -> tuple[arb, arb]:
    lam = as_arb(rate)
    sigma = lam + arb(1) / 2
    zeta_series = arb_series([sigma, arb(1)], 3).zeta()
    zeta_log_derivative = zeta_series[1] / zeta_series[0]
    zeta_log_derivative_prime = (
        2 * zeta_series[2] / zeta_series[0] - zeta_log_derivative**2
    )
    gamma_argument = arb(1) / 4 + lam / 2
    value = (
        1 / (lam + arb(1) / 2)
        + 1 / (lam - arb(1) / 2)
        - arb.pi().log() / 2
        + gamma_argument.digamma() / 2
        + zeta_log_derivative
    )
    derivative = (
        -1 / (lam + arb(1) / 2) ** 2
        - 1 / (lam - arb(1) / 2) ** 2
        + arb.zeta(arb(2), gamma_argument) / 4
        + zeta_log_derivative_prime
    )
    return value, derivative


def loewner_matrix(rates: tuple[Fraction, ...]) -> list[list[arb]]:
    data = [completed_log_derivative(rate) for rate in rates]
    matrix: list[list[arb]] = []
    for row, left_rate in enumerate(rates):
        left = as_arb(left_rate)
        entries: list[arb] = []
        for column, right_rate in enumerate(rates):
            right = as_arb(right_rate)
            if row == column:
                entry = 2 * (data[row][0] / left - data[row][1])
            else:
                entry = 4 * (
                    right * data[row][0] - left * data[column][0]
                ) / (right**2 - left**2)
            entries.append(entry)
        matrix.append(entries)
    return matrix


def midpoint_preconditioner(matrix: list[list[arb]]) -> arb_mat:
    digits = max(300, int(ctx.prec * 0.31))
    mp.mp.dps = digits
    midpoint = mp.matrix(
        [
            [
                mp.mpf(matrix[row][column].mid().str(digits, radius=False))
                for column in range(len(matrix))
            ]
            for row in range(len(matrix))
        ]
    )
    inverse_cholesky = mp.cholesky(midpoint) ** -1
    return arb_mat(
        [
            [
                arb(mp.nstr(inverse_cholesky[row, column], digits))
                for column in range(len(matrix))
            ]
            for row in range(len(matrix))
        ]
    )
