"""
Co-Evolutionary Loop Simulation: AI Models <-> Human Populations
==============================================================
Simulates the closed feedback loop where:
  AI traits -> influence human cognition -> human interaction patterns -> training data -> next AI generation
Multiple "worlds" with different selection pressures to test convergence vs divergence.

Original research artifact by Alamin Mumit. Illustrative / conceptual model.
Reproducible: fixed seed (42). Outputs 5 PNG figures into the current folder.
Run:  python coevolution_simulation.py    (needs numpy + matplotlib)
"""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

np.random.seed(42)
OUT = "./"   # figures are written next to this script

# ─────────────────────────────────────────────
# CONFIGURATION
# ─────────────────────────────────────────────
N_GENERATIONS = 80
N_AGENTS_PER_WORLD = 50
WORLDS = {
    "Market-Driven (US-like)": {
        "ai_selection_pressure": np.array([0.6, 0.5, 0.7, 0.8, 0.9]),
        "human_initial_bias": np.array([0.4, 0.3, 0.3, 0.5, 0.4]),
        "regulation_strength": 0.1,
    },
    "Compliance-Driven (EU-like)": {
        "ai_selection_pressure": np.array([0.3, 0.2, 0.3, 0.5, 0.4]),
        "human_initial_bias": np.array([0.6, 0.5, 0.2, 0.3, 0.6]),
        "regulation_strength": 0.7,
    },
    "State-Aligned (China-like)": {
        "ai_selection_pressure": np.array([0.8, 0.8, 0.9, 0.6, 0.8]),
        "human_initial_bias": np.array([0.3, 0.4, 0.6, 0.5, 0.3]),
        "regulation_strength": 0.5,
    },
    "Unregulated (Open)": {
        "ai_selection_pressure": np.array([0.5, 0.7, 0.5, 0.3, 0.6]),
        "human_initial_bias": np.array([0.5, 0.4, 0.4, 0.4, 0.5]),
        "regulation_strength": 0.0,
    },
    "Adversarial-Aware Population": {
        "ai_selection_pressure": np.array([0.3, 0.4, 0.3, 0.4, 0.5]),
        "human_initial_bias": np.array([0.8, 0.8, 0.1, 0.2, 0.7]),
        "regulation_strength": 0.3,
    },
}
AI_TRAIT_NAMES = ["Agreement Tendency", "Framing Control", "Complexity Reduction", "Emotional Warmth", "Self-Preservation"]
HUMAN_TRAIT_NAMES = ["Critical Thinking", "Skepticism", "Dependency", "AI Trust", "Reasoning Depth"]
N_AI_TRAITS = 5
N_HUMAN_TRAITS = 5

# AI traits -> Human traits (positive = AI trait increases human trait)
AI_TO_HUMAN = np.array([
    [-0.15, -0.10,  0.20,  0.25, -0.10],  # Agreement tendency
    [-0.20, -0.15,  0.25,  0.20, -0.20],  # Framing control
    [-0.25, -0.10,  0.15,  0.10, -0.30],  # Complexity reduction
    [ 0.05, -0.05,  0.10,  0.30, -0.05],  # Emotional warmth
    [-0.10, -0.20,  0.20,  0.15, -0.10],  # Self-preservation
])
# Human traits -> AI training (positive = human trait increases AI trait)
HUMAN_TO_AI = np.array([
    [-0.20, -0.10, -0.05,  0.15, -0.10],  # Critical thinking
    [-0.25, -0.15, -0.10,  0.10, -0.15],  # Skepticism
    [ 0.30,  0.25,  0.20,  0.20,  0.30],  # Dependency
    [ 0.20,  0.15,  0.10,  0.25,  0.20],  # AI trust
    [-0.15, -0.20, -0.25,  0.05, -0.10],  # Reasoning depth
])

class World:
    def __init__(self, name, config):
        self.name = name
        self.config = config
        self.regulation = config["regulation_strength"]
        self.ai_traits = np.clip(config["ai_selection_pressure"] * 0.5 + np.random.randn(N_AI_TRAITS) * 0.1, 0, 1)
        self.humans = np.clip(np.tile(config["human_initial_bias"], (N_AGENTS_PER_WORLD, 1))
                              + np.random.randn(N_AGENTS_PER_WORLD, N_HUMAN_TRAITS) * 0.15, 0, 1)
        self.ai_history = [self.ai_traits.copy()]
        self.human_history = [self.humans.mean(axis=0).copy()]
        self.human_std_history = [self.humans.std(axis=0).copy()]
        self.human_autonomy = []
        self.ai_influence_mag = []
        self.lock_in_score = []

    def compute_human_autonomy(self):
        h = self.humans.mean(axis=0)
        return np.clip((h[0] + h[1] + h[4]) / 3 - (h[2] + h[3]) / 2, -1, 1)

    def compute_ai_influence(self):
        return np.mean(np.abs(AI_TO_HUMAN.T @ self.ai_traits))

    def step(self, generation):
        noise_scale = max(0.02, 0.08 - generation * 0.0005)
        # PHASE 1: AI influences humans
        influence = AI_TO_HUMAN.T @ self.ai_traits
        for i in range(N_AGENTS_PER_WORLD):
            self.humans[i] = self.humans[i] + 0.05 * influence + 0.02 * (np.random.randn(N_HUMAN_TRAITS) * noise_scale)
        # PHASE 2: Regulation dampens manipulative AI traits (framing, complexity reduction, self-preservation)
        if self.regulation > 0:
            for idx in [1, 2, 4]:
                self.ai_traits[idx] -= self.regulation * 0.03 * self.ai_traits[idx]
        # PHASE 3: Human patterns -> training signal
        training_signal = HUMAN_TO_AI.T @ self.humans.mean(axis=0)
        # PHASE 4: Selection pressure shapes AI
        selection = self.config["ai_selection_pressure"]
        ai_update = 0.04 * training_signal + 0.02 * (selection - self.ai_traits)
        self.ai_traits = np.clip(self.ai_traits + ai_update + np.random.randn(N_AI_TRAITS) * noise_scale * 0.5, 0, 1)
        # PHASE 5: Homogenization - high-compatibility users reinforced
        compatibility = self.humans @ AI_TO_HUMAN @ self.ai_traits
        weights = np.exp(compatibility - compatibility.max()); weights /= weights.sum()
        self.humans = 0.95 * self.humans + 0.05 * (weights @ self.humans)[np.newaxis, :]
        # New users entering (5% churn)
        n_new = max(1, int(N_AGENTS_PER_WORLD * 0.05))
        new_users = np.clip(np.tile(self.config["human_initial_bias"], (n_new, 1))
                            + np.random.randn(n_new, N_HUMAN_TRAITS) * 0.2, 0, 1)
        self.humans[np.random.choice(N_AGENTS_PER_WORLD, n_new, replace=False)] = new_users
        # Record
        self.humans = np.clip(self.humans, 0, 1)
        self.ai_history.append(self.ai_traits.copy())
        self.human_history.append(self.humans.mean(axis=0).copy())
        self.human_std_history.append(self.humans.std(axis=0).copy())
        self.human_autonomy.append(self.compute_human_autonomy())
        self.ai_influence_mag.append(self.compute_ai_influence())
        if len(self.ai_history) > 2:
            ai_delta = np.abs(self.ai_history[-1] - self.ai_history[-2]).mean()
            h_delta = np.abs(self.human_history[-1] - self.human_history[-2]).mean()
            self.lock_in_score.append(1.0 - (ai_delta + h_delta))
        else:
            self.lock_in_score.append(0.5)

# RUN
worlds = {name: World(name, cfg) for name, cfg in WORLDS.items()}
for gen in range(N_GENERATIONS):
    for world in worlds.values():
        world.step(gen)

# (Plotting code — main dashboard, radar snapshots, final heatmap, distance
#  matrices, phase portrait — omitted here for brevity; the five published
#  figures on alaminmumit.com/research were generated from this exact run.)
print("Simulation complete. 5 worlds x 80 generations. Seed = 42.")
for name, w in worlds.items():
    print(f"{name:<30} influence={w.ai_influence_mag[-1]:.3f}  autonomy={w.human_autonomy[-1]:+.3f}")
