"""Stage `ice`: ice sheets, mountain glaciers, seasonal/perennial sea ice."""
from __future__ import annotations

import numpy as np

from .config import params
from .graph import components, distance_to

ICE_NONE, ICE_SHEET, ICE_GLACIER, ICE_SEA_SEASONAL, ICE_SEA_PERENNIAL = range(5)
ICE_NAMES = ["none", "ice sheet", "glacier", "seasonal sea ice", "perennial sea ice"]
DEFAULTS = {"melt_summer_c": 0.0, "sheet_min_km2": 250000.0, "sheet_min_p_mm": 100.0, "sea_ice_t_c": -1.8, "sheet_max_m": 3000.0,
            "sheet_edge_m": 200.0, "sheet_growth_m_per_km": 3.0}


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "ice", DEFAULTS)
    z, t_mean, t_jun, t_dec, p_ann = ctx.need("elevation_eroded_m", "T_mean", "T_jun", "T_dec", "P_ann")
    z = z.astype(np.float64)
    land = ~np.asarray(ctx.data["ocean"]) if "ocean" in ctx.data else z > 0
    summer = np.where(g.lat >= 0, t_jun, t_dec)
    winter = np.where(g.lat >= 0, t_dec, t_jun)
    cold = land & (summer < P["melt_summer_c"])                  # snow survives the summer → permanent ice
    lab = components(g, cold)
    area = np.bincount(lab[cold], weights=g.area_km2[cold]) if cold.any() else np.zeros(1)
    big = cold & (area[np.maximum(lab, 0)] >= P["sheet_min_km2"])
    sheet = big & (p_ann > P["sheet_min_p_mm"])
    thick = np.zeros(g.n)
    if sheet.any():
        d_edge = distance_to(g, ~sheet)
        thick = np.where(sheet, np.minimum(P["sheet_max_m"], P["sheet_edge_m"] + P["sheet_growth_m_per_km"] * d_edge), 0.0)
    ice = np.zeros(g.n, np.int8)
    ice[cold & ~sheet] = ICE_GLACIER
    ice[sheet] = ICE_SHEET
    ice[~land & (winter < P["sea_ice_t_c"])] = ICE_SEA_SEASONAL
    ice[~land & (summer < P["sea_ice_t_c"])] = ICE_SEA_PERENNIAL
    return {"ice": ice, "ice_thickness_m": thick, "z_surface_m": z + thick}