from osgeo import gdal
from PIL import Image
import numpy as np
import os
import datetime

def print_with_timestamp(message):
    # 获取当前时间
    current_time = datetime.datetime.now()
    # 格式化时间戳
    timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S")
    # 打印带有时间戳的消息
    print(f"[{timestamp}] {message}")

def reproject_and_resample_tiff(tiff_path, output_tiff_path, target_epsg, factor):
    # 打开TIFF文件
    dataset = gdal.Open(tiff_path, gdal.GA_ReadOnly)
    if dataset is None:
        print_with_timestamp(f"无法打开TIFF文件: {tiff_path}")
        return

    # 获取TIFF文件的宽度和高度
    width = dataset.RasterXSize
    height = dataset.RasterYSize

    # 计算重采样后的宽度和高度
    new_width = int(width / factor)
    new_height = int(height / factor)

    # 设置重投影和重采样参数
    options = gdal.WarpOptions(
        format='GTiff',
        width=new_width,
        height=new_height,
        dstSRS=f'EPSG:{target_epsg}',
        resampleAlg=gdal.GRA_Bilinear
    )

    filename = os.path.basename(tiff_path)
    output_path = os.path.join(output_tiff_path,filename)

    # 执行重投影和重采样
    gdal.Warp(output_path, dataset, options=options)
    print_with_timestamp(f"已将TIFF文件重投影并重采样到: {output_path}")
    return output_path

def tiff_to_png(tiff_path, png_path, waterFlag, imageFlag):
    # 打开重投影和重采样后的TIFF文件
    dataset = gdal.Open(tiff_path, gdal.GA_ReadOnly)
    if dataset is None:
        print_with_timestamp(f"无法打开TIFF文件:{tiff_path}")
        return

    # 获取TIFF文件的宽度和高度
    width = dataset.RasterXSize
    height = dataset.RasterYSize

    # 获取TIFF文件的高度数据
    band = dataset.GetRasterBand(1)
    heights = band.ReadAsArray()
    # mask_a =  heights > -32766
    # heights = mask_a * heights

    # 最高高度,要比测试数据的最高处21高
    maxHeight = 23 

    # 创建一个空白的PNG图像
    image = Image.new("RGB", (width, height))
    # image = Image.new("L", (width, height))

    # 创建二维数组
    array = np.zeros((width, height), dtype="<i2")

    # 将每个像素的RGB值设置为高度的灰度值
    for y in range(height):
        for x in range(width):
            # 获取该位置的高度值
            height_value = heights[y, x]
            alpha = 255
            if height_value <= 0:
                height_value = 0  if waterFlag else maxHeight
            elif np.isnan(height_value):
                height_value = 0  if waterFlag else maxHeight
            else:
                height_value = height_value + 1.0  if waterFlag else height_value
            if imageFlag:
                # 将高度值映射到RGB值（灰度），最高值以30为准
                gray_value = int((height_value / maxHeight) * 255) 
                # 在PNG图像中设置像素的RGB值
                # image.putpixel((x, y), (gray_value, gray_value, gray_value, alpha))
                image.putpixel((x, y), (gray_value, gray_value, gray_value))
                # image.putpixel((x, y), gray_value)
            else:
                # 将数值精确到厘米，保存成二进制的二维数组, 高程及水面的厘米表示要控制在65535之内
                height_value = int(height_value * 100) # 精确到小数后两位，厘米
                array[x][y] = height_value

    # 保存PNG图像
    if imageFlag:
        image.save(png_path)
        print_with_timestamp(f"已将TIFF文件转换为PNG: {png_path}")
    else:
        # 将二维数组保存为二进制文件
        # 使用 'w' 模式写入文件，'b' 表示以二进制形式
        with open(png_path, 'wb') as f:
            # tobytes() 方法将数组元素按内存中的顺序转换成一个字节串
            f.write(array.tobytes())
        print_with_timestamp(f"已将TIFF文件转换为BIN: {png_path}")

def list_tif_files(directory,tif_files):

    # 遍历指定目录中的所有文件和子目录
    for root, dirs, files in os.walk(directory):
        for file in files:
            # 检查文件是否以.tif结尾
            if file.lower().endswith('.tif'):
                # 将完整路径添加到列表中
                path = os.path.join(root,file)
                print_with_timestamp(f"识别到tiff文件：{path}")
                tif_files.append(os.path.join(root, file))
    
    return tif_files 

def main(imageFlag):
    # 设置输入TIFF文件路径
    input_tiff_path = ".\\tiff28"
    # 设置重投影和重采样后的TIFF文件路径
    output_tiff_path = ".\\resample"
    # 设置目标EPSG代码
    target_epsg = 4326  # WGS84
    # 设置重采样因子
    factor = 10

    # 设置输出含水面高度图PNG文件路径
    output_water_png_path = ".\\waterImage"
    # 设置输出不含水面的地形、建筑高度图PNG文件路径
    output_terrain_png_path = ".\\terrainImage"

    # 创建一个空列表来存储所有的.tif文件
    tif_files = []

    # 遍历文件夹中的tiff
    list_tif_files(input_tiff_path, tif_files)

    for file in tif_files:
        print_with_timestamp(f"{file} 开始重投影")

        # 执行重投影和重采样
        out_path = reproject_and_resample_tiff(file, output_tiff_path, target_epsg, factor)

        print_with_timestamp(f"{out_path} 开始转换为png")
        # 使用os.path.basename方法获取路径中的文件名
        file_name = os.path.basename(out_path)
        # 使用os.path.splitext方法分离文件名和扩展名
        file_name_without_ext, _ = os.path.splitext(file_name)
        # 水面高度图路径
        out_water_png_file = os.path.join(output_water_png_path, file_name_without_ext)
        out_water_png_file = f"{out_water_png_file}.png" if imageFlag else f"{out_water_png_file}.bin"
        # 地形&建筑高度图路径
        out_terrain_png_file = os.path.join(output_terrain_png_path, file_name_without_ext)
        out_terrain_png_file = f"{out_terrain_png_file}.png" if imageFlag else f"{out_terrain_png_file}.bin"
        # 将重投影和重采样后的带水面的TIFF文件转换为高度图PNG
        tiff_to_png(out_path, out_water_png_file, True, imageFlag)
        # 将重投影和重采样后的不带水面的TIFF文件转换为高度图PNG
        # tiff_to_png(out_path, out_terrain_png_file, False, imageFlag)
    
if __name__ == "__main__":
    main(True)