import collections
import copy
from functools import total_ordering
import re
import string

from common import FrissError
import nml_blocks

class Node:

    def __init__(self, name, parent=None):
        self.name = name
        self.parent = parent
        self.children = []
        self.transitions = []

    def __eq__(self, other):
        return (self.name, self.parent) == (other.name, other.parent)

    def __hash__(self):
        return hash((self.name, self.parent))

    def add(self, parent, child, start_y, end_y, prob=None):
        for n in self.bfs():
            if n.name == parent:
                n.children.append(Node(child, parent))
                n.transitions.append(
                    {"start_y": start_y, "end_y": end_y, "prob": prob})
                return
        raise FrissError

    def bfs(self):
        yield self
        q = collections.deque(self.children)
        while q:
            node = q.popleft()
            yield node
            for c in node.children:
                q.append(c)


@total_ordering
class Label:
    metro = {
        "3RDR", "MTRO"
    }
    electrified = {
        "ELRL"
    }
    # Maybe do these later if needed
    broad = {}
    narrow = {}

    def __init__(self, label):
        self.label = label

    def __getitem__(self, i):
        return self.label[i]

    def __hash__(self):
        return hash(self.label)

    def __str__(self):
        return self.label

    def __lt__(self, other):
        if not self.isstd() and other.isstd():
            return True
        if self.isstd() and not other.isstd():
            return False
        if not self.isstd() and not other.isstd():
            return self.label < other.label
        if self.label[0] != other.label[0]:
            return self.label[0] < other.label[0]
        if self.label[3] != other.label[3]:
            return self.label[3] < other.label[3]
        if self.label[2] != other.label[2]:
            return self.label[2] < other.label[2]
        return self.label[1] < other.label[1]
    def __eq__(self, other):
        return self.label == other

    def isstd(self):
        """
        https://newgrf-specs.tt-wiki.net/wiki/Standardized_Railtype_Scheme
        """
        return bool(re.match(r"[SBNnDdML][A-Z][A-E][NE3Z4YTDdAa]", self.label))

    def isbasic(self):
        return not self.isele() and not self.ismetro()

    def isele(self):
        if self.isstd():
            return bool(re.match(r"[EZYTDdAa]", self[3]))
        return self in self.electrified

    def ismetro(self):
        if self.isstd():
            return bool(re.match(r"[3Z4Y]", self[3]))
        return self in self.metro

    def iscompatible(self, other):
        # Speed class other than A should not exist in trains
        if self.isstd() and self[1] > "A":
            return False
        # Maybe create logic for other railtypes later
        if not self.isstd() or not other.isstd():
            return True
        # Axle class
        if not self[2] <= other[2]:
            return False
        # Gauge
        if self[0] == "S":
            return bool(re.match(r"[SDd]", other[0]))
        if self[0] == "B":
            return bool(re.match(r"[BD]", other[0]))
        if self[0] == "N":
            return bool(re.match(r"[Nd]", other[0]))
        # Assuming vehicles cant have multiple gauges
        return re.match(r"[SBNnML]", self[0]) and self[0] == other[0]

    def ispowered(self, other):
        if not self.iscompatible(other):
            return False
        if not self.isstd() or not other.isstd():
            if self.isbasic() \
                    or (self.isele() and other.isele()) \
                    or (self.ismetro() and other.ismetro()):
                return True
            return False
        # Power
        if self[3] == "N":
            return True
        if self[3] == "E":
            return bool(re.match(r"[EZY]", other[3]))
        if self[3] == "3":
            return bool(re.match(r"[3Z]", other[3]))
        if self[3] == "4":
            return bool(re.match(r"[4Y]", other[3]))
        # Assuming vehicles cant have multiple types of electric power
        return re.match(r"[E34TDdAa]", self[3]) and self[3] == other[3]

    def isalternative(self, other):
        # Speed class other than A should not exist in trains
        if other.isstd() and other[1] > "A":
            return False
        if not self.isstd() or not other.isstd():
            if other.isbasic() \
                    or self.isele() and other.isele() \
                    or self.ismetro() and other.ismetro():
                return True
            return False
        # Gauge
        if self[0] == "D" and not re.match(r"[SB]", other[0]):
            return False
        if self[0] == "d" and not re.match(r"[SN]", other[0]):
            return False
        if not re.match(r"[SBNnML]", other[0]) or not self[0] == other[0]:
            return False
        # Axle class
        if not self[2] >= other[2]:
            return False
        # Power
        if other[3] == "N":
            return True
        if self[3] == "Z" and re.match(r"[E3]", other[3]):
            return True
        if self[3] == "Y" and re.match(r"[E4]", other[3]):
            return True
        # Fallback for electricity, this works properly only with
        # frissrails since it only uses A for power
        if re.match(r"[TDdAa]", self[3]) and other[3] == "E":
            return True
        return re.match(r"[E34TDdAa]", other[3]) and self[3] == other[3]


class RailtypeDb:

    def __init__(self):
        self.railtypes = []
        # These trings will never be translated
        self.strings = {}

    def __getitem__(self, key):
        try:
            return self.railtypes[key]
        except KeyError:
            raise FrissError("No railtype with label {}".format(key))

    def __iter__(self):
        return iter(self.railtypes)

    def add_railtype(self, **kwargs):
        self.railtypes.append(Railtype(**kwargs))

    def blocks(self):
        return nml_blocks.switch_stack + nml_blocks.item_stack

    def finalize(self):
        for r in self.railtypes:
            r.add_other_railtypes(self)
        for r in self.railtypes:
            r.gen_item_block()


class Railtype:
    prop_strings = [
        "name", "toolbar_caption", "menu_text", "build_window_caption",
        "autoreplace_text", "new_engine_text"
    ]
    standardized_labels = {"SA{}{}".format(axle, pwr) \
        for axle in ["A", "B", "C", "D", "E"] \
        for pwr in ["N", "E", "3", "A"]}
    dataparsers = {
        "label": Label
    }

    def __init__(self, **kwargs):
        # Map parsing function to data if possible
        self.data = {k: self.dataparsers.get(k, lambda x: x)(kwargs[k]) for k in kwargs}

    def __getitem__(self, key):
        try:
            return self.data[key]
        except KeyError:
            raise FrissError("Railtype {} has no data with key {}".format(key))

    def get(self, key, default=None):
        try:
            return self.data[key]
        except KeyError:
            return default

    def add_other_railtypes(self, railtypedb):
        self.compatible = set()
        self.powered = set()
        self.alternative = set()
        others = set()
        others.update(rt["label"] for rt in railtypedb)
        others.update(Label(sl) for sl in self.standardized_labels)
        others.update(Label(ml) for ml in Label.metro)
        for o in others:
            if self["label"].iscompatible(o):
                self.compatible.add(o)
            if self["label"].ispowered(o):
                self.powered.add(o)
            if self["label"].isalternative(o):
                self.alternative.add(o)

    def gen_item_block(self):
        self.item = nml_blocks.Item(
            feature="FEAT_RAILTYPES", name=self["label"])
        mode_check = "if (param_mode == 0) {{\n{}\n}}"
        if not self["label"].ismetro():
            self.item.add_decor(mode_check)
        flags = []
        if self["catenary"]:
            flags.append("RAILTYPE_FLAG_CATENARY")
        if not self["lc_rural"]:
            flags.append("RAILTYPE_FLAG_NO_LEVEL_CROSSING")
        if "high_speed_param" in self["extra"]:
            dummy = nml_blocks.Item(feature="FEAT_RAILTYPES",
                                    name=self["label"])
            dummy.add_decor(
                "if (param_modern_speed_limit == 1) {{\n{}\n}}")
            dummy.add_decor(mode_check)
            dummy.add_prop("speed_limit", "{} km/h".format(self["speed"]))
            dummy = nml_blocks.Item(feature="FEAT_RAILTYPES",
                                    name=self["label"])
            dummy.add_decor(
                "if (param_modern_crossing == 0) {{\n{}\n}}")
            dummy.add_decor(mode_check)
            # Dont redefine catenary bitmask here,
            # otherwise it gets cancelled
            dflags = flags + ["RAILTYPE_FLAG_NO_LEVEL_CROSSING"]
            dummy.add_prop(
                "railtype_flags", "bitmask({})".format(','.join(dflags)))
        elif self["label"].ismetro():
            dummy = nml_blocks.Item(feature="FEAT_RAILTYPES",
                                    name=self["label"])
            dummy.add_decor(mode_check)
            dummy.add_decor(
                "if (param_metro_speed_limit == 1) {{\n{}\n}}")
            dummy.add_prop("speed_limit", "{} km/h".format(self["speed"]))
        else:
            self.item.add_prop("speed_limit", "{} km/h".format(self["speed"]))
        self.item.add_prop(
            "railtype_flags", "bitmask({})".format(','.join(flags)))
        self.item.add_prop("sort_order", self["sort"])
        self.item.add_prop("label", '"{}"'.format(self["label"]))
        self.item.add_prop("introduction_date", "date({},01,01)".format(
            self["available"]))
        for k in self.prop_strings:
            self.item.add_prop(
                k, "string({})".format(self.get(k, "STR_FIXME")))
        # Fill them elsewhere
        self.item.add_prop(
            "compatible_railtype_list", railtype_array(self.compatible))
        self.item.add_prop(
            "powered_railtype_list", railtype_array(self.powered))
        self.item.add_prop(
            "alternative_railtype_list", railtype_array(self.alternative))
        self.item.add_prop("curve_speed_multiplier", "0")
        self.item.add_prop("station_graphics", "RAILTYPE_STATION_NORMAL")
        self.item.add_prop("construction_cost", self["pur_cost"])
        self.item.add_prop("acceleration_model", "ACC_MODEL_RAIL")
        # Don't use because ottd map shows legend for only one color
        #self.item.add_prop("map_colour", self.get("color", self["sort"]))
        #self.item.add_prop("requires_railtype_list", something)
        if self["introduces"]:
            self.item.add_prop(
                "introduces_railtype_list", railtype_array(self["introduces"]))
        self.item.add_prop("maintenance_cost", self["maint_cost"])

        #Graphics block
        self.item.add_graphics("fences", "fences")
        if self["ballast"] or self["sleepers"]:
            self.item.add_graphics(
                "track_overlay", self.track_stack("overlay"))
            self.item.add_graphics("underlay", self.track_stack("underlay"))
            self.item.add_graphics(
                "tunnels", self.track_stack("tunnel_underlay"))
            self.item.add_graphics(
                "bridge_surfaces", self.track_stack("bridge"))
            self.item.add_graphics("tunnel_overlay", "tracks_tunnel_overlay")
        if self.data["catenary"]:
            self.item.add_graphics("catenary_pylons", self.catenary_stack())
            self.item.add_graphics("catenary_wire", "catenary_wires")
        if self.data["lc_rural"]:
            self.item.add_graphics("level_crossings", self.lc_stack())
        if self.data["signals"]:
            self.item.add_graphics("signals", self.signal_stack())

    def catenary_stack(self):
        nodes = copy.deepcopy(self.data["catenary"])
        for n in nodes.travel():
            n.name = "catenary_pylons_{}".format(n.name)
        return gen_switches(nodes)

    def lc_stack(self):
        lc_fmt = "lc_{}_{}"
        ballast_nodes = copy.deepcopy(self.data["ballast"])
        for bn in ballast_nodes.travel():
            rural_nodes = copy.deepcopy(self.data["lc_rural"])
            for n in rural_nodes.travel():
                n.name = lc_fmt.format(n.name, bn.name)
            city_nodes = copy.deepcopy(self.data["lc_city"])
            for n in city_nodes.travel():
                n.name = lc_fmt.format(n.name, bn.name)
            lc_root_rural = gen_switches(rural_nodes)
            lc_root_city = gen_switches(city_nodes)
            lc_switch = nml_blocks.Switch(expression="town_zone")
            lc_switch.add("TOWNZONE_EDGE", lc_root_rural)
            lc_switch.add("default", lc_root_city)
            bn.name = lc_switch
        return gen_switches(ballast_nodes)

    def signal_stack(self):
        nodes = copy.deepcopy(self.data["signals"])
        for n in nodes.travel():
            n.name = "signals_{}".format(n.name)
        stack = []
        stack.append(gen_switches(nodes))  # Map switches
        stack.append(gen_switches(nodes, phased=False))  # GUI switches
        stack.append(nml_blocks.Switch(  # Test if map or GUI
            expression="getbits(extra_callback_info1, 0, 8)"))
        stack[-1].add("0", stack[-3])
        stack[-1].add("0x10", stack[-2])
        if self["signal_group"]:
            stack.append(nml_blocks.Switch(
                expression="param_signal_group_{}_light".format(
                    self["signal_group"])))
            stack[-1].add("0", stack[-2])
            stack[-1].add("default", "signal_group_{}_light".format(
                self["signal_group"]))
            stack.append(nml_blocks.Switch(
                expression="getbits(extra_callback_info2, 8, 8)"))
            stack[-1].add("0", stack[-2])
            stack[-1].add("1", "signal_group_{}_sema".format(
                self["signal_group"]))
        return stack[-1]

    def track_stack(self, type_):
        # If no ballast or sleeper transition, then no need to generate
        # multiple transition switch stacks
        other = None
        if self["ballast"] and not self["sleepers"]:
            other = self.data["ballast"]
        elif not self["ballast"] and self["sleepers"]:
            other = self.data["sleepers"]
        if other:
            nodes = copy.deepcopy(other)
            for n in nodes.travel():
                n.name = "tracks_{}_{}".format(other.name, type_)
            return gen_switches(nodes)
        block_name = "tracks_{}_{}_{}"
        if len(list(self.data["ballast"].travel())) == 1:
            # Prepend all sleeper names with ballast name
            sleeper_nodes = copy.deepcopy(self.data["sleepers"])
            for n in sleeper_nodes.travel():
                n.name = block_name.format(self.data["ballast"].name,
                                           n.name, type_)
            return gen_switches(sleeper_nodes)
        if len(list(self.data["sleepers"].travel())) == 1:
            # Append all ballast names with sleeper name
            ballast_nodes = copy.deepcopy(self.data["ballast"])
            for n in ballast_nodes.travel():
                n.name = block_name.format(
                    n.name, self.data["sleepers"].name, type_)
            return gen_switches(ballast_nodes)
        # If both sleepers and ballast have transitions, then for each
        # ballast type, a corresponding transition switch stack for
        # sleepers must be generated
        ballast_nodes = copy.deepcopy(self.data["ballast"])
        for bn in ballast_nodes.travel():
            sleeper_nodes = copy.deepcopy(self.data["sleepers"])
            for sn in sleeper_nodes.travel():
                sn.name = block_name.format(bn.name, sn.name, type_)
            switches = gen_switches(sleeper_nodes)
            bn.name = switches
        return gen_switches(ballast_nodes)


def gen_switches(tree, phased=True):
    stack = {}
    for node in reversed(list(tree.travel())):
        if node.leaf:
            stack[node] = node.name
        elif not phased:
            assert len(node.children) == 1
            stack[node] = nml_blocks.Switch(expression="current_year")
            stack[node].add(
                "0..{}".format(node.transitions[0]["start_y"]-1),
                node.name)
            stack[node].add("default", stack[node.children[0]])
        else:
            prob_tot = 0
            random_stack = []
            for c, t in node:
                prob = t["prob"] // 25
                prob_tot += prob
                random_stack.append((prob, phased_transition(
                    start=node.name, end=stack[c],
                    start_y=t["start_y"], end_y=t["end_y"])))
            if prob_tot != 4:
                random_stack.append((4-prob_tot, node.name))
            # Only one 100% transition
            if len(random_stack) == 1:
                stack[node] = random_stack[0][1]
            else:
                stack[node] = nml_blocks.Switch(random=True)
                for p, s in random_stack:
                    stack[node].add(p, s)
    return stack[tree]  # "tree" is the root node


def even_year_ranges(start_y, end_y):
    if start_y == end_y:
        return []
    if 3 >= end_y - start_y >= 1:
        return [str(start_y + i) for i in range(end_y - start_y)]
    out = []
    quot, rem = divmod((end_y-start_y), 3)
    prev = start_y
    for _ in range(3):
        this_interval = quot
        if rem:
            this_interval += 1
            rem -= 1
        if this_interval > 1:
            out.append("{}..{}".format(prev, (prev+this_interval) - 1))
        else:
            out.append(str(prev))
        prev += this_interval
    return out


def phased_transition(start, end, start_y, end_y):
    years = even_year_ranges(start_y, end_y)
    if len(years) == 0:
        ret_sw = nml_blocks.Switch(expression="current_year")
        ret_sw.add("0..{}".format(end_y-1), start)
        ret_sw.add("default", end)
        return ret_sw
    if len(years) == 1:
        phases = [(2, 2)]
    elif len(years) == 2:
        phases = [(1, 3), (3, 1)]
    else:
        phases = [(1, 3), (2, 2), (3, 1)]
    temp = collections.deque()
    for y, p in zip(years, phases):
        temp.append(nml_blocks.Switch(random=True))
        temp[-1].add(p[0], end)
        temp[-1].add(p[1], start)
    ret_sw = nml_blocks.Switch(expression="current_year")
    ret_sw.add("0..{}".format(start_y-1), start)
    for y, p in zip(years, phases):
        ret_sw.add(y, temp.popleft())
    ret_sw.add("default", end)
    return ret_sw


def railtype_array(labels):
    quoted = ('"{}"'.format(r) for r in sorted(labels))
    return "[{}]".format(','.join(quoted))
