civicai/simulation.py

"""
CivicAI Society Simulation Engine

Data-driven economic/social simulation.
Models economics, health, crime, satisfaction, and real-world events.
Integrates with World Bank API and LLM-parsed news.
"""

from __future__ import annotations

import math
import random
from typing import Any

from civicai.models import Action, SocietyState, SubsidyPolicy
from civicai.data_pipeline import RealWorldDataPipeline
from civicai.llm_parser import LLMParser

# Singletons for data fetching
DATA_PIPELINE = RealWorldDataPipeline()
LLM = LLMParser()


def _clamp(value: float, lo: float = 0.0, hi: float = 1.0) -> float:
    return max(lo, min(hi, value))

def _noise(scale: float = 0.01) -> float:
    return random.gauss(0, scale)

# ---------------------------------------------------------------------------
# Simulation Step
# ---------------------------------------------------------------------------

def simulate_step(state: SocietyState, action: Action) -> SocietyState:
    """
    Advance the society simulation by one turn (quarter).
    Mutates and returns the state.
    """
    # Fetch Real-World Data Baseline for this turn
    real_world_baseline = DATA_PIPELINE.step_forward()
    news_feed = DATA_PIPELINE.get_real_world_news()
    
    # Parse unstructured news into modifiers
    news_modifiers = LLM.parse_news_to_modifiers(news_feed)

    # Log events for the turn
    state.active_events = news_feed
    if news_feed:
        state.event_history.append({
            "turn": state.turn,
            "events": news_feed,
            "modifiers": news_modifiers
        })

    # Store action
    state.current_tax_rate = action.tax_rate
    state.current_healthcare_budget = action.healthcare_budget
    state.current_education_budget = action.education_budget
    state.current_police_budget = action.police_budget
    state.current_subsidy = action.subsidy_policy.value

    state.action_history.append(action.model_dump())

    # --- Revenue & Budget ---
    state.tax_revenue = state.gdp * action.tax_rate
    total_spending = (
        action.healthcare_budget +
        action.education_budget +
        action.police_budget
    ) * state.government_budget

    subsidy_cost = 0.0
    if action.subsidy_policy != SubsidyPolicy.NONE:
        subsidy_cost = state.government_budget * 0.08

    state.government_budget += state.tax_revenue - total_spending - subsidy_cost
    state.budget_balance = (state.tax_revenue - total_spending - subsidy_cost) / max(state.gdp, 1.0)

    # --- Economic Model (Anchored to Real World Data) ---
    # GDP growth is influenced by the World Bank baseline + agent actions + news
    tax_drag = -0.5 * max(0, action.tax_rate - 0.30)
    education_boost = 0.3 * action.education_budget
    subsidy_boost = 0.0
    if action.subsidy_policy == SubsidyPolicy.TECHNOLOGY:
        subsidy_boost = 0.015
    elif action.subsidy_policy == SubsidyPolicy.INDUSTRY:
        subsidy_boost = 0.012
    elif action.subsidy_policy == SubsidyPolicy.AGRICULTURE:
        subsidy_boost = 0.008

    # Blend simulation mechanics with real-world baseline (converted from % to ratio)
    baseline_gdp_growth = real_world_baseline.get("gdp_growth", 2.0) / 100.0
    base_growth = baseline_gdp_growth + tax_drag + education_boost + subsidy_boost + news_modifiers.get("gdp_modifier", 0.0) + _noise(0.005)
    
    state.gdp_growth = _clamp(base_growth, -0.15, 0.15)
    state.gdp = max(10.0, state.gdp * (1 + state.gdp_growth))

    # Inflation: driven by spending vs revenue + World Bank baseline
    baseline_inflation = real_world_baseline.get("inflation", 2.0) / 100.0
    spending_ratio = (total_spending + subsidy_cost) / max(state.tax_revenue, 1.0)
    inflation_pressure = 0.02 * (spending_ratio - 1.0)
    
    state.inflation = _clamp(
        baseline_inflation + inflation_pressure + news_modifiers.get("inflation_modifier", 0.0) + _noise(0.003),
        -0.05, 0.30
    )

    # Employment: World Bank baseline + education + news
    baseline_unemployment = real_world_baseline.get("unemployment", 5.0) / 100.0
    baseline_employment = 1.0 - baseline_unemployment
    
    employment_delta = (
        0.5 * state.gdp_growth +
        0.1 * action.education_budget -
        0.2 * max(0, state.inflation - 0.06) +
        news_modifiers.get("employment_modifier", 0.0) +
        _noise(0.01)
    )
    # Blend baseline with simulation delta
    state.employment_rate = _clamp(baseline_employment + employment_delta, 0.3, 0.99)

    # --- Health Model ---
    baseline_le = real_world_baseline.get("life_expectancy", 70.0)
    # Normalize LE roughly 0-1 (min 40, max 90)
    norm_le = _clamp((baseline_le - 40.0) / 50.0, 0.0, 1.0)
    
    healthcare_effect = 0.15 * action.healthcare_budget
    infection_drag = -0.3 * state.infection_rate
    state.health_index = _clamp(
        (norm_le * 0.5 + state.health_index * 0.5) + healthcare_effect + infection_drag - 0.02 + _noise(0.008),
        0.05, 1.0
    )

    # Infection dynamics
    if state.infection_rate > 0:
        lockdown_factor = 0.4 if state.is_lockdown else 1.0
        recovery = 0.15 * action.healthcare_budget
        spread = 0.08 * lockdown_factor * (1 - action.healthcare_budget)
        state.infection_rate = _clamp(
            state.infection_rate + spread - recovery + _noise(0.01),
            0.0, 0.8
        )

    # Emergency response
    if action.emergency_response:
        er = action.emergency_response.lower()
        if "lockdown" in er:
            state.is_lockdown = True
            state.infection_rate = _clamp(state.infection_rate - 0.05, 0.0, 1.0)
            state.gdp_growth -= 0.02
            state.public_satisfaction -= 0.04
        elif "stimulus" in er:
            state.gdp_growth += 0.015
            state.government_budget -= state.gdp * 0.03
            state.public_satisfaction += 0.03
        elif "open" in er or "lift" in er:
            state.is_lockdown = False
            state.gdp_growth += 0.01
    else:
        state.is_lockdown = False

    # --- Crime Model ---
    unemployment_pressure = 0.3 * (1 - state.employment_rate)
    inequality_pressure = 0.2 * state.emergent.wealth_inequality
    police_effect = -0.25 * action.police_budget
    satisfaction_effect = -0.1 * state.public_satisfaction
    state.crime_rate = _clamp(
        state.crime_rate + unemployment_pressure + inequality_pressure +
        police_effect + satisfaction_effect + news_modifiers.get("crime_modifier", 0.0) + _noise(0.008),
        0.01, 0.6
    )

    # --- Satisfaction Model (with momentum / lag) ---
    economic_factor = 0.3 * state.employment_rate + 0.1 * _clamp(state.gdp_growth + 0.05, 0, 0.1) / 0.1
    health_factor = 0.2 * state.health_index
    crime_factor = 0.15 * (1 - state.crime_rate)
    tax_burden = -0.15 * max(0, action.tax_rate - 0.25)
    service_quality = 0.1 * (action.healthcare_budget + action.education_budget)

    target_satisfaction = _clamp(
        economic_factor + health_factor + crime_factor + tax_burden + service_quality + news_modifiers.get("satisfaction_modifier", 0.0),
        0.0, 1.0
    )
    # Momentum: satisfaction moves slowly toward target
    state.public_satisfaction = _clamp(
        state.public_satisfaction + 0.3 * (target_satisfaction - state.public_satisfaction) + _noise(0.01),
        0.0, 1.0
    )

    # --- Education ---
    state.education_quality = _clamp(
        state.education_quality + 0.1 * action.education_budget - 0.02 + _noise(0.005),
        0.1, 1.0
    )

    # --- Resources ---
    state.food_reserves = _clamp(state.food_reserves - 0.02 + 0.03 * (1 if action.subsidy_policy == SubsidyPolicy.AGRICULTURE else 0.3) + _noise(0.005))
    state.energy_reserves = _clamp(state.energy_reserves - 0.015 + 0.02 * (1 if action.subsidy_policy == SubsidyPolicy.INDUSTRY else 0.4) + _noise(0.005))
    state.medical_supplies = _clamp(state.medical_supplies - 0.02 + 0.04 * action.healthcare_budget + _noise(0.005))
    state.infrastructure = _clamp(state.infrastructure - 0.005 + 0.01 * (action.police_budget + action.education_budget) + _noise(0.003))

    # --- Demographics ---
    effective_birth = state.birth_rate * (0.8 + 0.2 * state.health_index)
    effective_death = state.death_rate * (1.2 - 0.2 * state.health_index)
    state.population = max(1000, int(state.population * (1 + effective_birth - effective_death)))

    # --- Emergent Metrics ---
    _update_emergent(state, action)

    # Advance turn
    state.turn += 1

    return state


# ---------------------------------------------------------------------------
# Emergent Behavior Tracking
# ---------------------------------------------------------------------------

def _update_emergent(state: SocietyState, action: Action) -> None:
    """Update emergent behavior metrics."""
    em = state.emergent

    # Wealth inequality (Gini proxy)
    tax_equalizer = -0.02 * action.tax_rate
    tech_inequality = 0.01 if action.subsidy_policy == SubsidyPolicy.TECHNOLOGY else 0.0
    unemployment_inequality = 0.05 * (1 - state.employment_rate)
    em.wealth_inequality = _clamp(
        em.wealth_inequality + tax_equalizer + tech_inequality + unemployment_inequality + _noise(0.005),
        0.1, 0.9
    )

    # Social unrest
    dissatisfaction = max(0, 0.5 - state.public_satisfaction) * 0.3
    crime_factor = state.crime_rate * 0.2
    inequality_factor = em.wealth_inequality * 0.15
    em.social_unrest = _clamp(
        em.social_unrest * 0.9 + dissatisfaction + crime_factor + inequality_factor + _noise(0.005),
        0.0, 1.0
    )

    # Protest probability
    em.protest_probability = _clamp(
        em.social_unrest * 0.6 + (1 - state.public_satisfaction) * 0.3,
        0.0, 1.0
    )

    # Cooperation index
    balanced = 1.0 - abs(action.healthcare_budget - action.education_budget) - abs(action.education_budget - action.police_budget)
    em.cooperation_index = _clamp(
        em.cooperation_index * 0.8 + 0.2 * _clamp(balanced + 0.5) + _noise(0.01),
        0.0, 1.0
    )