import pylab
import matplotlib
import networkx as nx
import scipy
import rpy

def plot_corrloads(R, pc1=0,pc2=1,s=20, c='b', zorder=5,expvar=None,ax=None,drawback=True, labels=None, **kwds):
    """ Correlation loading plot."""

    # background
    if ax==None or drawback==True:
        radius = 1
        center = (0,0)
        c100 = matplotlib.patches.Circle(center,radius=radius,
                                         facecolor='gray',
                                         alpha=.1,
                                         zorder=1)
        c50 = matplotlib.patches.Circle(center, radius=radius/2.0,
                                        facecolor='gray',
                                        alpha=.1,
                                        zorder=2)
        ax = pylab.gca()
        ax.add_patch(c100)
        ax.add_patch(c50)
        ax.axhline(lw=1.5,color='k')
        ax.axvline(lw=1.5,color='k')

    # corrloads
    ax.scatter(R[:,pc1], R[:,pc2], s=s, c=c,zorder=zorder, **kwds)
    ax.set_xlim([-1,1])
    ax.set_ylim([-1,1])
    if expvar!=None:
        xstring = "Comp: %d   expl.var:  %.1f " %(pc1+1, expvar[pc1])
        pylab.xlabel(xstring)
        ystring = "Comp: %d   expl.var.:  %.1f " %(pc2+1, expvar[pc2])
        pylab.ylabel(ystring)
    if labels!=None:
        assert(len(labels)==R.shape[0])
        for name, r in zip(labels, R):
            pylab.text(r[pc1], r[pc2], "  " + name)
    #pylab.show()


def dag(terms, ontology):
    rpy.r.library("GOstats")
    __parents = {'bp' : rpy.r.GOBPPARENTS,
                 'mf' : rpy.r.GOMFPARENTS,
                 'cc' : rpy.r.GOCCPARENTS}
    gograph = rpy.r.GOGraph(terms, __parents.get(ontology.lower()))
    dag = rpy.r.edges(gograph)
    #setattr(dag, "_ontology", ontology)
    return dag

def plot_dag(dag, node_color='b', node_size=30,with_labels=False,nodelist=None,pos=None):

    rpy.r.library("GOstats")
    
    dag_name = "GO-bp" 
    # networkx does not play well with colon in node names
    clean_edges = {}
    for head, neigb in dag.items():
        head = head.replace(":", "_")
        nei = [i.replace(":", "_") for i in neigb]
        clean_edges[head] = nei
    if pos==None:
        G = nx.from_dict_of_lists(clean_edges, nx.DiGraph(name=dag_name))
        pos = nx.pydot_layout(G, prog='dot')
        pos_new = {}
        for k, v in pos.items():
            x,y = v
            k = k.replace("_", ":")
            pos_new[k] = (x, -y) 
        pos = pos_new
    
    G = nx.from_dict_of_lists(dag, nx.Graph(name=dag_name))
    if len(node_color)>1:
        assert(len(node_color)==len(nodelist))
        
    nx.draw_networkx(G,pos, with_labels=with_labels, node_size=node_size, node_color=node_color, nodelist=nodelist)
    return pos


def plot_ZXcorr(gene_ids, term_ids, gene2go, X, D, scale=True):
    """ Plot correlation/covariance between genes as a function of
    semantic difference.

    input: X (n, p) data matrix
           D (p, p) gene-gene sematic similarity matrix
    """
    D = scipy.corrcoef(X)
    term2ind = dict(enumerate(term_ids))
    for i, gene_i in enumerate(gene_ids):
        for j, gene_j in enumerate(gene_ids):
            if j<i:
                r2 = D[i,j]
                terms_i = gene2go[gene_i]
                terms_j = gene2go[gene_j]
                for ti, term in enumerate(term_ids):
                    if term in terms_i:
                        pass


def clustering_index(T, Yg):
    pass

def draw_gene(gid, gene_ids, gene2go, Z, tmat, terms, G, pos):
    """Draw dags with marked go terms and distance to all terms.
    """
    sub_terms = gene2go[gid]
    sub_index = [i for i, tid in enumerate(terms) if tid in sub_terms]
    node_size = 70.*scipy.ones((len(terms),))
    node_size[sub_index] = 500
    gene_index = [i for i, gene_id in enumerate(gene_ids) if gene_id==gid]
    node_color = Z[:,gene_index].ravel()
    #1/0
    #node_size=200*node_color
    #node_color='g'
    pylab.figure()
    nx.draw_networkx(G, pos, node_color=node_color, node_size=node_size, with_labels=False, nodelist=terms)
    ax = pylab.gca()
    pylab.colorbar(ax.collections[0])

    for tid in sub_index:
        pylab.figure()
        node_color = tmat[tid,:]
        #node_size = 70*scipy.ones((len(terms),))
        node_size = 170*node_color
        node_size[tid] = 500
        nx.draw_networkx(G, pos, node_color=node_color, node_size=node_size, with_labels=False, nodelist=terms)
        pylab.title(terms[tid])
        ax = pylab.gca()
        pylab.colorbar(ax.collections[0])
    pylab.show()
    #nx.show()