# Flowpy -- takes a square interaction matrix, a corresponding set of point coordinates, and
# and creates a shapefile of flow lines from them. You can choose magnitudes as "raw" two way flow,
# net flow, or gross flow.

# Flowpy is written in Python 2.6 and requires ogr to run.
# Alan Glennon, 19 October 2009

def shapefilemaker(typeofcalculation,fulldirectorystring,outputfilename,fulldirODinput,fulldirPTinput,combotext):
        
        #import flowmapper
        
        import ogr, os, sys
        import math
        #directory location of output file
        #os.chdir(fulldirectorystring)
        #location of input file
        odmatrixfilename = fulldirODinput
        nodefilename = fulldirPTinput
        # type of flow calculation -- 1 is two way flow; 2 is gross flow; 3 is net flow

        try: 
            odmatrix = open(odmatrixfilename, 'r')
            nodes = open(nodefilename,'r')
        except IOError, e:
            print 'file open error:', e

        numberofnodes = 0
        for line in nodes:
                numberofnodes += 1
        nodes.close

        # interaction file is read. Output result is myodmatrix[rows][columns]
        rows = 0
        myodmatrix = []
        for icounter in xrange(numberofnodes):
                        myodmatrix.append([])
                        for jcounter in xrange(numberofnodes):
                                myodmatrix[icounter].append(icounter+jcounter)

        for eachLine in odmatrix:
                        separatestrings = eachLine.split()
                        columns = 0
                        while columns < numberofnodes:
                            onevalue = float(separatestrings[columns])
                            myodmatrix[rows][columns] = onevalue
                            columns += 1
                        rows += 1
        odmatrix.close

        # points file is read. Output result is mypoints[ptrows][ptcolumns]
        ptrows = 0
        mypoints = []
        for kcounter in xrange(numberofnodes):
                        mypoints.append([])
                        for lcounter in xrange(2):
                                mypoints[kcounter].append(kcounter+lcounter)

        nodesagain = open(nodefilename,'r')
        
        for eachLine2 in nodesagain:
                        separatestrings2 = eachLine2.split()
                        ptcolumns = 0
                        while ptcolumns < 2:
                                onevalue2 = float(separatestrings2[ptcolumns])
                                mypoints[ptrows][ptcolumns] = onevalue2
                                ptcolumns += 1
                        ptrows += 1
        nodesagain.close


        #START setup creation of shapefile

        # get the driver
        driver = ogr.GetDriverByName('ESRI Shapefile')
        # create a new data source and layer
        if os.path.exists(outputfilename):
            driver.DeleteDataSource(outputfilename)
        ds = driver.CreateDataSource(outputfilename)
        if ds is None:
            print 'Could not create file'
            sys.exit(1)
        layer = ds.CreateLayer('flow', geom_type=ogr.wkbLineString)
        fieldDefn = ogr.FieldDefn('magnitude', ogr.OFTReal)
        
        fieldDefn2 = ogr.FieldDefn('length_km', ogr.OFTReal)
        fieldDefn3 = ogr.FieldDefn('coord_x1', ogr.OFTReal)
        fieldDefn4 = ogr.FieldDefn('coord_y1', ogr.OFTReal)
        fieldDefn5 = ogr.FieldDefn('coord_x2', ogr.OFTReal)
        fieldDefn6 = ogr.FieldDefn('coord_y2', ogr.OFTReal)
        
        layer.CreateField(fieldDefn)
        
        layer.CreateField(fieldDefn2)
        layer.CreateField(fieldDefn3)
        layer.CreateField(fieldDefn4)
        layer.CreateField(fieldDefn5)
        layer.CreateField(fieldDefn6)
        
        #END setup creation of shapefile

        #START two way flow calculation. 
        if typeofcalculation == 1:
            counter1 = 0
            counter2 = 0
            while counter2 < numberofnodes:
                    while counter1 < numberofnodes:
                            linester = ogr.Geometry(ogr.wkbLineString)
                            linester.AddPoint(mypoints[counter2][0],mypoints[counter2][1])
                            linester.AddPoint(mypoints[counter1][0],mypoints[counter1][1])
                            featureDefn = layer.GetLayerDefn()
                            feature = ogr.Feature(featureDefn)
                            feature.SetGeometry(linester)
                            feature.SetField('magnitude',myodmatrix[counter2][counter1])
                            
                            x1 = mypoints[counter2][0]
                            y1 = mypoints[counter2][1]
                            x2 = mypoints[counter1][0]
                            y2 = mypoints[counter1][1]
                            
                            feature.SetField('coord_x1',x1)
                            feature.SetField('coord_y1',y1)
                            feature.SetField('coord_x2',x2)
                            feature.SetField('coord_y2',y2)                                    
                            
                            if combotext == "Cartesian":
                                length_m = (((x1-x2)**2)+((y1-y2)**2))**0.5
                                feature.SetField('length_km',length_m / (1000))
                            
                            else:
                                lon1 = float(x1)
                                lat1 = float(y1)
                                lon2 = float(x2)
                                lat2 = float(y2)
                                length_km = math.atan2(math.sqrt((math.cos((math.atan(1.0)/45.0)*lat2)*math.sin((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2 + (math.cos((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2)-math.sin((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2), math.sin((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2) + math.cos((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1)) * 6367.9375
                                feature.SetField('length_km',length_km)
                            
                            
                            layer.CreateFeature(feature)
                            counter1 = counter1 + 1
                    counter2 = counter2 + 1
                    counter1 = 0
        #FINISH two way flow calculation


        #START gross flow calculations
        if typeofcalculation == 2:
            g = 0
            h = 0
            while g < numberofnodes:
                while h < numberofnodes:
                            if (g <= h):
                                    linester = ogr.Geometry(ogr.wkbLineString)
                                    linester.AddPoint(mypoints[g][0], mypoints[g][1])
                                    linester.AddPoint(mypoints[h][0], mypoints[h][1])
                                    if h==g: 
                                        grossmagnitude = (myodmatrix[g][h] + myodmatrix[h][g])/2
                                    else: 
                                        grossmagnitude = (myodmatrix[g][h] + myodmatrix[h][g])
                                    featureDefn = layer.GetLayerDefn()
                                    feature = ogr.Feature(featureDefn)
                                    feature.SetGeometry(linester)
                                    feature.SetField('magnitude',grossmagnitude)

                                    x1 = mypoints[h][0]                                        
                                    y1 = mypoints[h][1]
                                    x2 = mypoints[g][0]                                        
                                    y2 = mypoints[g][1]                                        
                                    
                                    feature.SetField('coord_x1',x1)
                                    feature.SetField('coord_y1',y1)
                                    feature.SetField('coord_x2',x2)
                                    feature.SetField('coord_y2',y2)                                    
                                    
                                    if combotext == "Cartesian":
                                        length_m = (((x1-x2)**2)+((y1-y2)**2))**0.5
                                        feature.SetField('length_km',length_m / (1000))
                                    
                                    else:
                                        lon1 = float(x1)
                                        lat1 = float(y1)
                                        lon2 = float(x2)
                                        lat2 = float(y2)
                                        length_km = math.atan2(math.sqrt((math.cos((math.atan(1.0)/45.0)*lat2)*math.sin((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2 + (math.cos((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2)-math.sin((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2), math.sin((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2) + math.cos((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1)) * 6367.9375
                                        feature.SetField('length_km',length_km)
                                    
                                    layer.CreateFeature(feature)
                            h += 1
                h = 0
                g += 1
        #FINISH gross flow calculations

        #START net flow calculations
        if typeofcalculation == 3:
            g = 0
            h = 0
            while g < numberofnodes:
                while h < numberofnodes:
                            if (g <= h):
                                    if h==g: 
                                        netmagnitude = (myodmatrix[g][h] + myodmatrix[h][g])/2
                                    else: 
                                        netmagnitude = (myodmatrix[g][h] - myodmatrix[h][g])
                                    if netmagnitude < 0:
                                        linester = ogr.Geometry(ogr.wkbLineString)
                                        linester.AddPoint(mypoints[h][0], mypoints[h][1])
                                        linester.AddPoint(mypoints[g][0], mypoints[g][1])
                                        featureDefn = layer.GetLayerDefn()
                                        feature = ogr.Feature(featureDefn)
                                        feature.SetGeometry(linester)
                                        feature.SetField('magnitude',netmagnitude * (-1))

                                        x1 = mypoints[h][0]                                        
                                        y1 = mypoints[h][1]
                                        x2 = mypoints[g][0]                                        
                                        y2 = mypoints[g][1]                                        
                                        
                                        feature.SetField('coord_x1',x1)
                                        feature.SetField('coord_y1',y1)
                                        feature.SetField('coord_x2',x2)
                                        feature.SetField('coord_y2',y2)
                                        
                                        if combotext == "Cartesian":
                                            length_m = (((x1-x2)**2)+((y1-y2)**2))**0.5
                                            feature.SetField('length_km',length_m / (1000))
                                        
                                        else:
                                            lon1 = float(x1)
                                            lat1 = float(y1)
                                            lon2 = float(x2)
                                            lat2 = float(y2)
                                            length_km = math.atan2(math.sqrt((math.cos((math.atan(1.0)/45.0)*lat2)*math.sin((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2 + (math.cos((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2)-math.sin((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2), math.sin((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2) + math.cos((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1)) * 6367.9375
                                            feature.SetField('length_km',length_km)
                                        
                                        layer.CreateFeature(feature)
                                    
                                    else:
                                        linester = ogr.Geometry(ogr.wkbLineString)
                                        linester.AddPoint(mypoints[g][0], mypoints[g][1])
                                        linester.AddPoint(mypoints[h][0], mypoints[h][1])
                                        featureDefn = layer.GetLayerDefn()
                                        feature = ogr.Feature(featureDefn)
                                        feature.SetGeometry(linester)
                                        feature.SetField('magnitude',netmagnitude)
                                        
                                        x1 = mypoints[h][0]
                                        y1 = mypoints[h][1]
                                        x2 = mypoints[g][0]
                                        y2 = mypoints[g][1]
                                        
                                        feature.SetField('coord_x1',x1)
                                        feature.SetField('coord_y1',y1)
                                        feature.SetField('coord_x2',x2)
                                        feature.SetField('coord_y2',y2)
                                        
                                        if combotext == "Cartesian":
                                            length_m = (((x1-x2)**2)+((y1-y2)**2))**0.5
                                            feature.SetField('length_km',length_m / (1000))
                                            
                                        else:
                                            lon1 = float(x1)
                                            lat1 = float(y1)
                                            lon2 = float(x2)
                                            lat2 = float(y2)
                                            length_km = math.atan2(math.sqrt((math.cos((math.atan(1.0)/45.0)*lat2)*math.sin((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2 + (math.cos((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2)-math.sin((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1))**2), math.sin((math.atan(1.0)/45.0)*lat1)*math.sin((math.atan(1.0)/45.0)*lat2) + math.cos((math.atan(1.0)/45.0)*lat1)*math.cos((math.atan(1.0)/45.0)*lat2)*math.cos((math.atan(1.0)/45.0)*lon2-(math.atan(1.0)/45.0)*lon1)) * 6367.9375
                                            feature.SetField('length_km',length_km)

                                        layer.CreateFeature(feature)
                            h += 1
                h = 0
                g += 1
        #FINISH net flow calculations

        # shapefile cleanup
        # destroy the geometry and feature and close the data source
        linester.Destroy()
        feature.Destroy()
        ds.Destroy()

