"""Differential entropies of continuous probability distributions -- numberdb.org/T241

For a real random variable X with density f, this stores

    h(X) = - integral f(x) log(f(x)) dx,

in nats and in bits, for named continuous distributions in their standard
forms.

Run it with SageMath:

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

Values are computed from closed forms in real and complex ball arithmetic.
Each standard form is stated in the table's parameter labels.
"""

import os
import sys

import numberdb.sage as numberdb
from sage.rings.complex_arb import ComplexBallField
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField


# Bits of working precision beyond what the written digits need.
#
# Measured by dry_run.py over all 626 entries: the widest returned ball still
# carries more than the 100 decimal digits the table asks for.
WORKING_GUARD = 128

NO_SHAPE_DISTRIBUTIONS = (
    'normal',
    'cauchy',
    'laplace',
    'gumbel',
    'rayleigh',
    'maxwell',
    'half-normal',
    'arcsine',
    'semicircle',
    'levy',
    'triangular',
    'hypsecant',
    'hypsecant-unit-variance',
)

WEIBULL_SHAPES = (QQ(1) / 2, QQ(3) / 2, QQ(2), QQ(3), QQ(4), QQ(5))

NO_SHAPE_COMMENTS = {
    'normal': r'In nats, $h=\frac12\log(2\pi e)$.',
    'cauchy': r'In nats, $h=\log(4\pi)$.',
    'laplace': r'In nats, $h=1+\log2$.',
    'gumbel': r'In nats, $h=1+\gamma$.',
    'rayleigh': r'In nats, $h=1+\frac{\gamma-\log2}{2}$.',
    'maxwell': r'In nats, $h=\gamma+\frac12\log(2\pi)-\frac12$.',
    'half-normal': r'In nats, $h=\frac12(1+\log(\pi/2))$.',
    'arcsine': r'In nats, $h=\log(\pi/4)$.',
    'semicircle': r'In nats, $h=\log\pi-\frac12$.',
    'levy': r'In nats, $h=\frac12(1+3\gamma+\log(16\pi))$.',
    'triangular': r'In nats, $h=\frac12-\log2$.',
    'hypsecant': r'In nats, $h=\log(2\pi)$.',
    'hypsecant-unit-variance': r'In nats, $h=\log4$.',
}


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


def _real_field(digits):
    return RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))


def _complex_field(digits):
    return ComplexBallField(numberdb.bits(digits, losing=WORKING_GUARD))


def _q(text):
    return QQ(str(text))


def _log_beta(field, a, b):
    return field(a).log_gamma() + field(b).log_gamma() - field(a + b).log_gamma()


def _psi(field, value):
    return field(value).psi()


def _bessel_i(digits, order, argument):
    field = _complex_field(digits)
    value = field(argument).bessel_I(field(order))
    if not (value.real().is_finite() and value.imag().is_finite()):
        raise ArithmeticError('computed a non-finite Bessel ball')
    if not value.imag().contains_zero():
        raise ArithmeticError('expected I_%s(%s) to be real' % (order, argument))
    return value.real()


def _shape_values():
    for value in range(1, 31):
        yield 'student-t', str(value)
    for value in range(1, 31):
        yield 'chi', str(value)
    for value in range(1, 31):
        yield 'chi-square', str(value)
    for numerator in range(1, 31):
        yield 'gamma', str(QQ(numerator) / QQ(2))
    for value in WEIBULL_SHAPES:
        yield 'weibull', str(value)
    for value in range(2, 11):
        yield 'pareto', str(value)
    for value in range(1, 11):
        yield 'von-mises', str(value)
    beta_parameters = [QQ(numerator) / QQ(2) for numerator in range(1, 11)]
    for i, a in enumerate(beta_parameters):
        for b in beta_parameters[i:]:
            yield 'beta', '%s,%s' % (a, b)
    for n1 in range(1, 11):
        for n2 in range(1, 11):
            yield 'f', '%s,%s' % (n1, n2)


def _entropy_nats(distribution, shape, digits):
    field = _real_field(digits)
    one = field(1)
    two = field(2)
    half = field(QQ(1) / QQ(2))
    log2 = two.log()
    pi = field.pi()
    gamma = field.euler_constant()

    if distribution == 'normal':
        return (one + (two * pi).log()) / 2
    if distribution == 'cauchy':
        return (field(4) * pi).log()
    if distribution == 'laplace':
        return one + log2
    if distribution == 'gumbel':
        return one + gamma
    if distribution == 'rayleigh':
        return one - log2 / 2 + gamma / 2
    if distribution == 'maxwell':
        return gamma + (two * pi).log() / 2 - field(QQ(1) / QQ(2))
    if distribution == 'half-normal':
        return (one + (pi / 2).log()) / 2
    if distribution == 'arcsine':
        return pi.log() - two * log2
    if distribution == 'semicircle':
        return pi.log() - half
    if distribution == 'levy':
        return (one + 3 * gamma + (16 * pi).log()) / 2
    if distribution == 'triangular':
        return half - log2
    if distribution == 'hypsecant':
        return (2 * pi).log()
    if distribution == 'hypsecant-unit-variance':
        return field(4).log()

    if distribution == 'student-t':
        nu = _q(shape)
        a = nu / 2
        b = QQ(1) / 2
        return (field(nu).log() / 2 + _log_beta(field, a, b)
                + (field(nu) + 1) / 2
                * (_psi(field, (nu + 1) / 2) - _psi(field, nu / 2)))

    if distribution == 'chi':
        k = _q(shape)
        a = k / 2
        if k == 1:
            return (one + (pi / 2).log()) / 2
        return (field(a).log_gamma() + field(k) / 2 - log2 / 2
                - (field(k) - 1) * _psi(field, a) / 2)

    if distribution == 'chi-square':
        k = _q(shape)
        a = k / 2
        return (field(k) / 2 + log2 + field(a).log_gamma()
                + (1 - field(a)) * _psi(field, a))

    if distribution == 'gamma':
        a = _q(shape)
        if a == 1:
            return ZZ(1)
        return field(a) + field(a).log_gamma() + (1 - field(a)) * _psi(field, a)

    if distribution == 'weibull':
        k = _q(shape)
        return gamma * (1 - 1 / field(k)) + one - field(k).log()

    if distribution == 'pareto':
        alpha = _q(shape)
        return one + 1 / field(alpha) - field(alpha).log()

    if distribution == 'von-mises':
        kappa = _q(shape)
        i0 = _bessel_i(digits, 0, kappa)
        i1 = _bessel_i(digits, 1, kappa)
        return (2 * pi * i0).log() - field(kappa) * i1 / i0

    if distribution == 'beta':
        a_text, b_text = shape.split(',')
        a = _q(a_text)
        b = _q(b_text)
        if a == 1 and b == 1:
            return ZZ(0)
        return (_log_beta(field, a, b)
                - (field(a) - 1) * _psi(field, a)
                - (field(b) - 1) * _psi(field, b)
                + (field(a + b) - 2) * _psi(field, a + b))

    if distribution == 'f':
        n1_text, n2_text = shape.split(',')
        n1 = _q(n1_text)
        n2 = _q(n2_text)
        if n1 == 2:
            return QQ(1) + QQ(2) / n2
        a = n1 / 2
        b = n2 / 2
        return ((field(n2) / field(n1)).log() + _log_beta(field, a, b)
                - (field(a) - 1) * _psi(field, a)
                - (field(b) + 1) * _psi(field, b)
                + field(a + b) * _psi(field, a + b))

    raise ValueError('unknown distribution %r' % (distribution,))


def _to_unit(value, unit, digits):
    if unit == 'nats':
        return value
    if unit != 'bits':
        raise ValueError('unknown unit %r' % (unit,))
    if value == 0:
        return ZZ(0)
    field = _real_field(digits)
    return field(value) / field(2).log()


def _equals(distribution, shape, unit):
    if distribution == 'student-t' and shape == '1':
        return 'HREF{#cauchy,-,%s}' % (unit,)
    if distribution == 'chi':
        if shape == '1':
            return 'HREF{#half-normal,-,%s}' % (unit,)
        if shape == '2':
            return 'HREF{#rayleigh,-,%s}' % (unit,)
        if shape == '3':
            return 'HREF{#maxwell,-,%s}' % (unit,)
    if distribution == 'beta' and shape == '1/2,1/2':
        return 'HREF{#arcsine,-,%s}' % (unit,)
    return ''


def _comment(distribution, shape):
    if shape == '-':
        return NO_SHAPE_COMMENTS.get(distribution, '')
    if distribution == 'student-t' and shape == '1':
        return 'This is the Cauchy distribution.'
    if distribution == 'chi' and shape == '1':
        return 'This is the half-normal distribution.'
    if distribution == 'chi' and shape == '2':
        return 'This is the Rayleigh distribution.'
    if distribution == 'chi' and shape == '3':
        return 'This is the Maxwell-Boltzmann distribution.'
    if distribution == 'gamma' and shape == '1':
        return 'This is the standard exponential distribution.'
    if distribution == 'beta' and shape == '1/2,1/2':
        return 'This is the arcsine distribution.'
    if distribution == 'beta' and shape == '1,1':
        return 'This is the uniform distribution on $[0,1]$.'
    return ''


class DifferentialEntropies(numberdb.Generator):

    table = os.environ.get('NUMBERDB_TABLE', 'T241')
    parameters = ('distribution', 'shape', 'unit')
    type = 'R'
    digits = 100
    rigour = 'proven'

    def enumerate(self):
        for distribution in NO_SHAPE_DISTRIBUTIONS:
            for unit in ('nats', 'bits'):
                yield {'distribution': distribution, 'shape': '-', 'unit': unit}
        for distribution, shape in _shape_values():
            for unit in ('nats', 'bits'):
                yield {'distribution': distribution, 'shape': shape, 'unit': unit}

    def value(self, params, digits):
        distribution = str(params['distribution'])
        shape = str(params['shape'])
        unit = str(params['unit'])
        if distribution == 'hypsecant-unit-variance' and unit == 'bits':
            value = ZZ(2)
        else:
            value = _to_unit(_entropy_nats(distribution, shape, digits),
                             unit, digits)
        equals = _equals(distribution, shape, unit)
        comment = _comment(distribution, shape)
        if equals or comment:
            record = {'number': value}
            if equals:
                record['equals'] = equals
            if comment:
                record['comment'] = comment
            return record
        return value


if __name__ == '__main__':
    _key_from_stdin()
    generator = DifferentialEntropies()
    if os.environ.get('NUMBERDB_PUBLISH') == '1' or '--publish' in sys.argv:
        print(generator.publish(message='differential entropy values'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)