from util import *
import numpy as np
from osgeo import gdal
import time
import argparse


@timeit
def interpolate_kriging_pykrige(x, y, z, proj, spatial_resolution, output_path=None):
    # https://pygis.io/docs/e_interpolation.html
    from pykrige.ok import OrdinaryKriging

    x_min = x.min()
    x_max = x.max()
    y_min = y.min()
    y_max = y.max()

    # Generate a regular grid
    step_x = int(x_max - x_min / spatial_resolution)
    step_y = int(y_max - y_min / spatial_resolution)
    grid_x = np.linspace(x_min, x_max, step_x)
    grid_y = np.linspace(y_min, y_max, step_y)

    # Create ordinary kriging object:
    OK = OrdinaryKriging(
        x,
        y,
        z,
        variogram_model="spherical",
        verbose=False,
        enable_plotting=False,
        coordinates_type="euclidean",
    )

    # Execute on grid:
    print("执行插值")
    start1 = time.perf_counter()

    z1, ss1 = OK.execute("grid", grid_x, grid_y, backend="loop", n_closest_points=12)

    end1 = time.perf_counter()
    print("插值，运行时间：", end1 - start1, "秒")

    if output_path is not None:
        format_str = "GTiff"
        driver = gdal.GetDriverByName(format_str)
        # driver.Create()函数中RasterXSize代表影像的sample（列数），RasterYSize代表影像的line（行数）
        dataset_out = driver.Create(output_path, step_x, step_y, 1, gdal.GDT_Float64)

        dataset_out.SetProjection(proj)
        geotransform_out = (x_min, spatial_resolution, 0, y_max, 0, -spatial_resolution)
        dataset_out.SetGeoTransform(geotransform_out)

        dataset_out.GetRasterBand(1).WriteArray(z1)

        del dataset_out


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="此程序基于反演结果进行插值。")

    parser.add_argument('-i1', '--pos_content_path', type=str, required=True, help='输入含量文件的路径')
    parser.add_argument('-i2', '--img_path', type=str, required=True, help='输入影像文件的路径')
    parser.add_argument('-i3', '--spatial_resolution', type=float, default=1.0, help='输出控件分辨率')
    parser.add_argument('-o', '--outpath', required=True, type=str, help='输出文件的路径')

    parser.add_argument('-v', '--verbose', action='store_true', help='启用详细模式')

    args = parser.parse_args()

    pos_content = np.loadtxt(args.pos_content_path)

    dataset = gdal.Open(args.img_path)
    im_proj = dataset.GetProjection()  # 地图投影信息

    interpolate_kriging_pykrige(pos_content[:, 0], pos_content[:, 1], pos_content[:, 2], im_proj, args.spatial_resolution, args.outpath)