from sage.matroids.advanced import *

#Note, some of these are quick and dirty implementations
#In particular, trying to delta-Y on a triangle that is not coindependent will cause issues.

#Generalised DY for a LinearMatroid (constructs the representation)
def DY(M, segment, parallels=None):
    # First, construct the parallel extension of M
    if parallels is not None:
        for e in parallels:
            M = M.linear_extension(element=len(M), chain={e: 1})

    # Next, find a suitable representation of the segment:
    BS = M.augment([], segment)
    NBS = [e for e in segment if not e in BS]
    B = M.augment(BS)
    A, row_labels, col_labels = M.representation(reduced=True, B=B, labels=True)

    # construct the Theta matroid
    row_indices = [row_labels.index(e) for e in BS]
    col_indices = [col_labels.index(e) for e in NBS]
    k = len(segment)
    D = Matrix(A.base_ring(), nrows=k, ncols=k)
    D[0,1] = A[row_indices[1],col_indices[0]] / A[row_indices[0],col_indices[0]]
    D[1,0] = -D[0,1]
    for i in range(k-2):
        D[i+2,0] = A[row_indices[0],col_indices[i]]
        D[i+2,1] = A[row_indices[1],col_indices[i]]
        D[0,i+2] = A[row_indices[0],col_indices[i]]
        D[1,i+2] = A[row_indices[1],col_indices[i]]

    # construct the ground set of the generalized parallel connection
    tmpE = M.groundset_list()
    delset = []
    E = copy(row_labels)
    for i in row_indices:
        E[i] = newlabel(tmpE)
        tmpE.append(E[i])
        delset.append(E[i])
    E.extend(NBS)
    F = copy(col_labels)
    for i in col_indices:
        F[i] = newlabel(tmpE)
        tmpE.append(F[i])
        delset.append(F[i])
    F.extend(BS)
    E.extend(F)

    # construct the new representation
    AA = Matrix(A.base_ring(), nrows=A.nrows() + k - 2, ncols=A.ncols()+2)
    AA.set_block(0,0,A)
    AA[row_indices[0], A.ncols()+1] = D[0,1]
    AA[row_indices[1], A.ncols()] = D[1,0]
    for i in range(k-2):
        AA[A.nrows()+i, A.ncols()] = D[i+2,0]
        AA[A.nrows()+i, A.ncols()+1] = D[i+2,1]
    return Matroid(groundset=E,reduced_matrix=AA).delete(delset)

def YD(M, cosegment, series=None):
    return DY(M.dual(), segment=cosegment, parallels=series).dual()

#Abstract generalised Delta-Y on a coindependent segment (for any matroid, not necessarily a LinearMatroid)
#assumes no loops (3-connected even?)
def deltaY(matroid, seg):
    if len(seg) < 3:
        raise AttributeError("Can only delta-Y on a segment of size at least three")
    if matroid.rank(seg) != 2:
        raise AttributeError("Can only delta-Y on a segment (i.e. needs to have rank 2)")

    if isinstance(matroid, LinearMatroid):
        return DY(matroid, seg)
    
    E = list(matroid.groundset())
    newlabels = []
    elts = list(seg)
    eltps = []
    for elt in elts:
        eltps.append(newlabel(E + eltps))

    circuit_map = {}
    n = len(elts)
    for i in range(n):
        elt = elts[i]
        circuit_map[elt] = [elt] + eltps[0:i] + eltps[i+1:n]
    
    seg_size = len(seg)
    circ_rank = seg_size - 1
    new_cc_map = {circ_rank : circuit_map.values()}

    segment = frozenset(seg)
    for (r,ccs) in matroid.circuit_closures().iteritems():
        for cc in ccs:
            intersection = segment.intersection(cc)
            if len(intersection) == 0:
                add_to_ml(new_cc_map, r, list(cc))
            elif len(intersection) == 1:
                circ = circuit_map[list(intersection)[0]]
                add_to_ml(new_cc_map, r, list(cc))
                add_to_ml(new_cc_map, r + circ_rank - 1, list(cc) + list(circ))
            elif len(intersection) == seg_size:
                add_to_ml(new_cc_map, r, list(cc))
                add_to_ml(new_cc_map, r + circ_rank - 1, list(cc) + eltps)
            else:
                raise AssertionError("something busted")
        
    m = Matroid(groundset = E + eltps, circuit_closures=new_cc_map)
    return m.delete(elts)

def yDelta(M, cosegment):
    return deltaY(M.dual(), cosegment).dual()

def add_to_ml(ml, key, element):
    if key in ml:
        ml[key].append(element)
    else:
        ml[key] = [element]

#Find matroids with no triangles that we can obtain from given mats by repeatedly performing delta-Y exchanges
#Note, this method considers delta-Y equivalence only (no segment/cosegment exhanges) and doesn't handle the case where triangles that are not coindependent arise
def find_deltaY_suprema(mats):
    oldmats = list(mats)
    suprema = []
    while len(oldmats) > 0:
        newmats = []
        for oldmat in oldmats:
            #TODO coindependent triangles only (see also elsewhere.)
            #TODO circuit closures of rank 2. If segment, consider seg-coseg exchange(s) as well
            triangles = [c for c in oldmat.circuits() if len(c) == 3]
            if len(triangles) == 0:
                suprema.append(oldmat)
            else:
                newmats.extend([deltaY(oldmat, tri) for tri in triangles])
        oldmats = get_nonisomorphic_matroids(newmats)
    return suprema
    
#find delta-Y equivalent matroids to mat.  Assumes coindependent triangles only, e.g. will give unexpected results for something like U(2,3)
def get_deltaY_equivalent_matroids(mat,includeDuals=False):
    #first, find "least" matroids we can obtain by Y-Deltas
    seeds = [mat.dual()]
    if includeDuals:
        seeds.append(mat)
    infima = get_nonisomorphic_matroids([inf.dual() for inf in find_deltaY_suprema(seeds)])

    #next, find "greatest" matroids we can obtain by performing Delta-Ys starting from each of these
    suprema = find_deltaY_suprema(infima)
    
    #repeat until nothing new
    infima2 = [inf.dual() for inf in find_deltaY_suprema([sup.dual() for sup in suprema])]
    while len(infima) != len(infima2):
        infima = infima2
        suprema = find_deltaY_suprema(infima)
        infima2 = [inf.dual() for inf in find_deltaY_suprema([sup.dual() for sup in suprema])]
    infima = infima2

    #finally, find matroids we can obtain by Delta-Ys starting at the "least"
    oldmats = list(infima)
    allmats = list(infima)
    while len(oldmats) > 0:
        newmats = []
        for oldmat in oldmats:
            triangles = [c for c in oldmat.circuits() if len(c) == 3]
            newmats.extend([deltaY(oldmat, tri) for tri in triangles])
        newmats = get_nonisomorphic_matroids(newmats)
        allmats.extend(newmats)
        oldmats = newmats

    assert any(mat.is_isomorphic(N) for N in allmats)
    
    return get_nonisomorphic_matroids(allmats)

#returns a partition (as a list) of the matroids DY-equivalent to those given, partitioned by DY-class
def partitionIntoDeltaYClasses(matroidsList):
    deltaYclasses = []
    for mat in matroidsList:
        triangles = [c for c in mat.circuits() if len(c) == 3]
        triads = [c for c in mat.cocircuits() if len(c) == 3]
        if len(triangles) == 0 and len(triads) == 0:
            deltaYclasses.append([mat])
        #if not in some class already
        elif not any(mat.is_isomorphic(N) for N in flatten(deltaYclasses)):
            newclass = get_deltaY_equivalent_matroids(mat)
            deltaYclasses.append(newclass)
    return deltaYclasses

#dual closed version
def partitionIntoDCDeltaYClasses(matroidsList):
    deltaYclasses = []
    for mat in matroidsList:
        triangles = [c for c in mat.circuits() if len(c) == 3]
        triads = [c for c in mat.cocircuits() if len(c) == 3]
        #if not in some class already
        if not any(mat.is_isomorphic(N) for N in flatten(deltaYclasses)):
            if len(triangles) == 0 and len(triads) == 0:
                if mat.is_isomorphic(mat.dual()):
                    deltaYclasses.append([mat])
                else:
                    deltaYclasses.append([mat,mat.dual()])
            elif not any(mat.is_isomorphic(N) for N in flatten(deltaYclasses)):
                newclass = get_deltaY_equivalent_matroids(mat,includeDuals=True)
                deltaYclasses.append(newclass)
    return deltaYclasses