"""Translate symbolic representation to Python code."""
# Copyright 2017-2018 by California Institute of Technology
# All rights reserved. Licensed under 3-clause BSD.
#
from collections import defaultdict
import inspect
import pprint
import time

from omega.logic import bitvector as bv
from omega.symbolic import functions as fcn


def dumps_bdds_as_code(u, out_vars, aut, lang='python'):
    """Return code that evaluates BDD `u`.

    The returned code includes a function `step(state)`
    that takes as input a `dict` that maps unprimed
    variables to values, and returns a `dict` that maps
    `out_vars` to values.
    """
    out_bits = _list_bits(out_vars, aut.vars)
    outputs = fcn.make_functions(u, out_bits, aut.bdd)
    out_bdds = fcn.collect_functions(outputs)
    missing_bits = set(out_bits).difference(out_bdds)
    renaming = map_bits_to_bitvectors(aut.vars)
    code = dumps_bdd_as_code(
        out_bdds, aut.bdd, lang=lang, renaming=renaming)
    code = (4 * ' ').join(code.splitlines(True))
    code = template.format(
        code=code,
        vrs=pprint.pformat(aut.vars, indent=4),
        out_vars=pprint.pformat(out_vars, indent=4),
        missing_bits=repr(missing_bits),
        when=time.strftime('on %Y-%m-%d at %H:%M UTC%z'))
    # collect code of functions used here
    functions = [
        step, assign_bitvectors, int_to_bits,
        out_bits_to_ints, bv.bitfields_to_ints,
        bv._append_sign_bit, bv.twos_complement_to_int]
    lines = list()
    for func in functions:
        func_lines, _ = inspect.getsourcelines(func)
        lines.extend(func_lines)
        lines.append('\n\n')
    return code + ''.join(lines)


def step(state):
    """Return next values of controlled variables.

    @param state: `dict` that maps unprimed variable
        identifiers to values.
    """
    bitvectors = assign_bitvectors(state, vrs)
    out_bits = compute_bdds(bitvectors)
    out_var_values = out_bits_to_ints(out_bits, vrs)
    return out_var_values


template = '''\
# Code generated by the package `omega` {when}.
import logging


vrs = {vrs}
out_vars = {out_vars}
missing_bits = {missing_bits}
logger = logging.getLogger(__name__)


def compute_bdds(bitvectors):
    out_bits = dict()
    {code}
    return out_bits


# The code below is from `omega`, and:
#
# Copyright 2017-2018 by California Institute of Technology
# All rights reserved. Licensed under 3-clause BSD.

'''


def out_bits_to_ints(out_bits, vrs):
    """Convert output bit values to integer values."""
    new_state = dict(out_bits)
    # default value for don't care outputs
    other = {bit: False for bit in missing_bits}
    new_state.update(other)
    # restrict table to `out_vars`
    out_vars_table = {
        k: v for k, v in vrs.items()
        if k in out_vars}
    # next values of controlled variables.
    new_int_state = bitfields_to_ints(
        new_state, out_vars_table)
    return new_int_state


def _list_bits(vrs, table):
    """Return `list` of bits for `vrs`."""
    bits = list()
    for var in vrs:
        attr = table[var]
        typ = attr['type']
        if typ == 'bool':
            bits.append(var)
        else:
            bits.extend(attr['bitnames'])
    return bits


def map_bits_to_bitvectors(vrs):
    """Return `dict` that maps each bit to an array expression."""
    renaming = dict()
    for var, attr in vrs.items():
        typ = attr['type']
        if typ == 'bool':
            renaming[var] = var
            continue
        assert typ == 'int', typ
        bitnames = attr['bitnames']
        d = {
            name: 'bitvectors["{var}"][{i}]'.format(var=var, i=i)
            for i, name in enumerate(bitnames)}
        renaming.update(d)
    return renaming


def assign_bitvectors(state, vrs):
    """Return `dict` that maps each variable to bit values.

    @rtype: `list` of `bool`.
    """
    bitvectors = dict()
    for var, value in state.items():
        bitnames = vrs[var]['bitnames']
        width = len(bitnames)
        bitvectors[var] = int_to_bits(value, width)
    return bitvectors


def int_to_bits(x, width):
    """Return `list` of `bool` for integer value `x`."""
    # Adapted from the function
    #   `omega.logic.bitvector.int_to_twos_complement`.
    # The sign bit is not needed.
    n = x.bit_length()
    if x >= 0:
        y = x
    else:
        y = 2**n + x
    m = max(width, n, 1)  # if y == 0 then n == 0
    bits = bin(y).lstrip('-0b').zfill(m)
    bits = list(reversed(bits))
    bits = [bool(int(b)) for b in bits]
    return bits


languages = dict(
    python=dict(
        FALSE='False',
        TRUE='True',
        NOT='not',
        AND='and',
        OR='or',
        COMMENT='#',  # single-line comment delimiter
        SEP=''),  # statement separator
    c=dict(
        FALSE='false',
        TRUE='true',
        NOT='!',
        AND='&&',
        OR='||',
        COMMENT='//',
        SEP=';'))


def dumps_bdd_as_code(
        roots,
        bdd,
        lang='python',
        renaming=None):
    """Return code that computes root values from bits.

    @param roots: `dict` that maps each "output" bit to
        a BDD that depends on "input" bits
    @param bdd: BDD manager
    @param renaming: `dict` that maps each non-output
        bit in the support (unprimed bits for a Moore action)
        to an expression as `str`
    @return: code that computes the value of each output
        bit and stores it in a `dict` named `out_bits`,
        given the values of unprimed bits, as expressions
        defined in `renaming`.
    """
    if renaming is None:
        renaming = dict()
    syntax = languages[lang]
    layers, out_lines = _collect_layers(roots, syntax, bdd)
    # convert each layer to Python code
    lines = list()
    latches = set()  # names, in case we need to declare them
    # from bottom to top layers
    for level in sorted(layers, reverse=True):
        layer = layers[level]
        _comment_level(level, lines, syntax)
        _dumps_layer(
            layer, lines, latches, syntax, bdd, renaming)
    lines.extend(out_lines)
    gen = (_append_sep(line, syntax) for line in lines)
    code = '\n'.join(gen)
    return code


def _collect_layers(roots, syntax, bdd):
    """Return layers of nodes to dump, and root assignments."""
    layers = defaultdict(list)
    lines = list()
    for name, u in roots.items():
        _register_nodes(u, layers, bdd)
        # assign to `name` the latch value
        latch = _latch_ref(u, syntax)
        line = 'out_bits["{name}"] = {latch}'.format(
            name=name, latch=latch)
        lines.append(line)
    return layers, lines


def _comment_level(level, lines, syntax):
    """Comment indicating the level that follows."""
    line = '{C} level: {level}'.format(
        C=syntax['COMMENT'],
        level=level)
    lines.append(line)


def _append_sep(line, syntax):
    """Append statement separator to `line`.

    Omit separator if entire `line` is a comment.
    Raise `ValueError` if `line` contains code and comment.
    """
    com = syntax['COMMENT']
    if line.startswith(com):
        return line
    if com in line:
        raise ValueError(line)
    # line contains only code
    sep = syntax['SEP']
    return '{line}{SEP}'.format(line=line, SEP=sep)


def _dumps_layer(
        layer, lines, latches, syntax, bdd, renaming):
    """Format the nodes at a single layer."""
    # Nodes in each layer depend on only nodes from
    # layers at lower levels, so all latches are defined.
    for node_id in layer:
        _dumps_node(node_id, lines, latches, syntax, bdd, renaming)


def _dumps_node(
        node_id, lines, latches, syntax, bdd, renaming):
    """Format the assignment to a latch for an internal BDD node."""
    node = bdd._add_int(node_id)
    latch = _latch_name(node, syntax)
    # not terminal
    assert node.var is not None, node.var
    assert latch not in latches, latch
    latches.add(latch)
    #
    low_latch = _latch_ref(node.low, syntax)
    high_latch = _latch_ref(node.high, syntax)
    # ternary conditional for Boolean-valued variables
    bit = renaming.get(node.var, node.var)
    line = (
        '{latch} = (\n'
        '    ({bit} {AND} {high_latch}) {OR} \n'
        '    (({NOT} {bit}) {AND} {low_latch}))'
        ).format(
            latch=latch,
            bit=bit,
            low_latch=low_latch,
            high_latch=high_latch,
            AND=syntax['AND'],
            OR=syntax['OR'],
            NOT=syntax['NOT'])
    lines.append(line)


def _latch_ref(node, syntax):
    """Return latch reference for `node`.

    This function takes into account
    whether a reference is a complemented edge.
    """
    ref = _latch_name(node, syntax)
    if node.negated:
        ref = '({NOT} {ref})'.format(
            NOT=syntax['NOT'], ref=ref)
    return ref


def _latch_name(node, syntax):
    """Return latch name for `node`."""
    # terminal ?
    if node.var is None:
        latch = syntax['TRUE']
    else:
        # avoid identifiers containing `-`
        node_id = str(int(node)).replace('-', 'n')
        latch = 'latch_{i}'.format(i=node_id)
    return latch


def _register_nodes(u, layers, bdd):
    """Collect nodes to be dumped."""
    # terminal ?
    if u.var is None:
        return
    level, low, high = bdd.succ(u)
    # visited ?
    node_id = int(u)
    if node_id in layers[level]:
        return
    # use `int` to avoid creating a `Function`
    # instance for each node.
    layers[level].append(node_id)
    _register_nodes(low, layers, bdd)
    _register_nodes(high, layers, bdd)
