# -*- test-case-name: gamelib.tests.test_lab -*-

from random import random, choice, sample

import networkx

from gamelib import research, schematics
from gamelib.game_base import get_subclasses


class ScienceGraph(object):
    def __init__(self, all_science, known_science):
        self.graph = networkx.DiGraph()
        self.all_science = all_science
        self.known_science = known_science
        self.add_all_science()
        self.tag_known_science()
        assert networkx.is_directed_acyclic_graph(self.graph)

    def add_all_science(self):
        # Add level 0 of everything to the graph.
        for science in self.all_science:
            self.graph.add_node((science, 0), known=False)

        # Walk dependencies and fill in intermediate nodes.
        for science in self.all_science:
            for dep in science.PREREQUISITES:
                self.add_node_string(*dep)
                self.graph.add_edge(dep, (science, 0))

    def add_node_string(self, science, level):
        node = (science, level)
        if node in self.graph:
            return node

        # We prepopulate with level 0 of eveything.
        assert level >= 0

        self.graph.add_node(node, known=False)
        parent = self.add_node_string(science, level - 1)
        self.graph.add_edge(parent, node)
        return node

    def tag_known_science(self):
        for science in self.known_science:
            # We may know more of this than the graph has.
            self.add_node_string(type(science), science.points + 1)
            for i in range(science.points + 1):
                self.graph.node[(type(science), i)]['known'] = True

    def is_known(self, node):
        return self.graph.node[node]['known']

    def distances_to_known(self, science):
        nodes = set()
        nodes_to_check = [(science, 0)]

        while nodes_to_check:
            node = nodes_to_check.pop()
            if self.is_known(node) or node in nodes:
                continue
            nodes.add(node)
            nodes_to_check.extend(self.graph.predecessors(node))

        distances = {}
        for node in nodes:
            distances.setdefault(node[0], 0)
            distances[node[0]] += 1

        return distances

    def count_unknown(self, sciences):
        return len([s for s in sciences if not self.is_known((s, 0))])

    def find_prospects(self):
        prospects = {'research': [], 'schematic': []}
        for science in self.all_science:
            distances = self.distances_to_known(science)
            if not distances:
                # We already know this thing.
                continue
            # Remove the thing we're trying to get from the distance.
            distances.pop(science)
            if self.count_unknown(distances.keys()) > 0:
                # We only want direct breakthroughs.
                continue
            prospects[science.SCIENCE_TYPE].append(
                (sum(distances.values()), science, distances))
        return dict((k, sorted(v)) for k, v in prospects.items())

    def find_promising_areas(self, size=3):
        basic_science = False
        areas_for_research = set()
        areas_for_schematics = set()
        prospects = self.find_prospects()

        for points, target, distances in prospects['schematic']:
            if points > 0:
                # We need nonzero points in these things.
                areas_for_schematics.update(distances.keys())
            else:
                # Any of these things qualify us.
                areas_for_schematics.update(p for p, _ in target.PREREQUISITES)

        for points, target, distances in prospects['research']:
            if points == 0:
                basic_science = True
            else:
                areas_for_research.update(distances.keys())

        suggestions = []
        k = min(size, len(areas_for_schematics))
        suggestions.extend(sample(areas_for_schematics, k))
        if len(suggestions) < size:
            k = min(size - len(suggestions), len(areas_for_research))
            suggestions.extend(sample(areas_for_research, k))

        return basic_science, suggestions


class Lab(object):
    BASIC_RESEARCH_SUCCESS_RATE = 0.05
    BASIC_RESEARCH_SUCCESS_MULTIPLIER = 2

    def __init__(self, init_data=None):
        self.science = []
        self.new_research = get_subclasses(research.ResearchArea)
        self.new_schematics = get_subclasses(schematics.Schematic)
        self.all_science = [s for s in self.new_research + self.new_schematics]

        if init_data is not None:
            # Load stored state.
            self._load_data(init_data)
        else:
            # New game.
            self._choose_initial_science()

    def _load_data(self, init_data):
        sciences = init_data['science'].copy()
        for science in self.new_schematics + self.new_research:
            # Check if this science is one we should know.
            points = sciences.pop(science.save_name(), None)
            if points is not None:
                # It is! Learn it.
                self._gain_science(science(points))
        if sciences:
            # We're supposed to know an unknowable thing. :-(
            raise ValueError("Unknown science: %s" % (sciences.keys(),))

    def save_data(self):
        return {'science': dict(s.save_data() for s in self.science)}

    def _choose_initial_science(self):
        # We always get all starting schematics.
        for schematic in self.new_schematics[:]:
            if schematic.STARTING_PRODUCT:
                self._gain_science(schematic())

        # We start with Physics, because it's not Philately.
        physics = research.Physics()
        self._gain_science(physics)
        new_science = [physics]

        # We get two other random sciences with no prerequisites.
        for science in sample(self.find_new_research(), 2):
            science = science()
            self._gain_science(science)
            new_science.append(science)

        # Add a point to each of our sciences, and see if we get schematics.
        self.spend_points(new_science, 0)

    def _gain_science(self, science):
        self.science.append(science)
        if isinstance(science, research.ResearchArea):
            self.new_research.remove(type(science))
        elif isinstance(science, schematics.Schematic):
            self.new_schematics.remove(type(science))

    def spend_points(self, things, basic_research):
        breakthroughs = []

        # First, allocate the points.
        for thing in things:
            assert thing in self.science
            assert thing.can_spend(self, 1)
            thing.spend_point()

        # Next, check for schematic breakthroughs and upgrades
        breakthroughs.extend(self.apply_area_research([
                    thing for thing in things
                    if isinstance(thing, research.ResearchArea)]))

        # Finally, check for research breakthroughs.
        breakthroughs.extend(self.apply_basic_research(basic_research))

        return breakthroughs

    def _get_science(self, science_class):
        for science in self.science:
            if isinstance(science, science_class):
                return science
        return None

    def meet_requirements(self, science_class, extra=0):
        total_points = 0
        base_points = 0
        for science, level in science_class.PREREQUISITES:
            my_science = self._get_science(science)
            if my_science is None:
                return False
            if my_science.points < level:
                return False
            base_points += level
            total_points += my_science.points
        return total_points - base_points >= extra

    def find_new_schematics(self):
        available_schematics = []
        for schematic_class in self.new_schematics:
            if self.meet_requirements(schematic_class):
                available_schematics.append(schematic_class)
        return available_schematics

    def find_new_research(self):
        available_research = []
        for research_class in self.new_research:
            if self.meet_requirements(research_class):
                available_research.append(research_class)
        return available_research

    def apply_area_research(self, researches):
        options = [schema for schema in self.find_new_schematics()
                   if schema.depends_on(researches)]
        breakthroughs = [schematic for schematic in options
                         if random() < schematic.ACQUISITION_CHANCE]
        if breakthroughs:
            breakthrough = choice(breakthroughs)()
            self._gain_science(breakthrough)
            breakthroughs = [breakthrough]
        return breakthroughs

    def apply_basic_research(self, basic_research):
        if basic_research <= 0:
            return []

        options = self.find_new_research()
        success_chance = self.BASIC_RESEARCH_SUCCESS_RATE * (
            self.BASIC_RESEARCH_SUCCESS_MULTIPLIER ** basic_research)
        breakthroughs = [research for research in options
                         if random() < success_chance]
        if breakthroughs:
            breakthrough = choice(breakthroughs)(1)
            self._gain_science(breakthrough)
            breakthroughs = [breakthrough]
        return breakthroughs

    def suggest_research(self):
        """Suggest research areas to pursue.

        Return value is a tuple of (bool, list), where the first element
        indicates whether basic research might pay off and the second contains
        Science classes that can be profitably pursued.
        """
        graph = ScienceGraph(self.all_science, self.science)
        return graph.find_promising_areas()