import sys
sys.path.append('core')
import os
import torch
import numpy as np
from pathlib import Path
from osgeo import gdal  # 新增GDAL
from core.igev_stereo import IGEVStereo
from core.utils.utils import InputPadder
import argparse
from tqdm import tqdm
import imageio
DEVICE = 'cuda'
#from feature_match.eloftr.match_image import get_init_disp
#from eloftr.match_image import get_init_disp,init_matcher
from PIL import Image
import cv2
# 取外扩 50 个像素


def load_image_block(dataset, x_offset, y_offset, x_size, y_size):
    """从GDAL数据集中读取图像块"""
    band = dataset.GetRasterBand(1)
    img_block = band.ReadAsArray(x_offset, y_offset, x_size, y_size).astype(np.uint8)  # 读取指定位置的图像块
    if len(img_block.shape) == 2:  # 单通道转三通道
        img_block = np.stack([img_block] * 3, axis=-1)
    img_block = torch.from_numpy(img_block).permute(2, 0, 1).float()  # 转换为 PyTorch 张量
    return img_block[None].to(DEVICE)


def get_disp_img_pyramid(model, image1_path, image2_path,disp_init, output_directory):
    # 使用 GDAL 打开图像
    dataset1 = gdal.Open(image1_path)
    dataset2 = gdal.Open(image2_path)
    # 获取图像大小
    width = dataset1.RasterXSize
    height = dataset1.RasterYSize
    disp_map = np.zeros((height, width), dtype=np.float32)
    weight_map = np.zeros((height, width), dtype=np.float32)
    block_sizes = 512 #128 #512 #128
    

    w_block_number = width // block_sizes + 1 
    h_block_number = height // block_sizes + 1 
    
    for i in tqdm(range(h_block_number)):
        for j in range(w_block_number):
            #计算左图原始位置
            l_w_start = j*block_sizes
            l_w_end = min(width,l_w_start+block_sizes)
            l_h_start = i*block_sizes
            l_h_end = min(height,l_h_start+block_sizes)
            #计算右图的原始位置
            r_w_start = j*block_sizes
            r_w_end = min(width,r_w_start+block_sizes)
            r_h_start = i*block_sizes
            r_h_end = min(height,r_h_start+block_sizes)
            w_expand_pixels = 128 #512  #320
            h_expand_pixels = 32 
            # 原始 block_init_disp 计算
            block_init_disp = np.nanmean(disp_init[l_h_start:l_h_end, l_w_start:l_w_end])
            if abs(block_init_disp)>256:
                w_expand_pixels += int(abs(block_init_disp)/2)
                h_expand_pixels += 64
            # 如果整个区域都是 NaN，可以选择默认值，例如 0
            if np.isnan(block_init_disp):
                block_init_disp = 0
            else:
                block_init_disp = int(block_init_disp)
            #计算右图初始位置
            skew_r_w_start = r_w_start-block_init_disp
            skew_r_w_end = r_w_end-block_init_disp
            skew_disp = 0
            #如果超出范围，我就不管了
            if skew_r_w_start>0 and skew_r_w_end < width:
                r_w_start = skew_r_w_start
                r_w_end = skew_r_w_end
                skew_disp = block_init_disp
            

            
            
            l_w_start_adjusted = max(0, l_w_start - w_expand_pixels)
            l_w_end_adjusted = min(width, l_w_end + w_expand_pixels)
            l_h_start_adjusted = max(0, l_h_start - h_expand_pixels)
            l_h_end_adjusted = min(height, l_h_end + h_expand_pixels)
            
            r_w_start_adjusted = max(0, r_w_start - w_expand_pixels)
            r_w_end_adjusted = min(width, r_w_end + w_expand_pixels)
            r_h_start_adjusted = max(0, r_h_start - h_expand_pixels)
            r_h_end_adjusted = min(height, r_h_end + h_expand_pixels)
            
            #根据左图的外扩计算右图然后也得根据右图能不能外扩取计算左图
            map_left_expand = min(l_w_start-l_w_start_adjusted,r_w_start-r_w_start_adjusted)
            map_right_expand = min(l_w_end_adjusted - l_w_end,r_w_end_adjusted-r_w_end)
            map_up_expand = min(l_h_start-l_h_start_adjusted,r_h_start-r_h_start_adjusted)
            map_down_expand = min(l_h_end_adjusted - l_h_end,r_h_end_adjusted-r_h_end)
            
            
            l_w_start_adjusted = l_w_start-map_left_expand
            l_w_end_adjusted = l_w_end+map_right_expand
            l_h_start_adjusted = l_h_start-map_up_expand
            l_h_end_adjusted = l_h_end+map_down_expand
            
            r_w_start_adjusted = r_w_start - map_left_expand
            r_w_end_adjusted = r_w_end+map_right_expand
            r_h_start_adjusted = r_h_start - map_up_expand
            r_h_end_adjusted = r_h_end+map_down_expand

            
            image1_patch = load_image_block(dataset1, l_w_start_adjusted, l_h_start_adjusted, l_w_end_adjusted - l_w_start_adjusted, l_h_end_adjusted-l_h_start_adjusted)
            image2_patch = load_image_block(dataset2, r_w_start_adjusted, r_h_start_adjusted, r_w_end_adjusted - r_w_start_adjusted, r_h_end_adjusted-r_h_start_adjusted)


            padder = InputPadder(image1_patch.shape, divis_by=32)
            image1_patch, image2_patch = padder.pad(image1_patch, image2_patch)
            
            
            disp_patch = model(image1_patch, image2_patch, iters=12, test_mode=True)
            disp_patch = disp_patch.detach().cpu()
            disp_patch = padder.unpad(disp_patch).numpy()[0,0]

            #根据前面的外扩，再取出非外扩的值
            up_expand = l_h_start - l_h_start_adjusted
            down_expand = l_h_end_adjusted - l_h_end 
            left_expand = max(0,l_w_start - l_w_start_adjusted)
            right_expand = max(0,l_w_end_adjusted - l_w_end)
            
            # 去除外扩部分
            disp_path_non_expanded = None
            if(right_expand==0):
                if(down_expand==0):
                    disp_path_non_expanded = disp_patch[
                        up_expand:, 
                        left_expand:
                    ]
                    l_w_end = l_w_end_adjusted
                else:
                    disp_path_non_expanded = disp_patch[
                        up_expand:-down_expand, 
                        left_expand:
                    ]
                    l_w_end = l_w_end_adjusted

            else:
                if(down_expand==0):
                    disp_path_non_expanded = disp_patch[
                    up_expand:,
                    left_expand:-right_expand
                ]
                else:
                    disp_path_non_expanded = disp_patch[
                    up_expand:-down_expand, 
                    left_expand:-right_expand
                ]
            disp_map[l_h_start:l_h_end, l_w_start:l_w_end] += (disp_path_non_expanded+skew_disp)

            print(f"{l_h_start}_{l_h_end}_{l_w_start}_{l_w_end}")
    file_stem = image1_path.split('\\')[-2]
    filename = os.path.join(output_directory, f"{file_stem}_18.tif")
    imageio.imwrite(filename, disp_map.squeeze())
    # 当数据集不再需要时关闭它们
    dataset1 = None
    dataset2 = None        



def load_image(imfile):
        # 打开图像文件
    img = None
    img = Image.open(imfile)
    
    # 计算原始图像的宽度和高度
    original_width, original_height = img.size
    
    scale_size = 3096
    # 确定缩放比例
    scale = min(scale_size / original_width, scale_size / original_height)
    
    # 如果需要缩放，则进行缩放
    if scale < 1:
        # 计算新的尺寸
        new_width = int(original_width * scale)
        new_height = int(original_height * scale)
        # 等比例缩放图像
        img = img.resize((new_width, new_height))
    
    # 将图像转换为 numpy 数组
    
    img = np.array(img).astype(np.uint8)
    
    img = np.stack([img]*3, axis=-1)
    img = torch.from_numpy(img).permute(2, 0, 1).float()
    return img[None].to(DEVICE),scale

def get_init_disp(left,right,model):
    test = True
    if test:
        # "D:\jx\ai\stereo_match_dl\IGEV-main\IGEV-main\IGEV-Stereo\core\demo-output\test_patch\SV1-02_20190416_L1B0000976744_8041900187280009_01-PAN2_epi_right_half.tif_1.tif"
        disp_resized_path = r"D:\jx\ai\stereo_match_dl\IGEV-main\IGEV-main\IGEV-Stereo\core\demo-output\test_patch\result_resize.tif"
        disp_resized = imageio.imread(disp_resized_path)
        return disp_resized
    else:  
        image1,scale = load_image(left)
        image2,scale = load_image(right)
        padder = InputPadder(image1.shape, divis_by=32)
        image1, image2 = padder.pad(image1, image2)
        disp = model(image1, image2, iters=args.valid_iters, test_mode=True)
        disp = disp.detach().cpu().numpy()
        disp = padder.unpad(disp)
        #再根据scale还原回去
        disp_resized = cv2.resize(disp.squeeze(), (int(disp.shape[3] / scale), int(disp.shape[2] / scale)), interpolation=cv2.INTER_LINEAR)/scale
        return disp_resized

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--restore_ckpt', help="restore checkpoint", default=r"D:\jx\ai\stereo_match_dl\IGEV-main\IGEV-main\IGEV-Stereo\ckpt\41481_epoch_igev-stereo.pth")
    # parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default=r"D:\jx\ai\stereo_match_dl\test_image\bak\SV1-02_20190416_L1B0000976744_8041900187280009_01-PAN2_epi.tif")
    # parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default=r"D:\jx\ai\stereo_match_dl\test_image\bak\SV1-02_20190416_L1B0000976744_8041900187280009_01-PAN1_epi.tif")
    parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default=r"D:\jx\ai\stereo_match_dl\test_image\bak\SV1_PAN1_down_half_resized.tif")
    parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default=r"D:\jx\ai\stereo_match_dl\test_image\bak\SV1_PAN2_down_half_resized.tif")
    
    
    # parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default=r"D:\jx\ai\stereo_match_dl\test_image\right\3_6_right.tif")
    # parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default=r"D:\jx\ai\stereo_match_dl\test_image\left\3_6_left.tif")
    #parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default=r"D:\jx\ai\stereo_match_dl\train_image\experimental data\with ground truth\test\left\HY_left_0.tiff")
    #parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default=r"D:\jx\ai\stereo_match_dl\train_image\experimental data\with ground truth\test\right\HY_right_0.tiff")
    
    parser.add_argument('--output_directory', help="directory to save output", default=r"D:\jx\ai\stereo_match_dl\IGEV-main\IGEV-main\IGEV-Stereo\core\demo-output\test_patch")
    parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
    # Architecture choices
    parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions")
    parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation")
    parser.add_argument('--shared_backbone', action='store_true', help="use a single backbone for the context and feature encoders")
    parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
    parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
    parser.add_argument('--n_downsample', type=int, default=2, help="resolution of the disparity field (1/2^K)")
    parser.add_argument('--slow_fast_gru', action='store_true', help="iterate the low-res GRUs more frequently")
    parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
    parser.add_argument('--max_disp', type=int, default=192, help="max disp of geometry encoding volume")
    args = parser.parse_args()
    Path(args.output_directory).mkdir(exist_ok=True, parents=True)
    # 初始化模型
    model = torch.nn.DataParallel(IGEVStereo(args), device_ids=[0])
    model.load_state_dict(torch.load(args.restore_ckpt))
    model = model.module
    model.to(DEVICE)
    model.eval()
    result = get_init_disp(args.left_imgs,args.right_imgs,model)

    get_disp_img_pyramid(model, args.left_imgs,args.right_imgs, result ,args.output_directory)
    
    
我这个，拼缝，有点明显，你帮我修改一下，加一个权重矩阵？