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


def get_spectral_in_coor(imgpath, coorpath, outpath):
    coor_spectral = np.loadtxt(coorpath, delimiter="\t")

    dataset = gdal.Open(imgpath)
    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    num_bands = dataset.RasterCount  # 栅格矩阵的波段数
    geotransform = dataset.GetGeoTransform()  # 仿射矩阵
    im_proj = dataset.GetProjection()  # 地图投影信息

    new_columns = np.zeros((coor_spectral.shape[0], num_bands))
    coor_spectral = np.hstack((coor_spectral, new_columns))

    for i in range(coor_spectral.shape[0]):
        pixel_coor_x = int(coor_spectral[i, 2])
        pixel_coor_y = int(coor_spectral[i, 3])
        coor_spectral[i, 4:coor_spectral.shape[1]] = \
            dataset.ReadAsArray(int(pixel_coor_x), int(pixel_coor_y), 1, 1).reshape(num_bands)

    del dataset

    header = "image_coor_x\timage_coor_y\tpixel_coor_x\tpixel_coor_y\trest of colums are spectral"

    in_hdr_dict = spectral.envi.read_envi_header(get_hdr_file_path(imgpath))
    wavelengths = np.array(in_hdr_dict['wavelength']).astype('float64')
    new_row = np.zeros((1, coor_spectral.shape[1]))
    new_row[:, -num_bands:] = wavelengths
    coor_spectral = np.vstack((new_row, coor_spectral))

    np.savetxt(outpath, coor_spectral, fmt='%.4f', delimiter="\t", header=header)
    # position_content_restore = np.loadtxt(outpath)

    return coor_spectral


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="此程序获取给定坐标的光谱曲线。")

    parser.add_argument('-i1', '--imgpath', type=str, required=True, help='输入影像文件的路径')
    parser.add_argument('-i2', '--coorpath', type=str, required=True, help='输入坐标文件的路径')
    parser.add_argument('-o', '--output_path', type=str, required=True, help='输出csv文件的路径')

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

    args = parser.parse_args()

    tmp = get_spectral_in_coor(args.imgpath, args.coorpath, args.output_path)