# coding=utf-8
#author:        4N
#createtime:    2020/6/28
#email:         nheweijun@sina.com

from osgeo import ogr
from osgeo.ogr import DataSource,Driver,Layer,Feature


from osgeo.ogr import *
from osgeo import gdal,ogr
import uuid
import os
import math
import time

import threading
import copy
def get_info_from_sqlachemy_uri(uri):
    parts = uri.split(":")
    user = parts[1][2:]

    password_list = parts[2].split("@")
    if password_list.__len__() > 2:
        password = "@".join(password_list[:-1])
    else:
        password = parts[2].split("@")[0]
    host = parts[2].split("@")[-1]
    port = parts[3].split("/")[0]
    database = parts[3].split("/")[1]

    return user, password, host, port, database


def open_pg_data_source(iswrite, uri):
    """
    # 获取PostGIS数据源
    :return:
    """
    db_conn_tuple = get_info_from_sqlachemy_uri(uri)
    fn = "PG: user=%s password=%s host=%s port=%s dbname=%s " % db_conn_tuple
    driver = ogr.GetDriverByName("PostgreSQL")
    if driver is None:
        raise Exception("打开PostgreSQL驱动失败，可能是当前GDAL未支持PostgreSQL驱动！")
    ds = driver.Open(fn, iswrite)
    if ds is None:
        raise Exception("打开数据源失败！")
    return ds




def vacuate(extent: tuple, g,max_level,fill_dict,vacuate_layers, vacuate_layers_gridsize):

    for level in range(max_level):

        center: Geometry = g.Centroid()

        this_grid_len = vacuate_layers_gridsize[level]

        row = int((center.GetY() - extent[2]) / this_grid_len)
        col = int((center.GetX() - extent[0]) / this_grid_len)
        key = "{}.{}.{}".format(level,row, col)


        if not fill_dict.get(key):
            fill_dict[key] = 0
        if fill_dict[key] == 0:
            vacuate_layer = vacuate_layers.get(level)
            feat = ogr.Feature(vacuate_layer.GetLayerDefn())
            feat.SetGeometry(g)
            vacuate_layer.CreateFeature(feat)
            # if level==0:
            #     vacuate_layer.CreateFeature(feature)
        fill_dict[key] += 1


def change_geom_type(raw):
    if raw.__eq__(-2147483646):
        return 5
    if raw.__eq__(-2147483645) or raw.__eq__(-2147483642):
        return 6
    if raw==2 or raw ==3:
        return raw+3
    if raw==4:
        return 1
    return raw

def get_table_type(raw):
    if raw==4 or raw ==5 or raw ==6:
        return raw-3
    return raw

def change_geom(geo:Geometry,geom_type):
    '''
    转换空间对象的类型，以适应dmap只支持Multi类型
    :param geo:
    :param geom_type:
    :return: 转换后的空间对象
    '''

    if geom_type==5 or geom_type.__eq__(-2147483646):
        return ogr.ForceToMultiLineString(geo)
    if geom_type==6 or geom_type.__eq__(-2147483645):
        return ogr.ForceToMultiPolygon(geo)
    if geom_type==1:
        # 多点转单点，会有问题，只拿了第一个点
        xy = geo.GetPoint()
        point = ogr.Geometry(ogr.wkbPoint)
        point.AddPoint(xy[0], xy[1])
        return point
    return geo


def insert(data_path,new_layer_name,sqlalchemy_uri,out,num):

    work_dir =os.path.dirname(os.path.abspath(__file__))


    #打开图层
    driver: Driver = ogr.GetDriverByName("OpenFileGDB")
    if data_path.endswith("shp"):
        driver: Driver = ogr.GetDriverByName("ESRI Shapefile")
    ds: DataSource = driver.Open(data_path, 0)
    if not ds:
        raise Exception("打开数据失败！")
    layer: Layer = ds.GetLayer(0)


    #创建PG图层
    pg_ds:DataSource= open_pg_data_source(1,sqlalchemy_uri)
    geom_type =change_geom_type(layer.GetGeomType())
    print(geom_type)

    pg_layer: Layer = pg_ds.CreateLayer(new_layer_name, layer.GetSpatialRef(),geom_type,["GEOMETRY_NAME=g","OVERWRITE=YES"])
    pg_layer.CreateFields(layer.schema)



    #设置抽稀级别
    layer_count = layer.GetFeatureCount()
    max_level = 0
    while layer_count>50000:
        max_level+=1
        layer_count/=2
    print(max_level)

    pg_ds_dict={}
    #创建抽稀ds
    for l in range(max_level):
        pg_ds_l: DataSource = open_pg_data_source(1, sqlalchemy_uri)
        pg_ds_dict[l]=pg_ds_l

    #建立抽稀图层
    lc = layer.GetFeatureCount()
    extent = layer.GetExtent()
    vacuate_layers = {}
    vacuate_layers_gridsize = {}

    density=2

    for l in range(max_level):
        pg_ds_l=pg_ds_dict[l]
        vl = pg_ds_l.CreateLayer("{}_vacuate_{}".format(new_layer_name, l), layer.GetSpatialRef(),geom_type, ["GEOMETRY_NAME=g","OVERWRITE=YES"])
        # vl.CreateFields(layer.schema)
        vacuate_layers[l] = vl
        need_object = (lc * density / (2 ** (l + 1)))
        this_grid_ysize = math.sqrt(need_object / ((extent[1] - extent[0]) / (extent[3] - extent[2])))
        this_grid_len = (extent[3] - extent[2]) / this_grid_ysize
        print(this_grid_len)
        vacuate_layers_gridsize[l]=this_grid_len

    count=0
    fill_dict = dict()


    for feature in layer:
        count += 1
        if count==num:
            break

        out_feature: Feature = copy.copy(feature)
        out_geom=None
        geo: Geometry = feature.GetGeometryRef()
        # 如果是空对象不录入
        if geo is not None:
            if geo.IsEmpty():
                continue


        if geo is not None:
            out_geom: Geometry = change_geom(geo, geom_type)
            out_feature.SetGeometry(out_geom)

        pg_layer.CreateFeature(out_feature)

        if geo is not None:
            vacuate(extent, out_geom, max_level, fill_dict, vacuate_layers, vacuate_layers_gridsize)

        if count % 10000 == 0:
            print(count)
            with open(os.path.join(work_dir,out),"w") as fi:
                fi.write("已完成{}".format(count))

    pg_ds.Destroy()
    for pg in pg_ds_dict.values():
        pg.Destroy()
    ds.Destroy()


if __name__ == '__main__':
    data_path = r"E:\D2\桌面\FSSD_RES_PY_FWM\FSSD_RES_PY_FWM.shp"
    new_layer_name = "fs1"
    sqlalchemy_uri="postgresql://postgres:chinadci@172.26.99.173:5433/postgres"
    out="va.txt"
    begin=time.time()
    num=400
    insert(data_path,new_layer_name,sqlalchemy_uri,out,num)
    print(time.time()-begin)