# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np  # NOQA
import utool as ut
# import logging
import itertools as it
import copy
import six
import collections
from ibeis import constants as const
from ibeis.algo.graph import nx_dynamic_graph
# from ibeis.algo.graph import _dep_mixins
from ibeis.algo.graph import mixin_viz
from ibeis.algo.graph import mixin_helpers
from ibeis.algo.graph import mixin_dynamic
from ibeis.algo.graph import mixin_priority
from ibeis.algo.graph import mixin_loops
from ibeis.algo.graph import mixin_matching
from ibeis.algo.graph import mixin_groundtruth
from ibeis.algo.graph import mixin_simulation
from ibeis.algo.graph import mixin_ibeis
from ibeis.algo.graph import nx_utils as nxu
import pandas as pd
from ibeis.algo.graph.state import POSTV, NEGTV, INCMP, UNREV, UNKWN
from ibeis.algo.graph.state import UNINFERABLE
from ibeis.algo.graph.state import SAME, DIFF, NULL
import networkx as nx
import logging
print, rrr, profile = ut.inject2(__name__)


DEBUG_CC = False
# DEBUG_CC = True


def _rectify_decision(evidence_decision, meta_decision):
    """
    If evidence decision is not explicitly set, then meta decision is used to
    make a guess. Raises a ValueError if decisions are in incompatible states.
    """
    # Default to the decision based on the media evidence
    decision = evidence_decision
    # Overwrite the graph decision with the meta decision if necessary
    if meta_decision == SAME:
        if decision in UNINFERABLE:
            decision = POSTV
        elif decision == NEGTV:
            raise ValueError('evidence=negative and meta=same')
    elif meta_decision == DIFF:
        if decision in UNINFERABLE:
            decision = NEGTV
        elif decision == POSTV:
            raise ValueError('evidence=positive and meta=diff')
    return decision


class Feedback(object):
    def _check_edge(infr, edge):
        aid1, aid2 = edge
        if aid1 not in infr.aids_set:
            raise ValueError('aid1=%r is not part of the graph' % (aid1,))
        if aid2 not in infr.aids_set:
            raise ValueError('aid2=%r is not part of the graph' % (aid2,))

    def add_feedback_from(infr, items, verbose=None, **kwargs):
        if verbose is None:
            verbose = infr.verbose > 5
        if isinstance(items, pd.DataFrame):
            if list(items.index.names) == ['aid1', 'aid2']:
                for edge, data in items.iterrows():
                    infr.add_feedback(edge=edge, verbose=verbose, **data)
            else:
                raise ValueError(
                    'Cannot interpret pd.DataFrame without edge index')
        else:
            # Dangerous if item length > 3
            for item in items:
                args = []
                if len(item) == 1:
                    # Case where items=[edge1, edge2]
                    if isinstance(item[0], int) or len(item[0]) != 2:
                        raise ValueError('invalid edge')
                if len(item) == 2:
                    # Case where items=[(edge1, state), (edge2, state)]
                    if ut.isiterable(item[0]):
                        edge = item[0]
                        args = item[1:]
                    else:
                        edge = item
                else:
                    raise ValueError('invalid edge')
                    # Case where items=[(u, v, state), (u, v, state)]
                if len(item) > 3:
                    raise ValueError('pass in data as a dataframe or '
                                     'use kwargs')
                infr.add_feedback(edge, *args, verbose=verbose, **kwargs)

    def edge_decision(infr, edge):
        r"""
        Gets a decision on an edge, either explicitly or implicitly

        CommandLine:
            python -m ibeis.algo.graph.core edge_decision

        Doctest:
            >>> from ibeis.algo.graph.core import *  # NOQA
            >>> from ibeis.algo.graph import demo
            >>> infr = demo.demodata_infr(num_pccs=1, p_incon=1)
            >>> decision = infr.edge_decision((1, 2))
            >>> print('decision = %r' % (decision,))
            >>> assert decision == POSTV
            >>> decision = infr.edge_decision((199, 299))
            >>> print('decision = %r' % (decision,))
            >>> assert decision == UNREV
        """
        evidence_decision = infr.get_edge_attr(edge, 'evidence_decision',
                                               on_missing='default',
                                               default=UNREV)
        meta_decision = infr.get_edge_attr(edge, 'meta_decision',
                                           on_missing='default', default=NULL)
        decision = _rectify_decision(evidence_decision, meta_decision)
        return decision

    def edge_decision_from(infr, edges):
        r"""
        Gets a decision for multiple edges
        """
        edges = list(edges)
        evidence_decisions = infr.gen_edge_values(
            'evidence_decision', edges, on_missing='default', default=UNREV)
        meta_decisions = infr.gen_edge_values(
            'meta_decision', edges, on_missing='default', default=NULL)
        for ed, md in zip(evidence_decisions, meta_decisions):
            yield _rectify_decision(ed, md)

    def add_node_feedback(infr, aid, **attrs):
        infr.print('Writing annot aid=%r %s' % (aid, ut.repr2(attrs)))
        ibs = infr.ibs
        ibs.set_annot_quality_texts([aid], [attrs['quality_texts']])
        ibs.set_annot_viewpoint_code([aid], [attrs['viewpoint_code']])
        ibs.overwrite_annot_case_tags([aid], [attrs['case_tags']])
        ibs.set_annot_multiple([aid], [attrs['multiple']])

    @profile
    def add_feedback(infr, edge, evidence_decision=None, tags=None,
                     user_id=None, meta_decision=None, confidence=None,
                     timestamp_c1=None, timestamp_c2=None, timestamp_s1=None,
                     timestamp=None, verbose=None, priority=None):
        r"""
        Doctest:
            >>> from ibeis.algo.graph.core import *  # NOQA
            >>> infr = testdata_infr('testdb1')
            >>> infr.add_feedback((5, 6), POSTV)
            >>> infr.add_feedback((5, 6), NEGTV, tags=['photobomb'])
            >>> infr.add_feedback((1, 2), INCMP)
            >>> print(ut.repr2(infr.internal_feedback, nl=2))
            >>> assert len(infr.external_feedback) == 0
            >>> assert len(infr.internal_feedback) == 2
            >>> assert len(infr.internal_feedback[(5, 6)]) == 2
            >>> assert len(infr.internal_feedback[(1, 2)]) == 1
        """
        prev_verbose = infr.verbose
        if verbose is not None:
            infr.verbose = verbose
        edge = aid1, aid2 = nxu.e_(*edge)

        if not infr.has_edge(edge):
            if True:
                # Allow new aids
                if not infr.graph.has_node(aid1):
                    infr.add_aids([aid1])
                if not infr.graph.has_node(aid2):
                    infr.add_aids([aid2])
            infr._check_edge(edge)
            infr.graph.add_edge(aid1, aid2)

        if evidence_decision is None:
            evidence_decision = UNREV
        if meta_decision is None:
            meta_decision = const.META_DECISION.CODE.NULL
        if confidence is None:
            confidence = const.CONFIDENCE.CODE.UNKNOWN
        if timestamp is None:
            timestamp = ut.get_timestamp('int', isutc=True)

        msg = 'add_feedback ({}, {}), '.format(aid1, aid2)
        loc = locals()
        msg += ', '.join([
            str(val)
            # key + '=' + str(val)
            for key, val in (
                (key, loc[key])
                for key in ['evidence_decision', 'tags', 'user_id',
                            'confidence', 'meta_decision'])
            if val is not None
        ])
        infr.print(msg, 2, color='white')

        if meta_decision == NULL:
            # TODO: check previous meta_decision and use that if its consistent
            # with the evidence decision.
            pass

        decision = _rectify_decision(evidence_decision, meta_decision)

        if decision == UNREV:
            # Unreviewing an edge deletes anything not yet committed
            if edge in infr.external_feedback:
                raise ValueError('External edge reviews cannot be undone')
            if edge in infr.internal_feedback:
                del infr.internal_feedback[edge]

        # Remove the edge from the queue if it is in there.
        if infr.queue:
            if edge in infr.queue:
                del infr.queue[edge]

        # Keep track of sequential reviews and set properties on global graph
        num_reviews = infr.get_edge_attr(edge, 'num_reviews', default=0)

        review_id = next(infr.review_counter)
        feedback_item = {
            'tags': tags,
            'evidence_decision': evidence_decision,
            'meta_decision': meta_decision,
            'timestamp_c1': timestamp_c1,
            'timestamp_c2': timestamp_c2,
            'timestamp_s1': timestamp_s1,
            'timestamp': timestamp,
            'confidence': confidence,
            'user_id': user_id,
            'num_reviews': num_reviews + 1,
            'review_id': review_id,
        }
        infr.internal_feedback[edge].append(feedback_item)
        infr.set_edge_attr(edge, feedback_item)

        if infr.test_mode:
            prev_decision = infr._get_current_decision(edge)
            infr._dynamic_test_callback(edge, decision, prev_decision, user_id)

        # must happen after dynamic test callback
        infr.set_edge_attr(edge, {'decision': decision})

        if infr.params['inference.enabled']:
            assert infr.dirty is False, (
                'need to recompute before dynamic inference continues')
            # Update priority queue based on the new edge
            action = infr.add_review_edge(edge, decision)
            if infr.test_mode:
                infr.test_state['action'] = action
            if False:
                infr._print_debug_ccs()
        else:
            action = None
            infr.dirty = True
            infr._add_review_edge(edge, decision)

        if infr.params['inference.enabled'] and infr.refresh:
            # only add to criteria if this wasn't requested as a fix edge
            if priority is not None and priority <= 1.0:
                meaningful = bool({'merge', 'split'} & set(action))
                infr.refresh.add(meaningful, user_id, decision)

        if infr.test_mode:
            infr.metrics_list.append(infr.measure_metrics())

        infr.verbose = prev_verbose

    def _print_debug_ccs(infr):
        assert all([ut.allsame(infr.node_labels(*cc))
                    for cc in infr.positive_components()])
        sorted_ccs = sorted([
            set(cc) for cc in infr.pos_graph.connected_components()
        ])
        msg = '[' + ', '.join([
            repr(cc)
            if infr.is_consistent(cc) else
            ut.highlight_text(repr(cc), 'red')
            for cc in sorted_ccs]) + ']'
        print(msg)

    @ut.classproperty
    def feedback_keys(Infr):
        """ edge attribute keys used for feedback """
        return Infr.feedback_data_keys + ['num_reviews', 'review_id']

    @ut.classproperty
    def feedback_data_keys(Infr):
        """ edge attribute keys used for feedback """
        return [
            'evidence_decision', 'tags', 'user_id',
            'meta_decision', 'timestamp_c1', 'timestamp_c2',
            'timestamp_s1', 'timestamp', 'confidence'
        ]

    @profile
    def apply_feedback_edges(infr):
        r"""
        Transforms the feedback dictionaries into nx graph edge attributes

        CommandLine:
            python -m ibeis.algo.graph.core apply_feedback_edges

        Doctest:
            >>> from ibeis.algo.graph.core import *  # NOQA
            >>> infr = testdata_infr('testdb1')
            >>> infr.reset_feedback()
            >>> infr.params['inference.enabled'] = False
            >>> #infr.add_feedback((1, 2), 'unknown', tags=[])
            >>> infr.add_feedback((1, 2), INCMP, tags=[])
            >>> infr.apply_feedback_edges()
            >>> print('edges = ' + ut.repr4(dict(infr.graph.edges)))
            >>> result = str(infr)
            >>> print(result)
            <AnnotInference(nNodes=6, nEdges=3, nCCs=4)>
        """
        infr.print('apply_feedback_edges', 1)
        # Transforms dictionary feedback into numpy array
        edges = []
        attr_lists = {key: [] for key in infr.feedback_keys}
        for edge, vals in infr.all_feedback_items():
            # hack for feedback rectification
            feedback_item = infr._rectify_feedback_item(vals)
            feedback_item['review_id'] = next(infr.review_counter)
            feedback_item['num_reviews'] = len(vals)
            # if feedback_item['decision'] == 'unknown':
            #     continue
            set1 = set(feedback_item.keys())
            set2 = set(attr_lists.keys())
            if set1 != set2:
                raise AssertionError(
                    'Bad feedback keys: ' +
                    ut.repr2(ut.set_overlap_items(set1, set2, 'got', 'want'), nl=1)
                    # ut.repr2(sorted(feedback_item.keys()), sv=True) + ' ' +
                    # ut.repr2(sorted(attr_lists.keys()), sv=True)
                )
            for key, val in feedback_item.items():
                attr_lists[key].append(val)
            edges.append(edge)

        assert ut.allsame(list(map(len, attr_lists.values())))
        assert len(edges) == len(next(iter(attr_lists.values())))

        # Put pair orders in context of the graph
        infr.print('_set_feedback_edges(nEdges=%d)' % (len(edges),), 3)
        # Ensure edges exist
        for edge in edges:
            if not infr.graph.has_edge(*edge):
                infr.graph.add_edge(*edge)

        # take evidence_decision and meta_decision into account
        decisions = [
            _rectify_decision(ed, md) for ed, md in
            zip(attr_lists['evidence_decision'], attr_lists['meta_decision'])
        ]

        for state, es in ut.group_items(edges, decisions).items():
            infr._add_review_edges_from(es, state)

        for key, val_list in attr_lists.items():
            infr.set_edge_attrs(key, ut.dzip(edges, val_list))

        if infr.params['inference.enabled']:
            infr.apply_nondynamic_update()

    def _rectify_feedback(infr, feedback):
        return {edge: infr._rectify_feedback_item(vals)
                for edge, vals in feedback.items()}

    def _rectify_feedback_item(infr, vals):
        """ uses most recently use strategy """
        return vals[-1]

    def all_feedback_items(infr):
        for edge, vals in six.iteritems(infr.external_feedback):
            yield edge, vals
        for edge, vals in six.iteritems(infr.internal_feedback):
            yield edge, vals

    def all_feedback(infr):
        all_feedback = ut.ddict(list)
        all_feedback.update(infr.all_feedback_items())
        return all_feedback

    def clear_feedback(infr, edges=None):
        """ Delete all edges properties related to feedback """
        if edges is None:
            edges = infr.graph.edges()
        edges = list(edges)
        infr.print('clear_feedback len(edges) = %r' % (len(edges)), 2)
        infr.external_feedback = ut.ddict(list)
        infr.internal_feedback = ut.ddict(list)

        # Kill all feedback, remote edge labels, but leave graph edges alone
        keys = infr.feedback_keys + ['inferred_state']
        ut.nx_delete_edge_attr(infr.graph, keys, edges)

        # Move reviewed edges back into the unreviewed graph
        for key in (POSTV, NEGTV, INCMP):
            subgraph = infr.review_graphs[key]
            prev_edges = ut.compress(edges, list(subgraph.has_edges(edges)))
            subgraph.remove_edges_from(prev_edges)
            infr.review_graphs[UNREV].add_edges_from(prev_edges)

        infr.pos_redun_nids.clear()
        infr.neg_redun_metagraph.clear()
        infr.nid_to_errors.clear()

        if __debug__:
            infr.assert_disjoint_invariant()

    def clear_edges(infr):
        """
        Removes all edges from the graph
        """
        for graph in infr.review_graphs.values():
            graph.remove_edges_from(list(graph.edges()))
        infr.graph.remove_edges_from(list(infr.graph.edges()))
        infr.pos_redun_nids.clear()
        infr.neg_redun_metagraph.clear()
        infr.nid_to_errors.clear()

    def reset_feedback(infr, mode='annotmatch', apply=True):
        """ Resets feedback edges to state of the SQL annotmatch table """
        infr.print('reset_feedback mode=%r' % (mode,), 1)
        infr.clear_feedback()
        if mode == 'annotmatch':
            infr.external_feedback = infr.read_ibeis_annotmatch_feedback()
        elif mode == 'staging':
            infr.external_feedback = infr.read_ibeis_staging_feedback()
        else:
            raise ValueError('no mode=%r' % (mode,))
        infr.internal_feedback = ut.ddict(list)
        if apply:
            infr.apply_feedback_edges()

    def reset(infr, state='empty'):
        """
        Removes all edges from graph and resets name labels.

        Ignore:
            >>> from ibeis.algo.graph.core import *  # NOQA
            >>> from ibeis.algo.graph import demo
            >>> infr = demo.demodata_infr(num_pccs=5)
            >>> assert len(list(infr.edges())) > 0
            >>> infr.reset(state='empty')
            >>> assert len(list(infr.edges())) == 0
        """
        infr.clear_edges()
        infr.clear_feedback()
        if state == 'empty':
            # Remove all edges, and component names
            infr.clear_name_labels()
        elif state == 'orig':
            raise NotImplementedError('unused')
            infr.reset_name_labels()
        else:
            raise ValueError('Unknown state=%r' % (state,))

    def reset_name_labels(infr):
        """ Resets all annotation node name labels to their initial values """
        infr.print('reset_name_labels', 1)
        orig_names = infr.get_node_attrs('orig_name_label')
        infr.set_node_attrs('name_label', orig_names)

    def clear_name_labels(infr):
        """ Sets all annotation node name labels to be unknown """
        infr.print('clear_name_labels()', 1)
        # make distinct names for all nodes
        distinct_names = {
            node: -aid for node, aid in infr.get_node_attrs('aid').items()
        }
        infr.set_node_attrs('name_label', distinct_names)


class NameRelabel(object):

    def node_label(infr, aid):
        return infr.pos_graph.node_label(aid)

    def node_labels(infr, *aids):
        return infr.pos_graph.node_labels(*aids)

    def _next_nid(infr):
        if getattr(infr, 'nid_counter', None) is None:
            nids = nx.get_node_attributes(infr.graph, 'name_label')
            infr.nid_counter = max(nids)
        infr.nid_counter += 1
        new_nid = infr.nid_counter
        return new_nid

    def _rectify_names(infr, old_names, new_labels):
        """
        Finds the best assignment of old names based on the new groups each is
        assigned to.

        old_names  = [None, None, None, 1, 2, 3, 3, 4, 4, 4, 5, None]
        new_labels = [   1,    2,    2, 3, 4, 5, 5, 6, 3, 3, 7, 7]
        """
        infr.print('rectifying name lists', 3)
        from ibeis.scripts import name_recitifer
        newlabel_to_oldnames = ut.group_items(old_names, new_labels)
        unique_newlabels = list(newlabel_to_oldnames.keys())
        grouped_oldnames_ = ut.take(newlabel_to_oldnames, unique_newlabels)
        # Mark annots that are unknown and still grouped by themselves
        still_unknown = [len(g) == 1 and g[0] is None for g in grouped_oldnames_]
        # Remove nones for name rectifier
        grouped_oldnames = [
            [n for n in oldgroup if n is not None]
            for oldgroup in grouped_oldnames_]
        new_names = name_recitifer.find_consistent_labeling(
            grouped_oldnames, verbose=infr.verbose >= 3, extra_prefix=None)

        unknown_labels = ut.compress(unique_newlabels, still_unknown)

        new_flags = [n is None for n in new_names]
        #     isinstance(n, six.string_types) and n.startswith('_extra_name')
        #     for n in new_names
        # ]
        label_to_name = ut.dzip(unique_newlabels, new_names)
        needs_assign = ut.compress(unique_newlabels, new_flags)
        return label_to_name, needs_assign, unknown_labels

    def _rectified_relabel(infr, cc_subgraphs):
        """
        Reuses as many names as possible
        """
        # Determine which names can be reused
        from ibeis.scripts import name_recitifer
        infr.print('grouping names for rectification', 3)
        grouped_oldnames_ = [
            list(nx.get_node_attributes(subgraph, 'name_label').values())
            for count, subgraph in enumerate(cc_subgraphs)
        ]
        # Make sure negatives dont get priority
        grouped_oldnames = [
            [n for n in group if len(group) == 1 or n > 0]
            for group in grouped_oldnames_
        ]
        infr.print('begin rectification of %d grouped old names' % (
            len(grouped_oldnames)), 2)
        new_labels = name_recitifer.find_consistent_labeling(
            grouped_oldnames, verbose=infr.verbose >= 3)
        infr.print('done rectifying new names', 2)
        new_flags = [
            not isinstance(n, int) and n.startswith('_extra_name')
            for n in new_labels
        ]

        for idx in ut.where(new_flags):
            new_labels[idx] = infr._next_nid()

        for idx, label in enumerate(new_labels):
            if label < 0 and len(grouped_oldnames[idx]) > 1:
                # Remove negative ids for grouped items
                new_labels[idx] = infr._next_nid()
        return new_labels

    @profile
    def relabel_using_reviews(infr, graph=None, rectify=True):
        r"""
        Relabels nodes in graph based on positive connected components

        This will change all of the names on the nodes to be consistent while
        preserving any existing names as best as possible. If rectify=False,
        this will be faster, but the old names will not be preserved and each
        PCC will be assigned an arbitrary name.

        Note:
            if something messes up you can call infr.reset_labels_to_ibeis() to
            reset node labels to their original values --- this will almost
            always put the graph in an inconsistent state --- but then you can
            this with rectify=True to fix everything up.

        Args:
            graph (nx.Graph, optional): only edges in `graph` are relabeled
                defaults to current graph.
            rectify (bool, optional): if True names attempt to remain
                consistent otherwise there are no restrictions on name labels
                other than that they are distinct.
        """
        infr.print('relabel_using_reviews', 2)
        if graph is None:
            graph = infr.graph

        # Get subgraphs and check consistency
        cc_subgraphs = []
        num_inconsistent = 0
        for cc in infr.positive_components(graph=graph):
            cc_subgraphs.append(infr.graph.subgraph(cc))
            if not infr.is_consistent(cc):
                num_inconsistent += 1

        infr.print('num_inconsistent = %r' % (num_inconsistent,), 2)
        if infr.verbose >= 2:
            cc_sizes = list(map(len, cc_subgraphs))
            pcc_size_hist = ut.dict_hist(cc_sizes)
            pcc_size_stats = ut.get_stats(cc_sizes)
            if len(pcc_size_hist) < 8:
                infr.print('PCC size hist = %s' % (ut.repr2(pcc_size_hist),))
            infr.print('PCC size stats = %s' % (ut.repr2(pcc_size_stats),))

        if rectify:
            # Rectified relabeling, preserves grouping and labeling if possible
            new_labels = infr._rectified_relabel(cc_subgraphs)
        else:
            # Arbitrary relabeling, only preserves grouping
            if graph is infr.graph:
                # Use union find labels
                new_labels = {
                    count:
                    infr.node_label(next(iter(subgraph.nodes())))
                    for count, subgraph in enumerate(cc_subgraphs)
                }
            else:
                new_labels = {count: infr._next_nid()
                              for count, subgraph in enumerate(cc_subgraphs)}

        for count, subgraph in enumerate(cc_subgraphs):
            new_nid = new_labels[count]
            node_to_newlabel = ut.dzip(subgraph.nodes(), [new_nid])
            infr.set_node_attrs('name_label', node_to_newlabel)

        num_names = len(cc_subgraphs)
        infr.print('done relabeling', 3)
        return num_names, num_inconsistent

    def connected_component_status(infr):
        r"""
        Returns:
            dict: num_inconsistent, num_names_max

        CommandLine:
            python -m ibeis.algo.graph.core connected_component_status

        Example:
            >>> # DISABLE_DOCTEST
            >>> from ibeis.algo.graph.core import *  # NOQA
            >>> infr = testdata_infr('testdb1')
            >>> infr.add_feedback_from([(2, 3, NEGTV), (5, 6, NEGTV), (1, 2, POSTV)])
            >>> status = infr.connected_component_status()
            >>> print(ut.repr3(status))
        """
        infr.print('checking status', 3)

        num_inconsistent = len(infr.recovery_ccs)
        num_names_max = infr.pos_graph.number_of_components()

        status = dict(
            num_names_max=num_names_max,
            num_inconsistent=num_inconsistent,
        )
        infr.print('done checking status', 3)
        return status


class MiscHelpers(object):

    def _rectify_nids(infr, aids, nids):
        if nids is None:
            if infr.ibs is None:
                nids = [-aid for aid in aids]
            else:
                nids = infr.ibs.get_annot_nids(aids)
        elif ut.isscalar(nids):
            nids = [nids] * len(aids)
        return nids

    def remove_aids(infr, aids):
        """
        Remove annotations from the graph.
        Returns:
            dict: split: indicates which PCCs were split by this action.
        Note:
            This may cause unintended splits!

        Ignore:
            >>> from graphid import demo, util
            >>> infr = demo.demodata_infr(num_pccs=5, pos_redun=1)
            >>> infr.refresh_candidate_edges()
            >>> infr.pin_node_layout()
            >>> before = infr.copy()
            >>> aids = infr.aids[::5]