import numpy as np

class VanillaCFR:
    """Vanilla Counterfactual Regret Minimization."""

    def __init__(self, game):
        self.game = game
        # Cumulative regret per (infoset, action)
        self.cumulative_regret = {}
        # Cumulative strategy per (infoset, action)
        self.cumulative_strategy = {}
        self.iteration = 0

    def get_strategy(self, info_set: str, actions: list) -> np.ndarray:
        """Compute current strategy via regret matching."""
        regrets = np.array([
            max(self.cumulative_regret.get((info_set, a), 0.0), 0.0)
            for a in actions
        ])
        total = regrets.sum()
        if total > 0:
            return regrets / total
        else:
            return np.ones(len(actions)) / len(actions)

    def cfr(self, state, reach_probs: dict) -> dict:
        """Recursive CFR traversal."""
        if state.is_terminal():
            return state.payoffs()

        player = state.current_player()
        info_set = state.info_set()
        actions = state.legal_actions()
        strategy = self.get_strategy(info_set, actions)

        # Recurse for each action
        action_values = {}
        node_value = {p: 0.0 for p in self.game.players}

        for i, action in enumerate(actions):
            new_reach = dict(reach_probs)
            new_reach[player] *= strategy[i]
            child_state = state.apply_action(action)
            action_values[action] = self.cfr(child_state, new_reach)

            for p in self.game.players:
                node_value[p] += strategy[i] * action_values[action][p]

        # Compute counterfactual regret
        cf_reach = 1.0
        for p in self.game.players:
            if p != player:
                cf_reach *= reach_probs[p]

        for i, action in enumerate(actions):
            regret = cf_reach * (action_values[action][player] - node_value[player])
            key = (info_set, action)
            self.cumulative_regret[key] = self.cumulative_regret.get(key, 0.0) + regret

        # Accumulate strategy (weighted by reach probability)
        for i, action in enumerate(actions):
            key = (info_set, action)
            self.cumulative_strategy[key] = (
                self.cumulative_strategy.get(key, 0.0) + reach_probs[player] * strategy[i]
            )

        return node_value

    def solve(self, num_iterations: int):
        """Run CFR for num_iterations."""
        for t in range(num_iterations):
            self.iteration = t + 1
            initial_reach = {p: 1.0 for p in self.game.players}
            self.cfr(self.game.root(), initial_reach)

    def average_strategy(self, info_set: str, actions: list) -> np.ndarray:
        """Return the average strategy (converges to Nash)."""
        strat = np.array([
            self.cumulative_strategy.get((info_set, a), 0.0)
            for a in actions
        ])
        total = strat.sum()
        if total > 0:
            return strat / total
        return np.ones(len(actions)) / len(actions)

# AlphaEvolve configuration for game algorithm discovery
config = {
    "population_size": 100,
    "num_islands": 10,
    "llm": "gemini-2.5-pro",
    "evaluation_suite": [
        "kuhn_poker",
        "leduc_poker",
        "goofspiel_4",
        "liars_dice_4",
        "dark_hex_3x3",
    ],
    "iterations_per_eval": 1000,     # CFR/PSRO iterations per game
    "fitness": "geometric_mean_nashconv",
    "mutation_strategies": [
        "modify_regret_update",
        "modify_discount_weights",
        "modify_strategy_computation",
        "add_adaptive_parameter",
        "modify_meta_solver",
    ],
}

import numpy as np
from collections import defaultdict

class VADCFR:
    """Volatility-Adaptive Discounted CFR with Consistency-Enforced Optimism."""

    def __init__(self, game, alpha_base=0.5, eta=0.3, beta=0.1, gamma=0.01, k=10):
        self.game = game
        self.alpha_base = alpha_base  # Base discount factor
        self.eta = eta                # Volatility adaptation strength
        self.beta = beta              # Optimism bonus coefficient
        self.gamma = gamma            # Exploration bonus coefficient
        self.k = k                    # Consistency window size

        self.cumulative_regret = defaultdict(float)
        self.cumulative_strategy = defaultdict(float)
        self.prev_strategy = {}       # Previous iteration strategies
        self.regret_history = defaultdict(list)  # Recent regret signs
        self.action_counts = defaultdict(float)  # Cumulative action probabilities
        self.iteration = 0

    def compute_volatility(self, info_set, actions, current_strategy):
        """Compute strategy volatility for adaptive discounting."""
        if info_set not in self.prev_strategy:
            return 1.0  # Maximum volatility for first iteration
        prev = self.prev_strategy[info_set]
        diff = current_strategy - prev
        return np.linalg.norm(diff)

    def compute_consistency(self, info_set, action):
        """Compute action consistency for optimistic bonus."""
        history = self.regret_history.get((info_set, action), [])
        if len(history) == 0:
            return 0.0
        recent = history[-self.k:]
        return sum(1 for r in recent if r > 0) / len(recent)

    def get_strategy(self, info_set, actions):
        """Compute strategy with optimism and exploration bonuses."""
        # Base regret matching
        regrets = np.array([
            max(self.cumulative_regret[(info_set, a)], 0.0) for a in actions
        ])

        # Add consistency-enforced optimism
        for i, a in enumerate(actions):
            consistency = self.compute_consistency(info_set, a)
            regrets[i] += self.beta * consistency

        total = regrets.sum()
        if total > 0:
            strategy = regrets / total
        else:
            strategy = np.ones(len(actions)) / len(actions)

        # Add exploration bonus for rarely-played actions
        for i, a in enumerate(actions):
            count = self.action_counts.get((info_set, a), 0.0)
            strategy[i] += self.gamma / np.sqrt(count + 1)

        # Renormalize
        strategy = strategy / strategy.sum()
        return strategy

    def update_regret(self, info_set, actions, regrets, strategy):
        """Update regret with volatility-adaptive discounting."""
        volatility = self.compute_volatility(info_set, actions, strategy)
        vol_max = np.sqrt(2)  # Maximum L2 norm for simplex strategies

        # Adaptive discount: less discounting when volatile
        alpha = self.alpha_base + self.eta * (1 - volatility / vol_max)
        alpha = np.clip(alpha, 0.0, 1.0)

        for i, a in enumerate(actions):
            key = (info_set, a)
            # Discount old regret and add new
            self.cumulative_regret[key] = alpha * self.cumulative_regret[key] + regrets[i]
            # Track regret signs for consistency
            self.regret_history[key].append(regrets[i])
            if len(self.regret_history[key]) > self.k * 2:
                self.regret_history[key] = self.regret_history[key][-self.k:]
            # Update action counts
            self.action_counts[key] = self.action_counts.get(key, 0.0) + strategy[i]

        # Store strategy for volatility computation
        self.prev_strategy[info_set] = strategy.copy()

    def solve(self, num_iterations):
        """Run VAD-CFR for num_iterations."""
        for t in range(num_iterations):
            self.iteration = t + 1
            self._cfr_traversal(self.game.root(), {p: 1.0 for p in self.game.players})

    def _cfr_traversal(self, state, reach_probs):
        """Recursive CFR traversal with VAD modifications."""
        if state.is_terminal():
            return state.payoffs()

        player = state.current_player()
        info_set = state.info_set()
        actions = state.legal_actions()
        strategy = self.get_strategy(info_set, actions)

        # Standard recursive CFR
        action_values = {}
        node_value = {p: 0.0 for p in self.game.players}

        for i, action in enumerate(actions):
            new_reach = dict(reach_probs)
            new_reach[player] *= strategy[i]
            action_values[action] = self._cfr_traversal(
                state.apply_action(action), new_reach
            )
            for p in self.game.players:
                node_value[p] += strategy[i] * action_values[action][p]

        # Compute counterfactual regret
        cf_reach = np.prod([reach_probs[p] for p in self.game.players if p != player])
        regrets = np.array([
            cf_reach * (action_values[a][player] - node_value[player])
            for a in actions
        ])

        # Update with adaptive discounting
        self.update_regret(info_set, actions, regrets, strategy)

        # Accumulate average strategy
        for i, action in enumerate(actions):
            key = (info_set, action)
            self.cumulative_strategy[key] += reach_probs[player] * strategy[i]

        return node_value

import numpy as np
from scipy.optimize import linprog

class SHORPSRO:
    """Smoothed Hybrid Optimistic Regret PSRO."""

    def __init__(self, game, tau=0.1, lambda_novelty=0.3, p_br=0.7, mu=0.5):
        self.game = game
        self.tau = tau                # Smoothing temperature
        self.lambda_novelty = lambda_novelty  # Novelty bonus
        self.p_br = p_br              # Best response probability
        self.mu = mu                  # Optimistic correction strength

        self.population = {p: [] for p in game.players}
        self.meta_payoffs = None       # Meta-game payoff tensor
        self.meta_strategy = None      # Current meta-equilibrium
        self.prev_regrets = None       # For optimistic correction

    def smoothed_nash(self, payoff_matrix, player):
        """Compute entropy-regularized Nash equilibrium."""
        n = payoff_matrix.shape[0]
        # Iterative softmax best response with entropy bonus
        sigma = np.ones(n) / n  # Uniform start

        for _ in range(100):  # Fixed-point iteration
            utilities = payoff_matrix @ sigma if player == 0 else payoff_matrix.T @ sigma
            # Softmax with temperature tau
            exp_u = np.exp((utilities - utilities.max()) / self.tau)
            sigma_new = exp_u / exp_u.sum()
            if np.allclose(sigma, sigma_new, atol=1e-8):
                break
            sigma = sigma_new

        return sigma

    def compute_diverse_response(self, meta_strategy, player):
        """Compute a response that balances utility and novelty."""
        # Train a policy against meta_strategy mixture
        base_response = self.compute_best_response(meta_strategy, player)

        # Compute distances to existing population
        if len(self.population[player]) == 0:
            return base_response

        min_distance = min(
            self.policy_distance(base_response, existing)
            for existing in self.population[player]
        )

        # If already novel enough, return best response
        if min_distance > 0.1:
            return base_response

        # Otherwise, add perturbation for diversity
        perturbed = base_response.add_noise(scale=0.1)
        return perturbed

    def update_meta_strategy(self):
        """Update meta-strategy using optimistic regret matching."""
        payoffs = self.meta_payoffs
        n = len(self.population[0])

        # Compute current regrets
        current_regrets = np.zeros(n)
        for j in range(n):
            pure_strategy = np.zeros(n)
            pure_strategy[j] = 1.0
            current_regrets[j] = (
                pure_strategy @ payoffs @ self.meta_strategy[1]
                - self.meta_strategy[0] @ payoffs @ self.meta_strategy[1]
            )

        # Apply optimistic correction
        if self.prev_regrets is not None and len(self.prev_regrets) == n:
            optimistic_regrets = current_regrets + self.mu * (current_regrets - self.prev_regrets)
        else:
            optimistic_regrets = current_regrets

        self.prev_regrets = current_regrets.copy()

        # Regret matching with optimistic regrets
        positive_regrets = np.maximum(optimistic_regrets, 0)
        total = positive_regrets.sum()
        if total > 0:
            self.meta_strategy[0] = positive_regrets / total
        else:
            self.meta_strategy[0] = np.ones(n) / n

    def iterate(self):
        """Single PSRO iteration."""
        # 1. Update meta-strategy (smoothed + optimistic)
        self.update_meta_strategy()

        # 2. For each player, add new policy
        for player in self.game.players:
            if np.random.random() < self.p_br:
                # Best response
                new_policy = self.compute_best_response(
                    self.meta_strategy[1 - player], player
                )
            else:
                # Diverse response
                new_policy = self.compute_diverse_response(
                    self.meta_strategy[1 - player], player
                )
            self.population[player].append(new_policy)

        # 3. Update meta-game payoff matrix
        self.meta_payoffs = self.compute_meta_payoffs()

    def solve(self, num_iterations):
        """Run SHOR-PSRO for num_iterations."""
        # Initialize with random policies
        for player in self.game.players:
            self.population[player].append(self.random_policy(player))

        self.meta_payoffs = self.compute_meta_payoffs()
        n = len(self.population[0])
        self.meta_strategy = [np.ones(n) / n for _ in self.game.players]

        for t in range(num_iterations):
            self.iterate()
            nashconv = self.compute_nashconv()
            print(f"Iteration {t}: NashConv = {nashconv:.6f}, "
                  f"Population size = {len(self.population[0])}")

+-----------------------------------------------------------------------------------+
|                        Aletheia Agent Architecture                                 |
+-----------------------------------------------------------------------------------+
|                                                                                    |
|  +------------------+     +-------------------+     +-------------------------+    |
|  |   Problem Input  |     |   Gemini 3        |     |   Lean 4 Environment    |    |
|  |                  |---->|   Deep Think       |---->|                         |    |
|  |  Natural language |     |                   |     |  Proof checker          |    |
|  |  Lean 4 statement |     |  Extended thinking |     |  Type checker           |    |
|  +------------------+     |  Chain of thought  |     |  Tactic evaluation      |    |
|                           +-------------------+     +-------------------------+    |
|                                    |                         |                     |
|                                    v                         v                     |
|  +------------------------------------------------------------------------+        |
|  |                        Agent Loop                                       |        |
|  |                                                                         |        |
|  |  1. Understand problem statement                                        |        |
|  |  2. Formulate proof strategy (high-level)                               |        |
|  |  3. Decompose into lemmas                                               |        |
|  |  4. Write Lean 4 proof for each lemma                                   |        |
|  |  5. Submit to Lean 4 checker                                            |        |
|  |  6. If error: analyze error, revise strategy, go to step 3              |        |
|  |  7. If success: combine lemmas into complete proof                      |        |
|  |  8. Submit final proof                                                  |        |
|  +------------------------------------------------------------------------+        |
|                                    |                                               |
|                                    v                                               |
|  +------------------------------------------------------------------------+        |
|  |                     Strategy Selection                                  |        |
|  |                                                                         |        |
|  |  Direct Proof | Induction | Contradiction | Case Analysis | Counting   |        |
|  |                                                                         |        |
|  |  The agent selects and sometimes switches strategies based on           |        |
|  |  feedback from the Lean 4 proof checker.                                |        |
|  +------------------------------------------------------------------------+        |
+-----------------------------------------------------------------------------------+