import math

import ezdxf
from osgeo import gdal
import numpy as np


class Dem:
    def __init__(self, filepath):
        self._dataset = gdal.Open(filepath)

    def get_dem_info(self, if_print=False):
        """Get the information of DEM data.
        Parameters:
            dem_data <osgeo.gdal.Dataset> -- The data of DEM.
            if_print <bool> -- If print the information of DEM. Default is False (not print).
        Return:
            ... <...> -- The information and parameters of DEM.
        """
        dem_data = self._dataset
        dem_row = dem_data.RasterYSize  # height
        dem_col = dem_data.RasterXSize  # width
        dem_band = dem_data.RasterCount
        dem_gt = dem_data.GetGeoTransform()
        dem_proj = dem_data.GetProjection()

        if if_print:
            print("\nThe information of DEM:")
            print("The number of row (height) is: %d" % dem_row)
            print("The number of column (width) is: %d" % dem_col)
            print("The number of band is: %d" % dem_band)
            print("The 6 GeoTransform parameters are:\n", dem_gt)
            print("The GCS/PCS information is:\n", dem_proj)

        return dem_row, dem_col, dem_band, dem_gt, dem_proj

    def get_elevation(self, site_latlng):
        """Get the elevation of given locations from DEM in GCS.
        Parameters:
            dem_gcs <osgeo.gdal.Dataset> -- The input DEM (in GCS).
            site_latlng <numpy.ndarray> -- The latitude and longitude of given locations.
        Return:
            site_ele <numpy.ndarray> -- The elevation and other information of given locations.
        """
        gdal_data = self._dataset

        # gdal_band = gdal_data.GetRasterBand(1)
        # nodataval = gdal_band.GetNoDataValue()
        # if np.any(gdal_array == nodataval):
        #     gdal_array[gdal_array == nodataval] = np.nan

        gt = gdal_data.GetGeoTransform()
        # print("\nThe 6 GeoTransform parameters of DEM are:\n", gt)

        N_site = site_latlng.shape[0]
        Xgeo = site_latlng[:, 0]
        Ygeo = site_latlng[:, 1]
        site_ele = np.zeros(N_site)
        for i in range(N_site):
            # Note:
            #     Xgeo = gt[0] + Xpixel * gt[1] + Yline * gt[2]
            #     Ygeo = gt[3] + Xpixel * gt[4] + Yline * gt[5]
            #
            #     Xpixel - Pixel/column of DEM
            #     Yline - Line/row of DEM
            #
            #     Xgeo - Longitude
            #     Ygeo - Latitude
            #
            #     [0] = Longitude of left-top pixel
            #     [3] = Latitude of left-top pixel
            #
            #     [1] = + Pixel width
            #     [5] = - Pixel height
            #
            #     [2] = 0 for north up DEM
            #     [4] = 0 for north up DEM
            Xpixel = (Xgeo[i] - gt[0]) / gt[1]
            Yline = (Ygeo[i] - gt[3]) / gt[5]
            # 寻找左上，左下，右上，右下4个点
            lu_xy = np.array([math.floor(Xpixel), math.floor(Yline)])  # 左上
            ld_xy = np.array([math.floor(Xpixel), math.ceil(Yline)])  # 左下
            ru_xy = np.array([math.ceil(Xpixel), math.floor(Yline)])  # 右上
            rd_xy = np.array([math.ceil(Xpixel), math.ceil(Yline)])  # 右下
            lu_elevation = gdal_data.ReadAsArray(
                int(lu_xy[0]), int(lu_xy[1]), 1, 1
            ).astype(float)
            ld_elevation = gdal_data.ReadAsArray(
                int(ld_xy[0]), int(ld_xy[1]), 1, 1
            ).astype(float)
            ru_elevation = gdal_data.ReadAsArray(
                int(ru_xy[0]), int(ru_xy[1]), 1, 1
            ).astype(float)
            rd_elevation = gdal_data.ReadAsArray(
                int(rd_xy[0]), int(rd_xy[1]), 1, 1
            ).astype(float)
            # delta_x = (Xpixel - ld_xy[0]) / 1  # 距离是1
            # delta_y = (ld_xy[1]-Yline) / 1
            # site_ele[i] = (
            #     ld_elevation
            #     + (lu_elevation - ld_elevation) * delta_y
            #     + (rd_elevation - ld_elevation) * delta_x
            #     + (ld_elevation - lu_elevation + ru_elevation - rd_elevation)
            #     * delta_x
            #     * delta_y
            # )

            # 通过空间平面方程乃拟合Z值 参考《数字高程模型教程》 汤国安，李发源，刘学军编著 p89
            equation_a = np.array(
                [
                    [lu_xy[0], lu_xy[0] * lu_xy[1], lu_xy[1], 1],
                    [ld_xy[0], ld_xy[0] * ld_xy[1], ld_xy[1], 1],
                    [ru_xy[0], ru_xy[0] * ru_xy[1], ru_xy[1], 1],
                    [rd_xy[0], rd_xy[0] * rd_xy[1], rd_xy[1], 1],
                ]
            )
            equation_b = np.array(
                [lu_elevation[0], ld_elevation[0], ru_elevation[0], rd_elevation[0]]
            )
            equation = np.linalg.solve(equation_a, equation_b)
            point_z = (
                Xpixel * equation[0]
                + equation[1] * Xpixel * Yline
                + equation[2] * Yline
                + equation[3]
            )
            site_ele[i] = point_z
            print(f"row:{Yline} col:{Xpixel}  elev:{site_ele[i]}")
        return site_ele


def to_line_coordination(point_x_s, point_y_s, point_x_e, point_y_e):
    path_length = ((point_x_s - point_x_e) ** 2 + (point_y_s - point_y_e) ** 2) ** 0.5
    n = round(path_length / 1)
    center_point_x = np.linspace(point_x_s, point_x_e, n, endpoint=True)
    center_point_y = np.linspace(point_y_s, point_y_e, n, endpoint=True)
    # 计算左右边线
    # 计算角度
    side_width = 20  # 边线宽度
    _line_angel = math.atan((point_y_e - point_y_s) / (point_x_e - point_x_s))
    if point_x_e < point_x_s:
        line_angel = _line_angel + math.pi / 2
    else:
        line_angel = _line_angel
    left_offset_x = side_width * math.cos(line_angel + math.pi / 2)
    left_offset_y = side_width * math.sin(line_angel + math.pi / 2)
    left_offset_point_x = center_point_x + left_offset_x
    left_offset_point_y = center_point_y + left_offset_y
    right_offset_point_x = center_point_x - left_offset_x  # 向左偏移和向右偏移正好是相反的
    right_offset_point_y = center_point_y - left_offset_y
    # r = np.concatenate(
    #     (
    #         np.array(left_offset_point_x),
    #         np.array(left_offset_point_y),
    #         center_point_x,
    #         center_point_y,
    #         np.array(right_offset_point_x),
    #         np.array(right_offset_point_y),
    #     ),
    #     axis=0,
    # )
    r = np.array(
        [
            np.array(left_offset_point_x),
            np.array(left_offset_point_y),
            center_point_x,
            center_point_y,
            np.array(right_offset_point_x),
            np.array(right_offset_point_y),
        ]
    ).T
    return r


def distance(a, b):
    r = np.sum((a - b) ** 2) ** 0.5
    return r


if __name__ == "__main__":
    dem = Dem(r"d:\工程\金上线\DEM\四川-金上105-CGCS2000.tif")
    dem.get_dem_info(if_print=True)
    line_coordination = to_line_coordination(
        450077.359310936, 3304781.49970352, 451068.154964846, 3304604.32630216
    )
    left_elevation = dem.get_elevation(line_coordination[:, 0:2])
    center_elevation = dem.get_elevation(line_coordination[:, 2:4])
    right_elevation = dem.get_elevation(line_coordination[:, 4:6])
    doc = ezdxf.new(dxfversion="R2010")
    msp = doc.modelspace()
    x_axis = [0]
    cord_0 = line_coordination[0, :]
    for cord in line_coordination[1:]:
        x_axis.append(distance(cord, cord_0))
    msp.add_polyline2d(
        [(x_axis[i] / 5, left_elevation[i] * 2) for i in range(len(left_elevation))],
        dxfattribs={"color": 1},
    )  # 红色
    msp.add_polyline2d(
        [(x_axis[i] / 5, center_elevation[i] * 2) for i in range(len(center_elevation))]
    )
    msp.add_polyline2d(
        [(x_axis[i] / 5, right_elevation[i] * 2) for i in range(len(right_elevation))],
        dxfattribs={"color": 5},
    )  # 蓝色
    doc.saveas("d.dxf")
    print("Finished.")