#!/usr/bin/env python

from docopt import docopt
from jinja2 import Template
from sg_parser import load_sg
from lib_parser import load_lib
from lib_parser import merge_libs
from collections import defaultdict
from verilog_parser import load_verilog

import os
import re
import json
import math


usage = """generator.py

Usage:
  generator.py [options] <circuit.v>

Options:
  -o --output=<dir>         Write output files to directory.
  -s --spec=<spec.sg>       Generate spec model for verification.
  -t --template=<template>  Specify template [default: standard].
  -l --lib=<dir>            Load additional library file(s).

"""


def read_file(file):
    """Return content of file as a string."""
    with open(file, "r") as fid:
        return fid.read()


def write_file(content, file):
    """Write string 'content' to file."""
    with open(file, "w") as fid:
        fid.write(content)


def get_workcraft_state_ind(state):
    """Attempt to extract state index assuming state name is in the format
    generated by workcraft, e.g. s7_0101101 (returns 7 in this case).

    Return None is the state name cannot be matched.
    """

    reg1 = r"^s([0-9]+)_"
    pat1 = re.compile(reg1, flags=re.MULTILINE)
    matches = pat1.findall(state)
    return int(matches[0]) if len(matches) == 1 else None


def get_state_inds(spec):
    """Return a map: state -> index."""

    froms = [item[0] for item in spec["transitions"]]
    tos = [item[2] for item in spec["transitions"]]

    states = sorted(set(froms + tos))
    workcraft_inds = map(get_workcraft_state_ind, states)
    workcraft_valid = not any([ind is None for ind in workcraft_inds])
    inds = workcraft_inds if workcraft_valid else range(len(states))

    inds = {state: ind for ind, state in zip(inds, states)}
    return inds


def bit_size(n):
    """Return minimum number of bits required to represent n."""
    bits = math.log(n) / math.log(2)
    return int(math.ceil(bits))


def get_nond_groups(transitions):
    """Return two dicts, inds and counts, representing non-deterministic
    transitions choices.

    'counts' is (before, tr) -> count of (before, tr) in transitions
    'inds' is the index of (before, tr, after) within its group of (before, tr)
    """

    counts = defaultdict(lambda: 0)
    inds = dict()

    for before, tr, after in transitions:
        ckey = (before, tr)
        counts[ckey] += 1
        inds[(before, tr, after)] = counts[ckey]

    return inds, counts


def get_circuit_context(circuit, lib):
    """Return a Jinja context representing a circuit."""

    stateful = {
        inst: body
        for inst, body in circuit["modules"].iteritems()
        if not body.get("short_delay")
    }

    stateless = {
        inst: body
        for inst, body in circuit["modules"].iteritems()
        if body.get("short_delay")
    }

    def get_output_pin(mod):
        return lib[mod["type"]]["output"]

    def get_output_net(mod):
        return mod["connections"][get_output_pin(mod)]

    stateful_nets = sorted(map(get_output_net, stateful.values()))
    stateless_nets = sorted(map(get_output_net, stateless.values()))
    stateless_outs = set(circuit["outputs"]) & set(stateless_nets)
    inputs = sorted(circuit["inputs"])
    outputs = sorted(circuit["outputs"])
    nets = inputs + stateful_nets
    firebits = bit_size(len(nets))
    firing_indices = {net: nets.index(net) for net in nets}

    return {
        "lib": lib,
        "nets": nets,
        "inputs": inputs,
        "outputs": outputs,
        "firebits": firebits,
        "stateful": stateful,
        "stateless": stateless,  # dictionary of stateless modules
        "stateful_nets": stateful_nets,
        "stateless_outs": stateless_outs,
        "get_output_net": get_output_net,
        "get_output_pin": get_output_pin,
        "firing_indices": firing_indices,
        "initial_circuit": circuit["initial_state"]  # dict: signal -> state
    }


def get_spec_context(spec, circuit, lib):
    """Return a Jinja context representing a spec."""

    transitions = sorted(spec["transitions"])
    ndinds, ndcounts = get_nond_groups(spec["transitions"])
    ndbits = bit_size(max(ndcounts.values()))

    return {
        "ndinds": ndinds,
        "ndbits": ndbits,
        "ndcounts": ndcounts,
        "state_inds": get_state_inds(spec),  # state -> index
        "transitions": transitions,
        "initial_spec": spec["initial_state"]  # string
    }


def generate_circuit(circuit, lib, template):

    circuit_context = get_circuit_context(circuit, lib)
    template = Template(read_file(template))
    return template.render(circuit_context)


def generate_spec(spec, circuit, lib, template):

    spec_context = get_spec_context(spec, circuit, lib)
    circuit_context = get_circuit_context(circuit, lib)

    combined_context = dict()
    combined_context.update(spec_context)
    combined_context.update(circuit_context)

    template = Template(read_file(template))
    return template.render(combined_context)


def get_tool_path():
    """Return tool root directory."""
    return os.path.split(__file__)[0]


def main():

    args = docopt(usage, version="generator.py v0.1")

    built_in_libs = os.path.join(get_tool_path(), 'libraries', '*')

    templates_dir = os.path.join(get_tool_path(), 'templates',
                                 args["--template"])

    if args['--lib']:
        user_libs = os.path.join(args['--lib'], '*')
    else:
        user_libs = None

    lib_paths = [built_in_libs] + ([user_libs] if user_libs else [])

    lib = load_lib(*lib_paths)

    spec = load_sg(args["--spec"]) if args["--spec"] else None

    circuit = load_verilog(args["<circuit.v>"])

    def gen_circ(template_file):
        return generate_circuit(circuit, lib, template_file)

    def gen_spec(template_file):
        return generate_spec(spec, circuit, lib, template_file)

    if spec:
        template_files = [("circuit.v", gen_circ), ("spec.v", gen_spec)]
    else:
        template_files = [("circuit.v", gen_circ)]

    for file, generate in template_files:

        template_file = os.path.join(templates_dir, file)
        content = generate(template_file)

        if args['--output']:

            output_file = os.path.join(args['--output'], file)

            print "Generating %s ...." % file

            if not os.path.exists(args['--output']):
                os.makedirs(args['--output'])

            write_file(content, output_file)

        else:

            print content


if __name__ == '__main__':
    main()
