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[:, 1]  # longitude
        Ygeo = site_latlng[:, 0]  # latitude
        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 = int(round((Xgeo[i] - gt[0]) / gt[1]))
            Yline = int(round((Ygeo[i] - gt[3]) / gt[5]))
            site_ele[i] = gdal_data.ReadAsArray(Xpixel, Yline, 1, 1).astype(np.float)
        return site_ele