from scipy import *
from sets import Set
from copy import copy
import pysqlite2.dbapi2 as sqlite
from scipy.stats import random_integers
# -------------------------------------

def get_nwalk(A,nk=None,X=None):
    """U=get_nwalk(A,nk=None,X=None)
    
       For A being (n x n) adjacency matrix, returns n-vector U
       where U[i] is the matrix with rows U[i]={s_i_1,s_i_2,...,s_i_nk};
       s_i_k is the number of (generalised) walks from vertex i to vertex k. 
       If nk is not given, then nk = 2*n. Default value for X is E_n
       (classical walks). 
    """

    m,n=A.shape
    if not nk:
        nk=2*n
    U=zeros((n,n,nk))
    if not X:
        Ap=copy(A)
    else:
        Ap=copy(dot(A,X))
        
    for c in range(nk):
        for k in range(n):
            U[k,:,c]=Ap[:,k]
        Ap=dot(Ap,A)
    return U


def get_UV(A,B,nk=None,X=None):
    """U,V=get_UV(A,B,nk=None,X=None)

       For A,B - the (n x n) adjacency matrices, returns n-vectors U, V
       where U[i] is the matrix with rows U[i]={s_i_1,s_i_2,...,s_i_nk};
       s_i_k is the number of (generalised) walks from vertex i to vertex k
       in graph with adjacency matrix A. The same for vector V and matrix B. 
       If nk is not given, then nk = 2*n. Default value for X is E_n
       (classical walks). 
    
    """

    return get_nwalk(A,nk,X), get_nwalk(B,nk,X)


def local_complement(A,nv):
    """AL=local_complement(A,nv)

       Returns adjacency matrix of local complement of vertex nv
       (vertex numbers are 0,1,...) in graph with adjacency matrix A.  
    """
    
    Anv=copy(A)
    nzind=A[nv].nonzero()
    for r in nzind:
        for s in nzind:
            Anv[r,s]=1-A[r,s]
            Anv[s,s]=0
    return Anv


def check_iso(A,B,nk=None,X=None):
    """check_iso(A,B,nk=None,X=None)
    
    Isomorphism check of graphs with adjac. matrices A,B,
    using (generalised, if X is given) walks of length 1,...,nk.
    Returns False, if graphs are distinguishable, 
    otherwise returns the permutation P of vertices such that
    row- and column-permuted B could be A. 
    """
    
    n,m=A.shape
    if not nk:
        nk=2*n
    U,V=get_UV(A,B,nk)
    P=[-1]*n
    arr2d_tuple=lambda D: tuple([tuple(p) for p in D])
    
    UL=[Set(arr2d_tuple(p)) for p in U]
    VL=[Set(arr2d_tuple(p)) for p in V]

    for i in range(n):
        for j in range(n):
            if UL[i]==VL[j] and all(U[i][i]==V[j][j]) and not j in P:
                P[i]=j; break
        if P[i]==-1:
            return False
    return P


def check_LC(A,B,nk=3,chk_iso=0):
    """check_LC(A,B,nk=3,chk_iso=0)
    
    Isomorphism check of graphs with adjac. matrices A,B,
    using (generalised, if X is given) walks of length 1,...,nk.
    If chk_iso is nonzero, original graphs are first checked.
    In case they are indistinguishable, we form all possible
    tuples of local complements of A,B and make checks with them.
    
    Returns False, if graphs are distinguishable, 
    otherwise returns the permutation P of vertices such that
    row- and column-permuted B could be A. 
    """


    if chk_iso:
        if not check_iso(A,B,nk):
            return False
    m,n=A.shape
    P=[-1]*n
    for i in range(n):
        for j in range(n):
            A_loc,B_loc=local_complement(A,i), local_complement(B,j)
            if check_iso(A_loc,B_loc,nk) and not j in P:
                P[i]=j; break
        if P[i]==-1:
            return False 
    return P

def check_XW(A,B,nk=3,chk_iso=0):
    pass
    

def inv_perm(P):
    """IP=inv_perm(P)
    
       Returns inverse permutation IP to permutation P.
    """
    Pz=zip(P,range(len(P)))
    Pz.sort()
    return [p[1] for p in Pz]
    

def check_perm(A,B,P):
    """check_perm(A,B,P)
    
       Check, if the given permutation P applied to rows and columns
       of matrix B is the matrix A (returns True or False).
    """
    Bm=B[P]
    BP=Bm.transpose()[P].transpose()
    return all(BP==A)

# ------------ reading unlabelled graph from graphDB --------------
def read_unlab32(fname):           # works only for 32 bit machines
    S=file(fname,'r').read()
    D=fromstring(S,dtype=Int16)
    n=D[0]
    A=zeros((n,n))
    pos=1
    for v in range(n):
        dv=D[pos]
        if not dv:
            print "Isolated vertex %d" %v
        U=D[pos+1:pos+dv+1];  pos += (dv+1)
        for u in U:
            A[v,u]=1
    return A


def readword(S,pos):
    return ord(S[pos]) | (ord(S[pos+1]) << 8)
  

def read_unlab(fname):
    """A=read_unlab(fname)

       Read an unlabeled graph from file (in GraphDB format) with name fname.
       Returns the adjacency matrix A.
    """
    S=file(fname,'r').read()
    n=readword(S,0)

    pos=2
    A=zeros((n,n),dtype=Int32)

    for k in range(n):
        nv=readword(S,pos); pos += 2
        if not nv:
            print "Isolated vertex %d" %k
        else:
            for j in range(nv):
                v_node=readword(S,pos); pos += 2
                A[k,v_node]=1     
    return A


# ------------- reading regular graphs from database --------------

def get_conn(dbname='srg.db'):
    con=sqlite.connect(dbname)
    cur=con.cursor()
    return con,cur


# ----------------------- helper functions ------------------------
def get_gid(n,deg,n_adj,n_far,dbname='srg.db'):
    con,cur=get_conn(dbname)
    qr="select g_id from srg where n=%d and deg=%d "\
       "and n_adj=%d and n_far=%d" %(n,deg,n_adj,n_far)
    cur.execute(qr)  
    return cur.fetchone()[0]


def str_tomat(s,n):
    M1=array([int(c) for c in s])
    return M1.reshape(n,n)
# -----------------------------------------------------------------


def get_srgcount(n,deg,n_adj,n_far,dbname='srg.db'):
    con,cur=get_conn(dbname)
    g_id=get_gid(n,deg,n_adj,n_far,dbname)    
    qr="select count(*) from srg_mat where g_id=%d" %g_id
    cur.execute(qr)
    return cur.fetchone()[0] or 0


def get_avail_srg(g_count=None,dbname='srg.db'):
    """GA=get_avail_srg(g_count=None,dbname='srg.db')
    
       Returns the list of types (and counts, if g_count is nonzero)
       of strongly regular graphs contained in our database.
       The rows of GA are of the form: ((26, 10, 3, 4), 10).
    """
    con,cur=get_conn(dbname)
    qr_all='select n,deg,n_adj,n_far from srg'
    cur.execute(qr_all)
    res=cur.fetchall()
    res.sort()
    if g_count:
        for k in range(len(res)):
            n,deg,n_adj,n_far=res[k]
            ng=get_srgcount(n,deg,n_adj,n_far)
            res[k]=res[k],ng           
    return res

        
def get_srgmat(n,deg,n_adj,n_far,get_all=True,dbname='srg.db'):
    """LA=get_srgmat(n,deg,n_adj,n_far,get_all=True,dbname='srg.db')
    
       Gets list of adjac. matrices of SRG-s of type (n,deg,n_adj,n_far).
       If get_all is false, returns one of matrices, randomly choosen.     
    """
    con,cur=get_conn(dbname)    
    g_id=get_gid(n,deg,n_adj,n_far,dbname)
    qr="select adjmat from srg_mat where g_id=%d" %g_id
    cur.execute(qr)
    res=cur.fetchall()
    M_list=[]
    if get_all:
        for Ms, in res:
            M_list += [str_tomat(Ms,n)]        
        return M_list
    else:
        return str_tomat(res[random_integers(0,len(res)-1)][0],n)
    

iso_methods={'LC':check_LC,'XW':check_XW}

def check_srg_group(n,deg,n_adj,n_far,method='LC',dbname='srg.db'):
    """check_srg_group(n,deg,n_adj,n_far,method='LC',dbname='srg.db')

       Perform check of distinguashability of SRG-s of the given
       type (n,deg,n_adj,n_far) with 'Local Complement' checks
       if method='LC' or with generalised walks check, if method='XW'.
       Returns False, if all graphs in group are distinguishable,
       or the indices (n1,n2) of indistinguishable graphs
       (in the list returned by 'get_srgmat').
    """
    
    LA=get_srgmat(n,deg,n_adj,n_far,dbname)
    if not LA:
        raise ValueError, "No such group of SRG-s"

    ng=len(LA)
    check_fcn=iso_methods[method]
    res=False
    ind2=[(p,q) for p in range(ng) for q in range(ng) if p