import numpy as np
import os
import json
import itertools
import copy
import multiprocessing
import pandas as pd
import networkx as nx
import PIL.Image
from pathlib import Path
from typing import List
from nway.pairwise_matching import PairwiseMatching
from nway.schemas import NwayMatchingSchema, NwayMatchingOutputSchema
import nway.utils as utils
from nway import __version__ as nway_version
from argschema import ArgSchemaParser

from nway.diagnostics import NwaySummary


class NwayException(Exception):
    pass


def subgraphs(G):
    """returns connected component subgraphs. Replaces deprecated
    nx.connected_component_subgraphs()

    Parameters
    ----------
    G : :class:`networkx.Graph`

    Returns
    -------
    subgs : tuple
        tuple of connected subgraphs of G

    """
    subgs = (G.subgraph(c).copy() for c in nx.connected_components(G))
    return subgs


def greduce(glist, method='keepmin'):
    """recursively reduces a graph according to node attribute 'score'

    Parameters
    ----------
    glist : tuple
        tuple of subgraphs
    method : str
        'keepmin' reduce graph by finding the minimum score of a subraph
            and eliminating its neighbors
        'popmax' reduce graph by finding the maximum score and eliminating it

    Returns
    -------
    results : list
        list of nodes that survive the pruning

    """
    results = []
    for g in glist:
        gnodes = list(g.nodes.keys())
        if len(gnodes) == 1:
            results += [gnodes[0]]
        else:
            if method == 'popmax':
                maxind = np.argmax([i['score'] for i in g.nodes.values()])
                g.remove_node(gnodes[maxind])
            elif method == 'keepmin':
                minind = np.argmin([i['score'] for i in g.nodes.values()])
                neighbors = list(g.neighbors(gnodes[minind]))
                for n in neighbors:
                    g.remove_node(n)
            results += greduce(subgraphs(g), method=method)

    return results


def prune_matching_table(table, score, method='keepmin'):
    """eliminates match conflicts by comparing scores

    Parameters
    ----------
    table : list
        list (len = n_match_candidates) of numpy arrays (size = n_experiment)
        containing either cellIDs or -1
    score : :class:`numpy.ndarray`
        size = n_match_candidates, a float score value for each candidat
    method : str
        'keepmin' reduce graph by finding the minimum score of a subraph
            and eliminating its neighbors
        'popmax' reduce graph by finding the maximum score and eliminating it

    Returns
    -------
    pruned : list
        pruned version of the input table

    """
    G = nx.Graph()
    for i in range(len(table)):
        G.add_node(i, score=score[i])

    # if any 2 lines share an entry, make an edge between them
    for (i0, line0), (i1, line1) in itertools.combinations(
            enumerate(table), 2):
        nz = (line0 != -1) & (line1 != -1)
        if np.any(line0[nz] == line1[nz]):
            G.add_edge(i0, i1)

    inds = greduce(subgraphs(G), method=method)
    pruned = [table[i] for i in inds]

    return pruned


def pair_match_job(pair_args):
    pair_match = PairwiseMatching(input_data=pair_args, args=[])
    pair_match.run()
    delattr(pair_match, 'logger')
    return pair_match


def check_image_sizes(image_paths: List[Path]):
    sizes = dict()
    for image_path in image_paths:
        with open(image_path, 'rb') as fp:
            im = PIL.Image.open(fp)
            sizes.update({str(image_path): im.size})
    try:
        sz0 = next(iter(sizes.values()))
        assert all([i == sz0 for i in sizes.values()])
    except AssertionError:
        estr = "\n".join([f"{v}: {k}"
                          for k, v in sizes.items()])
        raise NwayException("not all experiments have the same size "
                            "average projection images, which nway "
                            "matching currently expects. sizes and image "
                            f"paths are\n{estr}")


class NwayMatching(ArgSchemaParser):
    default_schema = NwayMatchingSchema
    default_output_schema = NwayMatchingOutputSchema

    def make_masks_from_dicts(self):
        """read the input json, populate the experiments
        with nice masks.

        Parameters
        ----------
        input_json : str
            path to json file containing nway matching input

        """
        self.experiments = []
        image_paths = []
        for exp in self.args['experiment_containers']['ophys_experiments']:
            self.experiments.append(
                    utils.create_nice_mask(
                        exp,
                        self.args['output_directory']))
            image_paths.append(
                    Path(exp['ophys_average_intensity_projection_image']))

        check_image_sizes(image_paths)

        if len(self.experiments) < 2:
            raise NwayException("Need at least 2 experiements from input")

    def gen_nway_table_with_redundancy(self):
        """Combine all pairwise matches into one table
        regardless of conflicts.

        """
        matching_frame = pd.DataFrame(
            columns=[e['id'] for e in self.experiments])

        for pair in self.pair_matches:
            with open(pair.args['output_json'], 'r') as f:
                pairj = json.load(f)
            pairframe = pd.DataFrame(
                    [[i['fixed'], i['moving']] for i in pairj['matches']],
                    columns=[
                        pairj['fixed_experiment'],
                        pairj['moving_experiment']])
            matching_frame = matching_frame.append(pairframe)

        return matching_frame

    @staticmethod
    def remove_nway_table_redundancy(frame):
        """Remove any matching lines or subsets in a table

        Parameters
        ----------
        frame : :class:`pandas.DataFrame` or :class:`numpy.ndarray`
            table of cellIDs or zeros

        Returns
        -------
        table_new : list
            list of lines without redundancy or subsets

        """

        table = np.array(frame).astype('float')
        table[np.isnan(table)] = -1
        table = table.astype('int')

        linenum, expnum = np.shape(table)
        stoptag = 0

        while stoptag == 0:
            table_new = [table[0]]

            for i, oline in enumerate(table[1:], 1):
                for j, nline in enumerate(table_new):
                    mergetag = 1
                    omiss = oline == -1
                    nmiss = nline == -1

                    # matches not in any of the same experiments
                    if np.all(omiss | nmiss):
                        mergetag = 0

                    # matches in the same experiment(s), but no match
                    if np.any(
                            (nline != oline) &
                            np.invert(omiss) &
                            np.invert(nmiss)):
                        mergetag = 0

                    if mergetag == 1:
                        for k in range(expnum):
                            if oline[k] != -1:
                                nline[k] = oline[k]
                        break
                        # break for loop j
                        # only merge to one previous record

                # old line unrelated to any new lines, add it
                if mergetag == 0:
                    table_new.append(oline)

            linenum_new = np.shape(table_new)[0]

            if linenum == linenum_new:
                stoptag = 1
            else:
                table = table_new
                linenum = linenum_new

        return table_new

    def prune_matching_graph(self, matching_table_nway):
        """remove matching conflicts by scoring match candidates
        and pruning according to that score

        Parameters
        ----------
        matching_table_nway : list
            list of match arrays

        Returns
        -------
        matching_table_nway_new : list
            pruned list of match arrays

        """

        nline = np.shape(matching_table_nway)[0]

        exids = [exp['id'] for exp in self.experiments]

        # NOTE
        # assign a score to each line in the matching table
        # the score is the sum of costs for each pair in the
        # line. The costs are generally in [0, 2] when the
        # pair was within max_distance and 1000 when the pair
        # was not.
        # the scores are used to choose between conflicting lines
        # for n experiments, there are np = (n)(n - 1) / 2 pairs
        # example 1:
        # line1 = [a, b, -1, -1], line2 = [a, c, -1, -1]
        # score1 = cost(a, b) + 5000, score2 = cost(a, c) + 5000
        # if cost(a, b) < cost(a, c), line1 is preferred
        # example 2:
        # line1 = [a, b, c, -1], line2 = [a, c, -1, -1]
        # score1 = cost(a, b) + cost(b, c) + cost(a, c) + 3000
        # score2 = cost(a, c) + 5000
        # score1 in range [3000, 3006], score2 in range [5000, 5002]
        # line1 is preferred

        score = np.zeros(nline)
        for pair in self.pair_matches:
            with open(pair.args['output_json'], 'r') as f:
                j = json.load(f)
            # all possible candidates within max_distance
            candidates = {
                    (m['fixed'], m['moving']): m['cost']
                    for k in ['matches', 'rejected']
                    for m in j[k]}
            prow = exids.index(pair.args['fixed']['id'])
            pcol = exids.index(pair.args['moving']['id'])
            for i in range(nline):
                col = matching_table_nway[i][pcol]
                row = matching_table_nway[i][prow]
                if (row, col) in candidates:
                    pscore = candidates[(row, col)]
                else:
                    pscore = 1000
                score[i] += pscore

        matching_table_nway_new = prune_matching_table(
                matching_table_nway, score, self.args['pruning_method'])

        return matching_table_nway_new

    def add_remaining_cells(self, matching_table_nway):
        """add any remaining cells that had no match

        Parameters
        ----------
        matching_table_nway : list
            list of match arrays

        Returns
        -------
        matching_table_nway_new : list
            with added unmatched cells

        """

        linenum = np.shape(matching_table_nway)[0]
        cnt = linenum
        matching_table_nway_new = matching_table_nway
        label_remain = []

        for j in range(len(self.experiments)):
            labels = [v['id'] for v in self.experiments[j]['cell_rois']]
            for i in range(linenum):
                labels = np.setdiff1d(labels, [matching_table_nway[i][j]])
            label_remain.append(labels)

        for i in range(len(self.experiments)):
            num = len(label_remain[i])
            for j in range(num):
                cnt = cnt + 1
                this_record = np.zeros(len(self.experiments)) - 1
                this_record[i] = label_remain[i][j]
                matching_table_nway_new.append(this_record.astype(int))

        return matching_table_nway_new

    def create_output_dict(self):
        """write the nway matching results to a dict

        Returns
        -------
        matchingdata: dict
           "nway_matches": list of sets of cellID matches
           "pairwise_results": list of results from the pairwise matching
        """
        cellnum = np.shape(self.matching_table_nway)[0]
        matchingdata = dict()
        matchingdata["nway_matches"] = []
        prob = np.zeros(len(self.experiments))

        for i in range(cellnum):
            matching_exp_num = \
                    len(self.experiments) - 1 - \
                    np.count_nonzero(self.matching_table_nway[i] == -1)
            prob[matching_exp_num] = prob[matching_exp_num] + 1
            self.logger.debug(self.matching_table_nway[i])

        for i in range(cellnum):
            thisrgn = [v for v in self.matching_table_nway[i] if v != -1]

            if len(thisrgn) > 0:
                matchingdata["nway_matches"].append(thisrgn)

        matchingdata["pairwise_results"] = []
        for k in self.pair_matches:
            with open(k.args['output_json'], 'r') as f:
                j = json.load(f)
            matchingdata['pairwise_results'].append(j)
            if not self.args['save_pairwise_results']:
                os.remove(k.args['output_json'])

        return matchingdata

    def run(self):
        """Nway cell matching by calling pairwise
           matching and then combining the results
        """
        # log this ENV variable, if present
        commit = os.environ.get("NWAY_COMMIT_SHA", None)
        self.logger.name = type(self).__name__
        self.logger.info(f"NWAY_COMMIT_SHA {commit}")
        self.logger.info(f"Nway matching version {nway_version}")

        self.make_masks_from_dicts()

        # pair-wise matching
        self.pair_matches = []
        pair_arg_list = []
        for fixed, moving in itertools.combinations(self.experiments, 2):
            pair_args = copy.deepcopy(self.args)
            # marshmallow 3.0.0rc6 is less forgiving about extra keys around
            # so, pop out the unused shared keys here
            for popkey in [
                    "pruning_method",
                    "experiment_containers",
                    "save_pairwise_results",
                    "parallel_workers"]:
                pair_args.pop(popkey)
            pair_args["fixed"] = fixed
            pair_args["moving"] = moving
            pair_args["output_json"] = os.path.join(
                    self.args["output_directory"],
                    "{}_to_{}_output.json".format(moving['id'], fixed['id']))
            pair_arg_list.append(pair_args)

        if self.args['parallel_workers'] == 1:
            self.pair_matches = [pair_match_job(i) for i in pair_arg_list]
        else:
            with multiprocessing.Pool(self.args['parallel_workers']) as pool:
                self.pair_matches = pool.map(pair_match_job, pair_arg_list)

        self.logger.name = type(self).__name__
        df, id_map = utils.summarize_registration_success(
                [p.args['output_json'] for p in self.pair_matches])
        self.logger.info(f"registration success(1) or failure (0):\n{df}\n"
                         f"id map{json.dumps(id_map, indent=2)}")
        if not df.all(axis=None):
            raise NwayException("not all pairwise registrations succeeded.")

        # generate N-way matching table
        matching_frame = self.gen_nway_table_with_redundancy()
        matching_table_nway_tmp = self.remove_nway_table_redundancy(
                matching_frame)
        matching_table_nway_tmp = self.prune_matching_graph(
                matching_table_nway_tmp)
        self.matching_table_nway = self.add_remaining_cells(
                matching_table_nway_tmp)
        self.logger.info("Nway matching is done!")

        # Generate first part of output dict (necessary for summary plots)
        output_dict = self.create_output_dict()

        # Generate nway match summary plots
        self.logger.info("Creating match summary plots")
        nway_summarizer = NwaySummary(
            input_data={'nway_input': self.args,
                        'nway_output': output_dict,
                        'output_directory': self.args["output_directory"]},
            args=[])
        plot_save_paths: dict = nway_summarizer.run()

        output_dict.update(plot_save_paths)
        self.output(output_dict, indent=2)
        self.logger.info("wrote {}".format(self.args['output_json']))


if __name__ == "__main__":  # pragma: no cover
    nmod = NwayMatching()
    nmod.run()
