import numpy as np
import pandas as pd

# =========================================================
# 0. 参数设置
# =========================================================
RES = 0.5
RANDOM_SEED = 42
OUTPUT_CSV = "global_grassland_grazing_n2o_effect_0p5_final_9class.csv"

np.random.seed(RANDOM_SEED)

# =========================================================
# 1. 构建全球0.5°规则格网中心点
# =========================================================
lons = np.arange(-180 + RES / 2, 180, RES)   # -179.75 ~ 179.75
lats = np.arange(-90 + RES / 2, 90, RES)     # -89.75 ~ 89.75

grid = pd.DataFrame(
    [(lon, lat) for lat in lats for lon in lons],
    columns=["lon", "lat"]
)

# =========================================================
# 2. 草地区域划分（近似）
# =========================================================
def assign_grassland_region(lon, lat):
    # 北美高纬草地/苔原草地
    if (-170 <= lon <= -55) and (50 <= lat <= 72):
        return "NorthAmerica_HighLat"

    # 北美中西部/大平原/西部草地
    if (-125 <= lon <= -90) and (30 <= lat <= 55):
        return "NorthAmerica_MidGrass"

    # 南美安第斯高山草地
    if (-80 <= lon <= -65) and (-35 <= lat <= -5):
        return "SouthAmerica_Andes"

    # 南美南部草地（潘帕斯/巴塔哥尼亚）
    if (-75 <= lon <= -50) and (-55 <= lat <= -30):
        return "SouthAmerica_South"

    # 巴西中部/塞拉多稀树草原
    if (-60 <= lon <= -45) and (-25 <= lat <= -5):
        return "SouthAmerica_Cerrado"

    # 萨赫勒草地带
    if (-18 <= lon <= 40) and (8 <= lat <= 18):
        return "Africa_Sahel"

    # 东非草地
    if (25 <= lon <= 45) and (-10 <= lat <= 12):
        return "Africa_East"

    # 南部非洲草地
    if (10 <= lon <= 35) and (-35 <= lat <= -10):
        return "Africa_South"

    # 马达加斯加局部草地
    if (43 <= lon <= 50) and (-25 <= lat <= -12):
        return "Madagascar"

    # 欧亚高纬草地/寒带草地
    if (30 <= lon <= 180) and (55 <= lat <= 72):
        return "Eurasia_HighLat"

    # 中亚-蒙古-内亚温带草地
    if (45 <= lon <= 125) and (35 <= lat <= 55):
        return "CentralAsia_Mongolia"

    # 青藏高原草地
    if (75 <= lon <= 105) and (28 <= lat <= 38):
        return "TibetanPlateau"

    # 西亚-中东局部草地
    if (30 <= lon <= 60) and (20 <= lat <= 40):
        return "WestAsia_MiddleEast"

    # 澳大利亚草地
    if (112 <= lon <= 155) and (-40 <= lat <= -12):
        return "Australia"

    # 新西兰局部草地
    if (166 <= lon <= 179) and (-47 <= lat <= -34):
        return "NewZealand"

    return None

# =========================================================
# 3. 非草地区剔除
# =========================================================
def is_excluded_non_grassland(lon, lat, region):
    # 亚马逊核心雨林区
    if (-75 <= lon <= -45) and (-12 <= lat <= 5):
        return True

    # 撒哈拉核心沙漠区
    if (-15 <= lon <= 35) and (18 <= lat <= 30):
        return True

    # 刚果盆地核心热带雨林区
    if (10 <= lon <= 30) and (-8 <= lat <= 5):
        return True

    # 东南亚热带森林/岛屿区
    if (95 <= lon <= 130) and (-10 <= lat <= 20):
        return True

    # 中国东部湿润农田/森林区
    if (105 <= lon <= 125) and (25 <= lat <= 40):
        return True

    # 印度平原高密度农田区
    if (75 <= lon <= 90) and (20 <= lat <= 30):
        return True

    # 西欧农田与城市密集区
    if (-5 <= lon <= 20) and (40 <= lat <= 55):
        return True

    # 加拿大东南部湿润森林/农业区
    if (-90 <= lon <= -60) and (45 <= lat <= 55):
        return True

    # 美国东部湿润农业/森林区
    if (-90 <= lon <= -65) and (30 <= lat <= 45):
        return True

    # 澳大利亚东南沿海湿润带
    if (145 <= lon <= 155) and (-40 <= lat <= -25):
        return True

    # 日本及周边湿润山区
    if (130 <= lon <= 146) and (30 <= lat <= 45):
        return True

    return False

# =========================================================
# 4. 工具函数
# =========================================================
def clamp(x, low=0.0, high=1.0):
    return max(low, min(high, x))

# =========================================================
# 5. 温度适宜度（0~1）
# =========================================================
def temp_score(lat, region):
    abs_lat = abs(lat)
    base = 1.0 - abs_lat / 85.0

    if region in ["Eurasia_HighLat", "NorthAmerica_HighLat"]:
        base -= 0.20
    if region == "TibetanPlateau":
        base -= 0.25
    if region == "SouthAmerica_Andes":
        base -= 0.15
    if region == "Africa_East":
        base -= 0.05
    if region in ["Australia", "Africa_Sahel", "Africa_South"]:
        base += 0.05
    if region == "NewZealand":
        base += 0.03

    return clamp(base, 0, 1)

# =========================================================
# 6. 水分适宜度（0~1）
# =========================================================
def moisture_score(lat, lon, region):
    score = 0.50

    if region == "Africa_Sahel":
        score = 0.35
    elif region == "Africa_East":
        score = 0.55
    elif region == "Africa_South":
        score = 0.45
    elif region == "Australia":
        score = 0.30
    elif region == "CentralAsia_Mongolia":
        score = 0.35
    elif region == "WestAsia_MiddleEast":
        score = 0.25
    elif region == "TibetanPlateau":
        score = 0.30
    elif region == "Eurasia_HighLat":
        score = 0.40
    elif region == "NorthAmerica_HighLat":
        score = 0.45
    elif region == "NorthAmerica_MidGrass":
        score = 0.50
    elif region == "SouthAmerica_Andes":
        score = 0.40
    elif region == "SouthAmerica_South":
        score = 0.50
    elif region == "SouthAmerica_Cerrado":
        score = 0.55
    elif region == "Madagascar":
        score = 0.55
    elif region == "NewZealand":
        score = 0.65

    abs_lat = abs(lat)

    if abs_lat < 15:
        score += 0.05
    elif 15 <= abs_lat <= 45:
        score += 0.08
    elif abs_lat > 55:
        score -= 0.08

    score += 0.03 * np.sin(np.radians(lon * 2.0))

    return clamp(score, 0, 1)

# =========================================================
# 7. 水热综合适宜度（0~1）
# =========================================================
def hydrothermal_score(temp_s, moist_s):
    return clamp(np.sqrt(temp_s * moist_s), 0, 1)

# =========================================================
# 8. 冻融强度（0~1）
# =========================================================
def freeze_strength(lat, region):
    abs_lat = abs(lat)

    if abs_lat >= 65:
        strength = 0.95
    elif abs_lat >= 60:
        strength = 0.80
    elif abs_lat >= 55:
        strength = 0.65
    elif abs_lat >= 50:
        strength = 0.45
    else:
        strength = 0.15

    if region == "TibetanPlateau":
        strength = max(strength, 0.75)
    if region == "SouthAmerica_Andes":
        strength = max(strength, 0.65)
    if region in ["Eurasia_HighLat", "NorthAmerica_HighLat"]:
        strength = max(strength, 0.70)

    return clamp(strength, 0, 1)

# =========================================================
# 9. 草地筛选
# =========================================================
grid["region"] = grid.apply(lambda r: assign_grassland_region(r["lon"], r["lat"]), axis=1)
grass = grid.dropna(subset=["region"]).copy()

grass["exclude"] = grass.apply(
    lambda r: is_excluded_non_grassland(r["lon"], r["lat"], r["region"]),
    axis=1
)

grass = grass[grass["exclude"] == False].copy()
grass.drop(columns=["exclude"], inplace=True)
grass["grassland"] = 1

# =========================================================
# 10. 环境评分
# =========================================================
grass["temp_score"] = grass.apply(lambda r: temp_score(r["lat"], r["region"]), axis=1)
grass["moisture_score"] = grass.apply(lambda r: moisture_score(r["lat"], r["lon"], r["region"]), axis=1)
grass["hydrothermal_score"] = grass.apply(lambda r: hydrothermal_score(r["temp_score"], r["moisture_score"]), axis=1)
grass["freeze_strength"] = grass.apply(lambda r: freeze_strength(r["lat"], r["region"]), axis=1)

# =========================================================
# 11. 识别冻融区
# =========================================================
freeze_regions = [
    "Eurasia_HighLat",
    "NorthAmerica_HighLat",
    "TibetanPlateau",
    "SouthAmerica_Andes"
]

grass["is_freeze_region"] = grass["region"].isin(freeze_regions)

# 高纬区域也强制视为冻融区
grass.loc[grass["lat"].abs() >= 55, "is_freeze_region"] = True

# =========================================================
# 12. 生成 N2O 排放效应值（%）
# =========================================================
grass["n2o_effect_pct"] = np.nan

mask_nonfreeze = ~grass["is_freeze_region"]
mask_freeze = grass["is_freeze_region"]

# 非冻融区：负值 -25 ~ 0
grass.loc[mask_nonfreeze, "n2o_effect_pct"] = (
    -25.0 * grass.loc[mask_nonfreeze, "hydrothermal_score"]
)

# 冻融区：正值 0 ~ 15
grass.loc[mask_freeze, "n2o_effect_pct"] = (
    15.0 * grass.loc[mask_freeze, "freeze_strength"]
)

# 轻微扰动
noise = np.random.normal(loc=0.0, scale=0.8, size=len(grass))
grass["n2o_effect_pct"] = grass["n2o_effect_pct"] + noise

# 严格限制范围
grass.loc[mask_nonfreeze, "n2o_effect_pct"] = grass.loc[mask_nonfreeze, "n2o_effect_pct"].clip(-25.0, 0.0)
grass.loc[mask_freeze, "n2o_effect_pct"] = grass.loc[mask_freeze, "n2o_effect_pct"].clip(0.0, 15.0)

# =========================================================
# 13. 九级分类
# =========================================================
bins = [-25, -20, -15, -10, -5, 0, 3, 6, 10, 15.000001]
labels = [
    "1:-25~-20",
    "2:-20~-15",
    "3:-15~-10",
    "4:-10~-5",
    "5:-5~0",
    "6:0~3",
    "7:3~6",
    "8:6~10",
    "9:10~15"
]

grass["effect_class"] = pd.cut(
    grass["n2o_effect_pct"],
    bins=bins,
    labels=labels,
    include_lowest=True,
    right=True
)

class_map = {
    "1:-25~-20": 1,
    "2:-20~-15": 2,
    "3:-15~-10": 3,
    "4:-10~-5": 4,
    "5:-5~0": 5,
    "6:0~3": 6,
    "7:3~6": 7,
    "8:6~10": 8,
    "9:10~15": 9
}

grass["effect_class"] = grass["effect_class"].astype(str)
grass["effect_class_id"] = grass["effect_class"].map(class_map)

# 若仍有极少数空值，直接按数值补分
missing_mask = grass["effect_class_id"].isna()
if missing_mask.any():
    vals = grass.loc[missing_mask, "n2o_effect_pct"]
    grass.loc[missing_mask & (vals >= -25) & (vals <= -20), "effect_class_id"] = 1
    grass.loc[missing_mask & (vals > -20) & (vals <= -15), "effect_class_id"] = 2
    grass.loc[missing_mask & (vals > -15) & (vals <= -10), "effect_class_id"] = 3
    grass.loc[missing_mask & (vals > -10) & (vals <= -5), "effect_class_id"] = 4
    grass.loc[missing_mask & (vals > -5) & (vals <= 0), "effect_class_id"] = 5
    grass.loc[missing_mask & (vals > 0) & (vals <= 3), "effect_class_id"] = 6
    grass.loc[missing_mask & (vals > 3) & (vals <= 6), "effect_class_id"] = 7
    grass.loc[missing_mask & (vals > 6) & (vals <= 10), "effect_class_id"] = 8
    grass.loc[missing_mask & (vals > 10) & (vals <= 15), "effect_class_id"] = 9

    reverse_map = {v: k for k, v in class_map.items()}
    grass.loc[missing_mask, "effect_class"] = grass.loc[missing_mask, "effect_class_id"].map(reverse_map)

grass["effect_class_id"] = grass["effect_class_id"].astype(int)

# =========================================================
# 14. 输出字段整理
# =========================================================
grass = grass.reset_index(drop=True)
grass["pixel_id"] = np.arange(1, len(grass) + 1)
grass["grid_res_deg"] = RES

output_cols = [
    "pixel_id",
    "lon",
    "lat",
    "grid_res_deg",
    "region",
    "grassland",
    "temp_score",
    "moisture_score",
    "hydrothermal_score",
    "freeze_strength",
    "is_freeze_region",
    "n2o_effect_pct",
    "effect_class",
    "effect_class_id"
]

grass_out = grass[output_cols].copy()

# 数值保留小数
grass_out["lon"] = grass_out["lon"].round(2)
grass_out["lat"] = grass_out["lat"].round(2)
grass_out["grid_res_deg"] = grass_out["grid_res_deg"].round(2)

for col in ["temp_score", "moisture_score", "hydrothermal_score", "freeze_strength"]:
    grass_out[col] = grass_out[col].round(3)

grass_out["n2o_effect_pct"] = grass_out["n2o_effect_pct"].round(2)

# =========================================================
# 15. 导出 CSV
# =========================================================
grass_out.to_csv(OUTPUT_CSV, index=False, encoding="utf-8-sig")

# =========================================================
# 16. 结果检查
# =========================================================
print("=" * 60)
print("文件已生成：", OUTPUT_CSV)
print("全球0.5°总格点数：", len(grid))
print("最终草地像元数：", len(grass_out))
print("=" * 60)

print("\n经度间隔检查：", round(lons[1] - lons[0], 2))
print("纬度间隔检查：", round(lats[1] - lats[0], 2))

print("\nN2O效应值统计（%）：")
print(grass_out["n2o_effect_pct"].describe())

print("\n非冻融区效应值范围（应为 -25 ~ 0）：")
print(
    grass_out.loc[grass_out["is_freeze_region"] == False, "n2o_effect_pct"].min(),
    grass_out.loc[grass_out["is_freeze_region"] == False, "n2o_effect_pct"].max()
)

print("\n冻融区效应值范围（应为 0 ~ 15）：")
print(
    grass_out.loc[grass_out["is_freeze_region"] == True, "n2o_effect_pct"].min(),
    grass_out.loc[grass_out["is_freeze_region"] == True, "n2o_effect_pct"].max()
)

print("\n九个梯度数量：")
print(grass_out["effect_class"].value_counts().sort_index())

print("\n各区域N2O效应均值（%）：")
print(grass_out.groupby("region")["n2o_effect_pct"].mean().sort_values())