import numpy as np
from typing import List, Dict, Tuple

class AdSelector:
    """
    A simplified ad selection system using contextual bandits.

    In practice, this would be much more sophisticated with:
    - Real-time bidding integration
    - Budget pacing
    - Frequency capping
    - Multiple objectives
    """

    def __init__(self, n_features: int, alpha: float = 1.0):
        self.n_features = n_features
        self.alpha = alpha
        # Each ad gets its own LinUCB model
        self.ad_models: Dict[str, 'LinUCBModel'] = {}

    def select_ad(
        self,
        user_context: np.ndarray,
        candidate_ads: List[Dict]
    ) -> str:
        """
        Select the best ad for this user from candidates.

        Args:
            user_context: Feature vector describing the user
            candidate_ads: List of ad dictionaries with 'id' and 'features'

        Returns:
            ID of selected ad
        """
        best_ad_id = None
        best_ucb = -float('inf')

        for ad in candidate_ads:
            ad_id = ad['id']

            # Create model if new ad
            if ad_id not in self.ad_models:
                self.ad_models[ad_id] = LinUCBModel(
                    self.n_features, self.alpha
                )

            # Combine user and ad features
            combined_context = np.concatenate([
                user_context,
                ad['features']
            ])

            # Get UCB value
            ucb = self.ad_models[ad_id].get_ucb(combined_context)

            if ucb > best_ucb:
                best_ucb = ucb
                best_ad_id = ad_id

        return best_ad_id

    def update(
        self,
        ad_id: str,
        context: np.ndarray,
        clicked: bool
    ):
        """Update model after observing click/no-click."""
        if ad_id in self.ad_models:
            reward = 1.0 if clicked else 0.0
            self.ad_models[ad_id].update(context, reward)


class LinUCBModel:
    """Simple LinUCB model for a single arm."""

    def __init__(self, d: int, alpha: float = 1.0):
        self.A = np.eye(d)
        self.b = np.zeros(d)
        self.alpha = alpha

    def get_ucb(self, x: np.ndarray) -> float:
        A_inv = np.linalg.inv(self.A)
        theta = A_inv @ self.b
        mean = x @ theta
        std = np.sqrt(x @ A_inv @ x)
        return mean + self.alpha * std

    def update(self, x: np.ndarray, reward: float):
        self.A += np.outer(x, x)
        self.b += reward * x

import numpy as np
from typing import Dict, List, Tuple

class BanditABTest:
    """
    A/B testing using Thompson Sampling.

    Automatically balances exploration and exploitation,
    sending more traffic to winning variants over time.
    """

    def __init__(self, variant_names: List[str]):
        self.variants = variant_names
        # Beta priors for each variant (for Bernoulli rewards)
        self.alpha = {name: 1.0 for name in variant_names}
        self.beta = {name: 1.0 for name in variant_names}
        self.impressions = {name: 0 for name in variant_names}
        self.conversions = {name: 0 for name in variant_names}

    def select_variant(self) -> str:
        """Select a variant using Thompson Sampling."""
        samples = {
            name: np.random.beta(self.alpha[name], self.beta[name])
            for name in self.variants
        }
        return max(samples, key=samples.get)

    def record_outcome(self, variant: str, converted: bool):
        """Record the outcome of showing a variant."""
        self.impressions[variant] += 1
        if converted:
            self.conversions[variant] += 1
            self.alpha[variant] += 1
        else:
            self.beta[variant] += 1

    def get_statistics(self) -> Dict:
        """Return current statistics for all variants."""
        stats = {}
        for name in self.variants:
            mean = self.alpha[name] / (self.alpha[name] + self.beta[name])
            # 95% credible interval
            from scipy import stats
            low = stats.beta.ppf(0.025, self.alpha[name], self.beta[name])
            high = stats.beta.ppf(0.975, self.alpha[name], self.beta[name])
            stats[name] = {
                'impressions': self.impressions[name],
                'conversions': self.conversions[name],
                'rate': self.conversions[name] / max(1, self.impressions[name]),
                'posterior_mean': mean,
                'credible_interval': (low, high)
            }
        return stats

    def get_probability_best(self, n_samples: int = 10000) -> Dict[str, float]:
        """Estimate probability each variant is actually best."""
        samples = {
            name: np.random.beta(
                self.alpha[name], self.beta[name], n_samples
            )
            for name in self.variants
        }

        # Stack and find argmax for each sample
        all_samples = np.stack([samples[name] for name in self.variants])
        best_indices = np.argmax(all_samples, axis=0)

        probs = {}
        for i, name in enumerate(self.variants):
            probs[name] = np.mean(best_indices == i)
        return probs


# Example usage
test = BanditABTest(['control', 'variant_a', 'variant_b', 'variant_c'])

# Simulate 1000 users
true_rates = {'control': 0.03, 'variant_a': 0.035,
              'variant_b': 0.04, 'variant_c': 0.032}

for _ in range(1000):
    variant = test.select_variant()
    converted = np.random.random() < true_rates[variant]
    test.record_outcome(variant, converted)

# Check results
print("Traffic distribution:", {k: v['impressions'] for k, v in test.get_statistics().items()})
print("Probability best:", test.get_probability_best())

import numpy as np
import torch
import torch.nn as nn

class NeuralLinearBandit:
    """
    Neural contextual bandit using the Neural Linear approach.

    A neural network learns features, then LinUCB operates
    on the last layer for tractable uncertainty estimation.
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        n_arms: int,
        alpha: float = 1.0
    ):
        self.n_arms = n_arms
        self.alpha = alpha

        # Neural network for feature extraction
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )

        # LinUCB head for each arm
        self.feature_dim = hidden_dim
        self.A = [np.eye(hidden_dim) for _ in range(n_arms)]
        self.b = [np.zeros(hidden_dim) for _ in range(n_arms)]

        # Training data for neural net updates
        self.history = []

    def get_features(self, context: np.ndarray) -> np.ndarray:
        """Extract features using the neural network."""
        with torch.no_grad():
            x = torch.FloatTensor(context).unsqueeze(0)
            features = self.network(x).numpy().squeeze()
        return features

    def select_action(self, context: np.ndarray) -> int:
        """Select arm with highest UCB."""
        features = self.get_features(context)

        best_arm = 0
        best_ucb = -float('inf')

        for arm in range(self.n_arms):
            A_inv = np.linalg.inv(self.A[arm])
            theta = A_inv @ self.b[arm]

            mean = features @ theta
            std = np.sqrt(features @ A_inv @ features)
            ucb = mean + self.alpha * std

            if ucb > best_ucb:
                best_ucb = ucb
                best_arm = arm

        return best_arm

    def update(self, context: np.ndarray, action: int, reward: float):
        """Update LinUCB head for the selected arm."""
        features = self.get_features(context)

        self.A[action] += np.outer(features, features)
        self.b[action] += reward * features

        # Store for potential neural network retraining
        self.history.append((context, action, reward))