# -*- coding: utf-8 -*-

"""
@File    : 02_栅格HDF文件展示.py
@Author  : fungis@163.com
@notice  : https://zhuanlan.zhihu.com/p/536879741
"""
import os
import numpy as np
import matplotlib
from matplotlib import cm
import matplotlib.pylab as plt  # 画图模块

raster_path = './data-use/asc/waterbodyid.asc'  # 输入你的asc文件

if os.path.exists(raster_path) is False:
    raise Exception('[Errno 2] 该文件不存在: \'' + raster_path + '\'')

num = 6  # 读取前6行（如果是5行的数据，这里修改为5）
metadata_list = []  # 元数据列表
with open(raster_path, 'r') as rows_raster_path:
    for i in range(num):
        tem_row = str(rows_raster_path.readline()).strip()  # 读取两端剔除空格的字符串
        list_str = tem_row.split(" ")
        item_tem = list_str[-1]  # 取list的最后一个元素，即相应的头文件值
        metadata_list.append(item_tem)
# 根据读取结果取值
ncols = int(metadata_list[0])
nrows = int(metadata_list[1])
xllcorner = float(metadata_list[2])
yllcorner = float(metadata_list[3])
cellsize = float(metadata_list[4])  # 像元大小
NODATA_value = float(metadata_list[5])  # 空值默认值
xltcorner = xllcorner
yltcorner = yllcorner + cellsize * nrows
# 转换位置信息更改为asc文件制定的脚点
# 坐标信息
tran = (-180.0, 0.5, 0.0, 90.0, 0.0, -0.5)
new_tran = list(tran)
new_tran[0] = xltcorner
new_tran[3] = yltcorner
geo_trans_form = new_tran
lonmin = geo_trans_form[0]
latmin = geo_trans_form[3]
lonmax = geo_trans_form[0] + nrows * geo_trans_form[1] + ncols * geo_trans_form[2]
latmax = geo_trans_form[3] + nrows * geo_trans_form[4] + ncols * geo_trans_form[5]
extent = [lonmin, lonmax, latmin, latmax]

# 矩阵提取
location_array = np.loadtxt(raster_path, skiprows=6)
row = location_array.shape[0]  # 矩阵的行数
columns = location_array.shape[1]  # 矩阵的列数

plt.title('gdal-ASC-show')  # 设置图框标题为gdal-hdf-show

# 间隔色块
cmap = matplotlib.colors.ListedColormap(
    ['#81ecec', '#00FF00', '#00b894', '#6666FF', '#0984e3', '#0000FF', '#ffeaa7', '#FFFF00', '#fab1a0', '#ff7675',
     '#FF0000'])

newcmp = cm.get_cmap('rainbow', 6)

plt.imshow(location_array, extent=extent, interpolation='nearest', cmap=newcmp)

# plt.imshow(location_array, cmap='rainbow')  # 绘图，配色为rainbow
plt.colorbar(label='color bar')  # 绘制图例
plt.savefig("./results/20230314.01.png", dpi=300)
plt.show()