import os
import glob
import numpy as np
import xarray as xr
import rioxarray
from scipy import stats
import warnings

# 忽略数学运算产生的 RuntimeWarning
warnings.filterwarnings('ignore', category=RuntimeWarning)

# --- 1. 路径设置 ---
base_path = r"W:\tqq\drivers3"
paths = {
    "npp": os.path.join(base_path, "npp", "NPP_*_500m.tif"),
    "smrz": os.path.join(base_path, "SMrz", "ymSMrz_*_GLEAM_v4.2a_YR.tif"), # 替换为根区土壤水分路径
    "tmp": os.path.join(base_path, "tmp", "tmp_*.tif"),
    "pre": os.path.join(base_path, "pre", "pre_*.tif")
}


# --- 2. 数据加载与无效像元预处理 ---
def load_and_align(var_name, base_raster=None):
    files = sorted(glob.glob(paths[var_name]))
    if not files:
        raise FileNotFoundError(f"未找到变量 {var_name} 的文件，请检查路径。")
    data_list = []
    # 读取 2001-2022 年数据
    for f in files[:22]:
        ds = rioxarray.open_rasterio(f).squeeze()
        if base_raster is not None:
            # 空间重采样确保像元对齐
            ds = ds.rio.reproject_match(base_raster)
        data_list.append(ds)
    return xr.concat(data_list, dim="time").assign_coords(time=np.arange(2001, 2023))


print("正在读取并识别无效像元...")
npp_obs = load_and_align("npp")

# 创建全局有效像元掩膜：排除全为空值的背景像元
valid_mask = npp_obs.notnull().any(dim="time")

smrz_ds = load_and_align("smrz", npp_obs) # 加载根区土壤水分数据
tmp_ds = load_and_align("tmp", npp_obs)
pre_ds = load_and_align("pre", npp_obs)


# --- 3. 非线性回归预测 (气候响应 NPP_cli) ---
def pixel_nonlinear_fit(npp, f1, f2, f3):
    # 严格检查：有效样本点需满足非线性拟合要求
    mask = ~np.isnan(npp) & ~np.isnan(f1) & ~np.isnan(f2) & ~np.isnan(f3)
    if mask.sum() < 12:
        return np.full(npp.shape, np.nan)

    y = npp[mask]
    x1, x2, x3 = f1[mask], f2[mask], f3[mask]

    # 构建非线性矩阵：包含一次项、平方项及交互项
    X = np.column_stack([
        np.ones(len(y)),
        x1, x1 ** 2,           # SMrz 贡献
        x2, x2 ** 2,           # tmp 贡献
        x3, x3 ** 2,           # pre 贡献
        x1 * x2, x2 * x3, x1 * x3 # 因子间的交互作用
    ])

    try:
        coef, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
        full_X = np.column_stack([
            np.ones(len(npp)),
            f1, f1 ** 2,
            f2, f2 ** 2,
            f3, f3 ** 2,
            f1 * f2, f2 * f3, f1 * f3
        ])
        return np.dot(full_X, coef)
    except:
        return np.full(npp.shape, np.nan)


print("执行非线性回归计算...")
npp_cli = xr.apply_ufunc(
    pixel_nonlinear_fit, npp_obs, smrz_ds, tmp_ds, pre_ds,
    input_core_dims=[["time"], ["time"], ["time"], ["time"]],
    output_core_dims=[["time"]],
    vectorize=True,
    dask="parallelized",
    output_dtypes=[float]
)

# 计算残差并立即应用无效像元掩膜 [cite: 69]
npp_res = (npp_obs - npp_cli).where(valid_mask)


# --- 4. 趋势斜率分析 ---
def calculate_slope(data):
    mask = ~np.isnan(data)
    if mask.sum() < 10: return np.nan
    slope, _, _, _, _ = stats.linregress(np.arange(len(data))[mask], data[mask])
    return slope


print("计算各分量的趋势斜率...")
slope_npp = xr.apply_ufunc(calculate_slope, npp_obs, input_core_dims=[["time"]], vectorize=True).where(valid_mask)
slope_cli = xr.apply_ufunc(calculate_slope, npp_cli, input_core_dims=[["time"]], vectorize=True).where(valid_mask)
slope_res = xr.apply_ufunc(calculate_slope, npp_res, input_core_dims=[["time"]], vectorize=True).where(valid_mask)

# --- 5. 驱动因素划分  ---
# 仅对有效像元进行划分
cond_pos = (slope_npp > 0)
inc_drivers = xr.where(cond_pos & (slope_cli > 0) & (slope_res > 0), 3,
                       xr.where(cond_pos & (slope_cli > 0) & (slope_res <= 0), 4,
                                xr.where(cond_pos & (slope_cli <= 0) & (slope_res > 0), 6, 0)))

cond_neg = (slope_npp < 0)
dec_drivers = xr.where(cond_neg & (slope_cli < 0) & (slope_res < 0), -8,
                       xr.where(cond_neg & (slope_cli < 0) & (slope_res >= 0), -6,
                                xr.where(cond_neg & (slope_cli >= 0) & (slope_res < 0), -4, 0)))

# 合并结果并强制无效像元为 NoData (-9999) 以解决加载失败问题
npp_drivefactor = (inc_drivers + dec_drivers).where(valid_mask, -9999).astype("int16")

# --- 6. 相对贡献率计算  ---
# 防止分母为 0 导致溢出
safe_slope_npp = slope_npp.where(slope_npp != 0)
contrib_cli = (slope_cli / safe_slope_npp).clip(0, 1).where(valid_mask).fillna(0)
contrib_ha = (slope_res / safe_slope_npp).clip(0, 1).where(valid_mask).fillna(0)

# --- 7. 导出 GeoTIFF ---
print("正在导出结果文件...")


# 统一设置导出参数，确保兼容 ArcGIS
def export_raster(da, name, is_int=False):
    if is_int:
        # 写入 NoData 声明
        da.rio.write_nodata(-9999, inplace=True)
        da.rio.to_raster(name, dtype="int16")
    else:
        # 浮点数结果保存
        da.rio.to_raster(name)


export_raster(npp_drivefactor, "npp_drivefactor_smrz.tif", is_int=True)
export_raster(contrib_cli, "npp_climate_contribution_smrz.tif")
export_raster(contrib_ha, "npp_human_contribution_smrz.tif")

print("任务完成。已基于 SMrz 因子生成残差趋势分析结果。")