import math
from functools import reduce

import funcy as fn
from aiger_bv import atom, UnsignedBVExpr

from aiger_coins import utils
from aiger_coins.dice import Distribution


def coin(prob, name=None):
    prob = utils.to_frac(prob)
    return dist((prob, 1 - prob), name=name)[0] \
        .with_output(name)


def _normalize_probs(probs):
    if len(probs) == 1:
        assert not isinstance(probs[0], int)

    probs = list(probs)
    if all(isinstance(x, int) for x in probs):
        total = sum(probs)
        probs = [(x, total) for x in probs]
    probs = fn.lmap(utils.to_frac, probs)

    if len(probs) == 1:
        probs = [probs[0], 1 - probs[0]]

    assert sum(probs) == 1
    return probs


def dist(probs, name=None, keep_seperate=False):
    """Distribution generated by mutually exclusive coins.

    Encoded using the common denominator method.
    """
    is_coin = len(probs) == 1
    probs = _normalize_probs(probs)
    bots = [p.denominator for p in probs]
    lcm = reduce(utils.lcm, bots, 1)
    word_len = max(math.ceil(math.log2(lcm)), 1)
    weights = map(lambda p: p.numerator*(lcm // p.denominator), probs)

    bits = atom(word_len, name, signed=False)
    if name is not None:
        bits = bits.with_output(name)

    name = bits.output

    const_true = ~(bits @ 0)
    total, coins, max_val = 0, [], 2**word_len
    for i, weight in enumerate(weights):
        lb = const_true if total == 0 else (bits >= total)
        total += weight
        ub = const_true if total == max_val else (bits < total)
        coins.append(lb & ub)

    is_valid = const_true if lcm == max_val else bits < lcm

    if keep_seperate:
        return coins, is_valid

    coins = reduce(UnsignedBVExpr.concat, coins) \
        .with_output(name)

    _dist = Distribution(expr=coins, valid=is_valid)
    return _dist[0] if is_coin else _dist


def binomial(n):
    chain = utils.chain(n)
    expr = UnsignedBVExpr(chain.unroll(n, only_last_outputs=True))
    const_true = ~(expr @ 0)
    return Distribution(expr=expr, valid=const_true)
