"""

Description:
CoNetwork is the class for the pathway co-occurrence network.
Helper classes include Edge and Vertex.
CoNetwork represents a network of direct pathway-pathway relationships.

An example of how to use the CoNetwork to output a .csv network
(for visualization in Cytoscape) is provided in __main__.

Usage:
   python network.py test_network_file n_features output_file

(1) test_network_file: pathways identified in constructed features
                       (generated by pathway_coverage.R)
(2) n_features: number of features in the model
(3) output_file: the name of the output network file

Output:
A file (-output_file-) is created with the network for the pathway data
from a specified feature construction algorithm.

"""

import sys
import itertools
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random

import utils

# Usage: in `_pathway_feature_shuffle`. Sets an upper bound for the
# number of times we can randomly assign a pathway to a feature during the side-
# preserving model permutation.
MAX_ITERS = 30000

class Edge:

    def __init__(self, start, end, which_features=[]):
        self.edge = (start, end)
        self.weight = len(which_features)
        self.which_features = which_features

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        else:
            return sorted(self.edge) == sorted(other.edge)

    def __ne__(self, other):
        return not self.__eq__(other)
    
    def features_to_string(self):
        return " ".join(map(str, self.which_features))

    def connected_to(self, vid):
        """
        :param vid: vertex id
        :returns: the vertex that shares an edge with vid|None
        """
        if conn_to not in self.edge:
            return None
        conn_to = self.edge[1] if self.edge[0] == vid else self.edge[0]
        return conn_to


class Vertex:

    def __init__(self, vid):
        self._id = vid
        self.edges = {}

    def degree(self):
        # TODO: maybe just an instance variable
        return len(self.edges)

    def all_first_neighbor_ids(self):
        """
        :returns: all first neighbors of a vertex
        """
        return self.edges.keys()


class CoNetwork:

    def __init__(self, sig_pathways, n_features):
        self.vertices = {}  # vertex id (vid) -> vertex obj
        self.edges = {}  # (vid0, vid1) -> edge obj

        self.n_features = n_features
        self.features = set()
        
        self.n_pathways = 0
        self.pathways = {}  # pathway -> vid

        if isinstance(sig_pathways, str):
            self.sig_pathways_file = sig_pathways
            self._construct_from_file()
        else:
            # Some analyses generate network data dynamically
            # and store it in data structures.
            self.sig_pathways_file = None
            self._construct_from_permutation(sig_pathways)
        
    def __getitem__(self, item_id):
        for pw, vid in self.pathways.items():
            if item_id == vid:
                return pw
        return None

    def _construct_from_file(self):
        """ Builds the network using the significant pathways file.
        """
        feature_pathway_df = utils.read_significant_pathways_file(self.sig_pathways_file)

        network_dict = {}

        ### identify pairwise direct (same-side) relationships ###
        feature_side_grouping = feature_pathway_df.groupby(["feature", "side"])
        for (feature, _), group in feature_side_grouping:
            assoc_pathways = group["pathway"].tolist()
            pairings = list(itertools.combinations(assoc_pathways, 2))
            if pairings:
                self.features.add(feature)
                for pathway_pair in pairings:
                    v0 = self.add_pathway(pathway_pair[0])
                    v1 = self.add_pathway(pathway_pair[1])
                    new_edge = self.edge_tuple(v0, v1)
                    if new_edge not in network_dict:
                        network_dict[new_edge] = []
                    network_dict[new_edge].append(feature)

        self._augment_network(network_dict)
        self._total_edge_type_weight()

    def permute_pathways_across_features(self):
        """ Builds a permuted network from the significant pathways file.
        """
        feature_pathway_df = utils.read_significant_pathways_file(self.sig_pathways_file)
        feature_grouping = feature_pathway_df.groupby(["side", "feature"])

        pathway_features = {"pos": [], "neg": []}
        for (side, feature), group in feature_grouping:
            pathways = group["pathway"].tolist()
            for pathway in pathways:
                pathway_features[side].append((pathway, feature))
                
        pos_shuffle = self._pathway_feature_shuffle(pathway_features["pos"])
        while pos_shuffle is None:
            pos_shuffle = self._pathway_feature_shuffle(pathway_features["pos"])

        neg_shuffle = self._pathway_feature_shuffle(pathway_features["neg"])
        while neg_shuffle is None:
            neg_shuffle = self._pathway_feature_shuffle(pathway_features["neg"])

        pos_correct = self._shuffle_correctness(pos_shuffle,
                                                pathway_features["pos"])
        neg_correct = self._shuffle_correctness(neg_shuffle,
                                                pathway_features["neg"])
        
        assert pos_correct and neg_correct, "feature shuffle did not preserve the expected invariants."
        
        sig_pathways = {"pos": pos_shuffle, "neg": neg_shuffle}

        # TODO: Python referencing?

        # self-assignment
        self.permuted = CoNetwork(sig_pathways, self.n_features)
        self.permuted._total_edge_type_weight()

        # to return
        permutation = CoNetwork(sig_pathways, self.n_features)
        permutation._total_edge_type_weight()
        return permutation

    def divide(self):
        """ Specific to the permutation. Updates the
            network so that edge weights are based on the observed-to-expected
            ratio (original-to-permuted).
            Edges specific to the original network will maintain their edge
            weights, whereas those specific to the permuted network will be
            assigned their reciprocal weight.
        """
        self.features = set()
        mapping = self.aggregate_remap_pathways(self.permuted.pathways)
        for (v0, v1), edge in self.permuted.edges.items():
            edge_key = self.remapped_edge(mapping, v0, v1)
            if edge_key in self.edges:
                if self.edges[edge_key].which_features:
                    self.edges[edge_key].which_features = []
                observed_to_expected = (
                    float(self.edges[edge_key].weight) /
                    float(edge.weight))
                self.edges[edge_key].weight = obs_to_exp
            else:
                edge_obj = Edge(edge_key[0], edge_key[1], [])
                edge_obj.weight = 1. / edge.weight
                self.edges[edge_key] = edge_obj
                self._add_edge_to_vertex(edge_key[0], edge_obj)
                self._add_edge_to_vertex(edge_key[1], edge_obj)

    def aggregate(self):
        """ This is the merge function to combine multiple networks.
        :param merge: the CoNetwork object being merged into the current network
        :returns: None
        """
        self.sig_pathways_file = None
        self.n_features += merge.n_features

        mapping = self.aggregate_remap_pathways(merge.pathways)
        for (v0, v1), edge in merge.edges.items():
            edge_key = self.remapped_edge(mapping, v0, v1)
            if edge_key in self.edges:
                if self.edges[edge_key].which_features:
                    self.edges[edge_key].which_features = []
                self.edges[edge_key].weight += edge.weight
            else:
                edge_obj = Edge(edge_key[0], edge_key[1], [])
                edge_obj.weight = len(edge.which_features)
                self.edges[edge_key] = edge_obj
                self._add_edge_to_vertex(edge_key[0], edge_obj)
                self._add_edge_to_vertex(edge_key[1], edge_obj)

    def aggregate_remap_pathways(self, pathway_mapping):
        """ The following tasks are carried out in the remapping:
        (1) If self.pathways contains the pathway to be merged, map the merge
            vid to self vid.
        (2) If not, create a vertex in self and then add the key-val pair
            accordingly.
        :param pathway_mapping: the (pathway -> vid) map for the network being
                                merged with self
        :returns: a (merge.vid -> self.vid) mapping for pathway id conversion
        """
        mapping = {}
        for pathway, vid in pathway_mapping.items():
            if pathway in self.pathways:
                mapping[vid] = self.pathways[pathway]
            else:
                self_vid = self.add_pathway(pathway)
                self.vertices[self_vid] = Vertex(self_vid)
                mapping[vid] = self_vid
        return mapping

    def remapped_edge(self, remap, v0, v1):
        """ After aggregation, given an edge (v0, v1), return the corresponding
            edge name from the `remap` dictionary.
        """
        self_v0 = remap[v0]
        self_v1 = remap[v1]
        eid = self.edge_tuple(self_v0, self_v1)
        return eid

    def edge_tuple(self, v0, v1):
        """ Enforce an edge invariant: start ID always < end ID.
        """
        edge = [v0, v1]
        edge.sort()
        return tuple(edge)

    def add_pathway(self, pathway):
        if pathway not in self.pathways:
            self.pathways[pathway] = self.num_pathways
            self.n_pathways += 1
        return self.pathways[pathway]

    def get_edge_weight(self, pw0, pw1, etype):
        v0 = self.pathways[pw0]
        v1 = self.pathways[pw1]
        eid = self.edge_tuple(v0, v1)
        return self.edges[eid].weight

    def get_edge_pathways(self, edge_key):
        """ Return the pathways associated with an edge.
        """
        return (self.__getitem__(edge_key[0]), self.__getitem__(edge_key[1]))

    def get_vertex_obj(self, vid):
        return self.vertices[vid]

    def get_vertex_obj_from_pathway(self, pw):
        if pw in self.pathways:
            vid = self.pathways[pw]
            return self.vertices[vid]
        else:
            return None

    def get_pathway_from_vertex_id(self, vid):
        return self.__getitem__(vid)

    def network_dataframe(self, drop_weights_below=0, whitelist=None):
        """ Conversion of the network to a pandas dataframe. This allows for
            easy export to a CSV and visualization in Cytoscape.
        :param drop_weights_below: specify an edge weight threshold by which to
               filter our network.
        :param whitelist: list of edges to keep in the network (this is used
               during the permutation analysis).
        :returns: a pandas dataframe containing the network edge information.
             columns = [PW1, PW2, weight, features]
        """
        network_df = pd.DataFrame(columns=["PW1", "PW2",
                                           "weight", "features"])
        idx = 0
        for (v0, v1), edge_obj in self.edges.items():
            if (edge_obj.weight > drop_weights_below and
                    (whitelist is None or (v0, v1) in whitelist)):
                features = edge_obj.features_to_string()
                network_df.loc[idx] = [self.__getitem__(v0),
                                       self.__getitem__(v1),
                                       edge_obj.weight,
                                       features]
                idx += 1  # faster to append by index.
        network_df = network_df.sort_values(by=["weight"],
                                            ascending=False)
        print "This network contains " + str(self.n_pathways) + " pathways."
        return network_df

    def pathway_relations_count(self, pathway):
        """
        :param pathway: (str) pathway name
        :returns: (dict) for each etype, compute the # of features in which the
                  the pathway is involved
        """
        return self._pathway_relations_count_by_id(self.pathways[pathway])

    def pathway_neighbors_rel_count(self, pathway):
        """
        :param pathway: (str) pathway name
        :returns: { etype (key) -> dict (value) }
                  dict contains relation count information for all the first
                  neighbors of a given pathway
        """
        return self._pathway_neighbors_rel_count_by_id(self.pathways[pathway])

    def feature_counts_excluding_edge(self, pw0, pw1):
        """ Function used to collect counts for a hypergeometric test
        :param pw0: (str) pathway 0 of interest
        :param pw1: (str) pathway 1 of interest
        :returns: (int) # of -etype- edges in the network excluding those
                  connected to the two pathways of interest.
        """
        vobj0 = self.get_vertex_obj_from_pathway(pw0)
        vobj1 = self.get_vertex_obj_from_pathway(pw1)
        vobj0_relations = self._pathway_relations_by_id(vobj0, etype)
        vobj1_relations = self._pathway_relations_by_id(vobj1, etype)
        exclude = self.features[etype] - (vobj0_relations | vobj1_relations)
        return len(exclude)

    def _add_edge_to_vertex(self, vid, edge):
        conn_to = edge.connected_to(vid)
        if vid not in self.vertices:
            vert_obj = Vertex(vid)
            self.vertices[vid] = vert_obj
        self.vertices[vid].edges[conn_to] = edge.weight

    def _augment_network(self, edge_dict):
        for (s, e), feature_list in edge_dict.items():
            edge_obj = Edge(s, e, feature_list)
            self.edges[(s, e)] = edge_obj
            self._add_edge_to_vertex(s, edge_obj)
            self._add_edge_to_vertex(e, edge_obj)

    def _pathway_relations_count_by_id(self, vid):
        """
        :param vid: (int) vertex id
        :returns: { etype -> count }, where count is the # of -etype- edges
                  which contain the pathway with vertex id -vid-
        """
        vertex = self.vertices[vid]
        # TODO.
        return vertex.all_first_neighbor_ids
        rel_counts = {}

        for edge in vertex.edges.keys():
        for etype in vertex.edges.keys():
            rel_counts[etype] = self._pathway_relations_count_by_id_and_type(
                vertex, etype)
        return rel_counts

    def _pathway_relations_count_by_id_and_type(self, vertex_obj, etype):
        return len(self._pathway_relations_by_id_and_type(vertex_obj, etype))

    def _pathway_relations_by_id_and_type(self, vertex_obj):
        """ # of features in which a pathway show up with x type of relation.
            Get all of the first neighbors with that type. Add the features
            for which the relation appears between pw-neighbor to a set.
        """
        vneighbors = vertex_obj.first_neighbor_ids_by_type(etype)
        features = []
        for conn_to in vneighbors:
            eid = self.edge_tuple(vertex_obj.id, conn_to)
            features += self.edges[eid].which_features
        return set(features)

    def _pathway_neighbors_rel_count_by_id(self, vid):
        """ For each edge type, get the 1st neighbors of -vid-.
            For each 1st neighbor, determine the # of features where it has that
            edge type.
        :param vid: (int) vertex id
        :returns: { etype -> { neighbor -> count } }
        """
        vertex = self.vertices[vid]
        count_dicts = {}
        for etype in vertex.edges.keys():
            count_dicts[etype] = {}
            neighbors = vertex.first_neighbor_ids_by_type(etype)
            for conn_to in neighbors:
                pathway_name = self.__getitem__(conn_to)
                pathway_count = self._pathway_relations_count_by_id_and_type(
                    self.vertices[conn_to], etype)
                count_dicts[etype][pathway_name] = pathway_count
        return count_dicts

    def _total_edge_type_weight(self):
        """ Iterate over all edges in the network. Add the edge's weight
            to the corresponding edge type category.
        :returns: dictionary from etype -> total weight
        """
        totals = {}
        for edge_obj in self.edges.values():
            if edge_obj.etype not in totals:
                totals[edge_obj.etype] = 0
            totals[edge_obj.etype] += edge_obj.weight
            self.etype_weights = totals

    def _update_network_edge_info(self, to_update, weight, features, etype):
        to_update.append(weight)
        to_update.append(features)
        to_update.append(etype)
        return to_update

    ################################################################
    # PATHWAY-FEATURE PERMUTATION: RELATED HELPER FUNCTIONS.
    ################################################################

    def _unzip(self, tuple_list):
        return [list(t) for t in zip(*tuple_list)]

    def _pathway_feature_shuffle(self, pathway_feature_tuples):
        """ Shuffle the pathways across features of one side in the network.
        :param pathway_feature_tuples: [(pw, feature)] list from the significant
               pathways analysis of the constructed features.
        :returns: [(pw, feature)] permuted list
        """
        pathways, features = self._unzip(pathway_feature_tuples)
        random.shuffle(pathways)

        original_pws = pathways[:]
        feature_locs = {}
        i = 0
        while i < len(pathways):
            starting_index = i
            current_feature = features[i]
            pw_set = set()
            while i < len(pathways) and features[i] == current_feature:
                if pathways[i] not in pw_set:
                    pw_set.add(pathways[i])
                else:
                    k = 0
                    while True:
                        j = random.choice(range(0, len(pathways)))
                        rpathway = pathways[j]
                        rfeature = features[j]
                        if rpathway != pathways[i] and rpathway not in pw_set:
                            if rfeature not in feature_locs:
                                break
                            check = feature_locs[rfeature]
                            if pathways[i] not in pathways[check[0]:check[1]]:
                                break
                        k += 1
                        if k > MAX_ITERS:
                            print "Reached maximum number of iterations."
                            return None
                    pw_set.add(rpathway)
                    pathways[j] = pathways[i]
                    pathways[i] = rpathway
                i += 1
            ending_index = i
            feature_locs[current_feature] = (starting_index, ending_index)

        if original_pws == pathways:
            print "Pathway comparison: shuffle failed."
            return None

        return zip(pathways, features)

    def _shuffle_correctness(self, pathway_feature_tuples, original):
        if pathway_feature_tuples == original:
            return False
        feature_pw_dict = {}
        for (pw, feature) in pathway_feature_tuples:
            if feature not in feature_pw_dict:
                feature_pw_dict[feature] = set()
            if pw in feature_pw_dict[feature]:
                print "Duplicate detected."
                return False
            else:
                feature_pw_dict[feature].add(pw)
        return True

    def _create_direct_edges_p(self, feature_pw_dict):
        network_dict_direct = {}
        for feature, pw_list in feature_pw_dict.items():
            for i in xrange(len(pw_list)):
                for j in range(i+1, len(pw_list)):
                    vi = pw_list[i]
                    vj = pw_list[j]
                    new_edge = self.edge_tuple(vi, vj)
                    if new_edge not in network_dict_direct:
                        network_dict_direct[new_edge] = []
                    network_dict_direct[new_edge].append(feature)
        self._augment_network(network_dict_direct)

    def _collect_pathways_by_feature(self, pw_feature_tuples):
        feature_pw_dict = {}
        for (pw, feature) in pw_feature_tuples:
            v0 = self.add_pathway(pw)
            if feature not in feature_pw_dict:
                feature_pw_dict[feature] = []
            feature_pw_dict[feature].append(v0)
        return feature_pw_dict

    def _construct_from_permutation(self, sig_pathways):
        """ Takes in a dictionary of "pos" and "neg" tuple lists containing
            the pathway-feature mappings of a permuted ADAGE model.
        """
        # add all pathways:
        pos_tups = sig_pathways["pos"]
        neg_tups = sig_pathways["neg"]

        pos_dict = self._collect_pathways_by_feature(pos_tups)
        neg_dict = self._collect_pathways_by_feature(neg_tups)

        self._create_direct_edges_p(pos_dict)
        self._create_direct_edges_p(neg_dict)

if __name__ == "__main__":
    test_network_file = sys.argv[1]
    num_features = sys.argv[2]
    output_file = sys.argv[3]
    # edits for cor_thresh
    cor_file = sys.argv[4]
    output_prefix = sys.argv[5]

    test_network = InteractionNetwork(test_network_file, num_features)
    shuffled = test_network.adage_model_permutation()
    test_df = test_network.network_dataframe()
    test_df.to_csv(path_or_buf=output_file,
                   sep="\t",
                   index=False)