import os
import numpy as np
import geopandas as gpd
import rasterio
from rasterio.transform import Affine
from shapely.geometry import Point, LineString
from shapely.ops import unary_union, linemerge
from skimage.graph import route_through_array
import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.colors import LinearSegmentedColormap

# ===== 用户参数配置 =====
cost_raster = r"D:\desktop\数据分析\5.15边界模拟\成本面1.tif"
points_shp = r"D:\desktop\数据分析\5.15边界模拟\FIND POIND.shp"
out_dir = r"D:\desktop\数据分析\5.15边界模拟"
out_path = os.path.join(out_dir, "clockwise_boundary.shp")
img_path = os.path.join(out_dir, "boundary_plot.png")
os.makedirs(out_dir, exist_ok=True)

# ===== 1. 数据加载 =====
with rasterio.open(cost_raster) as src:
    cost_array = src.read(1).astype(float)
    transform = src.transform
    crs = src.crs
    nodata = src.nodata

cost_array[(~np.isfinite(cost_array)) | (cost_array == nodata)] = np.inf

gdf = gpd.read_file(points_shp).to_crs(crs)
valid_points, coords = [], []
for geom in gdf.geometry:
    if geom.geom_type == 'Point' and np.isfinite(geom.x) and np.isfinite(geom.y):
        coords.append((geom.x, geom.y))
        valid_points.append(Point(geom.x, geom.y))

if len(valid_points) < 3:
    raise RuntimeError("至少需要3个有效点用于形成闭环")

# ===== 2. 顺时针排序闭合 =====
pts_arr = np.array(coords)
center = pts_arr.mean(axis=0)
angles = np.arctan2(pts_arr[:, 1] - center[1], pts_arr[:, 0] - center[0])
order = np.argsort(angles)[::-1]
ordered = [valid_points[i] for i in order]
if ordered[0] != ordered[-1]:
    ordered.append(ordered[0])

# ===== 3. 最小成本路径连接 =====
def world_to_pixel(x, y):
    col, row = ~transform * (x, y)
    return int(round(row)), int(round(col))

path_segs = []
for i in range(len(ordered) - 1):
    p0, p1 = ordered[i], ordered[i + 1]
    r0, c0 = world_to_pixel(p0.x, p0.y)
    r1, c1 = world_to_pixel(p1.x, p1.y)
    if not (0 <= r0 < cost_array.shape[0] and 0 <= c0 < cost_array.shape[1] and
            0 <= r1 < cost_array.shape[0] and 0 <= c1 < cost_array.shape[1]):
        continue
    if cost_array[r0, c0] == np.inf or cost_array[r1, c1] == np.inf:
        continue
    if (r0, c0) == (r1, c1):
        continue
    try:
        inds, _ = route_through_array(cost_array, (r0, c0), (r1, c1), fully_connected=True)
        coords_line = [transform * (c + 0.5, r + 0.5) for r, c in inds]
        path_segs.append(LineString(coords_line))
    except:
        continue

if not path_segs:
    raise RuntimeError("未能计算出任何路径段")

boundary = linemerge(unary_union(path_segs))
gpd.GeoDataFrame(geometry=[boundary], crs=crs).to_file(out_path)

# ===== 4. 可视化与保存图像 =====
try:
    font_path = "C:/Windows/Fonts/simhei.ttf"
    font_prop = font_manager.FontProperties(fname=font_path)
    plt.rcParams['font.sans-serif'] = [font_prop.get_name()]
    plt.rcParams['axes.unicode_minus'] = False
except:
    font_prop = None

xmin = transform.c
xmax = transform.c + transform.a * cost_array.shape[1]
ymin = transform.f + transform.e * cost_array.shape[0]
ymax = transform.f

masked = np.where(np.isinf(cost_array), np.nan, cost_array)

fig, ax = plt.subplots(figsize=(12, 10))
cmap = LinearSegmentedColormap.from_list('custom', ['#f7fcf5', '#e5f5e0', '#c7e9c0',
                                                     '#a1d99b', '#74c476', '#41ab5d',
                                                     '#238b45', '#006d2c', '#00441b'])
im = ax.imshow(np.log1p(masked), extent=(xmin, xmax, ymin, ymax), cmap=cmap, alpha=0.8)

# 输入点
xs, ys = zip(*coords)
ax.scatter(xs, ys, c='red', s=40, label='输入点', zorder=5)

# 标注编号
for i, (x, y) in enumerate(coords):
    ax.annotate(str(order[i % len(order)]), (x, y), textcoords="offset points",
                xytext=(0, 10), ha='center', fontsize=8, color='red', zorder=6)

# 绘制边界
if not boundary.is_empty:
    if boundary.geom_type == 'LineString':
        x_b, y_b = boundary.xy
        ax.plot(x_b, y_b, c='darkblue', lw=2.5, label='闭合边界', zorder=4)
    elif boundary.geom_type == 'MultiLineString':
        for i, line in enumerate(boundary.geoms):
            x_b, y_b = line.xy
            ax.plot(x_b, y_b, c='darkblue', lw=2.5, label='闭合边界' if i == 0 else "", zorder=4)

ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)
ax.legend()
ax.set_title("最小成本闭合边界（点按顺时针顺序连接）", fontproperties=font_prop)
ax.set_xlabel("X坐标")
ax.set_ylabel("Y坐标")

# 添加颜色条
cbar = fig.colorbar(im, ax=ax, pad=0.02, aspect=30)
cbar.set_label('Log(成本值+1)', labelpad=15)
cbar.ax.tick_params(labelsize=9)

plt.tight_layout()
plt.savefig(img_path, dpi=300)
plt.show()
print(f"图像已保存至: {img_path}")
print("处理完成！")