"""Add MP-weighted support columns to right_wing_motions.

Adds centrist_support_mp, centrist_support_strict, center_right_support,
and left_support_mp — all computed as the fraction of individual MPs
within each party set who voted 'voor'.
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).parent.parent.parent.resolve()
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import duckdb

from analysis.config import CANONICAL_LEFT

CANONICAL_CENTRIST = frozenset({"VVD", "D66", "CDA", "NSC", "BBB", "CU"})
CANONICAL_CENTRIST_STRICT = frozenset({"D66", "CDA", "CU", "NSC"})
CANONICAL_CENTER_RIGHT = frozenset({"VVD", "BBB"})

COLUMNS = [
    ("centrist_support_mp", CANONICAL_CENTRIST),
    ("centrist_support_strict", CANONICAL_CENTRIST_STRICT),
    ("center_right_support", CANONICAL_CENTER_RIGHT),
    ("left_support_mp", CANONICAL_LEFT),
]


def compute_mp_support(
    votes: dict[str, dict[str, int]], parties: frozenset[str]
) -> float | None:
    total_voor = 0
    total_cast = 0
    for party, pv in votes.items():
        if party not in parties:
            continue
        voor = pv.get("voor", 0)
        tegen = pv.get("tegen", 0)
        tv = voor + tegen
        if tv == 0:
            continue
        total_voor += voor
        total_cast += tv
    if total_cast == 0:
        return None
    return total_voor / total_cast


def main(db_path: str = "data/motions.db"):
    db = Path(db_path)
    con = duckdb.connect(str(db))

    votemap: dict[int, dict[str, dict[str, int]]] = {}
    vote_rows = con.execute(
        """
        SELECT motion_id, party, vote, COUNT(*) as n
        FROM mp_votes
        WHERE party IS NOT NULL
        GROUP BY motion_id, party, vote
        """
    ).fetchall()

    for motion_id, party, vote, n in vote_rows:
        mv = votemap.setdefault(motion_id, {})
        pv = mv.setdefault(party, {"voor": 0, "tegen": 0, "afwezig": 0})
        pv[vote] = pv.get(vote, 0) + n

    # Add columns if missing
    for col_name, _party_set in COLUMNS:
        col_check = con.execute(
            "SELECT column_name FROM information_schema.columns "
            "WHERE table_name = 'right_wing_motions' AND column_name = ?",
            [col_name],
        ).fetchone()
        if col_check is None:
            con.execute(
                f"ALTER TABLE right_wing_motions ADD COLUMN {col_name} DOUBLE"
            )
            print(f"Added {col_name} column")

    # Update rows
    rows = con.execute(
        "SELECT motion_id FROM right_wing_motions"
    ).fetchall()

    updated = 0
    skipped = 0
    for (motion_id,) in rows:
        votes = votemap.get(motion_id)
        if votes is None:
            skipped += 1
            continue
        for col_name, party_set in COLUMNS:
            val = compute_mp_support(votes, party_set)
            con.execute(
                f"UPDATE right_wing_motions SET {col_name} = ? WHERE motion_id = ?",
                [val, motion_id],
            )
        updated += 1

    con.close()
    print(f"Updated {updated} rows, skipped {skipped}")


if __name__ == "__main__":
    main()