""" Generate grid from database dump """

from __future__ import print_function
import argparse
import prjxray.lib
import prjxray.util
import pyjson5 as json5
import multiprocessing
import progressbar
import os.path
import json
import datetime
import pickle
import sys

from utils import xjson


def get_tile_grid_info(fname):
    with open(fname, 'r') as f:
        tile = json5.load(f)

    return {
        tile['tile']: {
            'type': tile['type'],
            'grid_x': tile['x'],
            'grid_y': tile['y'],
            'sites': dict(
                (site['site'], site['type']) for site in tile['sites']),
            'wires': set((wire['wire'] for wire in tile['wires']))
        },
    }


def read_json5(fname):
    with open(fname, 'r') as f:
        return json5.load(f)


def is_edge_shared(edge1, edge2):
    """ Returns true if edge1 or edge2 overlap

  >>> is_edge_shared((0, 1), (0, 1))
  True
  >>> is_edge_shared((0, 2), (0, 1))
  True
  >>> is_edge_shared((0, 1), (0, 2))
  True
  >>> is_edge_shared((1, 2), (0, 3))
  True
  >>> is_edge_shared((0, 3), (1, 2))
  True
  >>> is_edge_shared((1, 2), (0, 2))
  True
  >>> is_edge_shared((0, 2), (1, 2))
  True
  >>> is_edge_shared((0, 2), (1, 3))
  True
  >>> is_edge_shared((1, 3), (0, 2))
  True
  >>> is_edge_shared((0, 1), (1, 2))
  False
  >>> is_edge_shared((1, 2), (0, 1))
  False
  >>> is_edge_shared((0, 1), (2, 3))
  False
  >>> is_edge_shared((2, 3), (0, 1))
  False
  """
    assert edge1[0] < edge1[1], edge1
    assert edge2[0] < edge2[1], edge2

    if edge1[0] <= edge2[0]:
        return edge2[0] < edge1[1]
    else:
        return edge1[0] < edge2[1]


def share_edge(a, b):
    """ Returns true if box defined by a and b share any edge.

  Box is defined as (x-min, y-min, x-max, y-max).

  >>> share_edge((0, 0, 1, 1), (1, 0, 2, 1))
  True
  >>> share_edge((1, 0, 2, 1), (0, 0, 1, 1))
  True
  >>> share_edge((0, 0, 1, 1), (0, 1, 1, 2))
  True
  >>> share_edge((0, 1, 1, 2), (0, 0, 1, 1))
  True
  >>> share_edge((0, 0, 1, 3), (1, 0, 2, 1))
  True
  >>> share_edge((1, 0, 2, 1), (0, 0, 1, 3))
  True
  >>> share_edge((0, 0, 3, 1), (0, 1, 1, 2))
  True
  >>> share_edge((0, 1, 1, 2), (0, 0, 3, 1))
  True
  >>> share_edge((0, 0, 1, 1), (1, 1, 2, 2))
  False
  >>> share_edge((1, 1, 2, 2), (0, 0, 1, 1))
  False
  >>> share_edge((0, 0, 1, 3), (1, 3, 2, 4))
  False
  >>> share_edge((0, 0, 1, 3), (1, 2, 2, 4))
  True
  """

    a_x_min, a_y_min, a_x_max, a_y_max = a
    b_x_min, b_y_min, b_x_max, b_y_max = b

    if a_x_min == b_x_max or a_x_max == b_x_min:
        return is_edge_shared((a_y_min, a_y_max), (b_y_min, b_y_max))
    if a_y_min == b_y_max or a_y_max == b_y_min:
        return is_edge_shared((a_x_min, a_x_max), (b_x_min, b_x_max))


def next_wire_in_dimension(
        wire1, tile1, wire2, tile2, tiles, x_wires, y_wires, wire_map,
        wires_in_node):
    """ next_wire_in_dimension returns true if tile1 and tile2 are in the same
  row and column, and must be adjcent.
  """
    tile1_info = tiles[tile1]
    tile2_info = tiles[tile2]

    tile1_x = tile1_info['grid_x']
    tile2_x = tile2_info['grid_x']
    tile1_y = tile1_info['grid_y']
    tile2_y = tile2_info['grid_y']

    # All wires are in the same row or column or if the each wire lies in its own
    # row or column.
    if len(y_wires) == 1 or len(x_wires) == len(wires_in_node) or abs(
            tile1_y - tile2_y) == 0:
        ordered_wires = sorted(x_wires.keys())

        idx1 = ordered_wires.index(tile1_x)
        idx2 = ordered_wires.index(tile2_x)

        if len(x_wires[tile1_x]) == 1 and len(x_wires[tile2_x]) == 1:
            return abs(idx1 - idx2) == 1

    if len(x_wires) == 1 or len(y_wires) == len(wires_in_node) or abs(
            tile1_x - tile2_x) == 0:
        ordered_wires = sorted(y_wires.keys())

        idx1 = ordered_wires.index(tile1_y)
        idx2 = ordered_wires.index(tile2_y)

        if len(y_wires[tile1_y]) == 1 and len(y_wires[tile2_y]) == 1:
            return abs(idx1 - idx2) == 1

    return None


def only_wire(tile1, tile2, tiles, x_wires, y_wires):
    """ only_wire returns true if tile1 and tile2 only have 1 wire in their respective x or y dimension.
  """
    tile1_info = tiles[tile1]
    tile2_info = tiles[tile2]

    tile1_x = tile1_info['grid_x']
    tile2_x = tile2_info['grid_x']

    tiles_x_adjacent = abs(tile1_x - tile2_x) == 1
    if tiles_x_adjacent and len(x_wires[tile1_x]) == 1 and len(
            x_wires[tile2_x]) == 1:
        return True

    tile1_y = tile1_info['grid_y']
    tile2_y = tile2_info['grid_y']

    tiles_y_adjacent = abs(tile1_y - tile2_y) == 1
    if tiles_y_adjacent and len(y_wires[tile1_y]) == 1 and len(
            y_wires[tile2_y]) == 1:
        return True

    return None


def is_directly_connected(node, node_tree, wire1, wire2):
    if 'wires' in node_tree:
        node_tree_wires = node_tree['wires']
    else:
        if len(node_tree['edges']) == 1 and len(node_tree['joins']) == 0:
            node_tree_wires = node_tree['edges'][0]
        else:
            return None

    if wire1 not in node_tree_wires:
        return None
    if wire2 not in node_tree_wires:
        return None

    # Is there than edge that has wire1 next to wire2?
    for edge in node_tree['edges']:
        idx1 = None
        idx2 = None
        try:
            idx1 = edge.index(wire1)
        except ValueError:
            pass

        try:
            idx2 = edge.index(wire2)
        except ValueError:
            pass

        if idx1 is not None and idx2 is not None:
            return abs(idx1 - idx2) == 1

        if idx1 is not None and (idx1 != 0 and idx1 != len(edge) - 1):
            return False

        if idx2 is not None and (idx2 != 0 and idx2 != len(edge) - 1):
            return False

    # Is there a join of nodes between wire1 and wire2?
    if wire1 in node_tree['joins']:
        return wire2 in node_tree['joins'][wire1]

    if wire2 in node_tree['joins']:
        assert wire1 not in node_tree['joins'][wire2]

    return None


def is_connected(
        wire1, tile1, wire2, tile2, node, wires_in_tiles, wire_map, node_tree,
        tiles, x_wires, y_wires, wires_in_node):
    """ Check if two wires are directly connected. """

    next_wire_in_dim = next_wire_in_dimension(
        wire1, tile1, wire2, tile2, tiles, x_wires, y_wires, wire_map,
        wires_in_node)
    if next_wire_in_dim is not None:
        return next_wire_in_dim

    # Because there are multiple possible wire connections between these two
    # tiles, consult the node_tree to determine if the two wires are actually connected.
    #
    # Warning: The node_tree is incomplete because it is not know how to extract
    # ordered wire information from the node.
    #
    # Example node CLK_BUFG_REBUF_X60Y142/CLK_BUFG_REBUF_R_CK_GCLK0_BOT
    # It does not appear to be possible to get ordered wire connection information
    # for the first two wires connected to PIP
    # CLK_BUFG_REBUF_X60Y117/CLK_BUFG_REBUF.CLK_BUFG_REBUF_R_CK_GCLK0_BOT<<->>CLK_BUFG_REBUF_R_CK_GCLK0_TOP
    #
    # However, it happens to be that theses wires are the only wires in their
    # tiles, so the earlier "only wires in tile" check will pass.

    connected = is_directly_connected(
        node['node'], node_tree[node['node']], wire1, wire2)
    if connected is not None:
        return connected

    is_only_wire = only_wire(tile1, tile2, tiles, x_wires, y_wires)
    if is_only_wire is not None:
        return is_only_wire

    # The node_tree didn't specify these wires, and the wires are not
    # unambiguously connected.
    return False


def process_node(tileconn, key_history, node, wire_map, node_tree, tiles):
    wires = [wire['wire'] for wire in node['wires']]

    wires_in_tiles = {}
    x_wires = {}
    y_wires = {}
    for wire in wires:
        wire_info = wire_map[wire]

        if wire_info['tile'] not in wires_in_tiles:
            wires_in_tiles[wire_info['tile']] = []
        wires_in_tiles[wire_info['tile']].append(wire)

        grid_x = tiles[wire_info['tile']]['grid_x']
        if grid_x not in x_wires:
            x_wires[grid_x] = []
        x_wires[grid_x].append(wire)

        grid_y = tiles[wire_info['tile']]['grid_y']
        if grid_y not in y_wires:
            y_wires[grid_y] = []
        y_wires[grid_y].append(wire)

    if len(wires) == 2:
        wire1 = wires[0]
        wire_info1 = wire_map[wire1]
        wire2 = wires[1]
        wire_info2 = wire_map[wire2]
        update_tile_conn(
            tileconn, key_history, wire1, wire_info1, wire2, wire_info2, tiles)
        return

    for idx, wire1 in enumerate(wires):
        wire_info1 = wire_map[wire1]
        for wire2 in wires[idx + 1:]:
            wire_info2 = wire_map[wire2]

            if not is_connected(wire1, wire_info1['tile'], wire2,
                                wire_info2['tile'], node, wires_in_tiles,
                                wire_map, node_tree, tiles, x_wires, y_wires,
                                wires):
                continue

            update_tile_conn(
                tileconn, key_history, wire1, wire_info1, wire2, wire_info2,
                tiles)


def update_tile_conn(
        tileconn, key_history, wirename1, wire1, wirename2, wire2, tiles):
    # Ensure that (wire1, wire2) is sorted, so we can easy check if a connection
    # already exists.

    tile1 = tiles[wire1['tile']]
    tile2 = tiles[wire2['tile']]
    if ((wire1['type'], wire1['shortname'], tile1['grid_x'], tile1['grid_y']) >
        (wire2['type'], wire2['shortname'], tile2['grid_x'], tile2['grid_y'])):
        wire1, tile1, wire2, tile2 = wire2, tile2, wire1, tile1

    tileconn.append(
        {
            "grid_deltas": [
                tile2['grid_x'] - tile1['grid_x'],
                tile2['grid_y'] - tile1['grid_y'],
            ],
            "tile_types": [
                tile1['type'],
                tile2['type'],
            ],
            "wire_pair": [
                wire1['shortname'],
                wire2['shortname'],
            ],
        })


def flatten_tile_conn(tileconn):
    """ Convert tileconn that is key'd to identify specific wire pairs between tiles
  key (tile1_type, wire1_name, tile2_type, wire2_name) to flat tile connect list
  that relates tile types and relative coordinates and a full list of wires to
  connect. """
    flat_tileconn = {}

    for conn in tileconn:
        key = (tuple(conn['tile_types']), tuple(conn['grid_deltas']))

        if key not in flat_tileconn:
            flat_tileconn[key] = {
                'tile_types': conn['tile_types'],
                'grid_deltas': conn['grid_deltas'],
                'wire_pairs': set()
            }

        flat_tileconn[key]['wire_pairs'].add(tuple(conn['wire_pair']))

    def inner():
        for output in flat_tileconn.values():
            yield {
                'tile_types': output['tile_types'],
                'grid_deltas': output['grid_deltas'],
                'wire_pairs': tuple(output['wire_pairs']),
            }

    return tuple(inner())


def is_tile_type(tiles, coord_to_tile, coord, tile_type):
    if coord not in coord_to_tile:
        return False

    target_tile = tiles[coord_to_tile[coord]]
    return target_tile['type'] == tile_type


def get_connections(wire, wire_info, conn, idx, coord_to_tile, tiles):
    """ Yields (tile_coord, wire) for each wire that should be connected to specified wire. """
    pair = conn['wire_pairs'][idx]
    wire_tile_type = wire_info['type']
    tile_types = conn['tile_types']
    shortname = wire_info['shortname']
    grid_deltas = conn['grid_deltas']

    wire1 = tile_types[0] == wire_tile_type and shortname == pair[0]
    wire2 = tile_types[1] == wire_tile_type and shortname == pair[1]
    assert wire1 or wire2, (wire, conn)

    tile_of_wire = wire_info['tile']
    start_coord_x = tiles[tile_of_wire]['grid_x']
    start_coord_y = tiles[tile_of_wire]['grid_y']
    if wire1:
        target_coord_x = start_coord_x + grid_deltas[0]
        target_coord_y = start_coord_y + grid_deltas[1]
        target_tile_type = tile_types[1]

        target_wire = pair[1]
        target_tile = (target_coord_x, target_coord_y)

        if is_tile_type(tiles, coord_to_tile, target_tile, target_tile_type):
            yield target_tile, target_wire

    if wire2:
        target_coord_x = start_coord_x - grid_deltas[0]
        target_coord_y = start_coord_y - grid_deltas[1]
        target_tile_type = tile_types[0]

        target_wire = pair[0]
        target_tile = (target_coord_x, target_coord_y)

        if is_tile_type(tiles, coord_to_tile, target_tile, target_tile_type):
            yield target_tile, target_wire


def make_connection(wire_nodes, wire1, wire2):
    if wire_nodes[wire1] is wire_nodes[wire2]:
        assert wire1 in wire_nodes[wire1]
        assert wire2 in wire_nodes[wire2]
        return

    new_node = wire_nodes[wire1] | wire_nodes[wire2]

    for wire in new_node:
        wire_nodes[wire] = new_node


def create_coord_to_tile(tiles):
    coord_to_tile = {}
    for tile, tileinfo in tiles.items():
        coord_to_tile[(tileinfo['grid_x'], tileinfo['grid_y'])] = tile

    return coord_to_tile


def connect_wires(tiles, tileconn, wire_map):
    """ Connect individual wires into groups of wires called nodes. """

    # Initialize all nodes to originally only contain the wire by itself.
    wire_nodes = {}
    for wire in wire_map:
        wire_nodes[wire] = set([wire])

    wire_connection_map = {}
    for conn in tileconn:
        for idx, (wire1, wire2) in enumerate(conn['wire_pairs']):
            key1 = (conn['tile_types'][0], wire1)
            if key1 not in wire_connection_map:
                wire_connection_map[key1] = []
            wire_connection_map[key1].append((conn, idx))

            key2 = (conn['tile_types'][1], wire2)
            if key2 not in wire_connection_map:
                wire_connection_map[key2] = []
            wire_connection_map[key2].append((conn, idx))

    coord_to_tile = create_coord_to_tile(tiles)

    for wire, wire_info in progressbar.progressbar(wire_map.items()):
        key = (wire_info['type'], wire_info['shortname'])
        if key not in wire_connection_map:
            continue

        for conn, idx in wire_connection_map[key]:
            for target_tile, target_wire in get_connections(
                    wire, wire_info, conn, idx, coord_to_tile, tiles):

                full_wire_name = coord_to_tile[target_tile] + '/' + target_wire
                assert wire_map[full_wire_name]['shortname'] == target_wire, (
                    target_tile, target_wire, wire, conn)
                assert wire_map[full_wire_name]['tile'] == coord_to_tile[
                    target_tile], (
                        wire_map[full_wire_name]['tile'],
                        coord_to_tile[target_tile])

                make_connection(wire_nodes, wire, full_wire_name)

    # Find unique nodes
    nodes = {}
    for node in wire_nodes.values():
        nodes[id(node)] = node

    # Flatten to list of lists.
    return tuple(tuple(node) for node in nodes.values())


def generate_tilegrid(pool, tiles):
    wire_map = {}

    grid = {}

    num_tiles = 0
    for tile_type in tiles:
        num_tiles += len(tiles[tile_type])

    idx = 0
    with progressbar.ProgressBar(max_value=num_tiles) as bar:
        for tile_type in tiles:
            for tile in pool.imap_unordered(
                    get_tile_grid_info,
                    tiles[tile_type],
                    chunksize=20,
            ):
                bar.update(idx)

                assert len(tile) == 1, tile
                tilename = tuple(tile.keys())[0]

                for wire in tile[tilename]['wires']:
                    assert wire not in wire_map, (wire, wire_map)
                    assert wire.startswith(tilename + '/'), (wire, tilename)

                    wire_map[wire] = {
                        'tile': tilename,
                        'type': tile[tilename]['type'],
                        'shortname': wire[len(tilename) + 1:],
                    }

                del tile[tilename]['wires']
                grid.update(tile)

                idx += 1
                bar.update(idx)

    return grid, wire_map


def generate_tileconn(pool, node_tree, nodes, wire_map, grid):
    tileconn = []
    key_history = {}
    raw_node_data = []
    with progressbar.ProgressBar(max_value=len(nodes)) as bar:
        for idx, node in enumerate(pool.imap_unordered(
                read_json5,
                nodes,
                chunksize=20,
        )):
            bar.update(idx)
            raw_node_data.append(node)
            process_node(
                tileconn, key_history, node, wire_map, node_tree, grid)
            bar.update(idx + 1)

    tileconn = flatten_tile_conn(tileconn)

    return tileconn, raw_node_data


def main():
    parser = argparse.ArgumentParser(
        description=
        "Reduces raw database dump into prototype tiles, grid, and connections."
    )
    parser.add_argument('--root_dir', required=True)
    parser.add_argument('--output_dir', required=True)
    parser.add_argument('--verify_only', action='store_true')
    parser.add_argument('--ignored_wires')

    args = parser.parse_args()

    tiles, nodes = prjxray.lib.read_root_csv(args.root_dir)

    processes = min(multiprocessing.cpu_count(), 10)
    print('{} Running {} processes'.format(datetime.datetime.now(), processes))
    pool = multiprocessing.Pool(processes=processes)

    node_tree_file = os.path.join(args.output_dir, 'node_tree.json')

    tileconn_file = os.path.join(args.output_dir, 'tileconn.json')
    wire_map_file = os.path.join(args.output_dir, 'wiremap.pickle')

    print('{} Reading tilegrid'.format(datetime.datetime.now()))
    with open(os.path.join(prjxray.util.get_db_root(), 'tilegrid.json')) as f:
        grid = json.load(f)

    if not args.verify_only:
        print('{} Creating tile map'.format(datetime.datetime.now()))
        grid2, wire_map = generate_tilegrid(pool, tiles)

        # Make sure tilegrid from 005-tilegrid matches tilegrid from
        # generate_tilegrid.
        db_grid_keys = set(grid.keys())
        generated_grid_keys = set(grid2.keys())
        assert db_grid_keys == generated_grid_keys, (
            db_grid_keys ^ generated_grid_keys)

        for tile in db_grid_keys:
            for k in grid2[tile]:
                assert k in grid[tile], k
                assert grid[tile][k] == grid2[tile][k], (
                    tile, k, grid[tile][k], grid2[tile][k])

        with open(wire_map_file, 'wb') as f:
            pickle.dump(wire_map, f)

        print('{} Reading node tree'.format(datetime.datetime.now()))
        with open(node_tree_file) as f:
            node_tree = json.load(f)

        print('{} Creating tile connections'.format(datetime.datetime.now()))
        tileconn, raw_node_data = generate_tileconn(
            pool, node_tree, nodes, wire_map, grid)

        print('{} Writing tileconn'.format(datetime.datetime.now()))
        with open(tileconn_file, 'w') as f:
            xjson.pprint(f, tileconn)
    else:
        with open(wire_map_file, 'rb') as f:
            wire_map = pickle.load(f)

        print('{} Reading raw_node_data'.format(datetime.datetime.now()))
        raw_node_data = []
        with progressbar.ProgressBar(max_value=len(nodes)) as bar:
            for idx, node in enumerate(pool.imap_unordered(
                    read_json5,
                    nodes,
                    chunksize=20,
            )):
                bar.update(idx)
                raw_node_data.append(node)
                bar.update(idx + 1)

        print('{} Reading tileconn'.format(datetime.datetime.now()))
        with open(tileconn_file) as f:
            tileconn = json.load(f)

    wire_nodes_file = os.path.join(args.output_dir, 'wire_nodes.pickle')
    if os.path.exists(wire_nodes_file) and args.verify_only:
        with open(wire_nodes_file, 'rb') as f:
            wire_nodes = pickle.load(f)
    else:
        print(
            "{} Connecting wires to verify tileconn".format(
                datetime.datetime.now()))
        wire_nodes = connect_wires(grid, tileconn, wire_map)
        with open(wire_nodes_file, 'wb') as f:
            pickle.dump(wire_nodes, f)

    print('{} Verifing tileconn'.format(datetime.datetime.now()))
    error_nodes = []
    prjxray.lib.verify_nodes(
        [
            (node['node'], tuple(wire['wire']
                                 for wire in node['wires']))
            for node in raw_node_data
        ], wire_nodes, error_nodes)

    if len(error_nodes) > 0:
        error_nodes_file = os.path.join(args.output_dir, 'error_nodes.json')
        with open(error_nodes_file, 'w') as f:
            xjson.pprint(f, error_nodes)

        ignored_wires = []
        ignored_wires_file = args.ignored_wires
        if os.path.exists(ignored_wires_file):
            with open(ignored_wires_file) as f:
                ignored_wires = set(l.strip() for l in f)

        if not prjxray.lib.check_errors(error_nodes, ignored_wires):
            print(
                '{} errors detected, see {} for details.'.format(
                    len(error_nodes), error_nodes_file))
            sys.exit(1)
        else:
            print(
                '{} errors ignored because of {}\nSee {} for details.'.format(
                    len(error_nodes), ignored_wires_file, error_nodes_file))


if __name__ == '__main__':
    main()