import sys
from .raster import Raster
import os
import numpy as np
from tqdm import tqdm
import json
import multiprocessing

from osgeo import gdal
from osgeo import ogr
from osgeo import osr
from osgeo import gdal_array
from osgeo import gdalconst

class Vector(Raster):
    def __init__(self, input_shp_path=''):

        print("# ---------------------------------------------------------------------------- #")
        print("#                                Vector Toolkit                                #")
        print("# ---------------------------------------------------------------------------- #")
        Raster.__init__(self)
        self.readshp(input_shp_path)
        self.ExportLayer = None

    def readshp(self, input_shp_path):
        if os.path.isfile(input_shp_path):
            filetype = input_shp_path.split(".")[-1]
            gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "NO")
            gdal.SetConfigOption("SHAPE_ENCODING", "")
            ogr.RegisterAll()
            self.DriveDict = {
                'shp': ogr.GetDriverByName('ESRI Shapefile'),
                'SHP': ogr.GetDriverByName('ESRI Shapefile'),
                'Geojson': ogr.GetDriverByName('GeoJSON'),
                'geojson': ogr.GetDriverByName('GeoJSON'),
                'pbf': ogr.GetDriverByName('MVT'),
                'mbtiles': ogr.GetDriverByName('MBTiles')
            }
            if filetype in self.DriveDict.keys():
                print("\033[0m# -----Valid vector format :\033[1;32m ", filetype)
                self.Driver = self.DriveDict[filetype]

            self.Input_path = input_shp_path
            self.DataSource = self.Driver.Open(self.Input_path)
            self.meta = self.DataSource.GetMetadata()
            
            print("\n\033[0m# ----------------------------- Meta Information ----------------------------- #")
            self.print_dict(self.meta)

            print("# ----------------------------- Meta Information ----------------------------- #\n")
            self.Description = self.DataSource.GetDescription()
            print('\033[0m# -----Description :\033[1;32m ', self.Description)
            assert self.DataSource is not None, '\n\n\nERROR-----' + \
                str(input_shp_path) + '  --- Invalid Input Shepefile\n\n\n'
            self.LayerCount = self.DataSource.GetLayerCount()
            print("\033[0m# -----LayerCount\033[1;32m :", self.LayerCount)
            self.LayerDict = {}
            for i in range(self.LayerCount):
                self.Layer = self.DataSource.GetLayer(i)
                print(
                    "\033[0m# -----Layer :\033[1;34m ",
                    i,
                    "\033[0m LayerName : \033[1;32m ",
                    self.Layer.GetName(),
                    self.Layer.GetGeometryColumn())
                self.LayerDict[self.Layer.GetName()] = self.Layer
            
            
            self.Srs = self.Layer.GetSpatialRef()
            self.Extent=self.Layer.GetExtent()
            print("\033[0m# -----Extent:\033[1;32m ",self.Extent)
            print("\033[0m# -----Alread Load:\033[1;32m ", input_shp_path,"\033[0m")
            print(
                "# -------------------------------- DEFINE DONE ------------------------------- #")

        else:
            print("# ----- \033[5;31m Warning: Vector Class init with wrong path or invalid vector file \nPath:",input_shp_path,"\033[0m\n")

    def getDefaultLayerbyName(self, name):
        """
        para:name string of layer name
        """
        self.defaultlayer = self.LayerDict[name]
        print("# -----Set Default Layer \033[1;32m|",name,"| : \033[0m",self.defaultlayer)
        return self.LayerDict[name]

    def Info(self):
        print(self.DataSource)
        print(self.Layer)
        print(self.Srs)

    def getbounding(self):
        adfGeoTransform = self.geotransform
        x1 = adfGeoTransform[0] + 0 * \
            adfGeoTransform[1] + 0 * adfGeoTransform[2]
        y1 = adfGeoTransform[3] + 0 * \
            adfGeoTransform[4] + 0 * adfGeoTransform[5]
        x2 = adfGeoTransform[0] + self.width * \
            adfGeoTransform[1] + self.height * adfGeoTransform[2]
        y2 = adfGeoTransform[3] + self.width * \
            adfGeoTransform[4] + self.height * adfGeoTransform[5]
        self.bounding = (x1, y1, x2, y2)
        return self.bounding

    def getdataset(self, tif_path):
        self.tif_path = tif_path
        self.dataset = gdal.Open(tif_path)
        print('\n# -----Dataset:', self.dataset)
        return self.dataset

    def SaveTo(self, name):
        filetype = name.split(".")[-1]
        if filetype in self.DriveDict.keys():
            out = self.DriveDict[filetype].CopyDataSource(
                self.DataSource, name)
            out = None

    def SaveVectorByLayerName(self, LayerName, outputname, format="GeoJSON"):
        """
        format:
            * ESRI Shapefile
            * GeoJSON
        """
        self.ExportLayer = self.LayerDict[LayerName]
        print("# ------------------Start Create Layer",
              LayerName, ' ---------------------- #')
        self.ExportDriver = ogr.GetDriverByName(format)

        if os.path.isfile(outputname):
            self.ExportDriver.DeleteDataSource(outputname)

        ExportResources = self.ExportDriver.CreateDataSource(outputname)
        print("-----ExportResources:", ExportResources)
        assert not ExportResources is None, "Invalid Export Resources ,Please check the path of output"
        ExportLayerTemp = ExportResources.CreateLayer(
            LayerName, geom_type=ogr.wkbPolygon)
        print("-----Data Resources Create Done")

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

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

        # Add features to the ouput Layer
        print("-----Start Loop Process")
        for i in tqdm(range(0, self.ExportLayer.GetFeatureCount())):
            # Get the input Feature
            inFeature = self.ExportLayer.GetFeature(i)
            # 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()
            inFeature = None
            centroid = geom.Centroid()
            outFeature.SetGeometry(centroid)
            # Add new feature to output Layer
            ExportLayerTemp.CreateFeature(outFeature)
            outFeature = None

        print("# -----LayerCopyDone,", LayerName)
        ExportResources.SetMetadata(self.meta)
        print("# -----Set Meta Done")
        ExportResources.SetDescription(self.Description)
        print("# -----Set Description Done")
        ExportResources.Destroy()
        print("# --------------------------------- Save Done -------------------------------- #")

        # Save and close DataSources
        inDataSource = None
        outDataSource = None

    def crop_layer_by_polygon(self, vector):
        self.ExportLayer.SetSpatialFilter()

    def crop_layer_by_rect(self, rect):
        assert self.ExportLayer is None, 'Invalid Export Layer'
        self.ExportLayer.SetSpatialFilterRect()

    
    def crop_default_layer_by_rect(self, rect):
        print("# -----Set filter Rect:\033[1;32m",rect,'\033[0m')
        self.defaultlayer.SetSpatialFilterRect(*rect)
    
    def print_dict(self,d,n=0):
        length=67
        for k,v in d.items():
            # print ('\t'*n)
            if type(v)==type({}):
                print("%s : {" % k)
                self.print_dict(v,n+1)
            else:
                
                strl=len(str(k))+len(str(v))
                space=length-strl
                if strl>length:                    
                    v=str(v)[:space]
                print("\033[0m# -----\033[1;34m %s :\033[1;32m  %s\033[0m" % (k,v)+" "*space+"#")
        if n!=0:
            print('\t'*(n-1)+ '}')
    
    def Rasterize(self, outputname, Nodata=0):

        targetDataSet = gdal.GetDriverByName('GTiff').Create(
            outputname, self.width, self.height, 1, gdal.GDT_Byte)
        targetDataSet.SetGeoTransform(self.geotransform)
        targetDataSet.SetProjection(self.projection)
        band = targetDataSet.GetRasterBand(1)
        band.SetNoDataValue(Nodata)
        band.FlushCache()
        gdal.RasterizeLayer(
            targetDataSet,
            [1],
            self.defaultlayer,
            )
        targetDataSet = None
        
    def generate(self,tiles,output_path="./label"):
        print('# ===== \033[5;33mStart Generate.....\033[0m')
        """
        修订  把Rasterize同样使用线程分解的可能性
        """

        self.labellist=[]
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        for tile in tqdm(tiles):
            self.readtif(tile)
            filename=tile.split('/')[-1]
            path=os.path.join(output_path,filename)
            self.labellist.append(path)
            self.Rasterize(path)
        return self.labellist
    
    def reset_layer(self):
        self.Layer.ResetReading()

    # def Json2Voclike(self, out_dir='./VocDataset'):
    #     """

    #     :param label_file:
    #     label.txt & the content like:
    #     __ignore__
    #     _background_
    #     Iron
    #     Mine
    #     None
    #     land

    #     The Content '__ignore__' maybe should add in file by manual
    #     :param in_dir:
    #     Dir input Data like :
    #     in_dir-
    #         |
    #         1.jpg
    #         1.json
    #         2.jpg
    #         2.json
    #         ....
    #     :param out_dir:
    #     output Voclike Dataset Path
    #     :return:
    #     """
    #     print('**********************strart voc-like  process******************')
    #     if os.path.exists(out_dir):
    #         os.removedirs(out_dir)
    #     os.makedirs(out_dir)
    #     os.makedirs(os.path.join(out_dir, 'JPEGImages'))
    #     os.makedirs(os.path.join(out_dir, 'SegmentationClass'))
    #     os.makedirs(os.path.join(out_dir, 'SegmentationClassPNG'))
    #     os.makedirs(os.path.join(out_dir, 'SegmentationClassVisualization'))

    #     class_names = []
    #     class_name_to_id = {}
    #     # for i, line in enumerate(open(labels_file).readlines()):
    #     #     class_id = i - 1  # starts with -1
    #     #     class_name = line.strip()
    #     #     print('current class name :', class_name)
    #     #     class_name_to_id[class_name] = class_id
    #     #     if class_id == -1:
    #     #         assert class_name == '__ignore__'
    #     #         continue
    #     #     elif class_id == 0:
    #     #         assert class_name == '_background_'
    #     #     class_names.append(class_name)
    #     assert self.TYPE is not None, 'None type list Please Set First'
    #     class_names = tuple(self.TYPE)
    #     self.classlist = class_names
    #     self.outputdir = out_dir
    #     print('class_names:', class_names)

    #     out_class_names_file = os.path.join(out_dir, 'class_names.txt')
    #     with open(out_class_names_file, 'w') as f:
    #         f.writelines('\n'.join(class_names))
    #     print('Saved class_names:', out_class_names_file)

    #     colormap = labelme.utils.label_colormap(255)
    #     in_dir = self.LabelmeJsonPath
    #     for label_file in glob.glob(os.path.join(in_dir, '*.json')):
    #         print('Generating dataset from:', label_file)
    #         with open(label_file) as f:
    #             base = os.path.splitext(os.path.basename(label_file))[0]
    #             out_img_file = os.path.join(
    #                 out_dir, 'JPEGImages', base + '.jpg')
    #             out_lbl_file = os.path.join(
    #                 out_dir, 'SegmentationClass', base + '.npy')
    #             out_png_file = os.path.join(
    #                 out_dir, 'SegmentationClassPNG', base + '.png')
    #             out_viz_file = os.path.join(
    #                 out_dir, 'SegmentationClassVisualization', base + '.jpg')

    #             data = json.load(f)

    #             img_file = os.path.join(os.path.dirname(label_file), data['imagePath'])
    #             img = np.asarray(PIL.Image.open(img_file))
    #             PIL.Image.fromarray(img).save(out_img_file)

    #             lbl = labelme.utils.shapes_to_label(
    #                 img_shape=img.shape,
    #                 shapes=data['shapes'],
    #                 label_name_to_value=class_name_to_id,
    #             )
    #             labelme.utils.lblsave(out_png_file, lbl)

    #             np.save(out_lbl_file, lbl)

    #             viz = labelme.utils.draw_label(
    #                 lbl, img, class_names, colormap=colormap)
    #             PIL.Image.fromarray(viz).save(out_viz_file)

   # def Shp2LabelmeJson(self, output_GeoJson=True,
    #                     Labelme_Json_path='./Segmentation/',
    #                     invert2Voc=True, FieldClassName='TYPE'):
    #     """
    #     invert the Shapefile to labelme_Json
    #     Flow I  : Read Tif & Shapefile -> Dataset & DataSource
    #     Flow II : Init Null labelme_json dict
    #     Flow III: Get Cords(x,y) & Feature and invert longtitude,latitude to image cols,rows
    #     Flow IV : Append image Cords list to dict
    #     Flow V  : Output Labelme Json
    #     Flow VI : Use System Command labelme_json_to_dataset / labelme2voc.py
    #             # It generates:
    #             #   - data_dataset_voc/JPEGImages
    #             #   - data_dataset_voc/SegmentationClass
    #             #   - data_dataset_voc/SegmentationClassVisualization
    #     :param shp_path:
    #     :param tif_path:
    #     :param Labelme_Json_path: Output Json&jpg path
    #     :param invert2Dataset: invert labelme json to Dataset
    #     :param  TYPE like (['Storge yard', 'Container', 'Oil Tank', 'Berth'])
    #     :return:
    #     """

    #     print('-----feature count :', self.Layer.GetFeatureCount())
    #     print("-----SpatialFilter :", self.Layer.GetSpatialFilter())

    #     print("-----", self.Layer.GetLayerDefn())

    #     print("-----", self.Layer.GetFeatureCount())

    #     print("-----", self.Layer.GetExtent())

    #     x1, y1, x2, y2 = self.Layer.GetExtent()
    #     x1, y1 = self.geo2imagexy(x1, y1)
    #     x2, y2 = self.geo2imagexy(x2, y2)
    #     print("-----Image:", x1, y1, x2, y2)
    #     width = abs(x2 - x1)
    #     height = abs(y2 - y1)
    #     print("-----Whole Image shape:", height, width)

    #     print('-----', self.Layer[0])

    #     print("-----", self.Layer.GetSpatialRef())
    #     exit(0)

    #     # for i in tqdm(range(self.Layer.GetFeatureCount())):
    #     #     in_Feature = self.Layer.GetFeature(i)
    #     #     geom = in_Feature.GetGeometryRef()
    #     #     print(type(geom))
    #     #     print(geom)
    #     # FieldValue=in_Feature.GetField(FieldClassName)
    #     # geodict=geom.ExportToJson()
    #     # geodict=

    #     # geodict=json.loads(geodict)

    #     # feature=self.initcordarrary(self.TYPE[FieldValue],imagecord)

    #     # shapes.append(feature)
    #     # print('\n\nshape is :',shapes,' process done \n\n')
    #     # return 0
    #     ####################################################################
    #     # shapes process done
    #     if(not os.path.exists(Labelme_Json_path)):
    #         os.system('mkdir ' + Labelme_Json_path)
    #     filename = filename.split('.')[-2]
    #     self.LabelmeJsonPath = Labelme_Json_path
    #     outputjsonname = Labelme_Json_path + filename + '.json'
    #     self.save(outputjsonname, shapes, filename + '.jpg', imageData=None,
    #               lineColor=(255, 0, 0, 128), fillColor=(0, 255, 0, 128))

    #     Newim = Image.fromarray(self.tif.image_nparray)
    #     Newim.save(Labelme_Json_path + filename + '.jpg')
    #     print('\n\n****json   &   image save done *****\n\n')

    #     if output_GeoJson:
    #         geojsonfilename = self.Input_path + '.geojson'
    #         formattransfrom_command = 'ogr2ogr -f GeoJSON ' + \
    #             geojsonfilename + ' ' + self.Input_path
    #         print(formattransfrom_command)
    #         os.system(formattransfrom_command)
    #         print(' output geojson transfrom done ')

    #     if invert2Voc:
    #         self.Json2Voclike()

    #         # command='labelme_json_to_dataset '+outputjsonname
    #         # os.system(command)
    #     #
    #     # if invertPascalVoc:
    #     #     imageDataPIL.save(Labelme_Json_path+filename+'.jpg')

 
def main():
    Vector = SHP("/workspace/data/Water/Beijing.geojson")
    # Vector=SHP("/workspace/SQCV/Data/IO/waterchina.geojson")

    # ------------------------------ Rasterize Demo ------------------------------ #

    # Vector=SHP("/home/winshare/Downloads/2017-07-03_asia_china.mbtiles").SaveTo("CHINA.shp")
    # Vector.SaveVectorByLayerName('water',"waterchina.shp",format="ESRI Shapefile")

    import glob

    tiflist = glob.glob("/workspace/data/Water/GoogleBeijing/*.tif")
    # tiflist.extend(glob.glob("/workspace/data/Water/BingBeijing/*.tif"))

    # Experiment--Geojson
    # ----------------------------------- Speed ---------------------------------- #
    #                               十万样本时间
    # CHINA             34          94  小时
    # Beijing           1/13        120 分钟
    # MBtilesFilter     3           8.3 小时

    # Experiment--Shapefile
    # ----------------------------------- Speed ---------------------------------- #

    print(len(tiflist))
    for file in tqdm(tiflist):
        Vector.readtif(file)
        Vector.getDefaultLayerbyName("bui")

        filename = file.split('/')[-1]
        path = '/workspace/data/Water/label/' + filename
        Vector.Rasterize(path)


if __name__ == '__main__':
    main()
