simulation.py

from __future__ import annotations

from collections import defaultdict

from economy import collect_bids, relationship, select_winning_bids
from llm import GossipLLM
from models import Bid, Channel, ConversationResult, Event, Evidence, Memory, Rumor, VillageState, Villager
from state import clamp

REPUTATION_WEIGHTS = {
    "financial": {"honesty": 0.7, "competence": 0.3},
    "moral": {"morality": 0.8, "honesty": 0.2},
    "leadership": {"leadership": 0.7, "competence": 0.3},
    "relationship": {"morality": 0.5, "kindness": 0.3, "honesty": 0.2},
    "competence": {"competence": 0.8, "leadership": 0.2},
    "kindness": {"kindness": 0.8, "morality": 0.2},
    "political": {"leadership": 0.5, "honesty": 0.3, "competence": 0.2},
}

EVIDENCE_TEMPLATES = {
    "financial": [
        ("Receipt found near the SACCO office", 0.72),
        ("Chama records conflict with the rumor", 0.84),
        ("Audit document mentioned by Cherono", 0.9),
    ],
    "moral": [
        ("Pastor privately knows missing context", 0.74),
        ("Witness contradicts the loudest version", 0.68),
    ],
    "political": [
        ("WhatsApp screenshot appears from a campaign group", 0.66),
        ("Someone saw a late-night meeting at the chief's office", 0.62),
    ],
    "competence": [
        ("School records clarify the timeline", 0.82),
        ("A teacher's note contradicts the rumor", 0.7),
    ],
    "kindness": [
        ("Clinic notes reveal private family pressure", 0.72),
        ("A neighbor explains who was actually helped", 0.65),
    ],
    "relationship": [
        ("A witness places the target somewhere else", 0.68),
        ("A private message changes the context", 0.6),
    ],
    "leadership": [
        ("Meeting minutes show who made the decision", 0.8),
        ("An elder remembers the earlier agreement", 0.7),
    ],
}


def active_rumor(state: VillageState) -> Rumor | None:
    if not state.rumors:
        return None
    return list(state.rumors.values())[-1]


def choose_channel(state: VillageState, speaker: Villager) -> Channel:
    best = sorted(
        state.channels.values(),
        key=lambda channel: (
            len(set(channel.typical_groups) & set(speaker.group_ids)),
            channel.reach,
        ),
        reverse=True,
    )
    return best[0]


def choose_listener(state: VillageState, speaker: Villager, rumor: Rumor, day_offset: int) -> Villager:
    candidates = [v for v in state.villagers.values() if v.id != speaker.id]
    ranked = sorted(
        candidates,
        key=lambda v: (
            len(set(v.group_ids) & set(speaker.group_ids)) * 30
            + speaker.relationships.get(v.id, 0)
            + (10 if v.id == rumor.target_id else 0)
            + ((state.day + day_offset + len(v.id)) % 7)
        ),
        reverse=True,
    )
    return ranked[0]


def belief_delta(speaker: Villager, listener: Villager, rumor: Rumor, channel: Channel, action: str) -> float:
    rel = relationship(speaker, listener.id)
    delta = (
        speaker.trust * 0.003
        + rel * 0.002
        + rumor.intensity * 0.001
        + channel.credibility_modifier * 0.1
    )
    if action == "distort":
        delta *= 0.7
    elif action == "verify":
        delta *= 1 + speaker.accuracy
    elif action == "accuse":
        delta *= 1.2
    elif action == "defend":
        delta *= -0.9
    elif action == "mediate":
        current = listener.private_beliefs.get(rumor.id, 0.5)
        delta = (0.5 - current) * 0.35
    elif action == "investigate":
        delta *= 0.4 if rumor.truth_status in {"false", "unknown"} else 1.1
    elif action == "deny":
        delta *= -0.6 if speaker.trust > 50 else 0.25
    elif action == "stay_silent":
        delta = 0
    return round(delta, 3)


def update_public_position(
    listener: Villager,
    rumor: Rumor,
    belief: float,
    action: str,
    state: VillageState,
) -> None:
    target_influence = state.villagers[rumor.target_id].influence if rumor.target_id else 0
    motive = listener.motive.lower()
    if "power" in motive and target_influence > 65 and belief > 0.38:
        listener.public_positions[rumor.id] = "accusing_target"
    elif "peace" in motive and belief < 0.7:
        listener.public_positions[rumor.id] = "defending_target"
    elif action == "defend":
        listener.public_positions[rumor.id] = "defending_target"
    elif action == "accuse":
        listener.public_positions[rumor.id] = "accusing_target"
    elif action in {"verify", "investigate"}:
        listener.public_positions[rumor.id] = "calling_for_investigation"
    elif belief > 0.66:
        listener.public_positions[rumor.id] = "believes"
    elif belief < 0.34:
        listener.public_positions[rumor.id] = "doubts"
    else:
        listener.public_positions[rumor.id] = "neutral"


def apply_reputation_change(
    state: VillageState,
    rumor: Rumor,
    speaker: Villager,
    listener_belief: float,
    channel: Channel,
    action: str,
) -> dict[str, float]:
    if not rumor.target_id:
        return {}
    target = state.villagers[rumor.target_id]
    reach_factor = max(0.25, channel.reach / 20)
    base_damage = 8
    damage = (
        rumor.intensity / 100
        * rumor.severity / 100
        * speaker.influence / 100
        * listener_belief
        * reach_factor
        * base_damage
    )
    if action == "distort":
        damage *= 1.35
    elif action == "accuse":
        damage *= 1.25
    elif action in {"defend", "mediate", "apologize"}:
        damage *= -0.45
    elif action in {"verify", "investigate"} and rumor.truth_status == "false":
        damage *= -0.2

    changes = {}
    for field, weight in REPUTATION_WEIGHTS.get(rumor.category, REPUTATION_WEIGHTS["moral"]).items():
        delta = round(-damage * weight, 2)
        current = getattr(target.reputation, field)
        setattr(target.reputation, field, clamp(current + delta))
        changes[field] = delta
    return changes


def apply_economy_changes(speaker: Villager, listener: Villager, rumor: Rumor, action: str) -> dict:
    before = {"trust": speaker.trust, "influence": speaker.influence, "attention": speaker.attention}
    if action in {"spread", "distort", "accuse"}:
        speaker.attention = clamp(speaker.attention + (9 if action == "distort" else 6))
        speaker.influence = clamp(speaker.influence + 2)
        speaker.trust = clamp(speaker.trust - (5 if action == "distort" else 2))
        rumor.intensity = clamp(rumor.intensity + (10 if action == "distort" else 6))
        rumor.spread_count += 1
    elif action in {"verify", "investigate"}:
        speaker.trust = clamp(speaker.trust + 4)
        speaker.influence = clamp(speaker.influence + 2)
        rumor.intensity = clamp(rumor.intensity - 3)
    elif action in {"defend", "mediate", "apologize"}:
        speaker.trust = clamp(speaker.trust + 3)
        rumor.intensity = clamp(rumor.intensity - 5)
    elif action == "deny":
        speaker.attention = clamp(speaker.attention + 4)
        rumor.intensity = clamp(rumor.intensity + (3 if speaker.trust < 55 else -4))

    speaker.relationships[listener.id] = clamp(speaker.relationships.get(listener.id, 0) + (3 if action in {"mediate", "defend", "verify"} else -2), -100, 100)
    after = {"trust": speaker.trust, "influence": speaker.influence, "attention": speaker.attention}
    return {
        "speaker_before": before,
        "speaker_after": after,
        "trust_delta": round(after["trust"] - before["trust"], 2),
        "influence_delta": round(after["influence"] - before["influence"], 2),
        "attention_delta": round(after["attention"] - before["attention"], 2),
    }


def record_strategy_memory(
    villager: Villager,
    day: int,
    rumor: Rumor,
    action: str,
    economy_changes: dict,
    reputation_delta: dict[str, float],
) -> Memory:
    trust_delta = economy_changes.get("trust_delta", 0)
    influence_delta = economy_changes.get("influence_delta", 0)
    attention_delta = economy_changes.get("attention_delta", 0)
    target_damage = -sum(value for value in reputation_delta.values() if value < 0)
    target_repair = sum(value for value in reputation_delta.values() if value > 0)
    payoff = round(attention_delta * 0.35 + influence_delta * 0.6 + trust_delta * 0.8 + target_damage * 0.18 - target_repair * 0.1, 2)
    if action in {"distort", "accuse", "spread"} and trust_delta < 0:
        payoff -= abs(trust_delta) * 0.45
    if action in {"verify", "investigate", "mediate", "defend"} and trust_delta > 0:
        payoff += trust_delta * 0.3
    payoff = round(payoff, 2)
    bias_before = villager.strategy_biases.get(action, 0)
    villager.strategy_biases[action] = round(clamp(bias_before + payoff * 0.22, -12, 18), 2)
    if payoff >= 4:
        verdict = "worked"
    elif payoff <= -2:
        verdict = "backfired"
    else:
        verdict = "had mixed results"
    summary = (
        f"Day {day}: {action.replace('_', ' ')} {verdict}; "
        f"trust {trust_delta:+.1f}, influence {influence_delta:+.1f}, attention {attention_delta:+.1f}."
    )
    memory = Memory(
        day=day,
        rumor_id=rumor.id,
        action=action,  # type: ignore[arg-type]
        summary=summary,
        trust_delta=trust_delta,
        influence_delta=influence_delta,
        attention_delta=attention_delta,
        payoff=payoff,
    )
    villager.memories.append(memory)
    villager.memories = villager.memories[-6:]
    return memory


def maybe_trigger_events(state: VillageState, rumor: Rumor) -> None:
    if not rumor.target_id:
        return
    target = state.villagers[rumor.target_id]
    existing = {event.type for event in state.events}
    if target.reputation.honesty < 30 and "public_confrontation" not in existing:
        state.events.append(Event(day=state.day, type="public_confrontation", title="Public Confrontation", description=f"{target.name} is confronted in public.", participants=[target.id], rumor_id=rumor.id))
    if target.reputation.leadership < 25 and "resignation_pressure" not in existing:
        state.events.append(Event(day=state.day, type="resignation_pressure", title="Resignation Pressure", description=f"Villagers begin asking whether {target.name} should step aside.", participants=[target.id], rumor_id=rumor.id))
    if rumor.intensity > 85 and "village_panic" not in existing:
        state.events.append(Event(day=state.day, type="village_panic", title="Village Panic", description="The rumor outruns facts and everyone starts picking sides.", participants=[], rumor_id=rumor.id))
    if rumor.public_belief > 0.75 and not public_evidence_for(state, rumor.id) and "mob_judgment" not in existing:
        state.events.append(Event(day=state.day, type="mob_judgment", title="Mob Judgment", description="Public belief hardens before strong evidence appears.", participants=[target.id], rumor_id=rumor.id))
    if contradictory_public_evidence(state, rumor) and total_reputation_damage(target) > 20 and "guilt_and_apologies" not in existing:
        state.events.append(Event(day=state.day, type="guilt_and_apologies", title="Guilt And Apologies", description=f"Evidence undercuts the rumor after {target.name}'s reputation has already been damaged.", participants=[target.id], rumor_id=rumor.id))
    mutai = state.villagers.get("mutai")
    if mutai and mutai.influence > 80 and "political_takeover" not in existing:
        state.events.append(Event(day=state.day, type="political_takeover", title="Political Takeover", description="Mutai turns the scandal into a power platform.", participants=["mutai"], rumor_id=rumor.id))


def public_evidence_for(state: VillageState, rumor_id: str) -> list[Evidence]:
    return [item for item in state.evidence.values() if item.rumor_id == rumor_id and item.public]


def strongest_evidence(state: VillageState, rumor_id: str) -> Evidence | None:
    items = public_evidence_for(state, rumor_id)
    if not items:
        return None
    return max(items, key=lambda item: item.reliability)


def contradictory_public_evidence(state: VillageState, rumor: Rumor) -> bool:
    if rumor.truth_status == "true":
        return False
    return any(not item.supports_rumor and item.reliability >= 0.65 for item in public_evidence_for(state, rumor.id))


def total_reputation_damage(villager: Villager) -> float:
    return sum(max(0, 70 - value) for value in villager.reputation.model_dump().values())


def evidence_supports_rumor(rumor: Rumor) -> bool:
    if rumor.truth_status == "true":
        return True
    if rumor.truth_status == "false":
        return False
    if rumor.truth_status == "partly_true":
        return rumor.spread_count % 2 == 0
    return rumor.intensity > 70


def discover_evidence(state: VillageState, rumor: Rumor, discoverer: Villager, action: str) -> Evidence | None:
    if action != "investigate" and not (state.day in {3, 5} and rumor.intensity > 60):
        return None
    if any(item.rumor_id == rumor.id and item.discovered_by == discoverer.id for item in state.evidence.values()):
        return None
    templates = EVIDENCE_TEMPLATES.get(rumor.category, EVIDENCE_TEMPLATES["moral"])
    description, reliability = templates[(len(state.evidence) + state.day) % len(templates)]
    evidence = Evidence(
        id=f"evidence_{len(state.evidence) + 1}",
        rumor_id=rumor.id,
        description=description,
        reliability=reliability,
        supports_rumor=evidence_supports_rumor(rumor),
        discovered_by=discoverer.id,
        public=action == "investigate" or reliability >= 0.8 or rumor.intensity > 78,
    )
    state.evidence[evidence.id] = evidence
    discoverer.known_evidence.append(evidence.id)
    state.trace.append(
        {
            "day": state.day,
            "villager": discoverer.name,
            "rumor_id": rumor.id,
            "chosen_action": "evidence_discovered",
            "evidence": evidence.model_dump(),
            "reason": f"{discoverer.name} uncovered evidence while the rumor market was moving.",
        }
    )
    if evidence.public:
        apply_public_evidence(state, rumor, evidence)
    return evidence


def apply_public_evidence(state: VillageState, rumor: Rumor, evidence: Evidence) -> None:
    direction = 1 if evidence.supports_rumor else -1
    strength = evidence.reliability * 0.24
    for villager in state.villagers.values():
        before = villager.private_beliefs.get(rumor.id, 0.5)
        after = clamp(before + direction * strength * villager.accuracy, 0, 1)
        villager.private_beliefs[rumor.id] = after
        update_public_position(villager, rumor, after, "investigate", state)
        predicted_correctly = (before >= 0.5) == evidence.supports_rumor
        if predicted_correctly:
            villager.trust = clamp(villager.trust + evidence.reliability * 3)
        else:
            villager.trust = clamp(villager.trust - evidence.reliability * 4)
    rumor.public_belief = average_public_belief(state, rumor.id)
    rumor.intensity = clamp(rumor.intensity + (8 if evidence.supports_rumor else -10))
    apply_evidence_reputation_effect(state, rumor, evidence)


def apply_evidence_reputation_effect(state: VillageState, rumor: Rumor, evidence: Evidence) -> None:
    if not rumor.target_id:
        return
    target = state.villagers[rumor.target_id]
    amount = evidence.reliability * rumor.severity / 100 * 10
    if evidence.supports_rumor:
        amount *= -1
    for field, weight in REPUTATION_WEIGHTS.get(rumor.category, REPUTATION_WEIGHTS["moral"]).items():
        current = getattr(target.reputation, field)
        setattr(target.reputation, field, clamp(current + amount * weight))


def make_trace(
    state: VillageState,
    bid: Bid,
    speaker: Villager,
    listener: Villager,
    channel: Channel,
    belief_before: float,
    belief_after: float,
    reputation_delta: dict[str, float],
    memory: Memory | None,
) -> dict:
    return {
        "day": state.day,
        "villager": speaker.name,
        "rumor_id": bid.rumor_id,
        "action_scores": bid.action_scores,
        "chosen_action": bid.action,
        "bid_score": bid.bid_score,
        "expected_gain": bid.expected_gain,
        "expected_risk": bid.expected_risk,
        "listener": listener.name,
        "channel": channel.name,
        "belief_before": round(belief_before, 3),
        "belief_after": round(belief_after, 3),
        "listener_public_position": listener.public_positions.get(bid.rumor_id, "neutral"),
        "public_evidence_count": len(public_evidence_for(state, bid.rumor_id)),
        "reputation_damage": reputation_delta,
        "economy_changes": {
            "trust_delta": memory.trust_delta if memory else 0,
            "influence_delta": memory.influence_delta if memory else 0,
            "attention_delta": memory.attention_delta if memory else 0,
            "payoff": memory.payoff if memory else 0,
        },
        "strategy_biases": speaker.strategy_biases,
        "memory": memory.model_dump() if memory else None,
        "reason": bid.reason,
    }


def simulate_day(state: VillageState, llm: GossipLLM | None = None) -> VillageState:
    rumor = active_rumor(state)
    if rumor is None:
        return state
    llm = llm or GossipLLM(mock_llm=True)

    state.latest_conversations = []
    bids = collect_bids(state, rumor)
    winners = select_winning_bids(bids, k=3)

    for index, bid in enumerate(winners):
        speaker = state.villagers[bid.villager_id]
        listener = choose_listener(state, speaker, rumor, index)
        channel = choose_channel(state, speaker)
        old_version = rumor.current_version
        before = listener.private_beliefs.get(rumor.id, 0.5)
        delta = belief_delta(speaker, listener, rumor, channel, bid.action)
        after = clamp(before + delta, 0, 1)
        listener.private_beliefs[rumor.id] = after
        update_public_position(listener, rumor, after, bid.action, state)
        target_name = state.villagers[rumor.target_id].name if rumor.target_id else "None"
        dialogue = llm.dialogue(
            state.day,
            state.tone,
            speaker,
            listener,
            rumor,
            bid.action,
            channel,
            relationship(speaker, listener.id),
            target_name,
        )
        new_version = llm.mutate_rumor(rumor, bid.action, channel, state.tone, speaker, dialogue)
        rumor.current_version = new_version
        reputation_delta = apply_reputation_change(state, rumor, speaker, after, channel, bid.action)
        economy_changes = apply_economy_changes(speaker, listener, rumor, bid.action)
        memory = record_strategy_memory(speaker, state.day, rumor, bid.action, economy_changes, reputation_delta)
        rumor.public_belief = average_public_belief(state, rumor.id)

        result = ConversationResult(
            day=state.day,
            speaker_id=speaker.id,
            listener_id=listener.id,
            rumor_id=rumor.id,
            action=bid.action,
            channel_id=channel.id,
            dialogue=dialogue,
            old_rumor_version=old_version,
            new_rumor_version=new_version,
            belief_delta=delta,
            target_reputation_delta=reputation_delta,
            economy_changes=economy_changes,
            event_summary=f"{speaker.name} used {bid.action.replace('_', ' ')} on {listener.name} at the {channel.name}.",
        )
        state.latest_conversations.append(result)
        state.trace.append(make_trace(state, bid, speaker, listener, channel, before, after, reputation_delta, memory))
        discover_evidence(state, rumor, speaker, bid.action)

    if not any(item.rumor_id == rumor.id and item.public for item in state.evidence.values()):
        top_investigator = max(state.villagers.values(), key=lambda v: v.accuracy + v.trust / 100)
        discover_evidence(state, rumor, top_investigator, "background")
    maybe_trigger_events(state, rumor)
    state.daily_summaries.append(llm.daily_summary(state.day, state, rumor))
    if state.day >= 7:
        state.baraza_output = generate_baraza(state, rumor)
    state.day += 1
    return state


def average_public_belief(state: VillageState, rumor_id: str) -> float:
    beliefs = [v.private_beliefs.get(rumor_id, 0.5) for v in state.villagers.values()]
    return round(sum(beliefs) / len(beliefs), 3)


def generate_baraza(state: VillageState, rumor: Rumor) -> str:
    action_counts = defaultdict(int)
    for row in state.trace:
        action_counts[(row["chosen_action"], row["villager"])] += 1
    accuser = top_actor(state, "accuse")
    defender = top_actor(state, "defend") or top_actor(state, "mediate")
    target = state.villagers.get(rumor.target_id) if rumor.target_id else None
    target_name = target.name if target else "the target"
    evidence = strongest_evidence(state, rumor.id)
    evidence_line = evidence.description if evidence else "No strong evidence reached the public."
    influential = max(state.villagers.values(), key=lambda villager: villager.influence)
    damaged_category = "reputation"
    outcome = "truth remains unresolved"
    if target:
        lowest = min(target.reputation.model_dump().items(), key=lambda item: item[1])
        damaged_category = lowest[0]
        if lowest[1] < 35:
            outcome = f"{target_name} is cleared by nobody and wounded by everybody"
        elif rumor.public_belief < 0.4:
            outcome = f"{target_name} survives the storm, but the village remembers the smoke"
    if evidence and not evidence.supports_rumor and rumor.public_belief < 0.5:
        outcome = f"{target_name} is partly cleared, but the apology arrives limping"
    elif evidence and evidence.supports_rumor and evidence.reliability > 0.75:
        outcome = f"{target_name} faces formal pressure because the evidence has teeth"
    if top_actor(state, "distort") == "Mutai":
        outcome = "the rumor mutates into political capital"
    return "\n".join(
        [
            "Chief Kiprono: We are here because this story has eaten more airtime than school fees.",
            f"{accuser or 'Mama Chebet'}: The village deserves answers about {rumor.current_version}",
            f"{defender or 'Pastor Langat'}: Answers, yes. Character assassination, no.",
            f"{target_name}: I have heard versions of myself that even I have not met.",
            f"Cherono: Evidence on the table: {evidence_line}",
            "Mercy: Let records speak louder than volume.",
            f"Chief Kiprono: The worst wound is {damaged_category}; the loudest voice is now {influential.name}.",
            "Kibet: Chief, even the jokes are now tired.",
            "Arap Rono: A rumor is a spear. Once thrown, it does not return clean.",
            f"Final outcome: {outcome}. Hidden reality: {rumor.hidden_reality}",
        ]
    )


def top_actor(state: VillageState, action: str) -> str | None:
    counts = defaultdict(int)
    for row in state.trace:
        if row["chosen_action"] == action:
            counts[row["villager"]] += 1
    if not counts:
        return None
    return max(counts.items(), key=lambda item: item[1])[0]