import numpy as np
from scipy.spatial.distance import cdist
from typing import Callable, Optional, Tuple

def compute_mmd(X: np.ndarray, 
                Y: np.ndarray, 
                kernel: Callable) -> float:
    """
    Compute Maximum Mean Discrepancy (MMD) between X and Y samples.
    
    As described in paper Section 4: "Sequential Two-Sample Kernel-MMD Test",
    MMD is computed as:
    MMD^2(P,Q) = E[k(X,X')] + E[k(Y,Y')] - 2E[k(X,Y)]
    where k is the kernel function.
    
    Args:
        X: First sample set 
        Y: Second sample set
        kernel: Kernel function k(x,y)
        
    Returns:
        float: MMD^2 value
    """
    # Compute kernel matrices
    Kxx = kernel(X)  # k(X,X)
    Kyy = kernel(Y)  # k(Y,Y)
    Kxy = kernel(X, Y)  # k(X,Y)
    
    # Compute expectations
    # Note: We exclude diagonal elements for Kxx and Kyy as they represent k(x,x)
    n = X.shape[0]
    m = Y.shape[0]
    
    xx = (Kxx.sum() - np.trace(Kxx)) / (n * (n-1))  # E[k(X,X')] 
    yy = (Kyy.sum() - np.trace(Kyy)) / (m * (m-1))  # E[k(Y,Y')]
    xy = Kxy.mean()  # E[k(X,Y)]
    
    # Compute MMD^2
    mmd2 = xx + yy - 2*xy
    
    return mmd2

def kernel_mmd_prediction(Xt: np.ndarray,
                         Yt: np.ndarray, 
                         kernel: Callable,
                         post_processing: str = 'clip') -> Tuple[float, dict]:
    """
    Compute MMD prediction for sequential testing using selected kernel.
    
    As discussed in paper Section 4, this implements the witness function
    associated with kernel-MMD for testing.
    
    Args:
        Xt: Current batch X samples
        Yt: Current batch Y samples
        kernel: Kernel function
        post_processing: How to constrain output ('clip' or 'tanh')
        
    Returns:
        tuple: (prediction, metadata)
    """
    # Compute centered kernel matrices
    Kxx = kernel(Xt)
    Kyy = kernel(Yt)
    Kxy = kernel(Xt, Yt)
    
    n = len(Xt)
    m = len(Yt)
    
    # Compute witness function values
    wx = (Kxx.sum(axis=1) - np.diag(Kxx)) / (n-1) 
    wy = Kxy.sum(axis=1) / m
    witness_x = (wx - wy) / n
    
    wy = (Kyy.sum(axis=1) - np.diag(Kyy)) / (m-1)
    wx = Kxy.sum(axis=0) / n  
    witness_y = (wy - wx) / m
    
    # Normalize predictions to [-1/2, 1/2] range as required by betting
    if post_processing == 'clip':
        prediction = np.clip(witness_x.mean() - witness_y.mean(), -0.5, 0.5)
    elif post_processing == 'tanh':
        # Use tanh for unbounded kernels as described in Section 2.3
        prediction = np.tanh(witness_x.mean() - witness_y.mean()) / 2
    else:
        raise ValueError(f"Unknown post_processing: {post_processing}")
    
    metadata = {
        'witness_x': witness_x,
        'witness_y': witness_y,
        'mmd2': compute_mmd(Xt, Yt, kernel)
    }
    
    return prediction, metadata

def get_optimal_kernel_bandwidth(X: np.ndarray, 
                               Y: np.ndarray) -> float:
    """
    Get optimal bandwidth for RBF kernel using median heuristic.
    
    As referenced in the paper's experiments section.
    
    Args:
        X: First sample set
        Y: Second sample set
        
    Returns:
        float: Optimal bandwidth
    """
    # Combine samples
    Z = np.vstack([X, Y])
    
    # Compute pairwise distances
    dists = cdist(Z, Z, 'euclidean')
    
    # Use median heuristic
    bandwidth = np.median(dists[np.triu_indices_from(dists, k=1)])
    
    return bandwidth

import numpy as np
from typing import Callable, Dict, Optional, Tuple
from dataclasses import dataclass

class SequentialKernelMMDTest:
    def __init__(self, 
                 kernel_func: Callable,
                 alpha: float = 0.05):
        """
        Sequential Kernel MMD Test as defined in paper Definition 8.
        
        Args:
            kernel_func: Kernel function k(x,y)
            alpha: Type I error rate
        """
        self.kernel = kernel_func
        self.alpha = alpha
        
        # Initialize test statistics
        self.K0 = 1.0
        self.Kt = 1.0  # Current wealth
        self.wealth_history = [1.0]
        
        # Initialize betting strategy
        self.lambda1 = 0.0
        self.lambda_t = 0.0
        self.a0 = 1.0  # For ONS strategy
        
        # Initialize prediction strategy
        self.g1 = None
        self.gt = None
        self.Mt = 0.0  # Sum of squared RKHS norms
        
    def _ons_betting_update(self, vt: float) -> float:
        """
        Online Newton Step betting strategy (Definition 5 in paper).
        """
        # Update at based on zt
        zt = vt / (1 - vt * self.lambda_t)
        self.a0 += zt * zt
        
        # Update lambda using ONS update rule
        lambda_next = self.lambda_t - (2 - np.log(3)) * zt / self.a0
        
        # Project to [-1/2, 1/2]
        lambda_next = np.clip(lambda_next, -0.5, 0.5)
        
        return lambda_next
    
    def _oga_prediction_update(self, 
                             Xt: np.ndarray, 
                             Yt: np.ndarray, 
                             kt_x: np.ndarray,
                             kt_y: np.ndarray) -> np.ndarray:
        """
        Online Gradient Ascent prediction strategy (Equation 17 in paper).
        """
        if self.gt is None:
            self.gt = np.zeros_like(kt_x)
            return self.gt
            
        # Compute RKHS norm of difference
        diff = kt_x - kt_y
        self.Mt += np.sum(diff * diff)
        
        # Update using OGA
        grad = diff / (2 * np.sqrt(self.Mt))
        gt_new = self.gt + grad
        
        # Project onto G using RKHS norm
        norm = np.sqrt(np.sum(gt_new * gt_new))
        if norm > 0:
            gt_new = gt_new / (2 * norm)  # ΠG operator
            
        return gt_new

    def update(self, 
               Xt: np.ndarray, 
               Yt: np.ndarray) -> Tuple[bool, str]:
        """
        Update test statistics for new observations.
        
        Args:
            Xt: New X observation
            Yt: New Y observation
            
        Returns:
            (should_stop, decision)
        """
        # Compute kernel evaluations
        kt_x = self.kernel(Xt)
        kt_y = self.kernel(Yt)
        
        # Update prediction gt+1 using OGA
        self.gt = self._oga_prediction_update(Xt, Yt, kt_x, kt_y)
        
        # Compute test statistic
        vt = np.dot(self.gt, (kt_x - kt_y))
        
        # Update wealth
        self.Kt *= (1 + self.lambda_t * vt)
        self.wealth_history.append(self.Kt)
        
        # Update betting fraction λt+1 using ONS
        self.lambda_t = self._ons_betting_update(vt)
        
        # Check stopping condition
        if self.Kt >= 1.0/self.alpha:
            return True, "Reject null hypothesis - Distributions differ significantly"
        
        return False, "Continue testing"
        
    def get_results(self) -> Dict:
        """Get test results."""
        return {
            'stopped': self.Kt >= 1.0/self.alpha,
            'wealth': self.Kt,
            'wealth_history': self.wealth_history,
            'lambda_final': self.lambda_t
        }

def run_sequential_mmd_test(X_stream: np.ndarray,
                           Y_stream: np.ndarray,
                           kernel_func: Callable,
                           alpha: float = 0.05,
                           max_samples: Optional[int] = None) -> Dict:
    """
    Run sequential kernel MMD test on two sample streams.
    
    Args:
        X_stream: First sample stream 
        Y_stream: Second sample stream
        kernel_func: Kernel function
        alpha: Type I error rate
        max_samples: Maximum samples to test (optional)
        
    Returns:
        Dict containing test results
    """
    # Initialize test
    test = SequentialKernelMMDTest(kernel_func, alpha)
    
    # Set max samples
    if max_samples is None:
        max_samples = len(X_stream)
    n = min(len(X_stream), len(Y_stream), max_samples)
    
    # Run sequential test
    for t in range(n):
        Xt = X_stream[t]
        Yt = Y_stream[t]
        
        should_stop, decision = test.update(Xt, Yt)
        if should_stop:
            break
            
    # Get results
    results = test.get_results()
    results.update({
        'samples_used': t + 1,
        'decision': decision
    })
    
    return results

import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass

@dataclass
class ROIMetrics:
    revenue: float
    spend: float
    
    def roi(self) -> float:
        return self.revenue / self.spend if self.spend > 0 else 0.0

class ROIKernelMMDTest:
    """
    Kernel MMD test for ROI comparisons using betting martingales.
    Combines ROI witness function with kernel MMD framework from paper.
    """
    def __init__(self, 
                kernel_bandwidth: float = 1.0,
                alpha: float = 0.05,
                min_effect: float = 0.10):
        self.kernel_bandwidth = kernel_bandwidth
        self.alpha = alpha
        self.min_effect = min_effect
        self.wealth = 1.0
        self.wealth_history = [1.0]

    def compute_kernel(self, x: ROIMetrics, y: ROIMetrics) -> float:
        """
        Compute kernel between two ROI values using RBF kernel.
        """
        roi_diff = x.roi() - y.roi()
        return np.exp(-(roi_diff**2) / (2 * self.kernel_bandwidth**2))

    def compute_mmd_squared(self, 
                          adaptive_samples: List[ROIMetrics],
                          fixed_samples: List[ROIMetrics]) -> float:
        """
        Compute squared MMD between adaptive and fixed strategy samples.
        
        MMD^2 = E[k(x,x')] + E[k(y,y')] - 2E[k(x,y)] 
        where k is the kernel function
        """
        n = len(adaptive_samples)
        m = len(fixed_samples)

        # Compute kernel matrices
        K_aa = np.zeros((n, n))
        K_ff = np.zeros((m, m))
        K_af = np.zeros((n, m))

        for i in range(n):
            for j in range(n):
                K_aa[i,j] = self.compute_kernel(adaptive_samples[i], adaptive_samples[j])
            for j in range(m):
                K_af[i,j] = self.compute_kernel(adaptive_samples[i], fixed_samples[j])

        for i in range(m):
            for j in range(m):
                K_ff[i,j] = self.compute_kernel(fixed_samples[i], fixed_samples[j])

        # Compute expectations (excluding diagonal elements where appropriate)
        E_aa = (K_aa.sum() - np.trace(K_aa)) / (n * (n-1))
        E_ff = (K_ff.sum() - np.trace(K_ff)) / (m * (m-1))
        E_af = K_af.mean()

        mmd_squared = E_aa + E_ff - 2*E_af
        return mmd_squared

    def compute_witness_value(self,
                            adaptive_metrics: ROIMetrics,
                            fixed_metrics: ROIMetrics,
                            adaptive_history: List[ROIMetrics],
                            fixed_history: List[ROIMetrics]) -> float:
        """
        Compute witness function value for sequential testing.
        Maps to [-1/2, 1/2] range as required by betting strategy.
        """
        # Compute ROI difference
        roi_diff = adaptive_metrics.roi() - fixed_metrics.roi()
        
        # Weight by kernel-based similarity to historical data
        kernel_weight = 0.0
        n_hist = len(adaptive_history)
        if n_hist > 0:
            for hist_a, hist_f in zip(adaptive_history, fixed_history):
                k_a = self.compute_kernel(adaptive_metrics, hist_a)
                k_f = self.compute_kernel(fixed_metrics, hist_f)
                kernel_weight += (k_a - k_f) / n_hist

        # Combine direct ROI difference with kernel-weighted historical influence
        witness = roi_diff + kernel_weight
        
        # Map to [-1/2, 1/2] using tanh
        return 0.5 * np.tanh(witness)

    def update(self,
              adaptive_metrics: ROIMetrics,
              fixed_metrics: ROIMetrics,
              adaptive_history: List[ROIMetrics],
              fixed_history: List[ROIMetrics]) -> Tuple[bool, str]:
        """
        Update test statistics and check stopping condition.
        
        Returns:
            Tuple of (should_stop, decision)
        """
        # Compute witness value
        witness = self.compute_witness_value(
            adaptive_metrics,
            fixed_metrics, 
            adaptive_history,
            fixed_history
        )
        
        # Compute optimal bet size (Kelly fraction)
        relative_improvement = (adaptive_metrics.roi() - fixed_metrics.roi()) / abs(fixed_metrics.roi()) if fixed_metrics.roi() != 0 else 0
        bet = (relative_improvement - self.min_effect) / 2
        bet = max(-0.5, min(0.5, bet))
        
        # Update wealth
        self.wealth *= (1 + bet * witness)
        self.wealth_history.append(self.wealth)

        # Check stopping conditions
        if self.wealth >= 1.0 / self.alpha:
            return True, "Reject null - Adaptive strategy shows significant improvement"
        elif self.wealth <= self.alpha:
            return True, "Accept null - Insufficient evidence for improvement"
        
        return False, "Continue testing"

def run_roi_mmd_experiment(
    adaptive_metrics: List[ROIMetrics],
    fixed_metrics: List[ROIMetrics],
    kernel_bandwidth: float = 1.0,
    alpha: float = 0.05,
    min_effect: float = 0.10
) -> Dict:
    """
    Run complete ROI MMD experiment.
    
    Args:
        adaptive_metrics: List of metrics from adaptive strategy
        fixed_metrics: List of metrics from fixed strategy
        kernel_bandwidth: Kernel sensitivity parameter
        alpha: Type I error rate
        min_effect: Minimum ROI improvement to detect
        
    Returns:
        Dict containing test results and metrics
    """
    test = ROIKernelMMDTest(
        kernel_bandwidth=kernel_bandwidth,
        alpha=alpha,
        min_effect=min_effect
    )
    
    n_days = len(adaptive_metrics)
    stopped = False
    stop_day = n_days
    decision = "Inconclusive"
    
    for day in range(n_days):
        # Get current day metrics
        adaptive = adaptive_metrics[day]
        fixed = fixed_metrics[day]
        
        # Get history up to current day
        adaptive_history = adaptive_metrics[:day]
        fixed_history = fixed_metrics[:day]
        
        # Update test
        should_stop, result = test.update(
            adaptive,
            fixed,
            adaptive_history,
            fixed_history
        )
        
        if should_stop:
            stopped = True
            stop_day = day + 1
            decision = result
            break
            
    # Compute final MMD
    final_mmd = test.compute_mmd_squared(adaptive_metrics, fixed_metrics)
    
    return {
        'stopped': stopped,
        'stop_day': stop_day,
        'decision': decision,
        'final_mmd': final_mmd,
        'wealth_history': test.wealth_history
    }