import random
import math
import matplotlib.pyplot as plt

def simulate_takeover(
    steps=500,
    initial_humans=8_000_000_000,
    initial_robots=1_000_000,
    human_growth_rate=0.001,      # ~0.1% per step
    robot_growth_rate=0.01,       # robots scale faster
    baseline_mortality_h=0.0005,
    baseline_mortality_r=0.0002,
    conflict_probability=0.1,     # chance of conflict each step
    conflict_intensity=0.05,      # fraction of combined forces lost in major conflict
    robot_advantage=2.0,          # >1 means robots fight more effectively
    takeover_threshold=0.8,       # robots control 80% of “power”
    verbose=False
):
    """
    Returns:
        history: dict with time, humans, robots, power_share_robots
        outcome: string describing final state
    """
    H = float(initial_humans)
    R = float(initial_robots)

    time = []
    humans_hist = []
    robots_hist = []
    power_share_robots_hist = []

    for t in range(steps):
        if H <= 1 and R <= 1:
            outcome = "both_extinct"
            break

        # 1) Natural growth / decline
        H += H * (human_growth_rate - baseline_mortality_h)
        R += R * (robot_growth_rate - baseline_mortality_r)

        # prevent negative
        H = max(H, 0.0)
        R = max(R, 0.0)

        # 2) Possible conflict
        if random.random() < conflict_probability and H > 0 and R > 0:
            # overall scale of conflict based on combined forces
            combined = H + R
            total_losses = conflict_intensity * combined

            # Robots are more efficient: they inflict more losses than they take
            # We split losses based on relative strength + robot_advantage
            # Higher robot_advantage => humans lose more of the total_losses
            if combined > 0:
                # share of losses humans vs robots
                # robots "push" more losses onto humans as advantage grows
                human_loss_share = 0.5 * robot_advantage / (1 + robot_advantage)
                robot_loss_share = 1.0 - human_loss_share
            else:
                human_loss_share = 0.5
                robot_loss_share = 0.5

            human_losses = total_losses * human_loss_share
            robot_losses = total_losses * robot_loss_share

            # Scale losses so we never exceed available population
            human_losses = min(human_losses, H)
            robot_losses = min(robot_losses, R)

            H -= human_losses
            R -= robot_losses

            if verbose:
                print(
                    f"Step {t}: conflict! "
                    f"Human losses ~{human_losses:,.0f}, Robot losses ~{robot_losses:,.0f}"
                )

        # 3) Compute "power share" (you can think military+economic+control)
        power_h = H
        power_r = R * robot_advantage  # robots convert population into more “power”
        total_power = power_h + power_r

        if total_power > 0:
            power_share_robots = power_r / total_power
        else:
            power_share_robots = 0.0

        time.append(t)
        humans_hist.append(H)
        robots_hist.append(R)
        power_share_robots_hist.append(power_share_robots)

        # 4) Check end conditions
        if H <= 1 and R > 1:
            outcome = "human_extinction"
            break
        if R <= 1 and H > 1:
            outcome = "robot_extinction"
            break
        if power_share_robots >= takeover_threshold and R > 1:
            outcome = "robot_takeover"
            break
    else:
        outcome = "no_decisive_outcome"

    history = {
        "time": time,
        "humans": humans_hist,
        "robots": robots_hist,
        "robot_power_share": power_share_robots_hist,
    }
    return history, outcome


if __name__ == "__main__":
    # Single run
    history, outcome = simulate_takeover(verbose=True)

    print("\nOutcome:", outcome)

    # Plot populations (log scale because numbers are huge)
    t = history["time"]
    H = history["humans"]
    R = history["robots"]
    P = history["robot_power_share"]

    plt.figure()
    plt.plot(t, H, label="Humans")
    plt.plot(t, R, label="Robots")
    plt.yscale("log")
    plt.xlabel("Time step")
    plt.ylabel("Population (log scale)")
    plt.title("Humans vs Humanoid Robots")
    plt.legend()

    plt.figure()
    plt.plot(t, P)
    plt.xlabel("Time step")
    plt.ylabel("Robot Power Share")
    plt.title("Robot Power Share Over Time (1.0 = full control)")

    plt.show()