import argparse
import json
from openpyxl import Workbook, utils
from prjxray.tile import OutPinTiming, InPinTiming
from prjxray.timing import Outpin, Inpin, Wire, Buffer, \
        PassTransistor, IntristicDelay, RcElement, PvtCorner
from prjxray.math_models import ExcelMathModel
from prjxray.db import Database

FAST = PvtCorner.FAST
SLOW = PvtCorner.SLOW


class TimingLookup(object):
    def __init__(self, db, nodes):
        self.db = db
        self.grid = db.grid()
        self.nodes = nodes

    def try_find_site_pin(self, site_pin_node, node_idx):
        site_pin_wire = self.nodes[site_pin_node]['wires'][node_idx]['name']
        tile, wire_in_tile = site_pin_wire.split('/')

        gridinfo = self.grid.gridinfo_at_tilename(tile)

        tile_type = self.db.get_tile_type(gridinfo.tile_type)

        for site in tile_type.get_sites():
            for site_pin in site.site_pins:
                if site_pin.wire == wire_in_tile:
                    return site_pin

        return None

    def find_site_pin(self, site_pin_node, node_idx):
        site_pin = self.try_find_site_pin(site_pin_node, node_idx)

        assert site_pin is not None, site_pin_node
        return site_pin

    def find_pip(self, pip_name):
        tile, pip = pip_name.split('/')

        gridinfo = self.grid.gridinfo_at_tilename(tile)

        tile_type = self.db.get_tile_type(gridinfo.tile_type)

        return tile_type.get_pip_by_name(pip)

    def find_wire(self, wire_name):
        tile, wire_in_tile = wire_name.split('/')

        gridinfo = self.grid.gridinfo_at_tilename(tile)

        tile_type = self.db.get_tile_type(gridinfo.tile_type)

        return tile_type.wires[wire_in_tile]


def delays_to_cells(ws, row, delays, cells):
    cells['FAST_MAX'] = 'E{}'.format(row)
    cells['FAST_MIN'] = 'F{}'.format(row)
    cells['SLOW_MAX'] = 'G{}'.format(row)
    cells['SLOW_MIN'] = 'H{}'.format(row)

    if delays is not None:
        ws[cells['FAST_MAX']] = delays[FAST].max
        ws[cells['FAST_MIN']] = delays[FAST].min
        ws[cells['SLOW_MAX']] = delays[SLOW].max
        ws[cells['SLOW_MIN']] = delays[SLOW].min
    else:
        ws[cells['FAST_MAX']] = 0
        ws[cells['FAST_MIN']] = 0
        ws[cells['SLOW_MAX']] = 0
        ws[cells['SLOW_MIN']] = 0


def cells_to_delays(cells):
    return {
        FAST: IntristicDelay(min=cells['FAST_MIN'], max=cells['FAST_MAX']),
        SLOW: IntristicDelay(min=cells['SLOW_MIN'], max=cells['SLOW_MAX']),
    }


class Net(object):
    def __init__(self, net):
        self.net = net
        self.ipin_nodes = {}
        self.row = None
        self.math = ExcelMathModel()
        self.models = {}

        for ipin in net['ipins']:
            self.ipin_nodes[ipin['node']] = ipin

        # Map of wire name to parent node
        self.wire_to_node = {}

        # Map of node name to node
        self.node_name_to_node = {}

        for node in net['nodes']:
            self.node_name_to_node[node['name']] = node
            for wire in node['wires']:
                self.wire_to_node[wire['name']] = node

        # Map of (src node, dst wire).
        self.pips = {}
        for pip in net['pips']:
            src_node = self.wire_to_node[pip['src_wire']]['name']
            dst_wire = pip['dst_wire'].split('/')[1]
            self.pips[(src_node, dst_wire)] = pip

            if not pip['is_directional']:
                dst_node = self.wire_to_node[pip['dst_wire']]['name']
                src_wire = pip['src_wire'].split('/')[1]
                self.pips[(dst_node, src_wire)] = pip

    def extend_rc_tree(self, ws, current_rc_root, timing_lookup, node):
        rc_elements = []
        for wire in node['wires']:
            wire_timing = timing_lookup.find_wire(wire['name'])
            ws['A{}'.format(self.row)] = wire['name']
            ws['B{}'.format(self.row)] = 'Part of wire'

            if wire_timing is not None:
                cells = {}
                cells['R'] = 'C{}'.format(self.row)
                cells['C'] = 'D{}'.format(self.row)
                ws[cells['R']] = wire_timing.resistance
                ws[cells['C']] = wire_timing.capacitance
                rc_elements.append(
                    RcElement(
                        resistance=cells['R'],
                        capacitance=cells['C'],
                    ))

            self.row += 1

        wire_rc_node = Wire(rc_elements=rc_elements, math=self.math)
        self.models[self.row - 1] = wire_rc_node

        current_rc_root.set_sink_wire(wire_rc_node)

        return wire_rc_node

    def descend_route(
            self,
            ws,
            timing_lookup,
            current_node,
            route,
            route_idx,
            current_rc_root,
            was_opin=False):
        """ Traverse the next pip, or recurse deeper. """

        # descend_route should've consumed this token
        assert route[route_idx] != '}'

        while route[route_idx] == '{':
            # Go deeper
            route_idx = self.descend_route(
                ws,
                timing_lookup,
                current_node,
                route,
                route_idx=route_idx + 1,
                current_rc_root=current_rc_root,
                was_opin=was_opin)

        next_edge = (current_node, route[route_idx])
        route_idx += 1
        assert next_edge in self.pips, (next_edge, self.pips.keys())

        pip = self.pips[next_edge]
        is_backward = self.wire_to_node[
            pip['dst_wire']]['name'] == current_node
        if not is_backward:
            assert self.wire_to_node[
                pip['src_wire']]['name'] == current_node, (current_node, pip)

        pip_info = timing_lookup.find_pip(pip['name'])
        if not is_backward:
            pip_timing = pip_info.timing
            current_node = self.wire_to_node[pip['dst_wire']]['name']
        else:
            pip_timing = pip_info.backward_timing
            current_node = self.wire_to_node[pip['src_wire']]['name']

        ws['A{}'.format(self.row)] = pip['name']

        cells = {}
        cells['R'] = 'C{}'.format(self.row)
        cells['C'] = 'D{}'.format(self.row)
        delays_to_cells(
            ws, row=self.row, delays=pip_timing.delays, cells=cells)
        delays = cells_to_delays(cells)

        if pip_info.is_pass_transistor:
            ws['B{}'.format(self.row)] = 'PassTransistor'
            ws[cells['R']] = pip_timing.drive_resistance

            pip_model = PassTransistor(
                drive_resistance=cells['R'],
                delays=delays,
            )
        else:
            ws['B{}'.format(self.row)] = 'Buffer'
            if pip_timing.drive_resistance is not None:
                ws[cells['R']] = pip_timing.drive_resistance
                if pip_timing.drive_resistance == 0 and was_opin:
                    new_site_pin = timing_lookup.try_find_site_pin(
                        current_node, node_idx=0)
                    if new_site_pin is not None:
                        ws[cells['R']] = new_site_pin.timing.drive_resistance
            else:
                ws[cells['R']] = 0

            if pip_timing.internal_capacitance is not None:
                ws[cells['C']] = pip_timing.internal_capacitance
            else:
                ws[cells['C']] = 0

            pip_model = Buffer(
                drive_resistance=cells['R'],
                internal_capacitance=cells['C'],
                delays=cells_to_delays(cells),
            )

        self.models[self.row] = pip_model

        current_rc_root.add_child(pip_model)
        self.row += 1

        if current_node in self.ipin_nodes:
            assert route[route_idx] == '}'
            route_idx += 1

        node = self.node_name_to_node[current_node]

        current_rc_root = self.extend_rc_tree(
            ws, pip_model, timing_lookup, node)

        if current_node in self.ipin_nodes:
            ipin = self.ipin_nodes[current_node]

            cells = {}
            name = '{} to {}'.format(self.net['opin']['name'], ipin['name'])
            ws['A{}'.format(self.row)] = ipin['name']
            ws['B{}'.format(self.row)] = 'Inpin'

            site_pin = timing_lookup.find_site_pin(ipin['node'], node_idx=-1)
            assert isinstance(site_pin.timing, InPinTiming)

            cells = {}
            cells['C'] = 'D{}'.format(self.row)
            delays_to_cells(
                ws, row=self.row, delays=site_pin.timing.delays, cells=cells)
            delays = cells_to_delays(cells)

            ws[cells['C']] = site_pin.timing.capacitance
            ipin_model = Inpin(
                capacitance=cells['C'], delays=delays, name=name)
            self.models[self.row] = ipin_model
            current_rc_root.add_child(ipin_model)
            self.row += 1

            #Sum delays only (sum*1000)
            #Sum delays + capacitive delay
            #
            #Total delay (from Vivado)
            ws['A{}'.format(self.row)] = '{}: {} sum delays'.format(
                self.net['net'], name)
            self.row += 1

            ws['A{}'.format(self.row)] = '{}: {}sum delays + cap delay'.format(
                self.net['net'], name)
            self.row += 2

            ws['A{}'.format(
                self.row)] = '{}: {}Total delay (from Vivado)'.format(
                    self.net['net'], name)

            ws['E{}'.format(self.row)] = ipin['ic_delays']['FAST_MAX']
            ws['F{}'.format(self.row)] = ipin['ic_delays']['FAST_MIN']
            ws['G{}'.format(self.row)] = ipin['ic_delays']['SLOW_MAX']
            ws['H{}'.format(self.row)] = ipin['ic_delays']['SLOW_MIN']

            self.row += 2

            return route_idx
        else:
            return self.descend_route(
                ws,
                timing_lookup,
                current_node,
                route,
                route_idx=route_idx,
                current_rc_root=current_rc_root)

    def walk_route(self, ws, timing_lookup):
        """ Walk route, creating rows in table.

        First row will always be the OPIN, followed by the node/wire connected
        to the OPIN.  After a node/wire is always 1 or more pips. After a pip
        is always a node/wire.  A terminal node/wire will then reach an IPIN.
        """

        self.row = 2
        ws['A{}'.format(self.row)] = self.net['opin']['wire']

        site_pin = timing_lookup.find_site_pin(
            self.net['opin']['node'], node_idx=0)
        assert isinstance(site_pin.timing, OutPinTiming)

        ws['B{}'.format(self.row)] = 'Outpin'
        ws['C{}'.format(self.row)] = site_pin.timing.drive_resistance

        cells = {}
        cells['R'] = 'C{}'.format(self.row)

        delays_to_cells(
            ws, row=self.row, delays=site_pin.timing.delays, cells=cells)

        model_root = Outpin(
            resistance=cells['R'], delays=cells_to_delays(cells))
        self.models[self.row] = model_root
        self.row += 1

        node = self.net['opin']['node']

        tile, first_wire = self.net['opin']['node'].split('/')

        route = [r for r in self.net['route'].strip().split(' ') if r != '']
        assert route[0] == '{'
        assert route[1] == first_wire

        node = self.node_name_to_node[node]
        current_rc_root = self.extend_rc_tree(
            ws, model_root, timing_lookup, node)

        self.descend_route(
            ws,
            timing_lookup,
            node['name'],
            route,
            route_idx=2,
            current_rc_root=current_rc_root,
            was_opin=True)

        model_root.propigate_delays(self.math)

        model_rows = {}

        for row, model in self.models.items():
            model_rows[id(model)] = row

        for row, model in self.models.items():
            rc_delay = model.get_rc_delay()
            if rc_delay is not None:
                ws['J{}'.format(row)] = self.math.eval(rc_delay)

            downstream_cap = model.get_downstream_cap()
            if downstream_cap is not None:
                ws['I{}'.format(row)] = self.math.eval(downstream_cap)

            if isinstance(model, Inpin):
                ipin_results = {
                    'Name': model.name,
                    'truth': {},
                    'computed': {},
                }

                ipin_delays = {}

                DELAY_COLS = (
                    ('E', 'FAST_MAX'),
                    ('F', 'FAST_MIN'),
                    ('G', 'SLOW_MAX'),
                    ('H', 'SLOW_MIN'),
                )

                for col, value in DELAY_COLS:
                    ipin_delays[value] = []
                    ipin_results['computed'][
                        value] = '{title}!{col}{row}'.format(
                            title=utils.quote_sheetname(ws.title),
                            col=col,
                            row=row + 2)
                    ipin_results['truth'][value] = '{title}!{col}{row}'.format(
                        title=utils.quote_sheetname(ws.title),
                        col=col,
                        row=row + 4)

                rc_delays = []

                for model in model.get_delays():
                    delays = model.get_intrinsic_delays()
                    if delays is not None:
                        ipin_delays['FAST_MAX'].append(delays[FAST].max)
                        ipin_delays['FAST_MIN'].append(delays[FAST].min)
                        ipin_delays['SLOW_MAX'].append(delays[SLOW].max)
                        ipin_delays['SLOW_MIN'].append(delays[SLOW].min)

                    if id(model) in model_rows:
                        rc_delays.append('J{}'.format(model_rows[id(model)]))

                ws['J{}'.format(row + 1)] = self.math.eval(
                    self.math.sum(rc_delays))

                for col, value in DELAY_COLS:
                    ws['{}{}'.format(col, row + 1)] = self.math.eval(
                        self.math.sum(ipin_delays[value]))
                    ws['{}{}'.format(
                        col, row + 2)] = '=1000*({col}{row} + J{row})'.format(
                            col=col, row=row + 1)

                yield ipin_results


def add_net(wb, net, timing_lookup):
    ws = wb.create_sheet(
        title="Net {}".format(net['net'].replace('[', '_').replace(']', '_')))

    # Header
    ws['A1'] = 'Name'
    ws['B1'] = 'Type'
    ws['C1'] = 'RES'
    ws['D1'] = 'CAP'
    ws['E1'] = 'FAST_MAX'
    ws['F1'] = 'FAST_MIN'
    ws['G1'] = 'SLOW_MAX'
    ws['H1'] = 'SLOW_MIN'
    ws['I1'] = 'Downstream C'
    ws['J1'] = 'Delay from cap'

    net_obj = Net(net)
    yield from net_obj.walk_route(ws, timing_lookup)


def main():
    parser = argparse.ArgumentParser(
        description="Create timing worksheet for 7-series timing analysis.")

    parser.add_argument('--timing_json', required=True)
    parser.add_argument('--db_root', required=True)
    parser.add_argument('--output_xlsx', required=True)

    args = parser.parse_args()

    with open(args.timing_json) as f:
        timing = json.load(f)

    db = Database(args.db_root)

    nodes = {}
    for net in timing:
        for node in net['nodes']:
            nodes[node['name']] = node

    timing_lookup = TimingLookup(db, nodes)

    wb = Workbook()
    summary_ws = wb.get_sheet_by_name(wb.sheetnames[0])
    summary_ws.title = 'Summary'

    summary_ws['A1'] = 'Name'

    cols = ['FAST_MAX', 'FAST_MIN', 'SLOW_MAX', 'SLOW_MIN']
    cur_col = 'B'
    for col in cols:
        summary_ws['{}1'.format(cur_col)] = col
        cur_col = chr(ord(cur_col) + 1)
        summary_ws['{}1'.format(cur_col)] = 'Computed ' + col
        cur_col = chr(ord(cur_col) + 3)

    summary_row = 2
    for net in timing:
        if '<' in net['route']:
            print(
                "WARNING: Skipping net {} because it has complicated route description."
                .format(net['net']))
            continue

        for summary_cells in add_net(wb, net, timing_lookup):
            summary_ws['A{}'.format(summary_row)] = summary_cells['Name']

            cur_col = 'B'
            for col in cols:
                truth_col = chr(ord(cur_col) + 0)
                computed_col = chr(ord(cur_col) + 1)
                error_col = chr(ord(cur_col) + 2)
                error_per_col = chr(ord(cur_col) + 3)
                summary_ws['{}{}'.format(
                    truth_col,
                    summary_row)] = '=' + summary_cells['truth'][col]
                summary_ws['{}{}'.format(
                    computed_col,
                    summary_row)] = '=' + summary_cells['computed'][col]
                summary_ws['{}{}'.format(
                    error_col,
                    summary_row)] = '={truth}{row}-{comp}{row}'.format(
                        truth=truth_col, comp=computed_col, row=summary_row)
                summary_ws['{}{}'.format(
                    error_per_col,
                    summary_row)] = '={error}{row}/{truth}{row}'.format(
                        error=error_col, truth=truth_col, row=summary_row)

                cur_col = chr(ord(cur_col) + 4)

            summary_row += 1

    wb.save(filename=args.output_xlsx)


if __name__ == "__main__":
    main()