from osm_parser import get_default_parser
from collections import defaultdict


class Node:
    def __init__(self, id, lat, lng):
        self.id = id
        self.lat = float(lat)
        self.lng = float(lng)
        self.neighbours = []

        self.parent = None
        self.size = 1

    def find_root(self):
        """ Recursively find the root of this node. """
        if not self.parent:
            return self
        else:
            return self.parent.find_root()

    def union(self, other):
        """ Mark this node and some other node as part of the same tree. """
        this = self.find_root()
        other = other.find_root()

        # If we share the same root we are already part of the same tree
        if this == other:
            return

        # Minimize tree size. This ensures optimal complexity.
        # We could also store the rank instead.
        if this.size < other.size:
            this, other = other, this

        # This is safe since other is a root node and currently doesn't have a
        # parent.
        other.parent = this
        this.size += other.size

    def coord_tuple(self):
        return self.lat, self.lng


parser = None  # Have a global reusable parser object


def add_neighbours(nodes):
    for way in parser.iter_ways():
        if 'highway' not in way['tags']:
            continue

        road = way['road']

        for i in range(len(road) - 1):
            node1 = road[i]
            node2 = road[i + 1]

            # FIXME: lots of repeated code here
            if suitable_bike(way):
                bike_nodes = nodes['bike']

                bike_nodes[node1].neighbours.append(bike_nodes[node2])
                bike_nodes[node2].neighbours.append(bike_nodes[node1])

                # These two are neighbours and should be part of the same tree.
                bike_nodes[node1].union(bike_nodes[node2])

            if suitable_car(way):
                car_nodes = nodes['car']

                car_nodes[node1].neighbours.append(car_nodes[node2])
                car_nodes[node2].neighbours.append(car_nodes[node1])

                # These two are neighbours and should be part of the same tree.
                car_nodes[node1].union(car_nodes[node2])

    return nodes


def extract_osm_nodes(f_name):
    global parser
    parser = get_default_parser(f_name)
    nodes = {
        'bike': dict(),
        'car': dict()
    }

    for node in parser.iter_nodes():
        # FIXME: this can probably be solved better
        nodes['bike'][node['id']] = Node(node['id'], node['lat'], node['lon'])
        nodes['car'][node['id']] = Node(node['id'], node['lat'], node['lon'])

    add_neighbours(nodes)

    # FIXME: this can probably be solved better
    # Remove nodes without neighbours
    for node_id, node in nodes['bike'].copy().items():
        if not node.neighbours:
            del nodes['bike'][node_id]

    for node_id, node in nodes['car'].copy().items():
        if not node.neighbours:
            del nodes['car'][node_id]

    nodes['bike'], unconnected_bike = make_forest(nodes['bike'])
    nodes['car'], unconnected_car = make_forest(nodes['car'])

    unconnected_nodes = {
        'bike': unconnected_bike,
        'car': unconnected_car
    }

    grids = {
        'bike': make_grid(nodes['bike']),
        'car': make_grid(nodes['car'])
    }

    return nodes, grids, unconnected_nodes


def make_forest(nodes):
    # Construct a forest of disjoint trees.
    # The forest is a { root: tree }-dict. If we construct a relation where
    # two nodes relate to each other if they are a part of the same tree,
    # this can be shown to be an equivalence relation where the equivalence
    # class is some node in the tree. In our case, the equivalence class
    # becomes the root of each node since it will be the same for all nodes
    # in the same tree due our use of union-find.
    forest = defaultdict(dict)
    for node_id, node in nodes.items():
        forest[node.find_root().id][node_id] = node

    # Find the largest disjoint tree.
    # We store the root so we can easily test if a node is part of this
    # tree or not.
    best_size = 0
    best_tree = None
    best_root = None
    for root in forest:
        tree = forest[root]
        size = len(tree)
        if size > best_size:
            best_size = size
            best_tree = tree
            best_root = root

    unconnected_nodes = []
    for _, node in nodes.items():
        if node.find_root().id != best_root:
            unconnected_nodes.append(node)

    nodes = best_tree

    return nodes, unconnected_nodes


def make_grid(nodes):
    grid = defaultdict()

    # create a "grid" by grouping nearby nodes.
    for node in nodes.copy().values():

        key = (int(round(node.lat, 3) * 1000), int(round(node.lng, 3)
                                                   * 1000))

        if key in grid.keys():
            grid[key].append(node)

        else:
            grid[key] = [node]

    return grid


def select_nodes_in_rectangle(nodes, min_lat, max_lat, min_long, max_long):
    return [node for node in nodes.values()
            if min_lat <= node.lat <= max_lat
            and min_long <= node.lng <= max_long]


def get_relevant_neighbours(nodes, transport_mode):
    if transport_mode == "bike":
        return nodes.neighbours_bike
    else:
        return nodes.neighbours_car


suitable_highway_types_bike = [
    # Special road types
    'living_street', 'service', 'pedestrian', 'track', 'road'
]

def suitable_bike(way):
    tags = way['tags']

    suitable_generic_type = tags['highway'] in suitable_highway_types_bike
    suitable_bike = False
    # This implies you can go by bike for some reason
    is_segregated = 'segregated' in tags

    if 'bicycle' in tags:
        bicycle_tag = tags['bicycle']
        suitable_bike = bicycle_tag == 'yes' or bicycle_tag == 'designated'

    return suitable_generic_type or suitable_bike or is_segregated


suitable_highway_types_car = [
    # Roads
    'motorway', 'trunk', 'primary', 'secondary',
    'tertiary', 'unclassified', 'residential',
    # Link roads
    'motorway_link', 'trunk_link', 'primary_link', 'secondary_link',
    'tertiary_link',
    # Special road types
    'living_street', 'service',
    'pedestrian', 'road' # FIXME: Handle predestrian and road differently?
]

def suitable_car(way):
    return way['tags']['highway'] in suitable_highway_types_car