from osgeo import gdal, ogr, osr
import numpy as np
import otb_numpy_proc
import otbApplication as otb
import os

def pixels_to_vector(selection_array, otb_pipeline, out_vector, out_key='out', var_dict=None):
    geo_transform = otb_pipeline.GetImageMetaData(out_key)['GeoTransform']
    epsg = int(otb_pipeline.GetImageMetaData(out_key)['ProjectionRef'].split('"')[-2])
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(epsg)

    drv = ogr.GetDriverByName('ESRI Shapefile')
    ds = drv.CreateDataSource(out_vector)
    ly = ds.CreateLayer(os.path.basename(os.path.splitext(out_vector)[0]), geom_type=ogr.wkbPoint, srs=srs)
    ly.CreateField(ogr.FieldDefn('id', ogr.OFTInteger64))
    ly.CreateField(ogr.FieldDefn('px', ogr.OFTInteger64))
    ly.CreateField(ogr.FieldDefn('py', ogr.OFTInteger64))
    if var_dict is not None:
        for v in var_dict:
            if 'float' in str(v['array'].dtype):
                typ = ogr.OFTReal
            else:
                typ = ogr.OFTInteger
            if v['array'].ndim == 2:
                for t in range(v['array'].shape[1]):
                    ly.CreateField(ogr.FieldDefn(v['prefix']+'_%d' % t, typ))
            else:
                ly.CreateField(ogr.FieldDefn(v['prefix'], typ))

    i = 0
    pix_pnts = np.where(selection_array)
    for p in zip(pix_pnts[0]+0.5,pix_pnts[1]+0.5):
        X = geo_transform[0] + p[1] * geo_transform[1]
        Y = geo_transform[3] + p[0] * geo_transform[5]
        gp = ogr.Geometry(ogr.wkbPoint)
        gp.AddPoint(X,Y)
        gf = ogr.Feature(ly.GetLayerDefn())
        gf.SetGeometry(gp)
        gf.SetField('id', i)
        gf.SetField('px', p[1])
        gf.SetField('py', p[0])
        if var_dict is not None:
            for v in var_dict:
                if v['array'].ndim == 2:
                    for t in range(v['array'].shape[1]):
                        gf.SetField(v['prefix']+'_%d' % t, float(v['array'][i,t]))
                else:
                    gf.SetField(v['prefix'], float(v['array'][i]))
        ly.CreateFeature(gf)
        i += 1

    ds = None

if __name__ == '__main__':
    p = otb_numpy_proc.to_otb_pipeline('./pippo.tif')
    p.Execute()
    arr = np.random.random((100, 100)) < 0.1
    drv = ogr.GetDriverByName('ESRI Shapefile')
    drv.DeleteDataSource('./pippuzzo.shp')
    pixels_to_vector(arr, p, 'pippuzzo.shp')