#!/usr/bin/env python3

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# LAMP License 
# 
# Author: chenzenghui 
# time: 2023.06.04
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# 自定义dem管理
# 1. 创建二进制文件 ENVI hdr
# 2. 读取文件构建.vrt .xml

import argparse
import isce
from ctypes import cdll, c_char_p, c_int, byref
from array import array
import struct
import zipfile
import os
import sys
import math
import urllib.request, urllib.parse, urllib.error
from isce import logging
from iscesys.Component.Component import Component
from isceobj.Image import createDemImage
from osgeo import gdal,osr,ogr
import xml.etree.ElementTree as ET
from html.parser import HTMLParser
import time



class DEM2ISCE(Component):

    def dem_merged(self,in_dem_path, out_dem_path):
        '''
        DEM重采样函数，默认坐标系为WGS84
        agrs:
            in_dem_path: 输入的DEM文件夹路径
            meta_file_path: 输入的xml元文件路径
            out_dem_path: 输出的DEM文件夹路径
        '''
        # 读取文件夹中所有的DEM
        dem_file_paths=in_dem_path#[os.path.join(in_dem_path,dem_name)  for dem_name in    os.listdir(in_dem_path)  if dem_name.find(".tif")>=0 and dem_name.find(".tif.")==-1]
        spatialreference=osr.SpatialReference()
        spatialreference.SetWellKnownGeogCS("WGS84") # 设置地理坐标，单位为度 degree # 设置投影坐标，单位为度 degree
        spatialproj=spatialreference.ExportToWkt() # 导出投影结果
        # 将DEM拼接成一张大图
        out_DEM=out_dem_path
        gdal.Warp(out_DEM,
                dem_file_paths,
                format="ENVI",
                dstSRS=spatialproj,
                dstNodata=self._NoDataValue,
                outputType=gdal.GDT_Float32)
        time.sleep(3)
        return out_DEM

    #this method also create an actual DeimImage object that is returned by the getImage() method
    def createXmlMetadata(self,outname):
        demImage = self.createImage(outname)
        demImage.renderHdr()

    def getDemWidth(self,outname):
        gdal.AllRegister()
        dataset=gdal.Open(outname)
        width=dataset.RasterXSize
        del dataset
        return width
    
    def getDemHeight(self,outname):
        gdal.AllRegister()
        dataset=gdal.Open(outname)
        height=dataset.RasterYSize
        del dataset
        return height        
    
    def getGeotransform(self,outname):
        gdal.AllRegister()
        dataset=gdal.Open(outname)
        geotransform = dataset.GetGeoTransform()
        del dataset
        return geotransform
    
    def createImage(self,outname):
        demImage = createDemImage()
        width = self.getDemWidth(outname)
        height=self.getDemHeight(outname)
        demImage.initImage(outname,'write',width,type="float")
        length = demImage.getLength()
        # 获取分辨率
        geotransform=self.getGeotransform(outname)

        dictProp = {'METADATA_LOCATION':outname+'.xml','REFERENCE':self._reference,'Coordinate1':{'size':width,'startingValue':geotransform[0],'delta':geotransform[1]},'Coordinate2':{'size':length,'startingValue':geotransform[3],'delta':geotransform[5]},'FILE_NAME':outname}
        #no need to pass the dictionaryOfFacilities since init will use the default one
        demImage.init(dictProp)
        self._image = demImage
        return demImage    

    def setFillingValue(self,val):
        self._fillingValue = val

    def setNoDataValue(self,val):
        self._NoDataValue = val


    def stitchDems(self,source, outname):
        import glob
        # 合并数据
        self.dem_merged(source, outname)
        self.createXmlMetadata(outname)
    family = 'DEM2ISCE'
    def __init__(self,family = '', name = ''):
        self._extension = '.tif'
        self._zip = '.zip'
        #to make it working with other urls, make sure that the second part of the url
        #it's /srtm/version2_1/SRTM(1,3)
        self._filters = {'region1':['Region'],'region3':['Africa','Australia','Eurasia','Islands','America'],'fileExtension':['.hgt.zip']}
        self._remove = ['.jpg']
        self._metadataFilename = 'fileDem.dem'
        self._createXmlMetadata = None
        self._createRscMetadata = None
        self._regionList = {'1':[],'3':[]}
        ##self._keepDems = False
        self._fillingFilename = 'filling.hgt' # synthetic tile to cover holes
        ##self._fillingValue = -32768 # fill the synthetic tile with this value
        ##self._noFilling = False
        self._failed = 'failed'
        self._succeded = 'succeded'
        self._image = None
        self._reference = 'EGM96' 
        
        super(DEM2ISCE, self).__init__(family if family else  self.__class__.family, name=name)
        # logger not defined until baseclass is called
        if not self.logger:
            self.logger = logging.getLogger('isce.contrib.demUtils.DEM2ISCE')    
    
    def getImage(self):
        return self._image


# DEM转换主流程
def processDEM2ISCE(name,source_path,target_path,fillvalue,noDataValue):
    ds = DEM2ISCE(name=name)
    # 构建
    ds.setFillingValue(fillvalue)
    ds.setNoDataValue(noDataValue)
    ds.stitchDems(source_path,target_path)
    

def main():
    #if not argument provided force the --help flag
    if(len(sys.argv) == 1):
        sys.argv.append('-h')
    # Use the epilog to add usage examples
    epilog = '将格式为tif 的DEM 转换为ISCE 支持的DEM格式:\n\n'
    epilog += 'Usage examples:\n\n'
    epilog += 'DEM2ISCE.py -s /mnt/d/codestorage/isce2/青海省.tif  -o  /mnt/d/codestorage/isce2/青海省_wgs84 -fillvalue  -9999 -Nodata -9999\n\n'

    #set the formatter_class=argparse.RawDescriptionHelpFormatter otherwise it splits the epilog lines with its own default format
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,epilog=epilog)
    parser.add_argument('-s', '--source', type = str, default ="/mnt/d/codestorage/isce2/青海省.tif", dest = 'source_path', help = '输入dem,格式为tif')
    parser.add_argument('-o', '--outpath', type = str, default = '/mnt/d/codestorage/isce2/青海省_wgs84', dest = 'outpath', help = '输出isce 支持的DEM ')
    parser.add_argument('-fillvalue', '--fillvalue', type = float, default = -9999, dest = 'fillvalue', help = '空值填充')
    parser.add_argument('-Nodata', '--Nodata', type = float, default = -9999, dest = 'Nodatavalue', help = '无效值填充')
    args = parser.parse_args()
    processDEM2ISCE("DEM2ISCE",args.source_path,args.outpath,args.fillvalue,args.Nodatavalue)
    return -1 
    
if __name__ == '__main__':
    sys.exit(main())
