import argparse

parser = argparse.ArgumentParser(description='Program to trim annotated iplane data into a single connected component.  Requires file generated by annotateIplaneData')
parser.add_argument('annotatedIplaneData', metavar='annotatedIplaneData', nargs=1, help='file containing annotated iplane data')
parser.add_argument('outFile', metavar = 'outputfile', nargs=1, help='file output trimed file to')

#parse arguments and open files
args = parser.parse_args()
annotatedIplaneDataFile= open(args.annotatedIplaneData[0], 'r')
outFile = open(args.outFile[0], 'w')

CUSTOMER_PROVIDER = -1
PEER_PEER = 0
PROVIDER_CUSTOMER = 1

asMap = {}

class AS:
    peers = []
    providers = []
    customers = []
    neighborMap = {}
    asn = ''
    def __init__(self, asn):
        self.asn = asn
        
    def addNeighbor(asNum, relationship):
        #if the asnumber is not a known neighbor then add it
        #the add it to the customer provider lists as appropriate
        if asNum not in neighborMap:
            neighborMap[asNum] = relationship
            if relationship == CUSTOMER_PROVIDER:
                providers.add(asNum)
            elif relationship == PEER_PEER:
                peers.add(asNum)
            elif relationship == PROVIDER_CUSTOMER:
                customer.add(asNum)


#function to parse data into an ASmap (integer - AS)
def parseIplane():
    #for each line in the iplane file
    for line in annotatedIplaneDataFile:
        split = line.split()
        AS1Num = int(split[0])
        AS2Num = int(split[1])
        relationship = int(split[2])

        if AS1Num not in asMap:
            #create as class
            AS1 = AS(AS1Num)
        else:
            AS1 = asMap[AS1Num]
        if AS2Num not in asMap:
            #create asclass
            AS2 = AS(AS2Num)
        else:
            AS2 = asMap[AS2Num]

        #add the AS relationships to AS
        if relationship == CUSTOMER_PROVIDER:
            AS2.addNeighbor(AS1Num, PROVIDER_CUSTOMER)
            AS1.addNeighbor(AS2Num, CUSTOMER_PROVIDER)
        if relationship == PEER_PEER:
            AS1.addNeighbor(AS2Num, PEER_PEER)
            AS2.addNeighbor(AS1Num, PEER_PEER)
        if relationship == PROVIDER_CUSTOMER: #included for redundancy sake, not contained in file
            AS1.addNeighbor(AS2Num, PROVIDER_CUSTOMER)
            AS2.addNeighbor(AS1Num, CUSTOMER_PROVIDER)

        #add to dictionary
        asMap[AS1Num] = AS1
        asMap[AS2Num] = AS2
            
