"""This script will contain many of the functions that I have created, which can be accessed
and reused for various purposes."""

from osgeo import ogr
import os

def ExtractFeaturesInGDB_ByAttribute_ToMemory(inGDB, inFileName, filterQuery, outLayerName):
    """
    This function will extract features from a feature class based on a SQL filter query and save the layer to
    memory. It will carry over any attributes in the input feature class.

    This function will only work with feature classes!!! The feature has to be in a geodatabase, and the driver
    MUST be 'FileGDB'. This function can be changed to use shapefiles by adjusting how the input is opened.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: input geodatabase (must end with .gdb)
        inFileName: the name of the input feature class
        filterQuery: SQL query that is in the format of 'column = string' ... string must be quoted and number should NOT be quoted
        outLayerName: the name of the output feature class

    Returns: output memory data source and memory layer
    """

    from osgeo import ogr

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 0)
    inLayer = inDataSource.GetLayer(inFileName)

    # query out the wanted fields
    inLayer.SetAttributeFilter(filterQuery)

    # create the output driver
    outDriver = ogr.GetDriverByName('MEMORY')

    print('Out driver set as ' + format(outDriver.GetName()))

    # create output shape file
    outDataSource = outDriver.CreateDataSource('memData_' + format(outLayerName))
    outMemory = outDataSource.CreateLayer(outLayerName, inLayer.GetSpatialRef(), inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outMemory.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outMemory.GetLayerDefn()

    # Add features to the output Layer
    for inFeature in inLayer:

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = inFeature.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            outFeature.SetField(i, inFeature.GetField(i))

        # Add new feature to output Layer
        outMemory.CreateFeature(outFeature)

    # Save and close DataSources
    del inFeature
    del inLayer
    del inDataSource
    del inDriver
    return outDataSource, outMemory

def ExtractFeaturesInGDB_ByAttribute_ToMemory(inGDB, inFileName, filterQuery, outLayerName):
    """
    This function will extract features from a feature class based on a SQL filter query and save the layer to
    memory. It will carry over any attributes in the input feature class.

    This function will only work with feature classes!!! The feature has to be in a geodatabase, and the driver
    MUST be 'FileGDB'. This function can be changed to use shapefiles by adjusting how the input is opened.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: input geodatabase (must end with .gdb)
        inFileName: the name of the input feature class
        filterQuery: SQL query that is in the format of 'column = string' ... string must be quoted and number should NOT be quoted
        outLayerName: the name of the output feature class

    Returns: output memory data source and memory layer
    """

    from osgeo import ogr

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 0)
    inLayer = inDataSource.GetLayer(inFileName)

    # query out the wanted fields
    inLayer.SetAttributeFilter(filterQuery)

    # create the output driver
    outDriver = ogr.GetDriverByName('MEMORY')

    print('Out driver set as ' + format(outDriver.GetName()))

    # create output shape file
    outDataSource = outDriver.CreateDataSource('memData_' + format(outLayerName))
    outMemory = outDataSource.CreateLayer(outLayerName, inLayer.GetSpatialRef(), inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outMemory.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outMemory.GetLayerDefn()

    # Add features to the output Layer
    for inFeature in inLayer:

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = inFeature.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            outFeature.SetField(i, inFeature.GetField(i))

        # Add new feature to output Layer
        outMemory.CreateFeature(outFeature)

    # Save and close DataSources
    del inFeature
    del inLayer
    del inDataSource
    del inDriver
    return outDataSource, outMemory

def ExtractFeaturesInMemory_ByAttribute_ToMemory(inDataSource, inMemory, filterQuery, outLayerName):
    """
    This function will extract features from an ogr memory data source based on a SQL filter query and save the layer to
    memory. It will carry over any attributes in the input feature. Both the data source and the memory and needed as
    inputs because the memory layer will not be able to be accessed without the data source.

    This function will only work with memory, and the driver MUST be 'Memory'.

    This function does NOT delete the input data source and layer.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inDataSource: input data source
        inMemory: input memory
        filterQuery: SQL query that is in the format of 'column = string' ... string must be quoted and number should NOT be quoted
        outLayerName: the name of the output

    Returns: input and output memory data source and memory layer
    """

    from osgeo import ogr

    # query out the wanted fields
    inMemory.SetAttributeFilter(filterQuery)

    # create the output driver
    outDriver = ogr.GetDriverByName('MEMORY')

    print('Out driver set as ' + format(outDriver.GetName()))

    # create output shape file
    outDataSource = outDriver.CreateDataSource('memData_' + format(outLayerName))
    outMemory = outDataSource.CreateLayer(outLayerName, inMemory.GetSpatialRef(), inMemory.GetGeomType())

    # Add input Layer Fields to the output Layer
    outMemory.CreateFields(inMemory.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outMemory.GetLayerDefn()

    # Add features to the output Layer
    for inFeature in inMemory:

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = inFeature.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            outFeature.SetField(i, inFeature.GetField(i))

        # Add new feature to output Layer
        outMemory.CreateFeature(outFeature)

    inMemory.ResetReading()
    del inFeature

    # return DataSources and layers
    return inDataSource, inMemory, outDataSource, outMemory

def ExtractFeaturesinMemory_ByAttribute_ToShapefile(inDataSource, inLayer, filterQuery, outShapefileName,
                                                    shapefileFolder):
    """
    This function will extract features from an ogr memory data source based on a SQL filter query and save the layer to
    memory. It will carry over any attributes in the input feature. Both the data source and the memory and needed as
    inputs because the memory layer will not be able to be accessed without the data source.

    This function does NOT delete the input data source and layer.

    Args:
        inDataSource: input data source
        inMemory: input layer
        filterQuery: SQL query that is in the format of 'column = string' ... string must be quoted and number should NOT be quoted
        outShapefileName: the name of the output
        shapefileFolder: path the the folder that will contain the output shapefile

    Returns: input memory data source and memory layer
    """

    from osgeo import ogr
    import os

    # query out the wanted fields
    inLayer.SetAttributeFilter(filterQuery)

    # create the output driver
    outDriver = ogr.GetDriverByName('ESRI Shapefile')

    # remove output shape file if it already exists
    if os.path.exists(shapefileFolder + '\\' + outShapefileName + '.shp'):
        outDriver.DeleteDataSource(shapefileFolder + '\\' + outShapefileName + '.shp')

    # create output shape file
    outDataSource = outDriver.CreateDataSource(shapefileFolder + '\\' + outShapefileName + '.shp')
    outLayer = outDataSource.CreateLayer(outShapefileName + '.shp', inLayer.GetSpatialRef(), inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outLayer.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outLayer.GetLayerDefn()

    # Add features to the output Layer
    for input_feat in inLayer:

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = input_feat.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            outFeature.SetField(i, input_feat.GetField(i))

        # Add new feature to output Layer
        outLayer.CreateFeature(outFeature)

    inLayer.ResetReading()

    del input_feat
    del outLayer
    del outDataSource
    del outDriver

    # return DataSources and layers
    return inDataSource, inLayer

def FeatureToFeature(inGDB, outGDB, inFileName, outFileName):
    """
    This tool will take a feature that is within a database and create an exact copy of it to another
    geodatabase. The geodatabases could be the same one.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: path to the geodatabase where the input feature class is located (must end with .gdb)
        outGDB: path to the geodatabase where the output feature class is to be placed (must end with .gdb)
        inFileName: nanme of the input feature class
        outFileName: name of the output feature class

    Returns: N/A

    """

    from osgeo import ogr
    import os

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 0)
    inLayer = inDataSource.GetLayer(inFileName)

    # grab the spatial reference system
    srs = inLayer.GetSpatialRef()

    # create the output driver
    outDriver = ogr.GetDriverByName('FileGDB')

    print('Out driver set as ' + format(outDriver.GetName()))

    # remove output shape file if it already exists
    if os.path.exists(outGDB + '\\' + outFileName):
        outDriver.DeleteDataSource(outGDB + '\\' + outFileName)

    # create output shape file
    outDataSource = outDriver.Open(outGDB, 1)
    outFile = outDataSource.CreateLayer(outFileName, srs, geom_type=inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outFile.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile.GetLayerDefn()

    # Add features to the output Layer
    for inFeature in inLayer:
        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = inFeature.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            outFeature.SetField(i, inFeature.GetField(i))

        # Add new feature to output Layer
        outFile.CreateFeature(outFeature)

    inLayer.ResetReading()

    # Save and close DataSources
    del inDataSource
    del outDataSource

def MemoryToFeature(inDataSource, inMemory, outGDB, outFileName):
    """
    This tool will take an input ogr memory layer and convert it to a feature class within a geodatabase.
    The memory data source and memory layer are both needed for inputs because the memory layer will not
    be able to be accessed without the data source.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inDataSource: in memory data source
        inMemory: in memory layer
        outGDB: path to the output geodatabase (must end with .gdb)
        outFileName: name of the output feature class

    Returns: N/A

    """

    from osgeo import ogr

    # create the output driver
    outDriver = ogr.GetDriverByName('FileGDB')

    print('Out driver set as ' + format(outDriver.GetName()))

    # create output shape file
    outDataSource = outDriver.Open(outGDB, 1)
    outFile = outDataSource.CreateLayer(outFileName, inMemory.GetSpatialRef(), inMemory.GetGeomType())

    # Add input Layer Fields to the output Layer
    inLayerDefn = inMemory.GetLayerDefn()
    for i in range(0, inLayerDefn.GetFieldCount()):
        fieldDefn = inLayerDefn.GetFieldDefn(i)
        outFile.CreateField(fieldDefn)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile.GetLayerDefn()

    # Add features to the output Layer
    for inFeature in inMemory:
        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            outFeature.SetField(outLayerDefn.GetFieldDefn(i).GetNameRef(),
                                inFeature.GetField(i))

        # Set geometry as centroid
        geom = inFeature.GetGeometryRef()
        outFeature.SetGeometry(geom.Clone())
        # Add new feature to output Layer
        outFile.CreateFeature(outFeature)
        outFeature = None

    # Save and close DataSources
    outDriver = None
    outDataSource = None
    outFile = None
    inDataSource = None
    inMemory = None

def CreateShapefile_Memory(name, srs, geom, schema=None):

    outDriver = ogr.GetDriverByName('MEMORY')
    # create output shape file
    outDataSource = outDriver.CreateDataSource('memData_' + format(name))
    outLyr = outDataSource.CreateLayer(name, srs, geom)

    if schema:
        # Add input Layer Fields to the output Layer
        for feat in schema:
            fld = ogr.FieldDefn(feat[0], feat[1])
            outLyr.CreateField(fld)

    del outLyr
    return outDataSource

def RasterToPoints(inRaster, srs, schema=None):

    from osgeo import gdal

    ds = gdal.Open(inRaster)
    band = ds.GetRasterBand(1)

    (upper_left_x, x_size, x_rotation, upper_left_y, y_rotation, y_size) = ds.GetGeoTransform()

    data = band.ReadAsArray()

    coords = []
    for i in range(0, data.shape[0]):
        for j in range(0, data.shape[1]):
            #         print(i, j)
            x_dist_i = upper_left_x + (x_size * i)
            #         print(x_dist_i)
            y_dist_j = upper_left_y + (y_size * j)
            #         print(y_dist_j)
            coords.append([x_dist_i, y_dist_j])

    # create point shapefile
    memLyr_ds = CreateShapefile_Memory('MemPoints', srs, ogr.wkbPoint, schema=schema)

    memLyr = memLyr_ds.GetLayer()

    memLyr_Defn = memLyr.GetLayerDefn()

    coord_fld = ogr.FieldDefn('Longitude', ogr.OFTReal)
    coord_fld.SetWidth(8)
    coord_fld.SetPrecision(3)
    memLyr.CreateField(coord_fld)
    coord_fld.SetName('Latitude')
    memLyr.CreateField(coord_fld)

    # loop through coordinates and create points based off of them
    for c in coords:
        outFeat = ogr.Feature(memLyr_Defn)
        point = ogr.Geometry(ogr.wkbPoint)
        point.AddPoint(c[0], c[1])
        outFeat.SetGeometry(point)
        outFeat.SetField('Longitude', c[0])
        outFeat.SetField('Latitude', c[1])
        memLyr.CreateFeature(outFeat)
        del outFeat

    del memLyr
    return memLyr_ds, coords

def SpatialFilter_Remove(inShp, filterShp, keep_within=True):

    if isinstance(inShp, str):
        driver = ogr.GetDriverByName('ESRI Shapefile')
        inShpDs = driver.Open(inShp, 0)
    else:
        inShpDs = inShp

    inLyr = inShpDs.GetLayer()

    if isinstance(filterShp, str):
        driver = ogr.GetDriverByName('ESRI Shapefile')
        filterDs = driver.Open(filterShp, 0)
    else:
        filterDs = filterShp

    filterLyr = filterDs.GetLayer()

    # get geomerty from filter
    featureCount = filterLyr.GetFeatureCount()
    if featureCount > 1:
        print('Feature count for filter shapfile is greater than 1. Spatial filter will only\n',
              'perform on the first feature.')
    filterFeat = filterLyr.GetNextFeature()
    geom = filterFeat.geometry().Clone()

    inLyr.SetSpatialFilter(geom)

    # ldefn = inLyr.GetLayerDefn()
    # for n in range(ldefn.GetFieldCount()):
    #     fdefn = ldefn.GetFieldDefn(n)
    #     print(fdefn.name)

    # gather feature FIDs that are left after the spatial filter
    FIDs = []
    for feature in inLyr:
        FIDs.append(feature.GetFID())

    inLyr.SetSpatialFilter(None)

    # if keep_within=True, then remove all features outside of the spatial filter
    # else, remove all features inside of the spatial filter
    for feature in inLyr:
        fid = feature.GetFID()
        if keep_within == True:
            if fid not in FIDs:
                inLyr.DeleteFeature(fid)
        elif keep_within == False:
            if fid in FIDs:
                inLyr.DeleteFeature(fid)
        else:
            print('Input for attribute keep_within not valid')

    return inShp

def MemoryToShapefile(inDataSource, inMemory, shapefileFolder, outFileName):
    """
    This tool will convert an ogr memory layer to a shapefile. The memory data source and memory layer are both
    needed for inputs because the memory layer will not be able to be accessed without the data source.

    Args:
        inDataSource: input memory data source
        inMemory: input memory layer
        shapefileFolder: path to the folder where the shapefile will be placed
        outFileName: name of the output shapefile (must end with .shp)

    Returns: N/A

    """

    from osgeo import ogr
    import os

    # create the output driver
    outDriver = ogr.GetDriverByName('ESRI Shapefile')

    # remove output shape file if it already exists
    if os.path.exists(shapefileFolder + '\\' + outFileName + '.shp'):
        outDriver.DeleteDataSource(shapefileFolder + '\\' + outFileName + '.shp')

    # create output shape file
    outDataSource = outDriver.CreateDataSource(shapefileFolder + '\\' + outFileName + '.shp')
    outFile = outDataSource.CreateLayer(outFileName + '.shp', inMemory.GetSpatialRef(), inMemory.GetGeomType())

    # Add input Layer Fields to the output Layer
    outFile.CreateFields(inMemory.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile.GetLayerDefn()

    inMemory.ResetReading()

    # Add features to the output Layer
    for input_feat in inMemory:

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = input_feat.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            outFeature.SetField(i, field_value)

        # Add new feature to output Layer
        outFile.CreateFeature(outFeature)

    # set the input data source and layer to none
    del inMemory
    del inDataSource
    del outFile
    del outDataSource
    del outDriver

def ShapefileToMemory(shapefileFolder, inFileName, outFileName):
    """
    This tool will convert a shapefile into an ogr memory layer and return the memory data source
    and memory layer. Both are returned because the memory layer cannot be opened without having the data
    source. So keep these together as long as the memory layer needs to be accessed.

    Args:
        shapefileFolder: path to the folder where the input shapefile is located
        inFileName: name of the input shapefile (must end with .shp)
        outFileName: name for the output memory layer

    Returns: memory data source and memory layer

    """

    from osgeo import ogr

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('ESRI Shapefile')
    inDataSource = inDriver.Open(shapefileFolder + '\\' + inFileName, 0)
    inLayer = inDataSource.GetLayer()

    # create the output driver
    outDriver = ogr.GetDriverByName('MEMORY')

    print('Out driver set as ' + format(outDriver.GetName()))

    # create output shape file
    outDataSource = outDriver.CreateDataSource('memData_' + format(outFileName))
    outFile = outDataSource.CreateLayer(outFileName, inLayer.GetSpatialRef(), inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outFile.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile.GetLayerDefn()

    inLayer.ResetReading()

    # Add features to the output Layer
    for input_feat in inLayer:

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = input_feat.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            outFeature.SetField(i, field_value)

        # Add new feature to output Layer
        outFile.CreateFeature(outFeature)

    # Save and close DataSources
    del input_feat
    del inLayer
    del inDataSource
    del inDriver

    return outDataSource, outFile

def ShapefileToMemory_ForceMultiPoint(shapefileFolder, inFileName, outFileName):

    from osgeo import ogr

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('ESRI Shapefile')
    inDataSource = inDriver.Open(shapefileFolder + '\\' + inFileName, 0)
    inLayer = inDataSource.GetLayer()

    # create the output driver
    outDriver = ogr.GetDriverByName('MEMORY')

    print('Out driver set as ' + format(outDriver.GetName()))

    # create output shape file
    outDataSource = outDriver.CreateDataSource('memData_' + format(outFileName))
    outFile = outDataSource.CreateLayer(outFileName, inLayer.GetSpatialRef(), ogr.wkbMultiPoint)

    # Add input Layer Fields to the output Layer
    outFile.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile.GetLayerDefn()

    inLayer.ResetReading()

    # Add features to the output Layer
    for input_feat in inLayer:
        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = input_feat.GetGeometryRef()
        geom_name = geom.GetGeometryName()
        if geom_name == 'POINT':
            geom = ogr.ForceToMultiPoint(geom)
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            outFeature.SetField(i, field_value)

        # Add new feature to output Layer
        outFile.CreateFeature(outFeature)

    # Save and close DataSources
    del inLayer
    del inDataSource
    del inDriver

    return outDataSource, outFile

def DissolveWithFiona_FromMemory(inDataSource, inMemory, inFileName, outFileName, shapefileFolder, dissolve_field):
    """
    This function uses a script found online that uses fiona to dissolve a shapefile.
    url: https://gis.stackexchange.com/questions/149959/dissolving-polygons-based-on-attributes-with-python-shapely-fiona

    This function takes an ogr memory layer as an input, converts its to a shapefile, dissolves, and
    converts the dissolved shapefile back into a memory layer.

    IMPORTANT NOTE!!!!!!: If the dissolve function fails the inFileName will be added to a list called 'incompleted_list'.
    This section will need to be taken out, unless used.

    Args:
        inDataSource:
        inMemory:
        inFileName:
        outFileName:
        shapefileFolder:
        dissolve_field:

    Returns: dissolved memory data source and layer

    """

    import fiona
    import itertools
    from shapely.geometry import shape, mapping
    from shapely.ops import unary_union

    MemoryToShapefile(inDataSource, inMemory, shapefileFolder, inFileName)

    """use fiona to dissolve"""
    try:
        with fiona.open(shapefileFolder + '\\' + inFileName + '.shp') as input_layer:
            # preserve the schema of the original shapefile, including the crs
            meta = input_layer.meta
            with fiona.open(shapefileFolder + '\\' + outFileName + '.shp', 'w', **meta) as output:
                # groupby clusters consecutive elements of an iterable which have the same key so you must first sort the features by the 'STATEFP' field
                e = sorted(input_layer, key=lambda k: k['properties'][dissolve_field])
                # group by the dissolve field
                for key, group in itertools.groupby(e, key=lambda x: x['properties'][dissolve_field]):
                    properties, geom = zip(*[(feature['properties'], shape(feature['geometry'])) for feature in group])
                    # write the feature, computing the unary_union of the elements in the group with the properties of the first element in the group
                    output.write({'geometry': mapping(unary_union(geom)), 'properties': properties[0]})

    except:
        print(inFileName + ' was not able to be dissolved.')

    outDataSource, outFile = ShapefileToMemory(shapefileFolder, outFileName + '.shp', outFileName)

    return outDataSource, outFile

def DissolveWithFiona(shapefileFolder, inFileName, outFileName, dissolveField):
    import fiona
    import itertools
    from shapely.geometry import shape, mapping
    from shapely.ops import unary_union
    import os
    from osgeo import ogr

    # remove output shape file if it already exists
    d_out = ogr.GetDriverByName('ESRI Shapefile')
    if os.path.exists(shapefileFolder + '\\' + outFileName + '.shp'):
        d_out.DeleteDataSource(shapefileFolder + '\\' + outFileName + '.shp')
    del d_out

    with fiona.open(shapefileFolder + '\\' + inFileName + '.shp') as input_layer:
        # preserve the schema of the original shapefile, including the crs
        meta = input_layer.meta
        with fiona.open(shapefileFolder + '\\' + outFileName + '.shp', 'w', **meta) as output:
            # groupby clusters consecutive elements of an iterable which have the same key so you must first sort the features by the 'STATEFP' field
            e = sorted(input_layer, key=lambda k: k['properties'][dissolveField])
            # group by the dissolve field
            for key, group in itertools.groupby(e, key=lambda x: x['properties'][dissolveField]):
                properties, geom = zip(
                    *[(fiona_feature['properties'], shape(fiona_feature['geometry'])) for fiona_feature in group])
                # write the feature, computing the unary_union of the elements in the group with the properties of the first element in the group
                output.write({'geometry': mapping(unary_union(geom)), 'properties': properties[0]})

def SortLayer_GDB(inGDB, layerName, field):
    """
    This function will sort a feature class by a certain field in ascending order. It works with numbered fields,
    but should work with string fields as well. This function will NOT create a new feature class.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: geodatabase where the input feature class is located (must end with .gdb)
        layer: name of the feature class
        field: Name of the field to sort by

    Returns: N/A

    """

    from osgeo import ogr

    # open input layer with driver
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 1)
    inLayer = inDataSource.GetLayer(layerName)
    # fids are unique, fids may be sorted or unsorted, fids may be consecutive or have gaps
    # don't care about semantics, don't touch fids and their order, reuse fids
    fids = []
    vals = []
    inLayer.ResetReading()
    for f in inLayer:
        fid = f.GetFID()
        fids.append(fid)
        vals.append((f.GetField(field), fid))
    vals.sort()
    # index as dict: {newfid: oldfid, ...}
    ix = {fids[i]: vals[i][1] for i in xrange(len(fids))}

    # swap features around in groups/rings
    for fidstart in ix.keys():
        if fidstart not in ix: continue
        ftmp = inLayer.GetFeature(fidstart)
        fiddst = fidstart
        while True:
            fidsrc = ix.pop(fiddst)
            if fidsrc == fidstart: break
            f = inLayer.GetFeature(fidsrc)
            f.SetFID(fiddst)
            inLayer.SetFeature(f)
            fiddst = fidsrc
        ftmp.SetFID(fiddst)
        inLayer.SetFeature(ftmp)

    inLayer.ResetReading()

    # Save and close DataSources
    del inLayer
    del inDataSource
    del inDriver

def GetFeatureFromShapefile(filePath, fid):
    """
    Use function to grab a certain feature from a shapefile
    based on the FID
    Args:
        file_path: path to shapefile (string)
        fid: fid number (integer)

    Returns: ogr data source and ogr feature object
    """

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('ESRI Shapefile')
    inDataSource = inDriver.Open(filePath, 0)
    inLayer = inDataSource.GetLayer()

    fid = int(fid)

    # grab feature from layer based on fid
    feat = inLayer.GetFeature(fid)

    return inDataSource, feat

def GetFieldFromShapefile(filePath, fid, fieldName):
    """
    Function to grab the value of a field from a specific
    features in a shapefile by fid

    Args:
        filePath: path to shapefile (string)
        fid: fid feature number (integer)
        fieldName: name of field (string)

    Returns: specified field value (string, integer, or float)

    """

    # open the inShapefile as the driver type
    inDriver = ogr.GetDriverByName('ESRI Shapefile')
    inDataSource = inDriver.Open(filePath, 0)
    inLayer = inDataSource.GetLayer()

    # if fid is "X", skip
    if fid.lower() == 'x':
        return None

    # grab feature from layer based on fid
    feat = inLayer.GetFeature(int(fid))
    # grab value of field
    val = feat.GetField(fieldName)

    return val

def Sort(inDataSource, inMemory, field):
    """
        This function will sort a memory layer by a certain field in ascending order. It works with numbered fields,
        but should work with string fields as well.

        Requirement: ogr

        Args:
            inDataSource: data source of the memory layer
            inMemory: layer source of the memory
            field: Name of the field to sort by

        Returns: inDataSource, inMemory

        """

    from osgeo import ogr

    # fids are unique, fids may be sorted or unsorted, fids may be consecutive or have gaps
    # don't care about semantics, don't touch fids and their order, reuse fids
    fids = []
    vals = []
    inMemory.ResetReading()
    for f in inMemory:
        fid = f.GetFID()
        fids.append(fid)
        vals.append((f.GetField(field), fid))
    vals.sort()
    # index as dict: {newfid: oldfid, ...}
    ix = {fids[i]: vals[i][1] for i in xrange(len(fids))}

    # swap features around in groups/rings
    for fidstart in ix.keys():
        if fidstart not in ix: continue
        ftmp = inMemory.GetFeature(fidstart)
        fiddst = fidstart
        while True:
            fidsrc = ix.pop(fiddst)
            if fidsrc == fidstart: break
            f = inMemory.GetFeature(fidsrc)
            f.SetFID(fiddst)
            inMemory.SetFeature(f)
            fiddst = fidsrc
        ftmp.SetFID(fiddst)
        inMemory.SetFeature(ftmp)

    inMemory.ResetReading()

    del f

    return inDataSource, inMemory

def GrabUniqueValuesFromField_GDB(inGDB, inFileName, field):
    """
    This function will create and return a list of the unique values or strings from a single
    column of the input feature class. This will only work with feature classes in a geodatabase!

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: path to the geodatabase where the input feature class is located (must end with .gdb)
        inFileName: name of the input feature class
        field: name of the field to grab the values or strings from

    Returns: unique values list

    """
    # open the feature
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 0)
    inLayer = inDataSource.GetLayer(inFileName)

    # reset reading
    inLayer.ResetReading()

    unique_values = []

    for feature in inLayer:
        value = feature.GetField(field)
        if value not in unique_values:
            unique_values.append(value)

    inLayer.ResetReading()

    del inLayer
    del inDataSource
    del inDriver

    return (unique_values)

def GrabUniqueValuesFromLayer(inDataSource, inMemory, field):
    """
        This function will create and return a list of the unique values or strings from a single
        column of the input memory layer.

        Requirement: ogr

        Args:
            inDatasource: data source that holds the input memory layer
            inFileName: input memory layer
            field: name of the field to grab the values or strings from

        Returns: unique values list, inDataSource, inMemory

        """

    from osgeo import ogr

    # reset reading
    inMemory.ResetReading()

    unique_values = []

    for feature in inMemory:
        value = feature.GetField(field)
        if value not in unique_values:
            unique_values.append(value)

    inMemory.ResetReading()

    return unique_values, inDataSource, inMemory

def AppendFeatureToFeature_GDB(inGDB, inFileName, appendingFileName, appendingFileGDB):
    """
    The function will append a feature class to another feature class. This will only work for feature classes
    in a geodatabase. At this point, it is unknown if this function will work with features that have different
    attributes. The appending file/GDB will be the feature that gets appended to the in file/GDB.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: path to the geodatabase that contains the main feature class
        inFileName: name of the main feature class
        appendingFileName: name of the feature class that will be appended to the main feature class
        appendingFileGDB: path to the geodatabase that contains the feature that is going to be appended.

    Returns: N/A

    """

    from osgeo import ogr

    # open in file which will be appended to
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 1)
    inLayer = inDataSource.GetLayer(inFileName)

    # open the file which will be appended to the other
    extraDriver = ogr.GetDriverByName('FileGDB')
    extraDataSource = extraDriver.Open(appendingFileGDB, 0)
    extraLayer = extraDataSource.GetLayer(appendingFileName)

    extraLayerDefn = extraLayer.GetLayerDefn()

    for input_feat in inLayer:

        # create new feature
        extraFeature = ogr.Feature(extraLayerDefn)

        # set geometry
        geom = input_feat.GetGeometryRef()
        extraFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, extraLayerDefn.GetFieldCount()):
            extraFeature.SetField(i, input_feat.GetField(i))

        # append extra layer to in layer
        inLayer.CreateFeature(extraFeature)

    # Save and close DataSources
    del extraFeature
    del inLayer
    del inDataSource
    del inDriver
    del extraLayer
    del extraDataSource
    del extraDriver

def AppendLayerToShapefile(inFolder, inFileName, appendingDataSource, appendingMemory):
    """
    This function will append an ogr memory layer to an existing feature class. As of now, it is uncertain
    if this will work with layers that have different attributes. Both the appending memory data source and layer
    are both needed because the memory layer cannot be opened without the data source.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: path to the geodatabase where the feature class is located
        inFileName: name of the feature class
        appendingDataSource: memory data source that will be appended to the feature class
        appendingMemory: memory layer that will be appended to the feature class

    Returns: N/A

    """

    from osgeo import ogr

    # open in file which will be appended to
    inDriver = ogr.GetDriverByName('ESRI Shapefile')
    inDataSource = inDriver.Open(inFolder + '\\' + inFileName + '.shp', 1)
    inLayer = inDataSource.GetLayer()

    appendingMemory.ResetReading()

    inLayerDefn = inLayer.GetLayerDefn()

    for input_feat in appendingMemory:

        # create new feature
        extraFeature = ogr.Feature(inLayerDefn)

        # Set geometry
        geom = input_feat.GetGeometryRef()
        extraFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, inLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            extraFeature.SetField(i, field_value)

        # append extra layer to in layer
        inLayer.CreateFeature(extraFeature)

    # Save and close DataSources
    del input_feat
    del inLayer
    del inDataSource
    del inDriver
    del appendingMemory
    del appendingDataSource

def AppendLayerToLayer(inDataSource, inMemory, appendingDataSource, appendingMemory):
    """
    This function will append a memory layer to another memory layer, return the memory data source
    and layer that was appended, and delete the appended memory data source and layer. The two layers must
    have the same field names!!!

    Args:
        inDataSource: input memory data source
        inMemory: input memory layer
        appendingDataSource: memory data source that will be appended to the input
        appendingMemory: memory layer that will be appended to the input

    Returns: input memory data source and layer that was appended to as well as the appended data source and memory

    """

    from osgeo import ogr

    inLayerDefn = inMemory.GetLayerDefn()

    appendingMemory.ResetReading()

    for input_feat in appendingMemory:

        extraFeature = ogr.Feature(inLayerDefn)

        # Set geometry
        geom = input_feat.GetGeometryRef()
        extraFeature.SetGeometry(geom.Clone())

        # Add field values from input Layer
        for i in range(0, inLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            extraFeature.SetField(i, field_value)

        # append extra layer to in layer
        inMemory.CreateFeature(extraFeature)

    del extraFeature

    inMemory.ResetReading()
    appendingMemory.ResetReading()

    return inDataSource, inMemory, appendingDataSource, appendingMemory

def AddFieldsFromFeatureToFeature_GDB(inGDB, inFileName, fieldsGDB, fieldsFeatureName):
    """
    This function will take the fields of a feature class and add them to another feature class. It will copy
    over all of the field names of the fields Feature.
    This function will only work with feature classes in a geodatabase.

    Requirement: 'FileGDB' driver for ogr

    Args:
        inGDB: path to geodatabase that contains the inFileName layer
        inFileName: name of the input feature class that will get the fields added to itself
        fieldsGDB: path to the geodatabase that contains the feature class with the wanted field names
        fieldsFeatureName: name of the feature class with the field names that are wanted for the in feature class

    Returns: N/A

    """

    from osgeo import ogr

    # open in feature using the driver
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 1)
    inLayer = inDataSource.GetLayer(inFileName)

    # open feature that has the fields wanted
    fieldsDriver = ogr.GetDriverByName('FileGDB')
    fieldsDataSource = fieldsDriver.Open(fieldsGDB, 0)
    fieldsLayer = fieldsDataSource.GetLayer(fieldsFeatureName)

    # add field layer fields to the inLayer
    inLayer.CreateFields(fieldsLayer.schema)

    # Save and close DataSources
    del inLayer
    del inDataSource
    del inDriver
    del fieldsLayer
    del fieldsDataSource
    del fieldsDriver

# def SeparateDissolveMerge(inDataSource, inMemory, inFeatureName, outFeaturePath, dissolveField, shapefileFolder):
# 	"""
#
# 	THIS FUNCTION DOESN'T WORK!!!!!!!!!!!!!!
#
# 	This function will take an input ogr memory layer, create a list of the unique values for a certain dissolve
# 	field, separate the memory layer based on unique dissolve field values, dissolve each unique dissolve field
# 	feature geometries into one, and merge all layers back together into one, fully dissolved, feature class.
#
# 	NOTE: This function requires the DissolveWithFiona_FromMemory function as well!!!!!!!! See the help for that
# 	for additional information.
#
# 	Requirement: 'FileGDB' driver for ogr
#
# 	Args:
# 		inDataSource: input memory data source
# 		inMemory: input memory layer
# 		inFeatureName: name of the input layer
# 		outFeaturePath: full path to the output feature class (must be in geodatabase)
# 		dissolveField: attribute field to be dissolved by
# 		shapefileFolder: folder that will be used to contain the shapefiles created with the
# 			DissolveWithFiona_FromMemory function.
#
# 	Returns: N/A
#
# 	"""
#
# 	print('Beginning dissolve function for ' + format(inFeatureName))
#
# 	# set inDataSource variable
# 	# inDataSource = inDataSource
#
# 	# from the input feature class, create a set of unique dissolve field values
# 	uniqueRARNUMs = set()
# 	for feature in inMemory:
# 		rarnum = feature.GetField(dissolveField)
# 		uniqueRARNUMs.add(rarnum)
#
# 	# create a list that the dissolved memory layers will be added to
# 	memory_list = []
#
# 	# for each unique RARNUM, separate into separate feature classes in the GDB
# 	for value in uniqueRARNUMs:
# 		print('Dissolving ' + format(value))
#
# 		query = '{} = {}'.format(dissolveField, value)
# 		dataSource, memory = ExtractFeaturesInMemory_ByAttribute_ToMemory(inDataSource, inMemory, query, format(inFeatureName) + '_separated_' + format(value))
#
# 		print('features for dissolve value of ' + format(value) + ' extracted from the input')
#
# 		inFileName = format(inFeatureName) + '_separated_' + format(value)
# 		outFileName = format(inFeatureName) + '_dissolved_' + format(value)
#
# 		dissolvedDataSource, dissolvedMemory = DissolveWithFiona_FromMemory(dataSource, memory, inFileName, outFileName,
# 																			shapefileFolder, dissolveField)
#
# 		print('dissolve completed for ' + format(value))
#
# 		# create a temporary list that will contain the previous data source and memory for the dissolved feature
# 		temp_list = [dissolvedDataSource, dissolvedMemory]
#
# 		# add dissolved memory layer to the memory list
# 		memory_list.append(temp_list)
#
# 		# delete memory to save space
# 		dissolvedDataSource = None
# 		dissolvedMemory = None
#
# 	# now the input data source and layer can be deleted
# 	inDataSource = None
# 	inMemory = None
#
# 	# now merge, if the new file doesn't exist, create a new one
# 	for pair in memory_list:
#
# 		if arcpy.Exists(outFeaturePath):
#
# 			# append to existing feature class
# 			AppendMemoryToFeature(outFeaturePath.rsplit('\\', 1)[0], outFeaturePath.rsplit('\\')[-1], pair[0],
# 								  pair[1])  # TODO: add data source input to function
#
# 		else:
#
# 			# create a new feature class
# 			MemoryToFeature(pair[0], pair[1], outFeaturePath.rsplit('\\', 1)[0], outFeaturePath.rsplit('\\')[-1])

def ListOGR_drivers():
    import ogr
    cnt = ogr.GetDriverCount()
    formatsList = []  # Empty List

    for i in range(cnt):
        driver = ogr.GetDriver(i)
        driverName = driver.GetName()
        if not driverName in formatsList:
            formatsList.append(driverName)

    formatsList.sort()  # Sorting the messy list of ogr drivers

    for i in formatsList:
        print(i)

def TableToCSV(inGDB, inTableName, outTablePath):
    from osgeo import ogr
    import csv

    # open table in geodatabase
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(inGDB, 0)
    inTable = inDataSource.GetLayer(inTableName)

    # pull header information
    header = []
    inTableDefn = inTable.GetLayerDefn()
    for i in range(0, inTableDefn.GetFieldCount()):
        header.append((inTableDefn.GetFieldDefn(i)).GetName())

    with open(outTablePath, 'wb') as csvFile:
        writer = csv.DictWriter(csvFile, fieldnames=header)

        writer.writeheader()

        for line in inTable:
            rowdict = {}
            for i in range(0, inTableDefn.GetFieldCount()):
                rowdict.update(
                    {inTableDefn.GetFieldDefn(i).GetNameRef(): line.GetField(inTableDefn.GetFieldDefn(i).GetNameRef())})

            line.Destroy()
            writer.writerow(rowdict)

        csvFile.close()

    del inTable
    del inDataSource
    del inDriver

def CheckFeatureCount_Shapefile(shapefileFolder, fileName):

    from osgeo import ogr

    d = ogr.GetDriverByName('ESRI Shapefile')
    ds = d.Open(shapefileFolder + '\\' + fileName + '.shp', 0)
    l = ds.GetLayer()
    feat_num = l.GetFeatureCount()
    del l
    del ds
    del d

    return feat_num

def Dissolve_ShapefileToShapefile(shapefileFolder, inFileName, outFileName, lengthFieldName, areaFieldName):
    """
    DOES NOT WORK YET!!!!

    Function to dissolve a shapefile. Will work for points, lines, or polygons. The function
    will calculate the new length and area for each feature, if applicable.

    Args:
        shapefileFolder: path to the folder where the shapefile is located
        inFileName: name of input shapefile
        outFileName: name of output shapefile
        lengthFieldName: name of the field in the shapefile that displays the
            length of a feature. If not applicable, enter empty string.
        areaFieldName: name of the field in the shapefile that displays the
            area of a feature. If not applicable, enter empty string.

    Returns: N/A

    """

    # get layer from data source
    d_in = ogr.GetDriverByName('ESRI Shapefile')
    ds_in = d_in.Open(os.path.join(shapefileFolder, inFileName + '.shp'), 0)
    l_in = ds_in.GetLayer()

    # check the geometry of the layer
    check_geom = l_in.GetGeomType()
    multi_geom = None
    if check_geom == 1:
        # crate multi point geometry
        multi_geom = ogr.Geometry(ogr.wkbMultiPoint)
    if check_geom == 2:
        # create multi line string geometry
        multi_geom = ogr.Geometry(ogr.wkbMultiLineString)
    if check_geom == 3:
        # create a multi polygon geometry
        multi_geom = ogr.Geometry(ogr.wkbMultiPolygon)

    # Keep track of length variable
    length = 0

    # loop through each feature until there are no more
    for input_feat in l_in:
        try:
            length += input_feat.GetField(lengthFieldName)
        except:
            pass
        # get geometry from feature
        g = input_feat.GetGeometryRef()

        # add geometry to multi geometry
        multi_geom.AddGeometry(g)

    l_in.ResetReading()

    # dissolve the multi geometry using union cascaded
    new_geom = multi_geom.UnionCascaded()

    d_out = ogr.GetDriverByName('ESRI Shapefile')
    outPath = os.path.join(shapefileFolder, outFileName + '.shp')

    # remove output shape file if it already exists
    if os.path.exists(outPath):
        d_out.DeleteDataSource(outPath)

    # open new shapefile
    ds_out = d_out.CreateDataSource(outPath)
    l_out = ds_out.CreateLayer(outFileName, l_in.GetSpatialRef(), multi_geom.GetGeometryType())

    # add field schema to out layer
    l_out.CreateFields(l_in.schema)

    defn = l_in.GetLayerDefn()

    # create a new feature
    newFeat = ogr.Feature(l_out.GetLayerDefn())
    # add geometry to the new feature
    newFeat.SetGeometry(new_geom)
    # add field values to the new feature
    for i in range(0, defn.GetFieldCount()):
        field_value = l_in.GetFeature(0).GetField(i)
        field_name = defn.GetFieldDefn(i).GetNameRef()
        # get geometry
        g = newFeat.GetGeometryRef()
        if field_name == lengthFieldName:
            # set the calculated length from above to the  field value
            field_value = length
        if field_name == areaFieldName:
            # calculate area
            field_value = g.Area()
        newFeat.SetField(i, field_value)
    # add new feature to the out layer
    l_out.CreateFeature(newFeat)

    # close data sources
    del l_in
    del ds_in
    del d_in
    del l_out
    del ds_out
    del d_out

def CheckFeatureCount_FeatureGDB(GDB, inFileName):
    """
    Use this function to calculate the number of features from a layer within a GDB.
    Args:
        GDB: path to the geodatabase
        inFileName: name of the layer

    Returns: object containing a number representing the number of features in the shapefile

    """

    # open in feature using the driver
    inDriver = ogr.GetDriverByName('FileGDB')
    inDataSource = inDriver.Open(GDB, 1)
    inLayer = inDataSource.GetLayer(inFileName)

    # get feature count from layer
    feat_num = inLayer.GetFeatureCount()
    del inLayer
    del inDataSource
    del inDriver

    return feat_num

def SplitIntoTwoLayers_Shapefile(shapefileFolder, shapefileName):
    """
    This function will take a shapefile and split it into 2 separate shapefiles. The first output will have the first half
    of the number of features, the second output will have the other half. The two outputs will have a '_A' and '_B'
    added to the names.

    Args:
        shapefileFolder:
        shapefileName:

    Returns: N/A

    """

    # Open input shapefile
    inDriver = ogr.GetDriverByName('ESRI Shapefile')
    inDataSource = inDriver.Open(shapefileFolder + '\\' + shapefileName + '.shp')
    inLayer = inDataSource.GetLayer()

    # Get feature count
    feat_count = inLayer.GetFeatureCount()

    # open first new shapefile
    outDriver_A = ogr.GetDriverByName('ESRI Shapefile')

    # remove output shape file if it already exists
    if os.path.exists(shapefileFolder + '\\' + shapefileName + '_A.shp'):
        outDriver_A.DeleteDataSource(shapefileFolder + '\\' + shapefileName + '_A.shp')

    # create output shape file
    outDataSource_A = outDriver_A.CreateDataSource(shapefileFolder + '\\' + shapefileName + '_A.shp')
    outFile_A = outDataSource_A.CreateLayer(shapefileName + '_A.shp', inLayer.GetSpatialRef(), inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outFile_A.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile_A.GetLayerDefn()

    inLayer.ResetReading()

    # now add first half of features from inLayer to shapefile A
    for i in range(0, int(feat_count / 2)):

        input_feat = inLayer.GetFeature(i)

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = input_feat.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            outFeature.SetField(i, field_value)

        # Add new feature to output Layer
        outFile_A.CreateFeature(outFeature)

    # close new shapefile
    del outFile_A
    del outDataSource_A
    del outDriver_A

    # reset reading
    inLayer.ResetReading()

    # open first new shapefile
    outDriver_B = ogr.GetDriverByName('ESRI Shapefile')

    # remove output shape file if it already exists
    if os.path.exists(shapefileFolder + '\\' + shapefileName + '_B.shp'):
        outDriver_B.DeleteDataSource(shapefileFolder + '\\' + shapefileName + '_B.shp')

    # create output shape file
    outDataSource_B = outDriver_B.CreateDataSource(shapefileFolder + '\\' + shapefileName + '_B.shp')
    outFile_B = outDataSource_B.CreateLayer(shapefileName + '_B.shp', inLayer.GetSpatialRef(), inLayer.GetGeomType())

    # Add input Layer Fields to the output Layer
    outFile_B.CreateFields(inLayer.schema)

    # Get the output Layer's Feature Definition
    outLayerDefn = outFile_B.GetLayerDefn()

    inLayer.ResetReading()

    # now add first half of features from inLayer to shapefile A
    for i in range(int(feat_count / 2), feat_count):

        input_feat = inLayer.GetFeature(i)

        # Create output Feature
        outFeature = ogr.Feature(outLayerDefn)

        # Set geometry as centroid
        geom = input_feat.GetGeometryRef()
        outFeature.SetGeometry(geom)

        # Add field values from input Layer
        for i in range(0, outLayerDefn.GetFieldCount()):
            field_value = input_feat.GetField(i)
            outFeature.SetField(i, field_value)

        # Add new feature to output Layer
        outFile_B.CreateFeature(outFeature)

    # close new shapefile
    del outFile_B
    del outDataSource_B
    del outDriver_B
    del inLayer
    del inDataSource
    del inDriver

def PullIndexFromFeature(GDB, inFileName, fieldName):
    """
    Simple function to grab the field index based off of a field name from a feature within a geodatabase.
    The field index of the first field matching the passed field name (case insensitively) is returned.

    Args:
        GDB: path to the geodatabase containing the feature (string)
        inFileName: name of the feature (string)
        fieldName: name of field for the index (string)

    Returns: field index (integer)

    """
    # open feature
    driver = ogr.GetDriverByName('FileGDB')
    dataSource = driver.Open(GDB, 0)
    layer = dataSource.GetLayer(inFileName)

    # get feature definition from the layer
    defn = layer.GetLayerDefn()

    # get field index based on field name
    index = defn.GetFieldIndex(fieldName)

    # delete objects
    del layer
    del dataSource
    del driver

    return index

def GetGeomType_OGR(folder, inFileName, driverType):
    """
    This function will take a feature within a geodatabase or a shapefile and return an integer indicating
    the geometry type of the layer.

    1: Point
    2: Polyline
    3: Polygon

    Args:
        folder: path to the folder or geodatabase
        inFileName: name of the layer
        driverType: OGR driver type, MUST BE 'ESRI Shapefile' or 'FileGDB'

    Requirement: 'FileGDB' driver for OGR, if using a feature within a geodatabase as the input

    Returns: integer indicating the geometry type of the input layer

    """
    # open layer
    if driverType == 'FileGDB':
        driver = ogr.GetDriverByName(driverType)
        dataSource = driver.Open(folder, 0)
        layer = dataSource.GetLayer(inFileName)
    if driverType == 'ESRI Shapefile':
        driver = ogr.GetDriverByName(driverType)
        dataSource = driver.Open(folder + '\\' + inFileName + '.shp', 0)
        layer = dataSource.GetLayer()

    # get the geometry type from the layer
    geometry_type = layer.GetGeomType()

    # return the geometry type integer
    return geometry_type

    # delete layers
    del layer
    del dataSource
    del driver

def CheckFeatureExistsInGDB(GDB, featureName):
    """
    This little function will check if a feature layer exists in a geodatabase using OGR.

    Args:
        GDB: path to the geodatabase
        featureName: name of the feature layer as a string

    Returns: True or False object

    """
    driver = ogr.GetDriverByName('FileGDB')
    dataSource = driver.Open(GDB, 0)
    if dataSource.GetLayer(featureName):
        feat = True
    else:
        feat = False

    return feat

def CSV_to_list(csv_path):
    import csv

    list = []
    with open(csv_path, 'r') as Record:
        reader = csv.reader(Record)
        for line in reader:
            list.append(line)

    return (list)

def CalculateAgeInYears(install_str, removal_str, format):
    """
    Function to calculate the number of years between two dates.
    Will return a year float to 2 decimal places.

    Args:
        install_str: string indicating install date
        removal_str: string indicating removal date
        format: date string format according to datetime module

    Returns: float indicating years

    """
    from datetime import datetime

    install = datetime.strptime(install_str, format)
    removal = datetime.strptime(removal_str, format)

    diff = removal - install

    years = round((diff.days + diff.seconds / 86400) / 365.2425, 1)

    return years

def ReturnAgeCategory(age):
    """
        This function is specific, as in it will return a string of the
    year range that an age falls between. These current year ranges
    were picked for the offshore infrastructure project to help
    predict the possible range of age at removal for platforms
    using machine learning. This should be adjusted as necessary.

    Args:
        age: float or integer indicating age

    Returns: string indicating age range category

    """

    if (age >= 0) and (age < 11):
        return '0-11'
    if (age >= 11) and (age <= 20):
        return '11-20'
    if (age >= 20) and (age < 30):
        return '20-30'
    if (age >= 30) and (age < 42):
        return '30-42'
    if (age >= 42):
        return '42-72'

def OpenShapefile_ReadOnly(path):
    """
    Open shapefile for reading only

    Args:
        path: path to shapefile

    Returns: ogr data source, ogr layer

    """

    d = ogr.GetDriverByName('ESRI Shapefile')
    ds = d.Open(path,0)
    l = ds.GetLayer()

    return ds, l

def OpenShapefile_ForEditing(path):
    """
    Open shapefile for reading only

    Args:
        path: path to shapefile

    Returns: ogr data source, ogr layer

    """

    d = ogr.GetDriverByName('ESRI Shapefile')
    ds = d.Open(path, 1)
    l = ds.GetLayer()

    return ds, l

def cloneFieldDefn(src_def):
    fdef = ogr.FieldDefn(src_def.GetName(), src_def.GetType())
    fdef.SetWidth(src_def.GetWidth())
    fdef.SetPrecision(src_def.GetPrecision())
    return fdef

def AlterFieldName(shapefilePath, oldFieldName, newFieldName):

    ds, l = OpenShapefile_ForEditing(shapefilePath)

    # get layer defn
    layerDefn = l.GetLayerDefn()

    # get field defn
    i = layerDefn.GetFieldIndex(oldFieldName)
    fieldDefn = layerDefn.GetFieldDefn(i)
    # clone field
    clone = cloneFieldDefn(fieldDefn)
    # rename field
    clone.SetName(newFieldName)
    # set new field to layer
    l.AlterFieldDefn(i, clone, ogr.ALTER_NAME_FLAG)

def RasMemToTiff(ds, output_ras):
    '''Converts raster datasource in memory to a geotiff'''

    # Get original no data value
    band1 = ds.GetRasterBand(1)
    NoDataVal = band1.GetNoDataValue()

    # Collect all raster information and delete ds from memory
    raster_data = ds.ReadAsArray()
    geo_transform = ds.GetGeoTransform()
    x_size = ds.RasterXSize  # Raster xsize
    y_size = ds.RasterYSize  # Raster ysize
    srs = ds.GetProjectionRef()  # Projection
    ds = None

    # Need a driver object. By default, we use GeoTIFF
    driver = gdal.GetDriverByName("gtiff")
    dataset_out = driver.Create(output_ras, x_size, y_size, 1, gdal.GDT_Float32)
    driver = None
    dataset_out.SetGeoTransform(geo_transform)
    dataset_out.SetProjection(srs)
    band_out = dataset_out.GetRasterBand(1)
    band_out.SetNoDataValue(NoDataVal)
    band_out.WriteArray(raster_data.astype(np.float32))
    band_out.FlushCache()
    band_out = None
    dataset_out = None

def ExtractValuesByLocations(rasDS, lats, lons):
    ''' Pull raster values from rasDS based on location (latitude, longitude), returns list of values'''

    values = []

    # Pull band & values of interest from raster
    band = rasDS.GetRasterBand(1)
    cols = rasDS.RasterXSize
    rows = rasDS.RasterYSize

    transform = rasDS.GetGeoTransform()

    xOrigin = transform[0]
    yOrigin = transform[3]
    pixelWidth = transform[1]
    pixelHeight = -transform[5]

    data = band.ReadAsArray(0, 0, cols, rows)

    # Loop through locations and pull raster values at locations
    l = 0
    for lat in lats:
        lon = lons[l]

        print(lat, lon)

        col = int((float(lat) - xOrigin) / pixelWidth)
        row = int((yOrigin - float(lon)) / pixelHeight)

        result = float(data[row][col])

        if result == band.GetNoDataValue():
            result = float('nan')

        values.append(result)

        l += 1

    return values

def ExtractValueByLocation(rasDS, lat, lon):
    ''' Pull raster values from rasDS based on location (latitude, longitude), returns list of values'''

    values = []

    # Pull band & values of interest from raster
    band = rasDS.GetRasterBand(1)
    cols = rasDS.RasterXSize
    rows = rasDS.RasterYSize

    transform = rasDS.GetGeoTransform()

    xOrigin = transform[0]
    yOrigin = transform[3]
    pixelWidth = transform[1]
    pixelHeight = -transform[5]

    data = band.ReadAsArray(0, 0, cols, rows)

    # Loop through locations and pull raster values at locations

    col = int((float(lat) - xOrigin) / pixelWidth)
    row = int((yOrigin - float(lon)) / pixelHeight) * -1

    result = float(data[row][col])

    if result == band.GetNoDataValue():
        result = float('nan')

    values.append(result)

    return values

def RasToMem(inRas):
    '''Copies raster into memory, returns raster datasource'''
    r_datasource = gdal.Open(inRas)

    mem_driver = gdal.GetDriverByName("MEM")
    ds = mem_driver.CreateCopy("", r_datasource)

    return ds

def ReprojectRaster_WGS84(inRaster):
    '''Reprojects raster (string path and name) to WGS84'''
    r_datasource = gdal.Open(inRaster)

    ds = gdal.Warp('', r_datasource, dstSRS='EPSG:4326', format='VRT',
                   outputType=gdal.GDT_CFloat32)  # , xRes=0.00892857142857143, yRes=0.0089285714285714

    return ds

def CheckSRS_WGS84(inData):
    '''Checks if input data is in WGS 84 geographic projection, returns True/False'''

    WGS84_srs = [
        'GEOGCS["WGS 84",\n    DATUM["WGS_1984",\n        SPHEROID["WGS 84",6378137,298.257223563,\n            AUTHORITY["EPSG","7030"]],\n        AUTHORITY["EPSG","6326"]],\n    PRIMEM["Greenwich",0,\n        AUTHORITY["EPSG","8901"]],\n    UNIT["degree",0.0174532925199433,\n        AUTHORITY["EPSG","9122"]],\n    AUTHORITY["EPSG","4326"]]',
        'GEOGCS["WGS 84",\n    DATUM["WGS_1984",\n        SPHEROID["WGS 84",6378137,298.257223563,\n            AUTHORITY["EPSG","7030"]],\n        AUTHORITY["EPSG","6326"]],\n    PRIMEM["Greenwich",0],\n    UNIT["degree",0.0174532925199433],\n    AUTHORITY["EPSG","4326"]]']

    if inData.split('.')[-1] == 'shp':
        driver = ogr.GetDriverByName("ESRI Shapefile")
        datasource = driver.Open(inData, 0)
        lyr = datasource.GetLayer()
        srs = str(lyr.GetSpatialRef())

    else:  # Assumed raster
        r_datasource = gdal.Open(inData)
        r_srs = r_datasource.GetProjection()
        srs = osr.SpatialReference(r_srs)

    if srs in WGS84_srs:
        return True
    else:
        return False

def CheckSRS(inData_List):
    ''' Check if all items in inData_List share the same srs or if there are differences '''
    SRS_list = []
    for d in inData_List:
        print(d)

        if d.split('.')[-1] == 'shp':
            driver = ogr.GetDriverByName("ESRI Shapefile")
            datasource = driver.Open(d, 0)
            lyr = datasource.GetLayer()
            SRS_list.append(str(lyr.GetSpatialRef()))

        else:  # Assumed raster
            r_datasource = gdal.Open(d)
            r_srs = r_datasource.GetProjection()
            SRS_list.append(osr.SpatialReference(r_srs))

    return all(x == SRS_list[0] for x in SRS_list)

def checkGeometry(inputShp):

    # get layer from data source
    d_in = ogr.GetDriverByName('ESRI Shapefile')
    ds_in = d_in.Open(inputShp, 0)
    l_in = ds_in.GetLayer()

    # check the geometry of the layer
    check_geom = l_in.GetGeomType()
    if check_geom == 1:
        return 'point'
    if check_geom == 3:
        return 'polygon'

def checkPointWithinMask(inputShp, maskFile):

    from osgeo import ogr
    import json

    def records(file):
        # generator
        reader = ogr.Open(file)
        layer = reader.GetLayer(0)
        for i in range(layer.GetFeatureCount()):
            feature = layer.GetFeature(i)
            yield json.loads(feature.ExportToJson())

    check = True

    points = [pt for pt in records(inputShp)]
    multipol = records(maskFile)
    multi = multipol.next()  # 1 feature
    for i, pt in enumerate(points):
        point = shape(pt['geometry'])
        if point.within(shape(multi['geometry'])):
            continue
        else:
            check = False

    return check

def getRasterValuesAtPoints(inputShp, inputRasters):

    outputValues = []

    # open shapefile
    shp_dr = ogr.GetDriverByName('ESRI Shapefile')
    shp_ds = shp_dr.Open(inputShp)
    shp_lyr = shp_ds.GetLayer()

    # loop through each point in the layer
    for p in shp_lyr:

        ptValues = []

        geom = p.GetGeometryRef()
        mx, my = geom.GetX(), geom.GetY()  # coord in map units

        # loop through each raster
        for raster in inputRasters:
            src_ds = gdal.Open(src_filename)
            geotrans = src_ds.GetGeoTransform()
            band = src_ds.GetRasterBand(1)

            # Convert from map to pixel coordinates.
            # Only works for geotransforms with no rotation.
            px = int((mx - geotrans[0]) / geotrans[1])  # x pixel
            py = int((my - geotrans[3]) / geotrans[5])  # y pixel

            intval = band.ReadAsArray(px, py, 1, 1)
            ptValues.append(intval[0])

        outputValues.append(ptValues)

    return outputValues


import os

import cv2
import keras
import numpy as np
import tensorflow as tf

###############################################################################
#   Here are three example model submissions.
#
#   Each model takes an OpenCV image and the name of the
#   input image (not the path). The model then performs
#   some operation on the image and saves to the designated
#   output path. The output paths are where the organizers
#   will receive the results of each model. An empty output
#   folder will signify there was no submission for that
#   that model. You must submit one model and have it write
#   its output to /root/output0. DO NOT MODIFY.
#
#   How each model handles its IO, either from input
#   source / to output or its method calls is up to the
#   discretion of the performer. See description in
#   __main__ for some IO optimization for Docker caching
#   recommendations and an example of it.
#
#   Absolute Requirements:
#   1.  A minimum of one model must be submitted
#   2.  The model must load all input images from /root/input
#   3.  The model must modify and save a new copy of all input images
#   4.  All image names must remain the same.
#       The image compression/extension can be changed to any
#       Ubuntu supported .jpeg, .jpg, .tiff, .tif, .png format.
###############################################################################

try:
    model = keras.models.load_model('superres_v3.0_2018-01-25 22:38:29.338372.h5', custom_objects={"tf":tf})
except:
    model = keras.models.load_model('/root/performer/superres_v3.0_2018-01-25 22:38:29.338372.h5', custom_objects={"tf": tf})

#   An example model that makes no changes to the input image
def model0(img, name):
    save = os.path.join('/root', 'output0', name)
    print(save)
    print(img.shape)
    print(np.expand_dims(img[:,:,0:], 0).shape)
    size = 32
    apron = 2
    new_img = np.zeros((img.shape[0]*2 + 2*apron, img.shape[1]*2 + 2*apron, 3))
    img = cv2.copyMakeBorder(img, apron, apron, apron, apron, cv2.BORDER_REFLECT)
    for x in np.arange(0, img.shape[0], size - 2*apron):
        for y in np.arange(0, img.shape[1], size - 2*apron):

            # print(model.predict(np.ones((1,160,160,1))))
            chunk = img[x:x + size, y:y + size]
            new_chunk0 = model.predict(np.stack([chunk[:,:,0:1],chunk[:,:, 1:2], chunk[:,:,2:]]))[0]
            # print(chunk.shape, new_chunk0.shape)
            # if new_chunk0.shape[0]>4 and new_chunk0.shape[1]>4:
            # if new_img[2 * x+apron*2:2 * x + 2 * size-apron*2, 2 * y +apron*2:2 * y + 2 * size -apron*2, 0].shape == new_chunk0[0, apron*2:-apron*2, apron*2:-apron*2, 0]:
            if True:
                new_img[2 * x+apron*2:2 * x + 2 * size-apron*2, 2 * y +apron*2:2 * y + 2 * size -apron*2, 0] = new_chunk0[0, apron*2:-apron*2, apron*2:-apron*2, 0]
                new_img[2 * x+apron*2:2 * x + 2 * size-apron*2, 2 * y +apron*2:2 * y + 2 * size -apron*2, 1] = new_chunk0[1, apron*2:-apron*2, apron*2:-apron*2, 0]
                new_img[2 * x+apron*2:2 * x + 2 * size-apron*2, 2 * y +apron*2:2 * y + 2 * size -apron*2, 2] = new_chunk0[2, apron*2:-apron*2, apron*2:-apron*2, 0]
    # print(new_img)
    #
    cv2.imwrite(save, new_img)

#   An example model that performs a median blur on the input image
# def model1(img, name):
#     save = os.path.join('', 'root', 'output1', name)
#
#     cv2.medianBlur(img, 3)
#
#     cv2.imwrite(save, img)

#   An example model that switches channel order
#   and changes image compression on the input image
# def model2(img, name):
#     save = os.path.join('', 'root', 'output2', name)
#
#     red = img[:,:,2].copy()
#     blue = img[:,:,0].copy()
#     img[:,:,0] = red
#     img[:,:,2] = blue
#
#     cv2.imwrite(save.split('.')[0]+'.tif', img)

if __name__ == '__main__':
    #   This is the path to where the input images will be when
    #   this container is executed by the organizers. Changing
    #   this path will fail to produce results. DO NOT CHANGE.
    input_ = '/root/input'
    # input_ = './root/input'

    paths = os.listdir(input_)

    #   In this example, the loop iterates over all of the input
    #   images and loads the image and then provides a copy
    #   to each model. Other image IO methods and model function
    #   calls can be used. E.g. pass image path to the model.
    #
    #   If a model can perform on one image at a time without
    #   effecting the performance, then this style is the optimal
    #   loading paradim for Docker caching rather than doing the
    #   entire input image set for a model, then repeating for
    #   subsequent models.
    for p in paths:
        img = cv2.imread(os.path.join(input_, p))

        #   Submit a minimum of one model, but a
        #   second and third model is optional.
        model0(np.copy(img), p)
        #model1(np.copy(img), p)   # uncomment to use more than one model
        #model2(np.copy(img), p)   # uncomment to use more than one model
    from keras import backend as K
    K.clear_session()


import os
import numpy as np
import pands as pd
import rasterio as rio
import rasterio
from rasterio.windows import Window

class SITSPatches:
  def __init__(self, sits_path, output_dir, patch_size=0, year=0):
    # self.__init__()
    self.sits_path = sits_path
    self.out_dir = output_dir
    self.patch_size = patch_size
    self.year = year

    dtypes = np.dtype([("Patch_path", str),("Rows", int),("Columns", int)])
    self.dataframe = pd.DataFrame(np.empty(0, dtype=dtypes))

  def patch_sits(self):
    out_year_dir = os.path.join(self.out_dir, self.year)
    if not os.path.exists(out_year_dir):
        os.makedirs(out_year_dir)
    patches = []
    rows = []
    cols = []

    with rio.open(self.sits_path) as src:
      for x in range(0, src.width, self.patch_size):
          for y in range(0, src.height, self.patch_size):
              window = Window(x, y, self.patch_size, self.patch_size)
              patch = src.read(window=window) #-- ??, out_shape=(src.count, self.patch_size, self.patch_size)
              output_file = f'patch_{x}_{y}.tif'
              patch_path = os.path.join(out_year_dir, output_file)
              patches.append(patch_path)
              rows.append(x)
              cols.append(y)
              self.dataframe=pd.DataFrame({'Patch_path': patches, 'Rows': rows, 'Columns': cols})

              with rio.open(patch_path,
                            'w',
                            driver='GTiff',
                            width=self.patch_size,
                            height=self.patch_size,
                            count=src.count,
                            dtype=patch.dtype,
                            crs=src.crs,
                            transform=src.window_transform(window)
                            ) as dst:
                  dst.write(patch)

    self.dataframe.to_csv(os.path.join(out_year_dir, f'df_patches_{self.year}.csv'), index=False)

project_dir = "/share/projects/erasmus/pasteis"
for y in ['data/2018/2018_subset.tif', 'data/2019/2019_subset.tif']:
    patches_ = SITSPatches(sits_path=os.path.join(project_dir, y),output_dir = os.path.join(project_dir,"/prepData/patches"), patch_size = 128, year=str(y))
    patches_.patch_sits()

# if __name__ == "__main__" :
#     project_dir = "/share/projects/erasmus/pasteis"
#     for y in ['data/2018/2018_subset.tif', 'data/2019/2019_subset.tif']:
#         patches_ = SITSPatches(sits_path=os.path.join(project_dir, y),output_dir = os.path.join(project_dir,"/prepData/patches"), patch_size = 32, year=str(y))
#         patches_.patch_sits()


import pandas as pd
from tqdm import tqdm
import numpy as np
import gc
import multiprocessing as mlp
from multiprocessing import Manager
from itertools import combinations,product
from tool import co_prob,condition_prob,condition_stat,num2bin
from process_worker import prob_feature_worker,tfidf_worker,conti_stat_feature_worker
from sklearn.linear_model import Ridge
from keras.callbacks import EarlyStopping

"""
tem_id                  - Ad id.
user_id                 - User id.
region                  - Ad region.
city                    - Ad city.
parent_category_name    - Top level ad category as classified by Avito's ad model.
category_name           - Fine grain ad category as classified by Avito's ad model.
param_1                 - Optional parameter from Avito's ad model.
param_2                 - Optional parameter from Avito's ad model.
param_3                 - Optional parameter from Avito's ad model.
title                   - Ad title.
description             - Ad description.
price                   - Ad price.
item_seq_number         - Ad sequential number for user.
activation_date         - Date ad was placed.
user_type               - User type.
image                   - Id code of image. Ties to a jpg file in train_jpg. Not every ad has an image.
image_top_1             - Avito's classification code for the image.
"""

def baseFeature(df):
    df['weekday'] = df.activation_date.dt.weekday
    df['month'] = df.activation_date.dt.month
    df['day'] = df.activation_date.dt.day
    df['week'] = df.activation_date.dt.week

    # length of description
    df['description_len'] = df['description'].fillna("").apply(lambda x: len(x.split()))

    # length of title
    df['title_len'] = df['title'].fillna("").apply(lambda x: len(x.split()))

    # param_combined and its length
    df['param_combined'] = df.apply(
        lambda row: ' '.join([str(row['param_1']), str(row['param_2']), str(row['param_3'])]), axis=1)
    df['param_combined_len'] = df['param_combined'].fillna("").apply(lambda x: len(x.split()))

    # charater len of text columns
    df['description_char'] = df['description'].fillna("").apply(len)
    df['title_char'] = df['title'].fillna("").apply(len)
    df['param_char'] = df['param_combined'].fillna("").apply(len)
    for col in ["description", "title"]:
        df[col+'_num_words'] = df[col].fillna("").apply(lambda comment: len(comment.split()))
        df[col+'_num_unique_words']=df[col].fillna("").apply(lambda comment:len(set(w for w in comment.split())))
        df[col+'_words_vs_unique']=df[col+'_num_unique_words']/df[col+'_num_words']*100

    return df

def get_geo_feature(df):
    df['city_region'] = df["city"] + ' ' + df["region"]
    city_region = pd.DataFrame()
    city_region['city_region'] = df['city_region'].drop_duplicates().reset_index(drop=True)
    city_region['latitude'] = np.nan
    city_region['longitude'] = np.nan

    from geopy import geocoders
    g = geocoders.GoogleV3()

    def get_geocode(row):
        try:
            geocode = g.geocode(row['city_region'],timeout=60,language='en')
            row['latitude'] = geocode.latitude
            row['longitude'] = geocode.longitude
        except:
            row['latitude'] = np.nan
            row['longitude'] = np.nan
        return row

    city_region = city_region.apply(get_geocode,axis=1)

    df = df.merge(city_region, on='city_region', how='left')
    df.drop(['city_region'],axis=1,inplace=True)
    return df

def tfidf_feat(train,test,features,n_components,analyzer,alg='pca',sparse=False,gram=3):

    print(features,n_components)

    pool = mlp.Pool(len(features))
    results = []
    for feat,k,anal in zip(features,n_components,analyzer):
        result = pool.apply_async(tfidf_worker,
                                  args=(feat,k,train[feat],test[feat],alg,sparse,anal,gram))
        results.append((feat,k,result))
        # tfidf_worker(feat,k,train[feat],test[feat],alg)
    pool.close()
    pool.join()

    for feat,k,result in tqdm(results):
        result.get()

def agg_time_feat(train,test):
    used_cols = ['item_id', 'user_id']

    train_sub = train[used_cols]
    train_active = pd.read_csv('./dataset/train_active.csv', usecols=used_cols)
    test_sub = test[used_cols]
    test_active = pd.read_csv('./dataset/test_active.csv', usecols=used_cols)

    train_periods = pd.read_csv('./dataset/periods_train.csv', parse_dates=['date_from', 'date_to'])
    test_periods = pd.read_csv('./dataset/periods_test.csv', parse_dates=['date_from', 'date_to'])

    all_samples = pd.concat([
        train_sub,
        train_active,
        test_sub,
        test_active
    ]).reset_index(drop=True)
    all_samples.drop_duplicates(['item_id'], inplace=True)

    all_periods = pd.concat([train_periods,test_periods])
    all_periods['days_up'] = all_periods['date_to'].dt.dayofyear - all_periods['date_from'].dt.dayofyear

    gp = all_periods.groupby(['item_id'])[['days_up']]

    gp_df = pd.DataFrame()
    gp_df['days_up_sum'] = gp.sum()['days_up']
    gp_df['times_put_up'] = gp.count()['days_up']
    gp_df.reset_index(inplace=True)
    gp_df.rename(index=str, columns={'index': 'item_id'})

    all_periods.drop_duplicates(['item_id'], inplace=True)
    all_periods = all_periods.merge(gp_df, on='item_id', how='left')

    all_periods = all_periods.merge(all_samples, on='item_id', how='left')

    gp = all_periods.groupby(['user_id'])[['days_up_sum', 'times_put_up']].mean().reset_index() \
        .rename(index=str, columns={
        'days_up_sum': 'avg_days_up_user',
        'times_put_up': 'avg_times_up_user'
    })

    n_user_items = all_samples.groupby(['user_id'])[['item_id']].count().reset_index() \
                        .rename(index=str, columns={'item_id': 'n_user_items'})
    gp = gp.merge(n_user_items, on='user_id', how='outer')

    train = train.merge(gp, on='user_id', how='left')
    test = test.merge(gp, on='user_id', how='left')

    return train,test

def get_img_info(df,datatype):
    from process_worker import img_info_worker
    if datatype=='tr':
        path = './dataset/train_img/'
    elif datatype=='te':
        path = './dataset/test_img/'
    else:
        raise RuntimeError("don't have this data")

    img_files = df['image'].values
    num_cpu = mlp.cpu_count()
    pool = mlp.Pool(num_cpu)
    num_task = 1+len(img_files)//num_cpu
    results = []

    for i in range(num_cpu):
        result = pool.apply_async(img_info_worker,
                        args=(img_files[i*num_task:(i+1)*num_task],path))
        results.append(result)
    pool.close()
    pool.join()

    img_feat = []
    for result in tqdm(results):
        img_feat += result.get()

    assert len(df) == img_feat
    columns = ['light_percent', 'dark_percent','average_color',
               'blurrness_score','dominant_color','average_pixel_width']
    img_feat = pd.DataFrame(img_feat,columns=columns)
    img_feat.to_csv('./dataset/'+datatype+'_info.csv',index=False)

def get_img_confi(df, datatype):
    from process_worker import img_confi_worker
    if datatype == 'tr':
        path = './dataset/train_img/'
    elif datatype == 'te':
        path = './dataset/test_img/'
    else:
        raise RuntimeError("don't have this data")

    img_files = df['image'].values
    num_cpu = mlp.cpu_count()
    pool = mlp.Pool(num_cpu)
    num_task = 1 + len(img_files) // num_cpu

    img_feat = pd.DataFrame()
    for model in ['resnet50','xception','inception_v3']:
        results = []
        for i in range(num_cpu):
            result = pool.apply_async(img_confi_worker,
                                      args=(img_files[i*num_task:(i + 1)*num_task],model,path))
            results.append(result)
        pool.close()
        pool.join()

        feat = []
        for result in tqdm(results):
            feat += result.get()

        assert len(df) == img_feat
        columns = [model+str(i) for i in range(3)]
        feat = pd.DataFrame(feat, columns=columns)
        img_feat = pd.concat([img_feat,feat],axis=1)

    img_feat.to_csv('./dataset/' + datatype + '_confi.csv', index=False)

def comb_feat():
    #  共现概率特征
    co_features = []
    co_prob_feat = [
        "category_name",
        "city",
        "param_1",
        "param_2",
        "param_3",
        "user_id",
        "weekday",
    ]
    for i in range(1, 4):
        co_features += list(combinations(co_prob_feat, i))

    # 条件概率特征
    from itertools import product
    condi_features = []

    area = ["city"]
    user_type = ["user_type"]
    user_id = ['user_id']
    date = ["activation_date", "weekday"]
    item = ["param_1", "category_name"]

    comb_feat = [area, user_type, date, item, user_id]
    for i in range(1, 2):
        for j in range(1, 3):
            for condi_varis in combinations(comb_feat, i):
                X = list(product(*condi_varis))
                res = [i for i in comb_feat if i not in condi_varis]
                for res_varis in combinations(res, j):
                    Y = list(product(*res_varis))
                    condi_features += list(product(X, Y))

    return co_features,condi_features

def prob_feature(train,test):

    def run_task(train,test,usecols,features,is_co):
        num_cpu = mlp.cpu_count()//2
        pool = mlp.Pool(num_cpu)
        num_task = 1 + len(features) // num_cpu
        results = []

        for i in range(num_cpu):
            result = pool.apply_async(prob_feature_worker,
                                      args=(usecols,features[i * num_task:(i + 1) * num_task],is_co))
            results.append(result)
        pool.close()
        pool.join()

        for result in tqdm(results):
            train_feat,test_feat = result.get()
            train = pd.concat([train,train_feat],axis=1)
            test = pd.concat([test,test_feat],axis=1)

        return train,test

    l_tr = len(train)
    l_te = len(test)
    usecols = [
        "item_id",
        "user_id",
        "region",
        "city",
        "parent_category_name",
        "category_name",
        "param_1",
        "param_2",
        "param_3",
        "activation_date",
        "user_type",
        "weekday"
    ]

    co_features,condi_features = comb_feat()
    train,test = run_task(train,test,usecols,co_features,True)

    train, test = run_task(train,test,usecols,condi_features,False)

    assert l_tr == len(train)
    assert l_te == len(test)

    return train,test

def user_item_matrix(train,test,dims):

    all_samples = pd.concat([
        train,
        test
    ]).reset_index(drop=True)

    def feat_encode(all_samples):

        from sklearn.preprocessing import LabelEncoder
        le = LabelEncoder()
        all_samples['city_id'] = le.fit_transform(all_samples['city'].values)
        le = LabelEncoder()
        all_samples['user_type_id'] = le.fit_transform(all_samples['user_type'].values)
        le = LabelEncoder()
        all_samples['category_name_id'] = le.fit_transform(all_samples['category_name'].values)

        all_samples['image_top_1_id'] = all_samples['image_top_1']
        all_samples['image_top_1_id'].fillna(all_samples['image_top_1'].max()+1, inplace=True)
        return all_samples

    all_samples = feat_encode(all_samples)

    i = len(all_samples['city_id'].unique())
    j = len(all_samples['user_type_id'].unique())
    assert all_samples['city_id'].max() == i-1
    assert all_samples['user_type_id'].max() == j-1
    num_user = i*j
    all_samples['u_id'] = i*all_samples['user_type_id'] + all_samples['city_id']
    all_samples['u_id'] = all_samples['u_id'].astype(int)

    i = int(all_samples['image_top_1_id'].max()+1)
    j = len(all_samples['category_name_id'].unique())
    all_samples['i_id'] = i * all_samples['category_name_id'] + all_samples['image_top_1_id']
    all_samples['i_id'] = all_samples['i_id'].astype(int)
    num_item = i*j

    ui_graph = all_samples[['u_id','i_id']].values

    u_i_matrix = np.zeros((num_user,num_item))
    print(u_i_matrix.shape)
    for i,j in ui_graph:
        u_i_matrix[i,j] +=1

    from scipy.sparse.linalg import svds
    [u,sigma,v] = svds(u_i_matrix,dims)
    v = v.T

    for i in range(dims):
        all_samples['user_vec_'+str(i)] = 0
        all_samples['item_vec_'+str(i)] = 0
    user_cols = ['user_vec_'+str(i) for i in range(dims)]
    item_cols = ['item_vec_'+str(i) for i in range(dims)]

    print('set vec')
    def set_vec(row):
        row[user_cols] = u[row['u_id']]
        row[item_cols] = v[row['i_id']]
        return row
    usecol = ['city', 'user_type', 'category_name', 'image_top_1']
    all_samples.drop_duplicates(usecol, inplace=True)
    all_samples = all_samples.apply(set_vec,axis=1)

    print('merge')
    all_samples = all_samples[usecol+user_cols+item_cols]
    train = train.merge(all_samples, on=usecol, how='left')
    test = test.merge(all_samples, on=usecol, how='left')

    for data,name in zip([train,test],['tr','te']):
        u = data[user_cols].values
        i = data[item_cols].values
        np.save('./dataset/'+name+'_u_id.csv',u)
        np.save('./dataset/'+name+'_i_id.csv',i)

def get_stat_features():

    cate_cols = [
        "city",
        "category_name",
        "param_1",
        "activation_date",
        "user_type",
        "user_id",
        "weekday",
    ]

    conti_cols = [
        "price",
        'item_seq_number',
    ]

    features = []
    for i in range(1, 3):
        for comb_feat in combinations(cate_cols, i):
            for conti_feat in conti_cols:
                features.append([conti_feat,list(comb_feat)])
    usecols = cate_cols + conti_cols + ['item_id']
    return usecols,features

def conti_stat_feature(train,test):

    l_train = len(train)
    l_test = len(test)

    usecols,features = get_stat_features()

    num_cpu = mlp.cpu_count()//2
    pool = mlp.Pool(num_cpu)
    num_task = 1+len(features)//num_cpu
    results = []

    for i in range(num_cpu):
        result = pool.apply_async(conti_stat_feature_worker,
                        args=(usecols,features[i*num_task:(i+1)*num_task]))
        results.append(result)
    pool.close()
    pool.join()

    for result in tqdm(results):
        train_feat,test_feat = result.get()
        train = pd.concat([train,train_feat],axis=1)
        test = pd.concat([test,test_feat],axis=1)

    assert len(train) == l_train
    assert len(test) == l_test

    return train,test

# def num_to_bin_feature(train,test):
#     samples = train[['price', 'item_seq_number']].append(test[['price', 'item_seq_number']]).reset_index(drop=True)
#
#     for num_bin,col in [()]
#     samples = num2bin(samples, ['price', 'item_seq_number'], 500)
#     train['price_bin'] = samples.loc[:len(train) - 1, 'price_bin']
#     test['price_bin'] = samples.loc[len(train):, 'price_bin']
#     train["item_seq_number_bin"] = samples.loc[:len(train) - 1, "item_seq_number_bin"]
#     test["item_seq_number_bin"] = samples.loc[len(train):, "item_seq_number_bin"]
#
#     return train, test

def target_encode(train,test):
    cate_cols = [
        "city",
        "category_name",
        "param_1",
        "activation_date",
        "user_type",
        "weekday",
        "image_top_1",
        # "price_bin",
        # "item_seq_number_bin"
    ]

    for col in tqdm(cate_cols):
        feat = condition_stat(train[[col,'deal_probability']],'deal_probability',col,col+'_TE_mean','mean')
        feat = feat[[col,col+'_TE_mean']].drop_duplicates([col]).reset_index(drop=True)
        train = train.merge(feat,how='left',on=[col])
        test = test.merge(feat, how='left', on=[col])

        feat = condition_stat(train[[col, 'deal_probability']], 'deal_probability', col, col +'_TE_std', 'std')
        feat = feat[[col,col+'_TE_std']].drop_duplicates([col]).reset_index(drop=True)
        train = train.merge(feat, how='left', on=[col])
        test = test.merge(feat, how='left', on=[col])
    return train,test

def oof_feature(train,test,alpha):
    NFOLDS = 5
    SEED = 42

    class SklearnWrapper(object):
        def __init__(self, clf, seed=0, params=None, seed_bool=True):
            if (seed_bool == True):
                params['random_state'] = seed
            self.clf = clf(**params)

        def train(self, x_train, y_train):
            self.clf.fit(x_train, y_train)

        def predict(self, x):
            return self.clf.predict(x)

    def get_oof(clf, x_train, y, x_test):

        from sklearn.cross_validation import KFold

        oof_train = np.zeros((len(y),))
        oof_test = np.zeros((len(test),))
        oof_test_skf = np.empty((NFOLDS,len(test)))

        kf = KFold(len(y), n_folds=NFOLDS, shuffle=True, random_state=SEED)
        for i, (train_index, test_index) in enumerate(kf):
            print('\nFold {}'.format(i))
            x_tr = x_train[train_index]
            y_tr = y[train_index]
            x_te = x_train[test_index]

            clf.train(x_tr, y_tr)

            oof_train[test_index] = clf.predict(x_te)
            oof_test_skf[i, :] = clf.predict(x_test)

        oof_test[:] = oof_test_skf.mean(axis=0)
        return oof_train.reshape(-1, 1), oof_test.reshape(-1, 1)

    print('load data')
    desc_matrix = np.load('./dataset/description.npy')
    titl_matrix = np.load('./dataset/title.npy')
    ready_df = np.concatenate((desc_matrix,titl_matrix),axis=1)

    print('train')
    ridge_params = {'alpha':alpha, 'fit_intercept': True, 'normalize': False, 'copy_X': True,
                    'max_iter': None, 'tol': 0.0001, 'solver': 'auto', 'random_state': SEED}
    y = train['deal_probability'].values
    ridge = SklearnWrapper(clf=Ridge, seed=SEED, params=ridge_params)
    ridge_oof_train, ridge_oof_test = get_oof(ridge, ready_df[:len(train)], y, ready_df[len(train):])

    from sklearn.metrics import mean_squared_error
    from math import sqrt
    rms = sqrt(mean_squared_error(y, ridge_oof_train))
    print('Ridge OOF RMSE: {}'.format(rms))

    train['ridge_preds'] = ridge_oof_train
    test['ridge_preds'] = ridge_oof_test

    return train,test

def generator(imgs,batchsize):
    from keras.preprocessing import image
    from keras.applications.vgg16 import preprocess_input

    iters = 1+len(imgs)//batchsize
    for i in tqdm(range(iters)):
        data = []
        for f in imgs[i*(batchsize):(i+1)*batchsize]:
            im = image.img_to_array(image.load_img(f,target_size=(256,256)))
            data.append(preprocess_input(im))

        yield {'input_1':np.array(data)}

def img_feat(train,test):
    import tensorflow as tf
    import keras.backend.tensorflow_backend as KTF
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True  # 不全部占满显存, 按需分配
    session = tf.Session(config=config)
    KTF.set_session(session)
    from keras.applications.vgg16 import VGG16
    from keras.models import Model
    from keras.layers import GlobalAveragePooling2D, GlobalMaxPool2D, concatenate
    from sklearn.decomposition import PCA

    out = VGG16(weights='imagenet',include_top=False)
    x = GlobalAveragePooling2D()(out.output)
    y = GlobalMaxPool2D()(out.output)
    z = concatenate([x,y])
    model = Model(inputs=out.input,outputs=z)

    train['image'].fillna('unk',inplace=True)
    train['image'] = './dataset/train_img/'+train['image']+'.jpg'
    test['image'].fillna('unk', inplace=True)
    test['image'] = './dataset/test_img/'+test['image']+'.jpg'
    train = train.append(test)
    imgs = train['image'].values
    del train,test
    print(len(imgs))
    print('predict')
    batchsize = 128
    steps = (len(imgs)+batchsize-1)//batchsize
    results = model.predict_generator(generator(imgs,batchsize=batchsize),steps=steps,verbose=1,use_multiprocessing=True)

    np.save('./dataset/img_feat.npy',results)

    pca = PCA(n_components=512)

    results = pca.fit_transform(results)
    np.save('./dataset/img_feat_demeni.npy', results)
    total_ratio = []
    ratio = 0
    for i in pca.explained_variance_ratio_:
        ratio += i
        total_ratio.append(ratio)
    print(total_ratio)

def get_word_vec(train,test,usecol):

    train = train[usecol]
    train_active = pd.read_csv('./dataset/train_active.csv', usecols=[usecol])
    test= test[usecol]
    test_active = pd.read_csv('./dataset/test_active.csv', usecols=[usecol])
    all_samples = pd.concat([
        train,
        train_active,
        test,
        test_active
    ]).fillna('unk').reset_index(drop=True)
    sentences = all_samples[usecol].values
    del train,train_active,test,test_active,all_samples
    gc.collect()

    from text_embedding import tokenize_word
    from gensim.models import Word2Vec
    from text_embedding import clean

    print('clean ')
    sentences = clean(sentences)
    num_seq = len(sentences)//4 + 1
    for i in range(4):
        sentences[i*num_seq:(i+1)*num_seq] = tokenize_word(sentences[i*num_seq:(i+1)*num_seq])

    print('train')
    model = Word2Vec(size=300,workers=32,iter=3,min_count=5)
    model.build_vocab(sentences)
    model.train(sentences,total_examples=model.corpus_count, epochs=3)  #
    del sentences
    print('save')
    model.save('./dataset/2'+usecol+'.model')
    del model
    print('success')
    gc.collect()

def pipeline():

    # train = pd.read_csv("./dataset/train.csv",parse_dates=["activation_date"])
    # test = pd.read_csv("./dataset/test.csv",parse_dates=["activation_date"])
    train = pd.read_csv("./dataset/train.csv",usecols=["description",'title'])
    test = pd.read_csv("./dataset/test.csv",usecols=["description",'title'])
    # train,test = baseFeature(train),baseFeature(test)
    # train,test = get_geo_feature(train),get_geo_feature(test)

    print('tfidf')
    # tfidf_feat(train, test, ['param_combined'],[1000,500],alg='ica')
    tfidf_feat(train,test,["description",'title'],[800,400],sparse=True,analyzer=['word','word'],gram=3)

    # get_word_vec(train,test,'description')
    # get_word_vec(train,test,'title')

    # img_feat(train,test)
    # train, test = agg_time_feat(train, test)

    # train,test = num_to_bin_feature(train,test)
    # train,test = target_encode(train,test)
    # get_img_info(train, 'tr')
    # get_img_info(test, 'te')
    # get_img_confi(train,'tr')
    # get_img_confi(test,'te')

    # train, test = oof_feature(train, test,20)

    # train, test = prob_feature(train, test)
    # train,test = conti_stat_feature(train,test)
    # user_item_matrix(train,test,100)

    # print('save')
    # train.to_csv("./dataset/train.csv",index=False)
    # test.to_csv("./dataset/test.csv",index=False)
    # print('save 2')
    # train.to_csv("./dataset/train.csv.gz", index=False,compression='gzip')
    # test.to_csv("./dataset/test.csv.gz", index=False,compression='gzip')

if __name__ == '__main__':
    pipeline()