"""Deterministic post-processing of the raw extraction results.

parse_results(name)        → per-image (or per-listing for text) DataFrame from results/raw/batches/<name>.results.jsonl
aggregate_images(df, K)    → one row per listing: cross-image aggregation rules (documented below), optionally truncated
                              to the first K images in seller order
reconcile(visual, text, admin) → B-level property tables (B0…B4) with per-attribute evidence status

Aggregation rules (fixed before results were seen; docs/experiment_plan.md):
  - Exterior attributes (building type, storeys, cladding, roof, garage, exterior condition) use exterior scenes only.
  - Categorical: confidence-weighted plurality vote among informative values (not 'not_visible'/'not_applicable');
    if no informative vote → 'not_visible'. Garage: any attached/detached/carport vote beats 'none_visible'.
  - Ordinal grades (kitchen, bathroom): weighted median of ordinal codes among informative values.
  - Conditions and finish quality: confidence-weighted mean of non-zero scores.
  - Renovation evidence: maximum of (none_apparent < partial < extensive), 'unsure' ignored.
  - Defects: union excluding 'none_visible'. Photo quality: mean. Staging share: share of interior images not 'none_apparent'.
  - Composite condition = mean(exterior_condition, interior_condition) over those available.
  - Composite finish grade (for the cost schedule) = weighted ordinal 0.5 kitchen + 0.3 bathroom + 0.2 (interior finish
    quality rescaled to 1–4), rounded to {basic, standard, upgraded, luxury}; None when no interior grade is visible.
"""
from __future__ import annotations

import json
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[2]
RAW = ROOT / "results" / "raw" / "batches"

EXTERIOR = {"exterior_front", "exterior_rear_or_side", "exterior_aerial_or_street"}
INTERIOR = {"kitchen", "bathroom", "living_or_dining", "bedroom", "basement"}
NONINFO = {"not_visible", "not_applicable", "not_stated", None}
GRADE_ORD = {"basic": 1, "standard": 2, "upgraded": 3, "luxury": 4}
ORD_GRADE = {v: k for k, v in GRADE_ORD.items()}
RENO_ORD = {"none_apparent": 0, "partial": 1, "extensive": 2}


def parse_cid(cid: str) -> tuple[str, str, str]:
    tier, rest = cid.split("-", 1)
    uid, idx = rest.rsplit("-", 1)
    return tier, uid.replace("__", ":"), idx


def parse_results(name: str) -> pd.DataFrame:
    f = RAW / f"{name}.results.jsonl"
    rows = []
    for line in open(f):
        r = json.loads(line)
        tier, uid, idx = parse_cid(r["custom_id"])
        base = {"tier": tier, "uid": uid, "image_index": None if idx == "T" else int(idx), "ok": r["type"] == "succeeded" and r.get("data") is not None,
                "error": r.get("error"), "error_message": r.get("error_message"), "stop_reason": r.get("stop_reason"),
                "input_tokens": (r.get("usage") or {}).get("input_tokens"), "output_tokens": (r.get("usage") or {}).get("output_tokens")}
        if r.get("data"):
            base.update(r["data"])
        rows.append(base)
    df = pd.DataFrame(rows)
    return df


def _wvote(values, weights, exclude=NONINFO):
    c = defaultdict(float)
    for v, w in zip(values, weights):
        if v in exclude:
            continue
        c[v] += max(w, 0.05)
    if not c:
        return None, 0.0
    v, w = max(c.items(), key=lambda kv: kv[1])
    return v, w / sum(c.values())


def _wmedian_ord(values, weights, ordmap):
    pairs = [(ordmap[v], max(w, 0.05)) for v, w in zip(values, weights) if v in ordmap]
    if not pairs:
        return None
    pairs.sort()
    tot = sum(w for _, w in pairs)
    acc = 0.0
    for v, w in pairs:
        acc += w
        if acc >= tot / 2:
            return v
    return pairs[-1][0]


def _wmean(values, weights):
    pairs = [(float(v), max(w, 0.05)) for v, w in zip(values, weights) if v is not None and not pd.isna(v) and float(v) > 0]
    if not pairs:
        return np.nan
    return sum(v * w for v, w in pairs) / sum(w for _, w in pairs)


def aggregate_images(img: pd.DataFrame, K: int | None = None) -> pd.DataFrame:
    df = img[img["ok"]].copy()
    if K is not None:
        df = df[df["image_index"] < K]
    out = []
    for uid, g in df.groupby("uid"):
        conf = g["confidence"].fillna(0.5).astype(float).tolist()
        ext = g[g["scene"].isin(EXTERIOR)]
        kit = g[g["scene"] == "kitchen"]
        bath = g[g["scene"] == "bathroom"]
        bsm = g[g["scene"] == "basement"]
        liv = g[g["scene"].isin({"living_or_dining", "bedroom", "kitchen"})]
        inte = g[g["scene"].isin(INTERIOR)]
        ec = ext["confidence"].fillna(0.5).astype(float).tolist()
        r = {"uid": uid, "n_images_ok": len(g), "n_exterior": len(ext), "n_kitchen": len(kit), "n_bathroom": len(bath),
             "n_basement": len(bsm), "n_living": len(liv), "n_unusable": int((g["scene"] == "unusable").sum()),
             "n_floorplan": int((g["scene"] == "floor_plan_or_document").sum()),
             "photo_quality_mean": g["photo_quality"].astype(float).mean(),
             "staging_share": float((inte["staging_or_rendering"] != "none_apparent").mean()) if len(inte) else np.nan,
             "confidence_mean": float(np.mean(conf)) if conf else np.nan}
        for col in ["building_type_visual", "storeys_visible", "roof_form", "roof_covering"]:
            v, s = _wvote(ext[col].tolist(), ec)
            r[col], r[col + "_share"] = (v or "not_visible"), s
        # cladding: votes over list elements
        cl_vals, cl_w = [], []
        for lst, w in zip(ext["cladding"].tolist(), ec):
            for m in (lst if isinstance(lst, (list, np.ndarray)) else []):
                cl_vals.append(m); cl_w.append(w)
        v, s = _wvote(cl_vals, cl_w)
        r["cladding_primary"], r["cladding_primary_share"] = (v or "not_visible"), s
        cnt = Counter(m for m in cl_vals if m not in NONINFO)
        r["cladding_set"] = sorted(cnt)
        # garage
        gv = ext["garage_visible"].tolist()
        pos = [x for x in gv if x in ("attached", "detached", "carport")]
        r["garage_visible"] = Counter(pos).most_common(1)[0][0] if pos else ("none_visible" if "none_visible" in gv else "not_visible")
        r["exterior_condition"] = _wmean(ext["exterior_condition"].tolist(), ec)
        r["interior_condition"] = _wmean(inte["interior_condition"].tolist(), inte["confidence"].fillna(0.5).tolist())
        r["interior_finish_quality"] = _wmean(inte["interior_finish_quality"].tolist(), inte["confidence"].fillna(0.5).tolist())
        kg = _wmedian_ord(kit["kitchen_grade"].tolist(), kit["confidence"].fillna(0.5).tolist(), GRADE_ORD)
        bg = _wmedian_ord(bath["bathroom_grade"].tolist(), bath["confidence"].fillna(0.5).tolist(), GRADE_ORD)
        r["kitchen_grade"] = ORD_GRADE.get(kg, "not_visible")
        r["bathroom_grade"] = ORD_GRADE.get(bg, "not_visible")
        for col, sub in [("kitchen_counters", kit), ("kitchen_cabinets", kit), ("basement_finish", bsm), ("flooring_main", liv)]:
            v, s = _wvote(sub[col].tolist(), sub["confidence"].fillna(0.5).tolist())
            r[col], r[col + "_share"] = (v or "not_visible"), s
        ren = [RENO_ORD[x] for x in g["renovation_evidence"].tolist() if x in RENO_ORD]
        r["renovation_evidence"] = {0: "none_apparent", 1: "partial", 2: "extensive"}[max(ren)] if ren else "unsure"
        defs = set()
        for lst in g["visible_defects"].tolist():
            for d in (lst if isinstance(lst, (list, np.ndarray)) else []):
                if d != "none_visible":
                    defs.add(d)
        r["defects"] = sorted(defs)
        r["n_defects"] = len(defs)
        conds = [c for c in (r["exterior_condition"], r["interior_condition"]) if not pd.isna(c)]
        r["condition_composite"] = float(np.mean(conds)) if conds else np.nan
        # composite finish grade
        parts, wts = [], []
        if kg: parts.append(kg); wts.append(0.5)
        if bg: parts.append(bg); wts.append(0.3)
        if not pd.isna(r["interior_finish_quality"]):
            parts.append(1 + 3 * (r["interior_finish_quality"] - 1) / 4); wts.append(0.2)
        r["grade_score"] = float(np.average(parts, weights=wts)) if parts else np.nan
        r["grade_composite"] = ORD_GRADE[int(min(4, max(1, round(r["grade_score"]))))] if parts else None
        out.append(r)
    return pd.DataFrame(out)


# ------------------------------------------------------------------------------------------------ reconciliation
def _status(vis, txt):
    v_inf = vis not in NONINFO
    t_inf = txt not in NONINFO
    if v_inf and t_inf:
        return "agree" if vis == txt else "contradict"
    if v_inf:
        return "visual_only"
    if t_inf:
        return "text_only"
    return "unknown"


def build_levels(frame: pd.DataFrame, vis: pd.DataFrame | None, txt: pd.DataFrame | None) -> dict[str, pd.DataFrame]:
    """Return {'B0','B1','B2','B3','B4'} property tables with the columns required by src/cost_model/cost.py."""
    base = frame.set_index("uid")
    attached_from_type = base["property_type"].map({"Maison": "detached", "Détachée": "detached", "Jumelé": "semi_detached",
                                                     "Maison de ville": "row_or_townhouse"}).fillna("detached")
    roll_storeys = base["roll_storeys"].map(lambda x: {1.0: "1", 1.5: "1.5", 2.0: "2", 2.5: "2.5", 3.0: "3"}.get(float(x) if x == x else 1.0, "1"))
    common = pd.DataFrame({"floor_area_m2": base["roll_floor_area_m2"], "age": base["age_2026"], "cma": base["cma"],
                           "roll_land_value": base["roll_land_value"], "roll_market_date": base["roll_market_date"],
                           "roll_building_value": base["roll_building_value"], "price": base["price"],
                           "weight": base["weight"] if "weight" in base else 1.0, "source": base["source"],
                           "age_band": base["age_band"].astype(str), "price_tier": base["price_tier"].astype(str),
                           "roll_municipalite": base["roll_municipalite"]})
    levels = {}
    # B0 — administrative fields + typical assumptions
    b0 = common.copy()
    b0["storeys"] = roll_storeys; b0["attached"] = attached_from_type
    b0["grade"] = "standard"; b0["condition"] = 3.0; b0["renovation"] = "none_apparent"
    b0["basement"] = "unfinished"; b0["garage"] = "none"; b0["pool"] = "none"
    levels["B0"] = b0
    if txt is not None:
        t = txt.set_index("uid").reindex(base.index)
        b1 = b0.copy()
        tg = t["kitchen_grade"].where(~t["kitchen_grade"].isin(NONINFO))
        bgd = t["bathroom_grade"].where(~t["bathroom_grade"].isin(NONINFO))
        g = pd.concat([tg.map(GRADE_ORD), bgd.map(GRADE_ORD)], axis=1).mean(axis=1)
        b1["grade"] = g.round().map(ORD_GRADE).fillna("standard")
        b1["condition"] = t["condition_claim"].map({"needs_work": 2.0, "average": 3.0, "well_maintained": 3.5,
                                                    "renovated_or_like_new": 4.5, "new_construction": 5.0}).fillna(3.0)
        b1["renovation"] = t["renovation_claim"].map({"none": "none_apparent", "partial": "partial", "extensive": "extensive"}).fillna("none_apparent")
        b1["basement"] = t["basement_finish"].map({"finished": "finished", "partially_finished": "partially_finished",
                                                   "unfinished": "unfinished", "no_basement": "unfinished"}).fillna("unfinished")
        b1["garage"] = t["garage"].map({"attached": "attached", "detached": "detached", "carport": "carport", "none": "none"}).fillna("none")
        b1["pool"] = t["pool"].where(~t["pool"].isin(NONINFO), "none")
        b1["storeys"] = t["storeys"].where(~t["storeys"].isin(NONINFO), roll_storeys)
        b1["attached"] = t["building_type"].where(~t["building_type"].isin(NONINFO), attached_from_type)
        levels["B1"] = b1
    if vis is not None:
        v = vis.set_index("uid").reindex(base.index)
        b2 = b0.copy()
        b2["grade"] = v["grade_composite"].fillna("standard")
        b2["grade_visible"] = v["grade_composite"].notna()
        b2["condition"] = v["condition_composite"].fillna(3.0)
        b2["condition_visible"] = v["condition_composite"].notna()
        b2["renovation"] = v["renovation_evidence"].replace({"unsure": "none_apparent"}).fillna("none_apparent")
        b2["basement"] = v["basement_finish"].where(~v["basement_finish"].isin(NONINFO), "unfinished")
        b2["garage"] = v["garage_visible"].map({"attached": "attached", "detached": "detached", "carport": "carport"}).fillna("none")
        b2["storeys"] = v["storeys_visible"].map({"1": "1", "1.5": "1.5", "2": "2", "2.5": "2.5", "3+": "3"}).fillna(roll_storeys)
        b2["attached"] = v["building_type_visual"].where(v["building_type_visual"].isin(["detached", "semi_detached", "row_or_townhouse"]), attached_from_type)
        levels["B2"] = b2
    if vis is not None and txt is not None:
        b3 = b2.copy()
        # visual first, text fills gaps; contradictions recorded
        for col, vcol, tcol in [("basement", "basement_finish", "basement_finish"), ("garage", "garage_visible", "garage"),
                                ("storeys", "storeys_visible", "storeys"), ("attached", "building_type_visual", "building_type")]:
            vv = v[vcol].replace({"3+": "3", "none_visible": "none"})
            tt = t[tcol].replace({"no_basement": "unfinished"})
            st = [_status(a, b) for a, b in zip(vv, tt)]
            b3[col + "_status"] = st
            b3[col] = np.where(~vv.isin(NONINFO), vv, np.where(~tt.isin(NONINFO), tt, b3[col]))
        # grade: visual composite, else text grade; condition: visual, else text claim
        b3["grade"] = v["grade_composite"].fillna(b1["grade"])
        b3["condition"] = v["condition_composite"].fillna(b1["condition"])
        b3["renovation"] = np.where(v["renovation_evidence"].isin(["partial", "extensive"]), v["renovation_evidence"], b1["renovation"])
        b3["pool"] = b1["pool"]
        b3["n_contradictions"] = sum((b3[c + "_status"] == "contradict").astype(int) for c in ["basement", "garage", "storeys", "attached"])
        levels["B3"] = b3
        # B4: administrative override (storeys from roll; attached from listing type when visual not visible) + flags for MC widening
        b4 = b3.copy()
        b4["storeys"] = roll_storeys
        b4["storeys_status"] = np.where(v["storeys_visible"].replace({"3+": "3"}) == roll_storeys, "agree",
                                        np.where(v["storeys_visible"].isin(NONINFO), "admin_only", "contradict"))
        b4["uncertain_grade"] = v["grade_composite"].isna()
        b4["uncertain_condition"] = v["condition_composite"].isna()
        levels["B4"] = b4
    return levels