# -*- coding: utf-8 -*-
"""Define nonogram solver that uses contradictions"""

import time
from collections import OrderedDict, defaultdict, deque
from itertools import product

from notetool.tool.log import logger
from six import iteritems
from six.moves import range

from notegame.games.nonogram.core import propagation
from notegame.games.nonogram.core.board import CellPosition, CellState
from notegame.games.nonogram.core.common import NonogramError
from notegame.games.nonogram.solver.base import cache_info
from notegame.games.nonogram.utils.iter import expand_generator
from notegame.games.nonogram.utils.priority_dict import PriorityDict

USE_CONTRADICTION_RESULTS = True

# the fewer possible colors for a cell
# the earlier it appears in a list of candidates
FEW_COLORS_FIRST = False

# use the position of a cell (number of neighbours and solution rate of row/column)
# to adjust its rate when choosing the next probe for DFS
ADJUST_RATE = True


class _SearchNode(object):
    def __init__(self, value):
        self.value = value
        self.children = OrderedDict()

    def to_dict(self):
        """Represent the given node as the dictionary of other nodes"""

        if not self.children:
            return self.value

        return {
            'value': self.value,
            'children': OrderedDict(
                (str(k), v.to_dict())
                for k, v in iteritems(self.children)
            )
        }


class Solver(object):
    """
    Solve the nonogram using contradictions and depth-first search
    """

    def __init__(self, board, max_solutions=None, timeout=None, max_depth=None):
        """
        :type board: Board
        """
        self.board = board

        self.max_solutions = max_solutions
        self.timeout = timeout
        if max_depth is None:
            # why 440?
            # I simply searched for some value that will be somehow bigger
            # than 434 (the maximum useful search depth reached on the test set so far),
            # but still reachable for recursion calls.
            #
            # https://webpbn.com/30654 can be solved only with max_depth > 429
            # and the maximum depth reached on this puzzle is 434.
            #
            # NB: in current implementation depth=444 fails with
            # 'RuntimeError: maximum recursion depth exceeded'
            self.max_depth = 440
        else:
            self.max_depth = max_depth

        self.depth_reached = 0
        self.start_time = None
        self.explored_paths = set()
        self.search_map = None

    def _add_search_result(self, path, score):
        if isinstance(score, float):
            score = round(score, 4)

        if not path:
            if not self.search_map:
                self.search_map = _SearchNode(score)
            else:
                self.search_map.value = score
            return

        current = self.search_map

        for i, node in enumerate(path):
            # node = (i + 1, node)
            if i == len(path) - 1:
                val = score
            else:
                val = None

            if node not in current.children:
                current.children[node] = _SearchNode(val)

            current = current.children[node]

        current.value = score

    def propagate_change(self, cell_state):
        """
        Set the given color to given cell
        and try to solve the board with that new info.
        :type cell_state: CellState
        """
        board = self.board
        logger.debug('Assume that (%i, %i) is %s',  *tuple(cell_state))

        board.set_color(cell_state)

        return propagation.solve(board,
                                 row_indexes=(cell_state.row_index,),
                                 column_indexes=(cell_state.column_index,),
                                 contradiction_mode=True)

    def probe(self, cell_state, rollback=True, force=False):
        """
        Try to find if the given cell can be in an assumed state.
        If the contradiction is found, set the cell
        in an inverted state and propagate the changes if needed.
        If `rollback` then the solved board will restore to the previous state
        after the assumption was made.
        If `force`, try to solve it anyway, even if the cell is already solved
        (to rerun the line solver).

        Return the pair `(is_contradiction, new_info)` where

        is_contradiction: whether the assumption led to a contradiction

        new_info:
          a) when contradiction is found, it contains the state of the board
          before any assumptions was made. It will help further to determine
          which cells has changed on that probe.
          b) if no contradiction found, but `rollback` is False,
          then we do not restore the board and return the previous state also.
          c) otherwise it contains the number of solved cells for the partially
          solved board (if the assumption made is true)
        """
        board = self.board

        pos = cell_state.position
        assumption = cell_state.color
        # already solved
        if board.is_cell_solved(pos):
            if not force:
                return False, None

        if assumption not in board.cell_colors(pos):
            logger.warning(
                "The probe is useless: color '%s' already unset" % (assumption))
            return False, None

        save = board.make_snapshot()

        try:
            solved_cells = self.propagate_change(cell_state)
        except NonogramError:
            logger.debug('Contradiction', exc_info=True)
            # rollback solved cells
            board.restore(save)

        else:
            if board.is_solved_full:
                self._add_solution()

            if rollback:
                board.restore(save)
                return False, solved_cells

            return False, save

        if USE_CONTRADICTION_RESULTS:
            before_contradiction = board.make_snapshot()
        else:
            before_contradiction = None

        pos = cell_state.position
        logger.info('Found contradiction at (%i, %i)' , *pos)
        try:
            board.unset_color(cell_state)
        except ValueError as ex:
            raise NonogramError(str(ex))

        # try to solve with additional info
        # solve with only one cell as new info
        propagation.solve(
            board,
            row_indexes=(pos.row_index,),
            column_indexes=(pos.column_index,))

        return True, before_contradiction

    def _new_jobs_from_solution(self, cell_state, previous_board, is_contradiction):
        board = self.board

        # evaluate generator
        changed = list(board.changed(previous_board))
        assumption = cell_state.color
        log_contradiction = '(not) ' if is_contradiction else ''
        logger.info('Changed %d cells with %s%s assumption' %
                    (len(changed), log_contradiction, assumption))

        # add the neighbours of the changed cells into jobs
        for pos in changed:
            for neighbour in board.unsolved_neighbours(pos):
                yield neighbour, 1

        # add the neighbours of the selected cell into jobs
        for neighbour in board.unsolved_neighbours(cell_state.position):
            yield neighbour, 0

    def _set_guess(self, state):
        board = self.board
        is_contradiction, prev_board = self.probe(
            state, rollback=False, force=True)

        if is_contradiction:
            raise NonogramError('Real contradiction was found: %s' % (state,))

        if prev_board is None:
            logger.warning(
                "The probe for state '%s' does not return anything new" % state)
            return ()

        if board.is_solved_full:
            self._add_solution()
            return ()

        return self._new_jobs_from_solution(state, prev_board, is_contradiction)

    def _solve_jobs(self, jobs, refill=False):
        """
        Given a board and a list of jobs try to solve that board
        using the jobs as probes.
        If `refill` specified, solve in several rounds,
        until all the contradictions disappear.

        Return the number of contradictions found and
        list of the best candidates for the tree-based search
        """

        counter, counter_found = 0, 0
        rates = dict()

        board = self.board

        processed_in_round = set()
        expired_assumptions = set()

        while jobs:
            state, priority = jobs.pop_smallest()
            counter += 1
            logger.info('Probe #%d: %s (%f)' % (counter, state, priority))

            # if the job is only coordinates
            # then try all the possible colors
            pos = state[:2]
            if len(state) == 2:
                assumptions = board.cell_colors(state)
            else:
                assumptions = (state.color,)

            for assumption in assumptions:
                state = CellState.from_position(pos, assumption)
                is_contradiction, info = self.probe(state)

                if info is None:
                    continue

                if is_contradiction:
                    counter_found += 1
                    if board.is_solved_full:
                        self._add_solution()
                        return counter_found, None

                    for new_job, priority in self._new_jobs_from_solution(
                            state, info, is_contradiction):
                        jobs[new_job] = priority

                    # All the jobs that was processed before the given
                    # contradiction considered expired. They will refill
                    # the active jobs queue when it gets empty.
                    expired_assumptions = set(processed_in_round)
                else:
                    rates[state] = (info, priority)

            if not refill:
                # only do one round of probes
                continue

            # we have work to do!
            if jobs:
                processed_in_round.add(pos)
            elif expired_assumptions:
                logger.warning('No more jobs. Refill all the jobs processed before '
                               'the last found contradiction (%s)' % (len(expired_assumptions)))
                refill_processed = self._get_all_unsolved_jobs(
                    choose_from_cells=expired_assumptions)
                for new_job, priority in iteritems(refill_processed):
                    jobs[new_job] = priority

                # they are no longer expired, as they active now!
                processed_in_round -= expired_assumptions
                expired_assumptions = set()
                counter = 0

        return counter_found, self._probes_from_rates(rates)

    def _add_solution(self):
        # force to check the board
        propagation.solve(self.board, contradiction_mode=True)
        self.board.add_solution()

    @expand_generator
    def _probes_from_rates(self, rates):
        jobs_with_rates = defaultdict(dict)

        for cell_state, (rate, priority) in iteritems(rates):
            pos = cell_state.position
            color = cell_state.color
            if self.board.is_cell_solved(pos):
                continue

            # the more priority the less desired that job
            # priority is in the range [0, 8]
            if ADJUST_RATE:
                rate += (10 - priority)
            # if rate > jobs_with_rates.get(job, 0):
            jobs_with_rates[pos][color] = rate

        max_rate = {pos: max(v.values())
                    for pos, v in iteritems(jobs_with_rates)}
        # the biggest rate appears first
        best = sorted(iteritems(max_rate), key=lambda x: x[1], reverse=True)
        if FEW_COLORS_FIRST:
            best = sorted(best, key=lambda x: len(jobs_with_rates[x[0]]))
        logger.debug('\n'.join(map(str, best)))

        for pos, max_rate in best:
            colors = sorted(
                iteritems(jobs_with_rates[pos]), key=lambda x: x[1], reverse=True)
            for color, rate in colors:
                yield CellState.from_position(pos, color)

    def _get_all_unsolved_jobs(self, choose_from_cells=None):
        board = self.board

        if choose_from_cells is None:
            # add every cell
            choose_from_cells = product(
                range(board.height), range(board.width))

        probe_jobs = PriorityDict()

        for pos in choose_from_cells:
            pos = CellPosition(*pos)
            if board.is_cell_solved(pos):
                continue

            no_unsolved = len(list(board.unsolved_neighbours(pos)))

            # if no_unsolved >= 4 and skip_low_rated:
            #     continue

            row_rate = board.row_solution_rate(pos.row_index)
            column_rate = board.column_solution_rate(pos.column_index)
            cell_rate = row_rate + column_rate

            probe_jobs[pos] = 4 - cell_rate + no_unsolved

        return probe_jobs

    def _solve_without_search(self, to_the_max=False):
        """
        Do the one round of solving with contradictions.
        Returns the number of contradictions found.

        Based on https://www.cs.bgu.ac.il/~benr/nonograms/
        """

        probe_jobs = self._get_all_unsolved_jobs()
        return self._solve_jobs(probe_jobs, refill=to_the_max)

    @classmethod
    def shrink_board(cls, board, candidates=()):
        """
        Skip the solved rows in a board before backtracking.
        Also fix candidate states.
        """
        solved_columns, solved_rows = board.reduce()

        if solved_rows and solved_rows[0]:
            skip_first_rows = len(solved_rows[0])

            for index, candidate in enumerate(candidates):  # type: CellState
                if candidate.row_index < skip_first_rows:
                    raise ValueError(
                        'Bad candidate: %r. The row skipped' % candidate)

                row_index, col_index, color = candidate
                new_candidate = CellState(
                    row_index - skip_first_rows, col_index, color)
                candidates[index] = new_candidate
                logger.info('Fixed candidate: %r -> %r' %
                            (candidate, new_candidate))

        if solved_columns and solved_columns[0]:
            skip_first_columns = len(solved_columns[0])

            for index, candidate in enumerate(candidates):  # type: CellState
                if candidate.column_index < skip_first_columns:
                    raise ValueError(
                        'Bad candidate: %r. The column skipped' % candidate)

                row_index, col_index, color = candidate
                new_candidate = CellState(
                    row_index, col_index - skip_first_columns, color)
                candidates[index] = new_candidate
                logger.info('Fixed candidate: %r -> %r',
                            candidate, new_candidate)

        return candidates

    @classmethod
    def _fix_candidates_colors(cls, candidates, color_mapping):
        for index, candidate in enumerate(candidates):  # type: CellState
            row_index, col_index, color = candidate
            if color not in color_mapping:
                logger.warning('Bad candidate %r', candidate)
                continue

            new_candidate = CellState(
                row_index, col_index, color_mapping[color])
            candidates[index] = new_candidate
            logger.info('Fixed candidate: %r -> %r', candidate, new_candidate)

        return candidates

    def solve(self, methods=None):
        """
        Solve the nonogram to the most with contradictions
        """

        board = self.board

        propagation.solve(board, methods=methods)
        if board.is_solved_full:
            board.set_finished()
            logger.info('No need to solve with contradictions')
            return

        logger.warning('Trying to solve using contradictions method')
        start = time.time()

        found_contradictions, best_candidates = self._solve_without_search(
            to_the_max=True)
        current_solution_rate = board.solution_rate
        logger.warning('Contradictions (found %d): %f',
                       found_contradictions, current_solution_rate)

        if current_solution_rate < 1:
            # if stalled with sophisticated selection of cells
            # do the brute force search
            logger.warning('Starting DFS (intelligent brute-force)')
            best_candidates = self.shrink_board(
                board, candidates=best_candidates)
            if board.is_colored:
                single_colored, color_mapping = board.reduce_to_single_color()
                if single_colored is not None:
                    logger.warning('Replacing colored board with the '
                                   'equivalent black and white: %r', color_mapping)

                    # from now we will search the black and white board
                    self.board = single_colored
                    best_candidates = self._fix_candidates_colors(
                        best_candidates, color_mapping)

            self.search(best_candidates)
            board.restore_reduced()

            current_solution_rate = board.solution_rate
            logger.warning('Search completed (depth reached: %d, solutions found: %d)',
                           self.depth_reached, len(board.solutions))

        if current_solution_rate != 1:
            logger.warning('The nonogram is not solved full (with contradictions). '
                           'The rate is %.4f', current_solution_rate)

        board.set_finished()
        logger.info('Full solution: %.6f sec', time.time() - start)
        for method, info in cache_info().items():
            size, hit_rate = info
            if size > 0:
                logger.warning('%s cache: size=%d, hit rate=%.4f%%',
                               method, size, hit_rate * 100.0)

    def _limits_reached(self, depth):
        """
        Whether we reached the defined limits:
        1) number of solutions found
        2) the maximum allowed run time
        """
        if self.max_solutions:
            solutions_number = len(self.board.solutions)
            if solutions_number >= self.max_solutions:
                if depth == 0:
                    # only show log on the top level
                    logger.warning('%d solutions is enough', solutions_number)

                return True

        if self.timeout and self.start_time:
            run_time = time.time() - self.start_time
            if run_time > self.timeout:
                if depth == 0:
                    # only show log on the top level
                    logger.warning('Searched too long: %.4fs', run_time)
                return True

        return False

    def _try_state(self, state, path):
        """
        Trying to search for solutions in the given direction.
        At first it set the given state and get a list of the
        further jobs for finding the contradictions.
        Later that jobs will be used as candidates for a deeper search.

        :param state: the next cell and color to set
        :param path: the path that already have done
        """
        board = self.board

        depth = len(path)
        full_path = path + (state,)

        # add every cell to the jobs queue
        probe_jobs = self._get_all_unsolved_jobs()
        try:
            # update with more prioritized cells
            for new_job, priority in self._set_guess(state):
                probe_jobs[new_job] = priority

            if self._limits_reached(depth):
                return True

            __, best_candidates = self._solve_jobs(probe_jobs)
        except NonogramError as ex:
            logger.warning('Dead end found (%s): %s', full_path[-1], str(ex))
            self._add_search_result(full_path, False)
            return False

        rate = board.solution_rate
        logger.info('Reached rate %.4f on %s path', rate, full_path)
        self._add_search_result(full_path, rate)

        if rate == 1 or self._limits_reached(depth):
            return True

        cells_left = round((1 - rate) * board.width * board.height)
        logger.info('Unsolved cells left: %d', cells_left)

        if best_candidates:
            return self.search(best_candidates, path=full_path)

        return True

    def _set_explored(self, path):
        self.explored_paths.add(tuple(sorted(path)))

    def _is_explored(self, path):
        return tuple(sorted(path)) in self.explored_paths

    def search(self, states, path=()):
        """
        Recursively search for solutions

        Return False if the given path is a dead end (no solutions can be found)
        """

        if self._is_explored(path):
            return True

        if self.start_time is None:
            self.start_time = time.time()

        board = self.board
        depth = len(path)

        if self._limits_reached(depth):
            return True

        if self.max_depth and depth >= self.max_depth:
            logger.warning('Next step on the depth %d is deeper than the max (%d)',
                           depth, self.max_depth)
            return True

        # going to dive deeper, so increment it (full_path's length)
        if depth + 1 > self.depth_reached:
            self.depth_reached = depth + 1

        search_directions = deque(states)

        unconditional = False
        search_counter = 0
        save = board.make_snapshot()
        try:
            while search_directions:
                total_number_of_directions = len(search_directions)
                state = search_directions.popleft()
                search_counter += 1

                if self._limits_reached(depth):
                    return True

                if state in path:
                    continue

                assumption = state.color
                pos = state.position
                cell_colors = board.cell_colors(pos)

                if assumption not in cell_colors:
                    logger.warning("The assumption '%s' is already expired. "
                                   "Possible colors for %s are %s",
                                   assumption, pos, cell_colors)
                    continue

                if len(cell_colors) == 1:
                    logger.warning('Only one color for cell %r left: %s. Solve it unconditionally',
                                   pos, assumption)
                    assert assumption == tuple(cell_colors)[0]
                    if unconditional:
                        logger.warning(
                            'The board does not change since the last unconditional solving, skip.')
                        continue

                    try:
                        self._solve_without_search()
                        unconditional = True
                    except NonogramError:
                        # the whole `path` branch of a search tree is a dead end
                        logger.warning(
                            "The last possible color '%s' for the cell '%s' "
                            "lead to the contradiction. "
                            "The path %s is invalid", assumption, pos, path)
                        # self._add_search_result(path, False)
                        return False

                    # rate = board.solution_rate
                    # self._add_search_result(path, rate)
                    if board.is_solved_full:
                        self._add_solution()
                        logger.warning(
                            "The only color '%s' for the cell '%s' lead to full solution. "
                            "No need to traverse the path %s anymore", assumption, pos, path)
                        return True
                    continue

                full_path = path + (state,)
                if self._is_explored(full_path):
                    logger.info('The path %s already explored', full_path)
                    continue

                unconditional = False
                rate = board.solution_rate
                guess_save = board.make_snapshot()
                try:
                    logger.warning('Trying state (%d/%d): %s (depth=%d, rate=%.4f, previous=%s)',
                                   search_counter, total_number_of_directions,
                                   state, depth, rate, path)
                    self._add_search_result(path, rate)
                    success = self._try_state(state, path)
                    # is_solved = board.is_solved_full
                finally:
                    board.restore(guess_save)
                    self._set_explored(full_path)

                if not success:
                    # TODO: add backjumping here
                    try:
                        logger.warning(
                            "Unset the color %s for cell '%s'. Solve it unconditionally",
                            assumption, pos)
                        board.unset_color(state)
                        self._solve_without_search()
                        unconditional = True
                    except ValueError:
                        # the whole `path` branch of a search tree is a dead end
                        logger.warning(
                            "The last possible color '%s' for the cell '%s' "
                            "lead to the contradiction. "
                            "The whole branch (depth=%d) is invalid. ", assumption, pos, depth)
                        # self._add_search_result(path, False)
                        return False

                    # rate = board.solution_rate
                    # self._add_search_result(path, rate)
                    if board.is_solved_full:
                        self._add_solution()
                        logger.warning(
                            "The negation of color '%s' for the cell '%s' lead to full solution. "
                            "No need to traverse the path %s anymore", assumption, pos, path)
                        return True

                if not success or board.is_solved_full:
                    # immediately try the other colors as well
                    # if all of them goes to the dead end,
                    # then the parent path is a dead end
                    states_to_try = []
                    for color in cell_colors:
                        if color == assumption:
                            continue

                        states_to_try.append(
                            CellState.from_position(pos, color))

                    # if all(self._is_explored(path + (state,)) for state in states_to_try):
                    #     LOG.warning('All other colors (%s) of cell %s already explored',
                    #                 states_to_try, cell)
                    #     return True

                    for state in states_to_try:
                        if state not in search_directions:
                            search_directions.appendleft(state)

        finally:
            # do not restore the solved cells on a root path - they are really solved!
            if path:
                board.restore(save)
                self._set_explored(path)

        return True
