''' Tool creating and manipulating normal distributions of random variable.

Features a composite class that treats the mean and standard deviation of
measurements as single entity.

'''

from math import hypot, sqrt, fabs, exp, erf, tau
from statistics import mean as fmean, stdev
import random

class NormalDist:
    'Normal distribution of a random variable'
    # https://en.wikipedia.org/wiki/Normal_distribution
    # https://en.wikipedia.org/wiki/Variance#Properties

    __slots__ = ('mu', 'sigma')

    def __init__(x1, mu, sigma=0.0):
        assert sigma >= 0.0
        x1.mu = mu
        x1.sigma = sigma

    @classmethod
    def from_samples(cls, data):
        'Make a normal distribution from a sample'
        data = list(data)
        xbar = fmean(data)
        return cls(xbar, stdev(data, xbar))

    @property
    def variance(self):
        'Square of the standard deviation'
        return self.sigma ** 2.0

    def examples(self, n, seed=None):
        'Generate *n* samples for given mean and standard deviation'
        gauss = random.gauss if seed is None else random.Random(seed).gauss
        mu, sigma = self.mu, self.sigma
        return [gauss(mu, sigma) for i in range(n)]

    def pdf(self, x):
        'Probability density function:  P(x <= X < x+dx) / dx'
        variance = self.sigma ** 2.0
        return exp((x - self.mu)**2.0 / (-2.0*variance)) / sqrt(tau * variance)

    def cdf(self, x):
        'Cumulative density function:  P(X <= x)'
        return 0.5 * (1.0 + erf((x - self.mu) / (self.sigma * sqrt(2.0))))

    def __repr__(self):
        return f'{type(self).__name__}(mu={self.mu!r}, sigma={self.sigma!r})'

    def __add__(x1, x2):
        if isinstance(x2, NormalDist):
            return NormalDist(x1.mu + x2.mu, hypot(x1.sigma, x2.sigma))
        return NormalDist(x1.mu + x2, x1.sigma)

    def __sub__(x1, x2):
        if isinstance(x2, NormalDist):
            return NormalDist(x1.mu - x2.mu, hypot(x1.sigma, x2.sigma))
        return NormalDist(x1.mu - x2, x1.sigma)

    def __mul__(x1, x2):
        return NormalDist(x1.mu * x2, x1.sigma * fabs(x2))

    def __truediv__(num, den):
        return NormalDist(num.mu / den, num.sigma / fabs(den))

    def __pos__(x1):
        return x1

    def __neg__(x1):
        return NormalDist(-x1.mu, x1.sigma)

    __radd__ = __add__

    def __rsub__(x1, x2):
        return -(x1 - x2)

    __rmul__ = __mul__


if __name__ == '__main__':

    from math import isclose
    from operator import add, sub, mul, truediv
    from itertools import repeat

    g1 = NormalDist(10, 20)
    g2 = NormalDist(-5, 25)

    # Test scaling by a constant
    assert (g1 * 5 / 5).mu == g1.mu
    assert (g1 * 5 / 5).sigma == g1.sigma

    n = 100_000
    G1 = g1.examples(n)
    G2 = g2.examples(n)

    for func in (add, sub):
        print(f'\nTest {func.__name__} with another NormalDist:')
        print(func(g1, g2))
        print(NormalDist.from_samples(map(func, G1, G2)))

    const = 11
    for func in (add, sub, mul, truediv):
        print(f'\nTest {func.__name__} with a constant:')
        print(func(g1, const))
        print(NormalDist.from_samples(map(func, G1, repeat(const))))

    const = 19
    for func in (add, sub, mul):
        print(f'\nTest constant with {func.__name__}:')
        print(func(const, g1))
        print(NormalDist.from_samples(map(func, repeat(const), G1)))

    def assert_close(G1, G2):
        assert isclose(G1.mu, G1.mu, rel_tol=0.01), (G1, G2)
        assert isclose(G1.sigma, G2.sigma, rel_tol=0.01), (G1, G2)

    X = NormalDist(-105, 73)
    Y = NormalDist(31, 47)
    s = 32.75
    n = 100_000

    S = NormalDist.from_samples([x + s for x in X.examples(n)])
    assert_close(X + s, S)

    S = NormalDist.from_samples([x - s for x in X.examples(n)])
    assert_close(X - s, S)

    S = NormalDist.from_samples([x * s for x in X.examples(n)])
    assert_close(X * s, S)

    S = NormalDist.from_samples([x / s for x in X.examples(n)])
    assert_close(X / s, S)

    S = NormalDist.from_samples([x + y for x, y in zip(X.examples(n),
                                                       Y.examples(n))])
    assert_close(X + Y, S)

    S = NormalDist.from_samples([x - y for x, y in zip(X.examples(n),
                                                       Y.examples(n))])
    assert_close(X - Y, S)