import os
import glob
import numpy as np
import pandas as pd
import xarray as xr
import rioxarray
from scipy import sparse
from esda.moran import Moran, Moran_Local
import libpysal
import warnings
import time
import gc  # 导入垃圾回收模块


# ==============================================================================
# 1. 核心修正：深度封装 WSP 对象
# ==============================================================================
class MockSummary:
    def __init__(self, s0, s1, s2):
        self.s0, self.s1, self.s2 = s0, s1, s2


class MockWSP:
    def __init__(self, wsp_obj):
        self.wsp = wsp_obj
        self.n = wsp_obj.n
        self.s0 = wsp_obj.s0
        self.sparse = wsp_obj.sparse

        # 计算分析矩，显式转换为 float64 避免溢出
        A = self.sparse
        AT = A.transpose()
        self.s1 = 0.5 * (A + AT).power(2).sum()

        r = np.array(A.sum(axis=1)).flatten()
        c = np.array(A.sum(axis=0)).flatten()
        self.s2 = np.sum((r + c) ** 2)
        self._summary_obj = MockSummary(self.s0, self.s1, self.s2)

    def lag(self, y):
        return self.sparse.dot(y)

    def transform(self, *args, **kwargs):
        return self

    def summary(self):
        return self._summary_obj

    def __getattr__(self, name):
        return getattr(self.wsp, name)


# ==============================================================================
# 2. 优化后的计算逻辑
# ==============================================================================
def build_fast_normalized_weights(h, w, mask):
    n = h * w
    idx = np.arange(n).reshape(h, w)
    rows, cols = [], []

    # Rook 邻域
    left, right = idx[:, :-1].flatten(), idx[:, 1:].flatten()
    rows.extend([left, right])
    cols.extend([right, left])
    top, bottom = idx[:-1, :].flatten(), idx[1:, :].flatten()
    rows.extend([top, bottom])
    cols.extend([bottom, top])

    r, c = np.concatenate(rows), np.concatenate(cols)

    # 【修改点 1】: 将 dtype=np.float32 改为 np.float64
    # 确保权重矩阵的基础类型为双精度，与输入数据保持一致，防止 Numba 报错
    adj = sparse.coo_matrix((np.ones(len(r), dtype=np.float64), (r, c)), shape=(n, n)).tocsr()

    # 裁剪有效像元并清理原始索引以释放内存
    valid_idx = np.where(mask.flatten())[0]
    adj_valid = adj[valid_idx, :][:, valid_idx]

    # 行标准化
    row_sums = np.array(adj_valid.sum(axis=1)).flatten()
    d_inv = np.power(row_sums, -1.0, where=row_sums != 0)
    adj_normalized = sparse.diags(d_inv).dot(adj_valid)

    return MockWSP(libpysal.weights.WSP(adj_normalized))


def process_single_file(input_tif, output_dir):
    print(f"\n--- 正在处理: {os.path.basename(input_tif)} ---")
    start_time = time.time()

    # 1. 加载数据
    da = rioxarray.open_rasterio(input_tif).squeeze()
    nodata = da.rio.nodata
    mask = (da.values != nodata) & (~np.isnan(da.values)) if nodata is not None else ~np.isnan(da.values)

    # 【关键点 2】: 确保数据转换为 float64
    # 这一步至关重要，因为 esda/numba 需要 float64 数据来匹配 float64 的权重
    valid_data = da.values[mask].astype(np.float64)

    # 构建权重
    w_mock = build_fast_normalized_weights(da.shape[0], da.shape[1], mask)

    # 2. 全局 Moran's I
    print("正在计算全局 Moran's I...")
    g_moran = Moran(valid_data, w=w_mock, permutations=999)
    print(f"全局结果: I={g_moran.I:.4f}, P={g_moran.p_sim:.4f}")

    # 3. 局部 Moran's I
    # n_jobs=1: 解决 WinError 1450 及 Numba 线程冲突问题
    print(f"正在计算局部 Moran's I (单线程模式以节省资源)...")
    l_moran = Moran_Local(valid_data, w=w_mock, permutations=999, n_jobs=1)

    # 4. 结果转换
    base_name = os.path.splitext(os.path.basename(input_tif))[0]

    # 结果可以使用 float32 保存以节省硬盘空间
    li_map = np.full(mask.shape, np.nan, dtype=np.float32)
    li_map[mask] = l_moran.Is.astype(np.float32)

    lz_map = np.full(mask.shape, np.nan, dtype=np.float32)
    lz_map[mask] = l_moran.z_sim.astype(np.float32)

    sig = l_moran.p_sim < 0.05
    quad = np.zeros_like(l_moran.q, dtype=np.int16)
    quad[sig] = l_moran.q[sig]
    q_map = np.full(mask.shape, -9999, dtype=np.int16)
    q_map[mask] = quad

    # 5. 导出 TIF
    results = [
        (li_map, "_local_I", "float32", np.nan),
        (lz_map, "_local_Z", "float32", np.nan),
        (q_map, "_lisa", "int16", -9999)
    ]
    for data, suffix, dtype, nd in results:
        out = xr.DataArray(data, coords=da.coords, dims=da.dims)
        out.rio.write_crs(da.rio.crs, inplace=True)
        out.rio.to_raster(os.path.join(output_dir, f"{base_name}{suffix}.tif"), dtype=dtype, nodata=nd)

    # 6. 强制内存清理
    del w_mock, l_moran, valid_data, mask, da, li_map, lz_map, q_map
    gc.collect()

    duration = (time.time() - start_time) / 60
    print(f"文件处理完成，耗时: {duration:.2f} 分钟")
    return {'File': base_name, 'MoranI': g_moran.I, 'P_val': g_moran.p_sim}


def main():
    # --- 用户配置路径 ---
    input_tif = r"D:\tqq\NPP_Albers Equal Area\NPP_2005_500m.tif"
    output_dir = r"D:\tqq\result\2005"
    # ------------------

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)
        print(f"已创建输出目录: {output_dir}")

    summary_list = []

    # 检查输入文件是否存在
    if not os.path.exists(input_tif):
        print(f"!!! 错误: 找不到输入文件: {input_tif}")
        return

    try:
        res = process_single_file(input_tif, output_dir)
        summary_list.append(res)
    except Exception as e:
        print(f"!!! 无法处理文件 {input_tif}: {str(e)}")
        import traceback
        traceback.print_exc()  # 打印详细错误栈，方便排查
        gc.collect()

    if summary_list:
        df = pd.DataFrame(summary_list)
        df.to_csv(os.path.join(output_dir, "moran_summary.csv"), index=False)
        print("\n=== 任务统计汇总 ===\n", df)


if __name__ == "__main__":
    warnings.filterwarnings('ignore')
    main()