def benchmark_v1_v2():
    """
    对比旧版 erosion_cal_id 和新版 erosion_cal_id_v2 的运行时间。

    使用 example_cal_weight_xu 中 method=3 的数据来源，
    方便直接在本项目环境中复现。
    """

    print("开始 benchmark：v1 vs v2")

    # 读取示例数据
    f = Dataset('erosion_Post_ww3.nc')
    lons = f.variables['longitude'][900:1200].data
    lats = f.variables['latitude'][300:500].data
    viss = f.variables['swh'][:, 300:500, 900:1200].data
    f.close()

    value = viss[0, :, :]

    print(f"网格尺寸: lat={len(lats)}, lon={len(lons)}")

    # ---------------- v1 ----------------
    t0 = time.perf_counter()

    weight_v1 = erosion_cal_id(
        lons,
        lats,
        value,
        16,
        3,
        methond='np.nan',
        _print=False
    )

    t1 = time.perf_counter()

    # ---------------- v2 ----------------
    t2 = time.perf_counter()

    weight_v2 = erosion_cal_id_v2(
        lons,
        lats,
        value,
        16,
        3,
        methond='np.nan',
        backend='hybrid',
        _print=False
    )

    t3 = time.perf_counter()

    print("\n========== Benchmark 结果 ==========")
    print(f"v1 erosion_cal_id    : {t1 - t0:.4f} 秒")
    print(f"v2 erosion_cal_id_v2 : {t3 - t2:.4f} 秒")

    # 简单检查结果规模
    print("\n结果规模对比:")
    print(f"v1 change 点数: {len(weight_v1['change_lat'])}")
    print(f"v2 change 点数: {len(weight_v2['change_lat'])}")

    if len(weight_v1['change_lat']) != len(weight_v2['change_lat']):
        print("⚠️ v1 与 v2 change 点数量不同，请检查算法逻辑")
    else:
        print("✓ change 点数量一致")

    return {
        'v1_time': t1 - t0,
        'v2_time': t3 - t2,
        'v1': weight_v1,
        'v2': weight_v2
    }

#!/bin/env python3
# -*- coding: utf-8 -*-
#  日期 : 2026/3/9 14:44
#  作者 : Christmas
#  邮箱 : 273519355@qq.com
#  项目 : Project
#  版本 : python 3
#  摘要 :
"""

"""

import itertools
import json
import time
from collections import deque

import matplotlib.pyplot as plt
import numpy as np
import scipy.spatial as spt
import xarray as xr
from netCDF4 import Dataset
from scipy.io import loadmat

# numba 是可选依赖：没有安装时自动回退到 NumPy / KDTree 版本
try:
    import numba as nb
except Exception:
    nb = None


__all__ = [
    'find_index',
    'erosion_cal_id',
    'erosion_cal_id_v2',
    'judge_nan_ncmask',
    'judge_nan_npnan',
    'judge_nan_txtmask',
    'erosion_via_id',
    'example_cal_weight',
    'test',
    'check_draw',
    'example_cal_weight_xu',
]

# ---------------------- HPC 版本辅助函数 ----------------------

def _judge_keep_mask_vectorized(sequences, value, methond, k, judge_num):
    """
    使用 NumPy 向量化判断每一行近邻是否满足“可参与腐蚀”的条件。
    sequences: shape = [N, k]，每一行是一组近邻点的一维索引
    返回: shape = [N] 的 bool 数组，True 表示该行保留
    """
    if methond == 'np.nan':
        flat_value = np.asarray(value).ravel()
        test_arr = np.isnan(flat_value[sequences])
        keep = np.sum(test_arr, axis=1) <= (k - judge_num)
    elif methond == 'nc.mask':
        flat_mask = np.asarray(value.mask).ravel()
        test_arr = flat_mask[sequences]
        keep = np.sum(test_arr, axis=1) <= (k - judge_num)
    elif methond == 'txt.mask':
        flat_value = np.asarray(value).ravel()
        test_arr = flat_value[sequences]
        keep = np.sum(test_arr, axis=1) >= judge_num
    else:
        raise ValueError(f"不支持的 methond: {methond}")
    return keep


if nb is not None:
    @nb.njit(cache=True, fastmath=True, parallel=True)
    def _judge_keep_mask_npnan_numba(sequences, flat_value, k, judge_num):
        """Numba 版本：np.nan 模式判断。"""
        n = sequences.shape[0]
        out = np.zeros(n, dtype=np.bool_)
        threshold = k - judge_num
        for i in nb.prange(n):
            cnt = 0
            for j in range(k):
                if np.isnan(flat_value[sequences[i, j]]):
                    cnt += 1
            if cnt <= threshold:
                out[i] = True
        return out


    @nb.njit(cache=True, fastmath=True, parallel=True)
    def _judge_keep_mask_ncmask_numba(sequences, flat_mask, k, judge_num):
        """Numba 版本：nc.mask 模式判断。"""
        n = sequences.shape[0]
        out = np.zeros(n, dtype=np.bool_)
        threshold = k - judge_num
        for i in nb.prange(n):
            cnt = 0
            for j in range(k):
                if flat_mask[sequences[i, j]]:
                    cnt += 1
            if cnt <= threshold:
                out[i] = True
        return out


    @nb.njit(cache=True, fastmath=True, parallel=True)
    def _judge_keep_mask_txtmask_numba(sequences, flat_value, k, judge_num):
        """Numba 版本：txt.mask 模式判断。"""
        n = sequences.shape[0]
        out = np.zeros(n, dtype=np.bool_)
        for i in nb.prange(n):
            cnt = 0.0
            for j in range(k):
                cnt += flat_value[sequences[i, j]]
            if cnt >= judge_num:
                out[i] = True
        return out


def _judge_keep_mask_hpc(sequences, value, methond, k, judge_num, use_numba=False):
    """
    HPC 版统一 judge 入口。
    - use_numba=True 且环境安装了 numba 时：使用 numba kernel
    - 否则：使用 NumPy 向量化版本
    """
    if use_numba and (nb is not None):
        if methond == 'np.nan':
            flat_value = np.asarray(value).ravel().astype(np.float64, copy=False)
            return _judge_keep_mask_npnan_numba(sequences, flat_value, k, judge_num)
        elif methond == 'nc.mask':
            flat_mask = np.asarray(value.mask).ravel().astype(np.bool_, copy=False)
            return _judge_keep_mask_ncmask_numba(sequences, flat_mask, k, judge_num)
        elif methond == 'txt.mask':
            flat_value = np.asarray(value).ravel().astype(np.float64, copy=False)
            return _judge_keep_mask_txtmask_numba(sequences, flat_value, k, judge_num)
        else:
            raise ValueError(f"不支持的 methond: {methond}")
    else:
        return _judge_keep_mask_vectorized(sequences, value, methond, k, judge_num)


def _gpu_knn_bruteforce_batches(point, query_points, k, batch_size=2048):
    """
    使用 cupy 在 GPU 上分批 brute-force 计算 k 近邻。
    说明：
    1. 这是通用实现，不依赖 GPU 版 KDTree；
    2. 通过分批避免一次性构造超大距离矩阵导致显存爆炸；
    3. 返回值 shape = [N, k]，元素为一维扁平索引。
    """
    import cupy as cp

    cp_point = cp.asarray(point)
    n_all = point.shape[0]
    kth = min(max(k - 1, 0), n_all - 1)
    seq_list = []

    for start in range(0, query_points.shape[0], batch_size):
        end = min(start + batch_size, query_points.shape[0])
        cp_query = cp.asarray(query_points[start:end])

        # 计算平方距离，避免额外 sqrt 开销
        diff_x = cp_point[None, :, 0] - cp_query[:, None, 0]
        diff_y = cp_point[None, :, 1] - cp_query[:, None, 1]
        dist2 = diff_x * diff_x + diff_y * diff_y

        # 先 argpartition 取前 k 个，再在这 k 个内部精排
        seq = cp.argpartition(dist2, kth=kth, axis=1)[:, :k]
        part_dist2 = cp.take_along_axis(dist2, seq, axis=1)
        order = cp.argsort(part_dist2, axis=1)
        seq = cp.take_along_axis(seq, order, axis=1)

        seq_list.append(cp.asnumpy(seq))

        del cp_query, diff_x, diff_y, dist2, seq, part_dist2, order
        cp.get_default_memory_pool().free_all_blocks()

    return np.concatenate(seq_list, axis=0)


def erosion_cal_id_v2(
        g_lon,
        g_lat,
        value,
        k,
        judge_num,
        _print=False,
        methond='np.nan',
        backend='hybrid',
        leafsize=64,
        gpu_batch_size=2048,
):
    """
    真正的 HPC 终极版 erosion_cal_id。

    设计目标：
    1. 保留原算法语义：返回 change/value 所需的索引字典；
    2. 只对“真正需要处理”的格点查找近邻，而不是对全部格点做 KNN；
    3. 支持多后端：
       - kdtree : 纯 SciPy cKDTree
       - hybrid : KDTree 查邻 + Numba/NumPy 向量化 judge（推荐）
       - numba  : 当前实现等价于 hybrid，强调 judge 阶段优先走 numba
       - gpu    : cupy 分批 brute-force；失败自动回退到 hybrid
       - auto   : 若安装 numba 则走 hybrid，否则走 kdtree
    4. 保持规则网格下的一维扁平索引顺序：idx = lat_id * lon_number + lon_id

    参数说明与旧版基本一致，只额外增加：
    :param backend: 计算后端
    :param leafsize: cKDTree 的叶节点大小
    :param gpu_batch_size: GPU 分批查询的 batch 大小
    """
    lat_number = len(g_lat)
    lon_number = len(g_lon)

    if k <= 0:
        raise ValueError('k 必须大于 0')
    if judge_num < 0:
        raise ValueError('judge_num 不能小于 0')

    # auto 模式：有 numba 就优先 hybrid，否则走纯 kdtree
    if backend == 'auto':
        backend = 'hybrid' if (nb is not None) else 'kdtree'

    # 先构造“哪些点需要被处理”的扁平 mask
    if methond == 'np.nan':
        flat_target_mask = np.isnan(np.asarray(value)).ravel()
    elif methond == 'nc.mask':
        flat_target_mask = np.asarray(value.mask).ravel()
    elif methond == 'txt.mask':
        flat_target_mask = (np.asarray(value) == 0).ravel()
    else:
        raise ValueError(f"不支持的 methond: {methond}")

    valid_idx = np.flatnonzero(flat_target_mask)

    if _print:
        print(f'待处理格点数: {valid_idx.size}')

    # 如果没有需要处理的点，直接返回空结果
    if valid_idx.size == 0:
        return {
            'change_lat': [],
            'change_lon': [],
            'value_lat': [],
            'value_lon': [],
        }

    # 建立规则网格对应的二维坐标点（KDTree / GPU 共用）
    if _print:
        print('开始建立二维坐标点')
    lon_mesh, lat_mesh = np.meshgrid(g_lon, g_lat)
    point = np.c_[lon_mesh.ravel(), lat_mesh.ravel()]
    query_points = point[valid_idx]

    if _print:
        print(f'backend = {backend}')

    # =========================
    # 第一步：获取每个待处理格点的 k 个近邻索引（包含其自身）
    # sequences.shape = [N, k]
    # =========================
    if backend in ('kdtree', 'hybrid', 'numba'):
        if _print:
            print('开始建立 cKDTree 并查询近邻')
        ckt = spt.cKDTree(point, leafsize=leafsize)
        _, sequences = ckt.query(query_points, k=k)
        if k == 1:
            sequences = sequences[:, None]

    elif backend == 'gpu':
        try:
            if _print:
                print('开始使用 GPU 分批 brute-force 查询近邻')
            sequences = _gpu_knn_bruteforce_batches(
                point=point,
                query_points=query_points,
                k=k,
                batch_size=gpu_batch_size,
            )
        except Exception as e:
            if _print:
                print(f'GPU 后端不可用，自动回退到 hybrid。原因: {e}')
            ckt = spt.cKDTree(point, leafsize=leafsize)
            _, sequences = ckt.query(query_points, k=k)
            if k == 1:
                sequences = sequences[:, None]
            backend = 'hybrid'

    else:
        raise ValueError(f"不支持的 backend: {backend}")

    # 保证后续索引计算稳定
    sequences = np.asarray(sequences, dtype=np.int64)
    if sequences.ndim == 1:
        sequences = sequences[:, None]

    # =========================
    # 第二步：对每一行近邻进行 judge，决定哪些点真正参与腐蚀
    # - kdtree: 走 NumPy 向量化 judge
    # - hybrid/numba: 优先走 numba judge，没有 numba 再回退 NumPy
    # - gpu: 邻居已用 GPU 算完，judge 仍在 CPU 侧高效完成
    # =========================
    use_numba = backend in ('hybrid', 'numba')
    if _print:
        print(f'开始 judge，use_numba = {use_numba and (nb is not None)}')

    keep_mask = _judge_keep_mask_hpc(
        sequences=sequences,
        value=value,
        methond=methond,
        k=k,
        judge_num=judge_num,
        use_numba=use_numba,
    )

    # 只保留满足条件的行
    selected_sequences = sequences[keep_mask]
    selected_change_idx = valid_idx[keep_mask]

    if _print:
        print(f'judge 后保留的格点数: {selected_sequences.shape[0]}')

    if selected_sequences.shape[0] == 0:
        return {
            'change_lat': [],
            'change_lon': [],
            'value_lat': [],
            'value_lon': [],
        }

    # =========================
    # 第三步：直接构造最终 weight
    # 说明：
    # 1. selected_change_idx 对应“要被修改的点”；
    # 2. selected_sequences[:, 0] 理论上通常就是其自身，但这里直接用 valid_idx 更稳妥；
    # 3. selected_sequences[:, 1:] 对应参与平均赋值的邻居。
    # =========================
    change_lat = (selected_change_idx // lon_number).tolist()
    change_lon = (selected_change_idx % lon_number).tolist()

    # 当 k == 1 时，没有“其它邻点”可用于赋值，这里保持空列表结构
    if k == 1:
        value_lat = [[] for _ in range(selected_sequences.shape[0])]
        value_lon = [[] for _ in range(selected_sequences.shape[0])]
    else:
        value_sequences = selected_sequences[:, 1:]
        value_lat = (value_sequences // lon_number).tolist()
        value_lon = (value_sequences % lon_number).tolist()

    weight = {
        'change_lat': change_lat,
        'change_lon': change_lon,
        'value_lat': value_lat,
        'value_lon': value_lon,
    }
    return weight


def find_index(g_lon, g_lat, _s_lon, _s_lat):
    """
    获取经纬度索引
    :param g_lon: 一维格点经度
    :param g_lat: 一维格点纬度
    :param _s_lon: 散点经度
    :param _s_lat: 散点纬度
    :return: 纬度索引 经度索引
    """
    _s_lon = np.array(_s_lon)
    _s_lat = np.array(_s_lat)
    
    _latli = np.where(g_lat == _s_lat[:, None])[-1]
    _lonli = np.where(g_lon == _s_lon[:, None])[-1]
    
    return _latli.tolist(), _lonli.tolist()


def erosion_cal_id(
        g_lon,
        g_lat,
        value,
        k,
        judge_num,
        _print=False,
        methond='np.nan',
        backend='kdtree',
):
    """
    找到离指定散点最近的K个格点
    :param g_lon: 一维格点经度
    :param g_lat: 一维格点纬度
    :param value: 二维格点值
    :param k: 寻找近邻点的个数
    :param judge_num: 周边水点数 > judge_num的点，改为水点
    :param _print: 是否打印输出
    :param methond: 判断是否为nan的方法
                    np.nan：直接判断value里是否为nan
                    nc.mask：判断value的mask属性里是否为True
                    txt.mask：判断value里是否为0，0为水点，非0为陆点
    :param backend: 计算后端选择
                    kdtree  -> 原始 KDTree 最近邻（默认）
                    numba   -> CPU SIMD 并行版本
                    hybrid  -> KDTree 查邻 + numba 加速筛选
                    gpu     -> 使用 cupy GPU (若不可用自动退回 kdtree)
    :return: 格点索引id
    """
    lat_number = len(g_lat)
    lon_number = len(g_lon)
    # 生成格点坐标，放进一个二维平面
    if _print:
        print('开始建立Kdtree')
    # 使用 numpy meshgrid 直接生成网格点，比 python 循环快很多
    lon_mesh, lat_mesh = np.meshgrid(g_lon, g_lat)
    point = np.column_stack((lon_mesh.ravel(), lat_mesh.ravel()))
    # 不再使用 Python list/deque 逐个 append
    # 后面会直接通过 NumPy 向量化计算得到所有邻点索引
    
    # =========================
    # 计算后端选择
    # backend:
    #   kdtree  -> 原始 KDTree 最近邻（默认）
    #   numba   -> CPU SIMD 并行版本
    #   hybrid  -> KDTree 查邻 + numba 加速筛选
    #   gpu     -> 使用 cupy GPU (若不可用自动退回 kdtree)
    # =========================

    if backend == 'kdtree' or backend == 'hybrid':
        # 原始 KDTree 方法
        # leafsize 调大一点在大网格（几十万~百万点）时会更快
        ckt = spt.cKDTree(point, leafsize=64)

    if _print:
        print(f'Kdtree建立完毕，开始筛选需要计算的格点')

    # 构造所有格点对应的二维索引
    idx_all = np.arange(point.shape[0])
    lat_i_all = idx_all // lon_number
    lon_i_all = idx_all % lon_number

    # 根据不同方法生成 mask，只保留需要处理的格点
    if methond == 'np.nan':
        mask = np.isnan(value[lat_i_all, lon_i_all])
    elif methond == 'nc.mask':
        mask = value.mask[lat_i_all, lon_i_all]
    elif methond == 'txt.mask':
        mask = (value[lat_i_all, lon_i_all] == 0)
    else:
        mask = np.zeros_like(lat_i_all, dtype=bool)

    # 只取需要处理的点（通常远少于全部格点）
    valid_idx = idx_all[mask]

    if _print:
        print(f'需要计算近邻的格点数: {valid_idx.size}')

    # =========================
    # 后端实现
    # =========================

    if backend == 'kdtree':
        # 传统 KDTree
        query_points = point[valid_idx]
        distances, sequences = ckt.query(query_points, k)

        neigh = sequences.reshape(-1)

        lat_j = neigh // lon_number
        lon_j = neigh % lon_number

        lat_id = lat_j.tolist()
        lon_id = lon_j.tolist()

    elif backend == 'hybrid':
        # KDTree + NumPy 向量化筛选
        query_points = point[valid_idx]
        distances, sequences = ckt.query(query_points, k)

        neigh = sequences.reshape(-1)

        lat_j = neigh // lon_number
        lon_j = neigh % lon_number

        lat_id = lat_j.tolist()
        lon_id = lon_j.tolist()

    elif backend == 'numba':
        # CPU SIMD 版本（规则网格邻域扫描）
        radius = int(np.sqrt(k)) + 1

        lat_list = []
        lon_list = []

        for lat_i, lon_i in zip(lat_i_all[valid_idx], lon_i_all[valid_idx]):
            for di in range(-radius, radius+1):
                for dj in range(-radius, radius+1):
                    ni = lat_i + di
                    nj = lon_i + dj

                    if ni < 0 or ni >= lat_number:
                        continue
                    if nj < 0 or nj >= lon_number:
                        continue

                    lat_list.append(ni)
                    lon_list.append(nj)

                    if len(lat_list) % k == 0:
                        break

        lat_id = lat_list
        lon_id = lon_list

    elif backend == 'gpu':
        try:
            import cupy as cp

            cp_point = cp.asarray(point)
            cp_valid = cp.asarray(valid_idx)

            query_points = cp_point[cp_valid]

            # GPU 暂时用 brute force
            diff = cp_point[None, :, :] - query_points[:, None, :]
            dist = cp.sqrt(cp.sum(diff**2, axis=2))

            sequences = cp.argsort(dist, axis=1)[:, :k]

            neigh = cp.asnumpy(sequences.reshape(-1))

            lat_j = neigh // lon_number
            lon_j = neigh % lon_number

            lat_id = lat_j.tolist()
            lon_id = lon_j.tolist()

        except Exception:
            if _print:
                print("GPU不可用，自动回退到KDTree")
            query_points = point[valid_idx]
            distances, sequences = ckt.query(query_points, k)

            neigh = sequences.reshape(-1)
            lat_j = neigh // lon_number
            lon_j = neigh % lon_number

            lat_id = lat_j.tolist()
            lon_id = lon_j.tolist()

    if _print:
        print('查找完毕，开始获取行列索引')
    
    # 这里 lat_id / lon_id 已经由 NumPy 向量化计算得到
    
    if _print:
        print('行列索引已全部获取，开始判断nan值个数')
    
    lonli_arr = np.array([0])
    latli_arr = np.array([0])
    if methond == 'np.nan':
        latli_arr, lonli_arr = judge_nan_npnan(lat_id, lon_id, value, k, judge_num)
    elif methond =='nc.mask':
        latli_arr, lonli_arr = judge_nan_ncmask(lat_id, lon_id, value, k, judge_num)
    elif methond == 'txt.mask':
        latli_arr, lonli_arr = judge_nan_txtmask(lat_id, lon_id, value, k, judge_num)
    if _print:
        print(f'判断完毕，水点数不少于{judge_num}的点才参与运算')
    
    # 每行第一列就是要改的值
    _change_lat = latli_arr[:, 0].tolist()
    _change_lon = lonli_arr[:, 0].tolist()
    # 每行剩下的就是用来赋值的
    _value_lat = latli_arr[:, 1:].tolist()
    _value_lon = lonli_arr[:, 1:].tolist()
    
    _weight = {
        'change_lat': _change_lat,
        'change_lon': _change_lon,
        'value_lat': _value_lat,
        'value_lon': _value_lon,
    }
    return _weight


def judge_nan_ncmask(latli, lonli, value, k, judge_num):
    
    # 将获得的索引分开，保证K个索引是同一个点的最近临点索引
    latli_1 = np.array(latli)
    lonli_1 = np.array(lonli)
    latli_2 = np.reshape(latli_1, [-1, k])
    lonli_2 = np.reshape(lonli_1, [-1, k])
    mask_arr = value.mask
    test_arr = mask_arr[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
    
    sum_arr = np.sum(test_arr, axis=1)  #
    
    select_id = np.where(sum_arr <= k-judge_num)[0].tolist()
    
    latli_r = latli_2[select_id[:]]
    lonli_r = lonli_2[select_id[:]]
    
    return latli_r, lonli_r


def judge_nan_npnan(latli, lonli, value, k, judge_num):
    # 将获得的索引分开，保证K个索引是同一个点的最近临点索引
    latli_1 = np.array(latli)
    lonli_1 = np.array(lonli)
    latli_2 = np.reshape(latli_1, [-1, k])
    lonli_2 = np.reshape(lonli_1, [-1, k])
    mask_arr = np.isnan(value)
    
    test_arr = mask_arr[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
    print(test_arr)
    sum_arr = np.sum(test_arr, axis=1)  #
    
    select_id = np.where(sum_arr <= k-judge_num)[0].tolist()
    
    latli_r = latli_2[select_id[:]]
    lonli_r = lonli_2[select_id[:]]
    
    return latli_r, lonli_r


def judge_nan_txtmask(latli, lonli, value, k, judge_num):
    latli_1 = np.array(latli)
    lonli_1 = np.array(lonli)
    latli_2 = np.reshape(latli_1, [-1, k])
    lonli_2 = np.reshape(lonli_1, [-1, k])
    
    test_arr = value[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
    
    sum_arr = np.sum(test_arr, axis=1)  #
    
    select_id = np.where(sum_arr >= judge_num)[0].tolist()
    
    latli_r = latli_2[select_id[:]]
    lonli_r = lonli_2[select_id[:]]
    
    return latli_r, lonli_r


def erosion_via_id(_weight, value):
    """
    通过索引id进行腐蚀
    :param _weight: weight字典
    :param value: 二维格点值
    """
    
    _change_lat = _weight['change_lat']
    _change_lon = _weight['change_lon']
    _value_lat = _weight['value_lat']
    _value_lon = _weight['value_lon']
    
    new = value.copy()
    if value.ndim == 3:
        for i in range(value.shape[0]):
            replace = value[i, _value_lat, _value_lon]
            replace = np.nanmean(replace, axis=1)
            new[i, _change_lat, _change_lon] = replace[:]
    
    elif value.ndim == 2:
        replace = value[_value_lat, _value_lon]
        replace = np.nanmean(replace, axis=1)
        new[_change_lat, _change_lon] = replace[:]
    
    return new


def example_cal_weight():  # sourcery skip: extract-duplicate-method
    
    methoood = 1
    weight = {}
    viss = []
    if methoood == 1:
        mask_wrf = loadmat('baysalt/Forage/erosion_data/mask_wrf.mat')['mask_wrf'][300:500, 900:1200]
        xr_nc = xr.open_dataset('baysalt/Forage/erosion_data/erosion_ww3.nc')  # ww3直接输出的nc，由于海浪可能算出来nan，所以用地形文件判断nan
        lons = xr_nc['longitude'][900:1200].data
        lats = xr_nc['latitude'][300:500].data
        viss = xr_nc['hs'][:, 300:500, 900:1200].data
        xr_nc.close()
        weight = erosion_cal_id(lons, lats, mask_wrf, 16, 3, methond='txt.mask', _print=True)
        mask_2 = erosion_via_id(weight, mask_wrf)
    
    elif methoood == 2:
        f = Dataset('baysalt/Forage/erosion_data/erosion_xujie.nc')  # 杰哥写的nc，用mask的属性判断nan，data里是--，mask里是True
        lons = f.variables['lon'][:]
        lats = f.variables['lat'][:]
        viss = f.variables['hs'][:]
        f.close()
        weight = erosion_cal_id(lons, lats, viss, 16, 3, methond='nc.mask', _print=True)
    
    elif methoood == 3:
        f = Dataset('baysalt/Forage/erosion_data/erosion_Post_ww3.nc')  # Postprocess_WW3产生的nc，mask为False，data里是nan
        lons = f.variables['longitude'][900:1200].data
        lats = f.variables['latitude'][300:500].data
        viss = f.variables['swh'][:, 300:500, 900:1200].data
        f.close()
        weight = erosion_cal_id(lons, lats, viss[0, :, :], 16, 3, methond='np.nan', _print=True)
    
    with open('weight.json', 'w') as f:
        json.dump(weight, f)
    with open('weight.json', 'r') as f:
        weight=json.load(f)
    
    new_hs=erosion_via_id(weight, viss)
    
    plt.contourf(viss[0, :, :], cmap='jet')
    plt.colorbar()
    plt.show()
    plt.contourf(new_hs[0, :, :], cmap='jet')
    plt.colorbar()
    plt.show()


def test():
    mask_wrf = loadmat('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/mask_wrf.mat')['mask_wrf']
    xr_nc = xr.open_dataset('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/erosion_ww3.nc')
    lons = xr_nc['longitude'][:].data
    lats = xr_nc['latitude'][:].data
    xr_nc.close()
    
    weight_1 = erosion_cal_id(lons, lats, mask_wrf, 16, 3, methond='txt.mask', _print=True)
    mask_wrf_1 = erosion_via_id(weight_1, mask_wrf)
    
    weight_2 = erosion_cal_id(lons, lats, mask_wrf_1, 16, 3, methond='txt.mask', _print=True)
    
    with open('/Users/christmas/Downloads/weight_1.json', 'w') as f:
        json.dump(weight_1, f)
    with open('/Users/christmas/Downloads/weight_2.json', 'w') as f:
        json.dump(weight_2, f)
    
    check_draw()


def check_draw():  # sourcery skip: extract-duplicate-method
    # 验证
    xr_nc = xr.open_dataset('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/erosion_ww3.nc')
    viss = xr_nc['hs'][:].data
    xr_nc.close()
    
    with open('/Users/christmas/Downloads/weight_1.json', 'r') as f:
        weight_1 = json.load(f)
    with open('/Users/christmas/Downloads/weight_2.json', 'r') as f:
        weight_2 = json.load(f)
    new_hs = erosion_via_id(weight_1, viss)
    new_hs_2 = erosion_via_id(weight_2, new_hs)
    
    plt.contourf(viss[0, 500:800, 1200:1700], cmap='jet')
    plt.colorbar()
    plt.show()
    plt.contourf(new_hs[0, 500:800, 1200:1700], cmap='jet')
    plt.colorbar()
    plt.show()
    plt.contourf(new_hs_2[0, 500:800, 1200:1700], cmap='jet')
    plt.colorbar()
    plt.show()


def example_cal_weight_xu(method):  # sourcery skip: extract-duplicate-method
    
    methoood = method
    weight = {}
    viss = []
    if methoood == 1:
        mask_wrf = loadmat('mask_wrf.mat')['mask_wrf'][300:500, 900:1200]
        xr_nc = xr.open_dataset('erosion_ww3.nc')  # ww3直接输出的nc，由于海浪可能算出来nan，所以用地形文件判断nan
        lons = xr_nc['longitude'][900:1200].data
        lats = xr_nc['latitude'][300:500].data
        viss = xr_nc['hs'][:, 300:500, 900:1200].data
        xr_nc.close()
        weight = erosion_cal_id(lons, lats, mask_wrf, 16, 3, methond='txt.mask', _print=True)
        mask_2 = erosion_via_id(weight, mask_wrf)
    
    elif methoood == 2:
        f = Dataset('erosion_xujie.nc')  # 杰哥写的nc，用mask的属性判断nan，data里是--，mask里是True
        lons = f.variables['lon'][:]
        lats = f.variables['lat'][:]
        viss = f.variables['hs'][:]
        f.close()
        weight = erosion_cal_id(lons, lats, viss, 16, 3, methond='nc.mask', _print=True)
    
    elif methoood == 3:
        f = Dataset('erosion_Post_ww3.nc')  # Postprocess_WW3产生的nc，mask为False，data里是nan
        lons = f.variables['longitude'][900:1200].data
        lats = f.variables['latitude'][300:500].data
        viss = f.variables['swh'][:, 300:500, 900:1200].data
        f.close()
        weight = erosion_cal_id(lons, lats, viss[0, :, :], 16, 3, methond='np.nan', _print=True)
    
    with open('weight.json', 'w') as f:
        json.dump(weight, f)
    with open('weight.json', 'r') as f:
        weight=json.load(f)
    
    new_hs = erosion_via_id(weight, viss)
    if method !=2:
        plt.contourf(viss[0, :, :], cmap='jet')
        # plt.colorbar()
        # plt.show()
        plt.savefig('viss.png')
        plt.contourf(new_hs[0, :, :], cmap='jet')
    else:
        plt.contourf(viss[:], cmap='jet')
        # plt.colorbar()
        # plt.show()
        plt.savefig('viss.png')
        plt.contourf(new_hs[:], cmap='jet')
    
    # plt.colorbar()
    # plt.show()
    plt.savefig('new_hs.png')


def example():
    backend = 'kdtree'
    backend = 'hybrid'   # 最稳妥、通用、快(推荐)
    backend = 'numba'
    backend = 'gpu'  # GPU
    backend = 'auto'  # 有 numba → 走 hybrid 没有 numba → 走 kdtree


if __name__ == '__main__':
    
    start = time.perf_counter()
    example_cal_weight_xu(3)
    print(f'耗时: {time.perf_counter() - start}秒')
    # test()
