#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#  日期 : 2026/3/9 14:44
#  作者 : Christmas
#  邮箱 : 273519355@qq.com
#  项目 : Project
#  版本 : python 3
#  摘要 :
"""
1. 先找出“需要被填补”的点（通常是海陆交界处的水点、空值点、mask 点）
2. 再为这些点找到周围最近的 K 个邻居
3. 过滤掉“周围有效邻居太少”的点
4. 把可用邻居的值做平均，回填到目标点中

这里的“erosion（侵蚀）”更接近一种海岸线附近的掩膜修正 / 邻域填补：
"""

import itertools
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

from Alaea.Blogging import Blog
from Alaea.tools.Read import read_json, write_json

__all__ = [
    'erosion_cal_id',
    'erosion_via_id',
    'LOGGER',
]

BLOG = Blog('erosion_coastline')
BLOG.setup_Blog(log_level=BLOG.DEBUG)
LOGGER = BLOG.logger


def erosion_cal_id(
        g_lon,
        g_lat,
        value,
        k,
        judge_num,
        method='np.nan',
):
    """
    计算“哪些点需要被回填”，以及“每个点应该参考哪些邻居点”。

    这个函数不直接修改 value，而是先生成一个索引说明书 `_weight`：
    - change_lat / change_lon：哪些点要被修改
    - value_lat / value_lon：每个待修改点应当使用哪些邻居点的值

    整体流程：
    1. 把规则网格展开成二维坐标点云
    2. 用 KDTree 为每个格点查找最近的 k 个邻居
    3. 根据 method 判断哪些点是“待处理点”
    4. 再检查它们的 k 个邻居里是否有足够多的有效点
    5. 只保留满足 judge_num 条件的点，返回后续填补所需索引

    参数说明：
    - k：每个目标点向周围找多少个最近邻
    - judge_num：最少需要多少个“有效邻居”才允许回填
    - method='np.nan'：把 NaN 看成待处理点
    - method='nc.mask'：把 masked array 中 mask=True 的点看成待处理点
    - method='txt.mask'：把值为 0 的点看成待处理点（常用于陆海 mask）

    :param g_lon: 一维格点经度
    :param g_lat: 一维格点纬度
    :param value: 二维格点值
    :param k: 寻找近邻点的个数
    :param judge_num: 至少需要多少个有效邻居点才保留该目标点
    :param method: 判断目标点类型的方法
                    np.nan：直接在 value 中用 NaN 表示无效点
                    nc.mask：NetCDF masked array，用 mask 属性判断无效点
                    txt.mask：文本/掩膜型数据，0 表示水点，非 0 表示陆点
    :return: 字典，描述待修改点及其邻居索引
    """
    # 网格大小（这里只是记录下来，后面真正高频使用的是 lon_number）
    lat_number = len(g_lat)
    # 一行有多少个经度格点，用于后面把一维展开索引还原成二维行列索引
    lon_number = len(g_lon)
    # 例如 g_lon.shape=(nlon,), g_lat.shape=(nlat,)
    # meshgrid 后得到两个 (nlat, nlon) 的二维数组，分别存放每个格点的经纬度
    LOGGER.debug(f'开始建立Kdtree，格点数量约为 {lat_number} * {lon_number}，每个点找 {k} 个邻居，至少需要 {judge_num} 个有效邻居')
    lon_mesh, lat_mesh = np.meshgrid(g_lon, g_lat)
    # 后面 KDTree 就是在这个二维坐标集合上做最近邻搜索
    point = np.column_stack((lon_mesh.ravel(), lat_mesh.ravel()))
    # 下面两个 deque 用来保存所有“候选邻居点”的二维索引
    fp_lon = deque([])
    fp_lat = deque([])

    # 建立 KDTree：适合大量最近邻查询
    # 这里 point 里装的是所有格点坐标，所以相当于“让每个格点去找离自己最近的 k 个格点”
    ckt = spt.cKDTree(point)  # 用C写的查找类，执行速度更快
    # 这里不再逐点循环，而是一次性把所有格点的近邻都查出来
    LOGGER.debug('Kdtree建立完毕，开始查找每个格点的近邻点')

    # 对所有格点一次性查询最近邻：
    # distances.shape 约为 (n_points, k)
    # sequences.shape 约为 (n_points, k)
    # sequences 里保存的是最近邻在 point 数组中的一维编号
    distances, sequences = ckt.query(point, k)

    # point 中每个元素对应一个展开后的一维索引 [0, 1, 2, ...]
    idx_all = np.arange(point.shape[0])
    # 把展开索引还原回二维索引：
    # 行号 = index // 每行列数
    lat_i_all = idx_all // lon_number
    # 列号 = index % 每行列数
    lon_i_all = idx_all % lon_number

    # 先找出哪些格点属于“待处理点”
    # 不同数据源中，水点/无效点/边界点的表达方式不同，所以这里分三种判断方法
    if method == 'np.nan':
        mask = np.isnan(value[lat_i_all, lon_i_all])
    elif method == 'nc.mask':
        mask = value.mask[lat_i_all, lon_i_all]
    elif method == 'txt.mask':
        mask = (value[lat_i_all, lon_i_all] == 0)
    else:
        mask = np.zeros_like(lat_i_all, dtype=bool)

    # 只保留那些“需要处理”的点在 point 里的编号
    valid_idx = idx_all[mask]

    # 取出这些目标点各自对应的 k 个最近邻编号
    neigh = sequences[valid_idx]

    # 展平后，后面统一换算成二维索引
    neigh = neigh.reshape(-1)

    # 邻居点的一维编号 -> 二维行列编号
    lat_j = neigh // lon_number
    lon_j = neigh % lon_number

    # 这里会把所有目标点的 k 个邻居顺序拼在一起
    # 因此后面需要再按 k reshape 回去，恢复“每 k 个属于同一个目标点”的结构
    fp_lat.extend(lat_j.tolist())
    fp_lon.extend(lon_j.tolist())
    LOGGER.debug('所有格点的近邻点查找完毕，开始获取行列索引')

    # 这里 fp_lat / fp_lon 已经是索引，不再需要 find_index
    lat_id = list(fp_lat)
    lon_id = list(fp_lon)

    LOGGER.debug('行列索引获取完毕，开始判断每个目标点的邻居里有效点数量')
    
    # 根据不同类型的数据，进一步判断：
    # 某个目标点的 k 个邻居里，是否有足够多的有效点可以拿来回填
    lonli_arr = np.array([0])
    latli_arr = np.array([0])
    if method == 'np.nan':
        latli_arr, lonli_arr = judge_nan_npnan(lat_id, lon_id, value, k, judge_num)
    elif method =='nc.mask':
        latli_arr, lonli_arr = judge_nan_ncmask(lat_id, lon_id, value, k, judge_num)
    elif method == 'txt.mask':
        latli_arr, lonli_arr = judge_nan_txtmask(lat_id, lon_id, value, k, judge_num)

    LOGGER.debug('有效邻居点数量判断完毕，开始生成索引')

    # reshape 之后，每一行对应“一个目标点 + 它的 k-1 或 k 个邻居”
    # 约定每行第 1 列是目标点本身，后面各列是参考邻居
    _change_lat = latli_arr[:, 0].tolist()
    _change_lon = lonli_arr[:, 0].tolist()
    # 每行剩下的就是用来赋值的
    _value_lat = latli_arr[:, 1:].tolist()
    _value_lon = lonli_arr[:, 1:].tolist()

    # 组织成统一的索引字典，供 erosion_via_id 直接使用
    _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):
    """
    针对 masked array（mask=True 表示无效点）筛选可参与回填的目标点。

    思路：
    - latli / lonli 当前是展平后的邻居索引序列
    - 每 k 个索引对应 1 个目标点及其 k 个最近邻
    - 如果这 k 个点里“未被 mask 的有效点”数量不少于 judge_num，
      就保留这一组索引
    """
    # 先转成数组，便于 reshape 成 [目标点数, k] 的二维结构
    latli_1 = np.array(latli)
    lonli_1 = np.array(lonli)
    # 每一行表示一个目标点对应的 k 个最近邻索引
    latli_2 = np.reshape(latli_1, [-1, k])
    lonli_2 = np.reshape(lonli_1, [-1, k])
    # mask=True 说明该位置是无效/被遮罩的点
    mask_arr = value.mask
    # 取出每组 k 个邻居的 mask 情况
    test_arr = mask_arr[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理

    # 统计每组中“无效点”的数量
    sum_arr = np.sum(test_arr, axis=1)  #

    # 无效点数 <= k - judge_num，等价于有效点数 >= judge_num
    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):
    """
    针对普通 ndarray 中 NaN 表示无效值的情况，筛选可参与回填的目标点。

    逻辑与 judge_nan_ncmask 一样，只不过这里通过 np.isnan(value)
    判断某个点是否无效。
    """
    # 每 k 个索引是一组，后面要恢复成二维结构
    latli_1 = np.array(latli)
    lonli_1 = np.array(lonli)
    # reshape 后每一行对应 1 个目标点的 k 个邻居
    latli_2 = np.reshape(latli_1, [-1, k])
    lonli_2 = np.reshape(lonli_1, [-1, k])
    # NaN 为 True，表示该点不能直接作为有效值参与平均
    mask_arr = np.isnan(value)

    # 检查每组 k 个邻居中哪些位置是 NaN
    test_arr = mask_arr[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
    sum_arr = np.sum(test_arr, axis=1)  #

    # 只有有效值数量 >= judge_num 的目标点，才保留下来
    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):
    """
    针对文本/掩膜型数据（0 表示水点，非 0 表示陆点）筛选可参与回填的目标点。

    这里与前两个函数不同：
    - 前两个函数统计的是“无效点个数”
    - 这里直接统计邻居值之和，本质上是在数一组邻居里有多少个陆点（值为 1）

    当一组邻居中陆点数量 >= 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 中通常 0=水，1=陆
    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):
    """
    根据 erosion_cal_id 生成的索引说明书，对目标点执行回填。

    核心思想：
    - 对每个待修改点，取出它对应邻居点的值
    - 对这些邻居值做平均（np.nanmean 会自动跳过 NaN）
    - 再把平均值写回目标点

    支持两种输入：
    - 2D: [lat, lon]
    - 3D: [time, lat, lon]

    :param _weight: erosion_cal_id 返回的字典
    :param value: 二维或三维格点值
    :return: 回填后的新数组
    """

    # 要被修改的目标点索引
    _change_lat = _weight['change_lat']
    _change_lon = _weight['change_lon']
    _value_lat = _weight['value_lat']
    _value_lon = _weight['value_lon']

    # 拷贝一份，避免直接改原数组
    new = value.copy()
    # 3D 情况：常见于 [time, lat, lon]
    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[:]

    # 2D 情况：直接对一个平面进行处理
    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(methoood):  # sourcery skip: extract-duplicate-method
    """
    示例函数：演示如何针对不同类型的数据源计算 weight，并执行一次回填。

    methoood 含义：
    1 -> txt.mask：外部 mask 文件中 0/1 表示海陆
    2 -> nc.mask：NetCDF masked array，用 mask 属性判断无效点
    3 -> np.nan：直接在数据中用 NaN 表示无效点
    """

    if methoood == 1:
        mask_wrf = loadmat('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/mask_wrf.mat')['mask_wrf']
        ds = xr.open_dataset('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/erosion_ww3.nc')  # ww3直接输出的nc，由于海浪可能算出来nan，所以用地形文件判断nan
        lons = ds['longitude'][:].data
        lats = ds['latitude'][:].data
        viss = ds['hs'][:].data
        ds.close()
        weight_1 = erosion_cal_id(lons, lats, mask_wrf, 16, 3, method='txt.mask')
        mask_wrf_1 = erosion_via_id(weight_1, mask_wrf)
        weight_2 = erosion_cal_id(lons, lats, mask_wrf_1, 24, 3, method='txt.mask')
        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()

    elif methoood == 2:
        f = Dataset('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/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, method='nc.mask')
        
        new_hs = erosion_via_id(weight, viss)
        
        plt.contourf(viss[:, :], cmap='jet')
        plt.colorbar()
        plt.show()
        plt.contourf(new_hs[:, :], cmap='jet')
        plt.colorbar()
        plt.show()

    elif methoood == 3:
        # f = Dataset('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/erosion_Post_ww3.nc')  # Postprocess_WW3产生的nc，mask为False，data里是nan
        ds = xr.open_dataset('/Users/christmas/Documents/Code/Project/A-pypi/Alaea/Alaea/baysalt/Forage/erosion_data/erosion_Post_ww3.nc')
        lons = ds['longitude'][900:1200].data
        lats = ds['latitude'][300:500].data
        viss = ds['swh'][:, 300:500, 900:1200].data
        ds.close()
        # f.close()
        weight = erosion_cal_id(lons, lats, viss[0, :, :], 16, 3, method='np.nan')
        
        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()
        
    else:
        raise ValueError()


if __name__ == '__main__':
    # 简单计时，看看一次示例运行要多久
    start = time.perf_counter()
    example_cal_weight(1)
    example_cal_weight(2)
    example_cal_weight(3)
    print(f'耗时: {time.perf_counter() - start}秒')
    # test()
