import numpy as np
from modelling.math_utils import natural_log
from modelling.row_schema import ModeledRowsPayload


def compute_v1_team_surprise(event):
    """Return per-team V1 surprise contribution for a single event."""
    if not isinstance(event, dict):
        return None, None
    if event.get("is_marker_only_action"):
        return None, None
    if event.get("exclude_from_surprise"):
        return None, None

    dice_roller_raw = event.get("dice_roller")
    if dice_roller_raw is None:
        dice_roller_raw = event.get("team_id")
    try:
        team_id = int(str(dice_roller_raw))
    except (TypeError, ValueError):
        return None, None

    if team_id not in (0, 1):
        return None, None

    p_success = event.get("probability_success")
    p_fail = event.get("probability_fail")
    if not isinstance(p_success, (int, float)):
        return None, None
    if isinstance(p_fail, (int, float)):
        p_fail_value = float(p_fail)
    else:
        p_fail_value = None

    surprise = 0.0
    classification = str(
        event.get(
            "report_result_classification",
            event.get("result_classification", ""),
        )
    ).strip().lower()

    if classification == "success":
        if float(p_success) <= 0 or float(p_success) >= 1.0:
            return None, None
        surprise = np.log(1.0 / float(p_success))
    elif classification == "fail":
        if p_fail_value is None or p_fail_value <= 0:
            return None, None
        surprise = np.log(p_fail_value)
    elif classification in ("neutral", "unknown", ""):
        surprise = 0.0
    else:
        result_value = event.get("result_value")
        if not isinstance(result_value, (int, float)):
            result_value = 0
        if result_value == 1:
            if float(p_success) <= 0 or float(p_success) >= 1.0:
                return None, None
            surprise = np.log(1.0 / float(p_success))
        elif result_value == -1:
            if p_fail_value is None or p_fail_value <= 0:
                return None, None
            surprise = np.log(p_fail_value)

    if team_id == 0:
        return float(surprise), 0.0
    return 0.0, float(surprise)


def expected_v1_surprise_scalar(event):
    """Return expected per-event V1 surprise E[X] for current probabilities."""
    if not isinstance(event, dict):
        return None
    if event.get("is_marker_only_action"):
        return None
    if event.get("exclude_from_surprise"):
        return None

    p_success = event.get("probability_success")
    p_fail = event.get("probability_fail")
    if not isinstance(p_success, (int, float)) or not isinstance(p_fail, (int, float)):
        return None

    expected = 0.0
    if p_success > 0:
        expected += float(p_success) * float(natural_log(1.0 / float(p_success)))
    if p_fail > 0:
        expected += float(p_fail) * float(natural_log(float(p_fail)))
    return float(expected)


def expected_v1_surprise_by_team(event):
    """Return expected per-team V1 surprise allocated to the dice-rolling team."""
    expected = expected_v1_surprise_scalar(event)
    if expected is None:
        return None, None

    dice_roller_raw = event.get("dice_roller")
    if dice_roller_raw is None:
        dice_roller_raw = event.get("team_id")

    try:
        team_id = int(str(dice_roller_raw))
    except (TypeError, ValueError):
        return None, None

    if team_id == 0:
        return expected, 0.0
    if team_id == 1:
        return 0.0, expected
    return None, None


def materialize_modeled_rows(modeled_sections) -> ModeledRowsPayload:
    """Flatten modeled sections into row lists used by reports and stats."""
    rows = []
    kickoff_rows = []

    for turn_info in modeled_sections:
        turn_key = turn_info.get("turn_key")
        for section in turn_info.get("sections", []):
            for kickoff_row in section.get("kickoff_report_rows") or []:
                if isinstance(kickoff_row, dict):
                    kickoff_row_copy = dict(kickoff_row)
                    if kickoff_row_copy.get("game_turn") is None:
                        kickoff_row_copy["game_turn"] = turn_key
                    kickoff_rows.append(kickoff_row_copy)
                    rows.append(kickoff_row_copy)
            for event in section.get("events") or []:
                if isinstance(event, dict):
                    event_copy = dict(event)
                    event_copy["game_turn"] = turn_key
                    rows.append(event_copy)

    return {
        "rows": rows,
        "kickoff_rows": kickoff_rows,
        "sections": modeled_sections,
    }