from __future__ import annotations

import math
from collections import defaultdict
from decimal import Decimal, ROUND_HALF_UP

from .excel_parse import class_sort_key

LAYER_SIZE = 20
MISSING_AS_SCORE = 0.0


def infer_track_group(track: str) -> str:
    t = track or ""
    if "历史" in t:
        return "历史"
    if "物理" in t:
        return "物理"
    return t or "未分方向"


def default_line_settings(groups: list[str]) -> dict[str, dict[str, int]]:
    out = {}
    for g in groups:
        if g == "物理":
            out[g] = {"top": 80, "control": 262, "layer_size": 20}
        elif g == "历史":
            # 历史样本较小；basedata.xlsx 的历史方向使用 5 人一层。
            out[g] = {"top": 10, "control": 45, "layer_size": 5}
        else:
            out[g] = {"top": 20, "control": 70, "layer_size": 20}
    return out


def select_students(parsed: dict, settings: dict) -> list[dict]:
    include_mingyuan = settings.get("include_mingyuan", True)
    include_other = settings.get("include_other", True)
    exclude_moves = bool(settings.get("exclude_moves", False))
    exams = settings.get("exams") or parsed["exams"]
    selected = []
    for s in parsed["students"]:
        move_flag = bool((s.get("move") or "").strip()) or s.get("sheet_kind") in ("art", "leave")
        if exclude_moves and move_flag:
            continue
        st = s.get("status") or "在籍"
        if st == "明远" and not include_mingyuan:
            continue
        if st not in ("在籍", "明远", "") and not include_other:
            continue
        row = dict(s)
        row["move_flag"] = move_flag
        row["scores"] = {e: s["scores"].get(e) for e in exams}
        row["computed_scores"] = {
            e: MISSING_AS_SCORE if v is None else float(v) for e, v in row["scores"].items()
        }
        row["missing_exams"] = [e for e, v in row["scores"].items() if v is None]
        row["group"] = infer_track_group(s.get("track") or "")
        selected.append(row)
    return selected


def competition_ranks(values: list[float]) -> list[int]:
    """1 = highest; all equal scores receive the same rank and later ranks are skipped."""
    return [1 + sum(1 for v in values if v > x) for x in values]


def layer_count(n: int, size: int = LAYER_SIZE) -> int:
    return max(1, math.ceil(n / size)) if n else 1


def rank_layer(rank: int, size: int = LAYER_SIZE) -> int:
    """A tie at a boundary stays in the upper layer because the next distinct rank is skipped."""
    return max(1, (int(rank) - 1) // size + 1)


def layer_label(layer: int, n: int, size: int = LAYER_SIZE) -> str:
    start = (layer - 1) * size + 1
    if layer * size >= n:
        return f"前{start}名及以后"
    return f"前{start}-{layer * size}名"


def _sigmoid(x: float) -> float:
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-min(x, 40.0)))
    z = math.exp(max(x, -40.0))
    return z / (1.0 + z)


def _log_likelihood(beta: list[float], rows: list[tuple[float, float, float]]) -> float:
    total = 0.0
    for t, n, k in rows:
        p = _sigmoid(beta[0] - beta[1] * t)
        p = min(max(p, 1e-10), 1 - 1e-10)
        total += k * math.log(p) + (n - k) * math.log(1 - p)
    return total


def fit_retention_probabilities(
    counts: list[dict], round_digits: int = 6
) -> list[float]:
    """Weighted binomial logistic fit: p=1/(1+exp(-(a-b*layer)))."""
    rows = [(float(x["layer"]), float(x["previous_n"]), float(x["hit_n"])) for x in counts]
    if not rows:
        return []

    beta = [0.3, 0.15]
    best = _log_likelihood(beta, rows)
    for _ in range(100):
        g1 = g2 = 0.0
        h11 = h12 = h22 = 0.0
        for t, n, k in rows:
            p = _sigmoid(beta[0] - beta[1] * t)
            w = n * p * (1 - p)
            r = k - n * p
            xt1, xt2 = 1.0, -t
            g1 += xt1 * r
            g2 += xt2 * r
            h11 += w * xt1 * xt1
            h12 += w * xt1 * xt2
            h22 += w * xt2 * xt2
        # A tiny ridge prevents singular matrices when a layer is perfectly separated.
        h11 += 1e-7
        h22 += 1e-7
        det = h11 * h22 - h12 * h12
        if abs(det) < 1e-14:
            break
        d1 = (h22 * g1 - h12 * g2) / det
        d2 = (h11 * g2 - h12 * g1) / det
        step = 1.0
        improved = False
        for _ls in range(20):
            cand = [beta[0] + step * d1, beta[1] + step * d2]
            # Keep the curve monotone decreasing and numerically bounded.
            cand[0] = min(max(cand[0], -6.0), 14.0)
            cand[1] = min(max(cand[1], 0.0), 10.0)
            value = _log_likelihood(cand, rows)
            if value + 1e-9 >= best:
                beta, best = cand, value
                improved = True
                break
            step *= 0.5
        if not improved or math.sqrt(d1 * d1 + d2 * d2) < 1e-8:
            break

    raw = [_sigmoid(beta[0] - beta[1] * t) for t, _, _ in rows]

    # basedata.xlsx 的约定：第 1 层不使用 logistic 估计值，而是保留合成概率。
    # 这样小样本首层的边界观测不会被模型收缩。
    if counts and counts[0]["previous_n"]:
        raw[0] = counts[0]["hit_n"] / counts[0]["previous_n"]

    # basedata.xlsx 的尾部约定：第一次出现命中为 0 的层及其后，拟合概率取 0。
    first_zero = next((i for i, item in enumerate(counts) if item["hit_n"] == 0), None)
    if first_zero is not None:
        raw[first_zero:] = [0.0] * (len(raw) - first_zero)

    # Force monotone decreasing, even when sparse layers reverse the fitted trend.
    for i in range(1, len(raw)):
        raw[i] = min(raw[i], raw[i - 1])
    return [round(max(0.0, min(1.0, p)), round_digits) for p in raw]


def round_half_up(value: float) -> int:
    return int(Decimal(str(float(value))).quantize(Decimal("1"), rounding=ROUND_HALF_UP))


def _pair_counts(
    members: list[dict], e0: str, e1: str, threshold: int, size: int = LAYER_SIZE
) -> list[dict]:
    total_layers = layer_count(len(members), size)
    buckets = [
        {"layer": i, "previous_n": 0, "hit_n": 0} for i in range(1, total_layers + 1)
    ]
    for s in members:
        layer = rank_layer(s["ranks"][e0], size)
        buckets[layer - 1]["previous_n"] += 1
        if s["ranks"][e1] <= threshold:
            buckets[layer - 1]["hit_n"] += 1
    return buckets


def _merge_counts(parts: list[list[dict]]) -> list[dict]:
    n_layers = max(len(p) for p in parts)
    out = [{"layer": i, "previous_n": 0, "hit_n": 0} for i in range(1, n_layers + 1)]
    for part in parts:
        for item in part:
            out[item["layer"] - 1]["previous_n"] += item["previous_n"]
            out[item["layer"] - 1]["hit_n"] += item["hit_n"]
    return out


def analyze(parsed: dict, settings: dict) -> dict:
    exams = settings.get("exams") or parsed["exams"]
    if len(exams) < 2:
        raise ValueError("相邻层次法至少需要两次考试")
    selected = select_students(parsed, settings)
    if not selected:
        raise ValueError("没有学生进入测算")
    groups = sorted({s["group"] for s in selected})
    line_settings = settings.get("line_settings") or default_line_settings(groups)

    # 用户口径：勾选后的所有相邻两次都要计算，最后一次也进入相邻转移样本。
    probability_exams = exams

    # 排名与层次划分；缺考已按 0 分参加排名。
    for s in selected:
        s["ranks"] = {}
        s["layers"] = {}
    for e in exams:
        vals = [s["computed_scores"][e] for s in selected]
        # Rankings are global within each direction, so calculate by group.
        for g in groups:
            gm = [s for s in selected if s["group"] == g]
            gvals = [s["computed_scores"][e] for s in gm]
            ranks = competition_ranks(gvals)
            for s, rank in zip(gm, ranks):
                s["ranks"][e] = rank
                s["layers"][e] = rank_layer(rank)

    results = []
    all_views = []
    for g in groups:
        members = [s for s in selected if s["group"] == g]
        n = len(members)
        lines = line_settings.get(g) or default_line_settings([g])[g]
        top_k = max(1, min(int(lines.get("top", 1)), n))
        control_k = max(top_k, min(int(lines.get("control", top_k)), n))
        line_defs = [
            {"id": "top", "name": "尖子线", "k": top_k},
            {"id": "control", "name": "特控线", "k": control_k},
        ]
        layer_size = max(1, min(int(lines.get("layer_size", LAYER_SIZE)), max(1, n)))
        n_layers = layer_count(n, layer_size)
        labels = [layer_label(i, n, layer_size) for i in range(1, n_layers + 1)]
        pair_blocks = []
        summary_lines = {}
        target_blocks = []
        for line in line_defs:
            pair_parts = []
            pair_details = []
            pair_labels = [layer_label(i, n, layer_size) for i in range(1, n_layers + 1)]
            for a, b in zip(probability_exams, probability_exams[1:]):
                part = _pair_counts(members, a, b, line["k"], layer_size)
                pair_parts.append(part)
                pair_details.append(
                    {
                        "from_exam": a,
                        "to_exam": b,
                        "layers": [
                            {
                                **x,
                                "label": pair_labels[x["layer"] - 1],
                                "probability": (
                                    x["hit_n"] / x["previous_n"]
                                    if x["previous_n"]
                                    else None
                                ),
                            }
                            for x in part
                        ],
                    }
                )
            merged = _merge_counts(pair_parts)
            for item, label in zip(merged, labels):
                item["label"] = label
                item["probability"] = (
                    item["hit_n"] / item["previous_n"]
                    if item["previous_n"]
                    else None
                )
            fitted = fit_retention_probabilities(merged)
            for item, p in zip(merged, fitted):
                item["fitted_probability"] = p

            classes = sorted({s["class_name"] for s in members}, key=class_sort_key)
            class_rows = []
            for c in classes:
                cm = [s for s in members if s["class_name"] == c]
                layer_counts = [0] * n_layers
                for s in cm:
                    layer_counts[rank_layer(s["ranks"][exams[-1]], layer_size) - 1] += 1
                expected = sum(a * p for a, p in zip(layer_counts, fitted))
                current = sum(1 for s in cm if s["ranks"][exams[-1]] <= line["k"])
                class_rows.append(
                    {
                        "class_name": c,
                        "n": len(cm),
                        "current": current,
                        "layer_counts": layer_counts,
                        "expected": round(expected, 4),
                        "assigned": round_half_up(expected),
                        "delta": round_half_up(expected) - current,
                    }
                )
            summary_lines[line["id"]] = {
                "layers": merged,
                "labels": labels,
                "total_previous": sum(x["previous_n"] for x in merged),
                "total_hit": sum(x["hit_n"] for x in merged),
                "fitted": fitted,
            }
            target_blocks.append(
                {
                    **line,
                    "classes": class_rows,
                    "layers": merged,
                    "labels": labels,
                    "fitted": fitted,
                    "sum_expected": round(sum(r["expected"] for r in class_rows), 4),
                    "sum_assigned": sum(r["assigned"] for r in class_rows),
                    "sum_current": sum(r["current"] for r in class_rows),
                }
            )
            pair_blocks.append({"line": line["name"], "pairs": pair_details})

        results.append(
            {
                "group": g,
                "n": n,
                "layer_size": layer_size,
                "n_layers": n_layers,
                "line_settings": {"top": top_k, "control": control_k},
                "pairs": pair_blocks,
                "summary": summary_lines,
                "targets": target_blocks,
            }
        )

    for s in selected:
        all_views.append(
            {
                "sid": s["sid"],
                "name": s["name"],
                "class_name": s["class_name"],
                "group": s["group"],
                "track": s.get("track") or "",
                "subjects": s.get("subjects") or "",
                "status": s.get("status") or "",
                "move": s.get("move") or "",
                "has_move": bool(s.get("move_flag", False)),
                "scores": s["scores"],
                "computed_scores": s["computed_scores"],
                "missing_exams": s["missing_exams"],
                "ranks": s["ranks"],
                "layers": s["layers"],
            }
        )

    return {
        "algorithm": "adjacent-layer-markov-v2",
        "exams": exams,
        "layer_size": LAYER_SIZE,
        "pair_count": len(probability_exams) - 1 if len(probability_exams) > 1 else 0,
        "probability_exams": probability_exams,
        "missing_as_score": MISSING_AS_SCORE,
        "exclude_moves": bool(settings.get("exclude_moves", False)),
        "n_selected": len(selected),
        "n_total": len(parsed["students"]),
        "groups": results,
        "students": all_views,
    }


def class_comment(row: dict, target_name: str) -> str:
    assigned = row["assigned"]
    current = row["current"]
    expected = row["expected"]
    if assigned > current:
        tail = f"比最后一次到线人数多 {assigned - current} 人。"
    elif assigned < current:
        tail = f"比最后一次到线人数少 {current - assigned} 人。"
    else:
        tail = "与最后一次到线人数持平。"
    return f"{row['class_name']}班{target_name}拟合摊派 {assigned} 人（加权计算 {expected:.2f}）。{tail}"