import numpy as np
import time
import sys

# Define parameters
NDIM = 10
N_POINTS = 10000  # "alpha"
N_MEANS = 10000   # "alpha"
N_REPEAT = 10      # Number of repeats for timing

# Generate random data
np.random.seed(42)
points = np.random.randn(N_POINTS, NDIM)
means = np.random.randn(N_MEANS, NDIM)
# Create a valid positive-definite covariance matrix
A = np.random.randn(NDIM, NDIM)
cov = np.dot(A, A.T)
inv_cov = np.linalg.inv(cov)
norm_term = 1.0 / np.sqrt((2 * np.pi) ** NDIM * np.linalg.det(cov))
proposal_coeffs = np.ones(N_MEANS)

# --- Original Implementation (from parismc/utils.py) ---
def original_weighting(points_array, means_array, inv_covariance, norm_term, proposalcoeff_array):
    weights = np.zeros((points_array.shape[0]))
    for k in range(means_array.shape[0]):
        diff = points_array - means_array[k]
        # Note: Original code uses 'ij,jk,ik->i' which implies (N, D) @ (D, D) @ (N, D)^T diagonal
        exponent = -0.5 * np.einsum('ij,jk,ik->i', diff, inv_covariance, diff)
        # exponent[exponent == 0] = ... (omitted for speed test, minor logic)
        weights += norm_term * np.exp(exponent) * proposalcoeff_array[k]
    return weights

# --- Optimized Implementation (Vectorized) ---
def optimized_weighting(points_array, means_array, inv_covariance, norm_term, proposalcoeff_array):
    # 1. Compute term 1: x^T * inv_cov * x (Shape: N)
    # x_inv = points @ inv_cov
    # term1 = np.sum(x_inv * points, axis=1)
    # Optimized einsum for term 1:
    term1 = np.einsum('ij,jk,ik->i', points_array, inv_covariance, points_array)
    
    # 2. Compute term 2: mu^T * inv_cov * mu (Shape: M)
    term2 = np.einsum('ij,jk,ik->i', means_array, inv_covariance, means_array)
    
    # 3. Compute cross term: x^T * inv_cov * mu (Shape: N x M)
    # This is the heavy part: (N, D) @ (D, D) @ (D, M) -> (N, M)
    # Let Y = points @ inv_cov (N, D)
    # Cross = Y @ means.T (N, M)
    Y = points_array @ inv_covariance
    cross_term = Y @ means_array.T
    
    # Combine: dist_sq = term1[:, None] + term2[None, :] - 2 * cross_term
    # Shape: (N, M)
    dist_sq = term1[:, None] + term2[None, :] - 2 * cross_term
    
    # Compute weights
    # exponent = -0.5 * dist_sq
    # weights_matrix = norm_term * np.exp(exponent) * proposalcoeff_array[None, :]
    # Final weights: sum over means (axis 1)
    
    # Use log-sum-exp trick or direct sum depending on numerical stability needs
    # Original code does direct sum of exp. We match that.
    weights = np.sum(norm_term * np.exp(-0.5 * dist_sq) * proposalcoeff_array[None, :], axis=1)
    
    return weights

# --- Verification ---
print("Verifying correctness...")
t0 = time.time()
w_orig = original_weighting(points[:100], means[:100], inv_cov, norm_term, proposal_coeffs[:100])
print(f"Original (small batch) time: {time.time()-t0:.4f}s")

t0 = time.time()
w_opt = optimized_weighting(points[:100], means[:100], inv_cov, norm_term, proposal_coeffs[:100])
print(f"Optimized (small batch) time: {time.time()-t0:.4f}s")

# Check alignment
if np.allclose(w_orig, w_opt, atol=1e-10):
    print("SUCCESS: Optimized implementation matches original results.")
else:
    print("FAILURE: Implementations differ!")
    print("Max diff:", np.max(np.abs(w_orig - w_opt)))
    sys.exit(1)

# --- Benchmarking ---
print(f"\nBenchmarking with N={N_POINTS}, Alpha={N_MEANS}, repeat={N_REPEAT}...")

# Original
start_time = time.time()
for _ in range(N_REPEAT):
    # Running on subset to save time if N is huge, but here we requested 10k
    # 10k x 10k loop in Python is SLOW. We might need to reduce N_REPEAT or reduce N for the SLOW one.
    # Let's run just ONCE for the slow one to estimate.
    if _ == 0:
        _ = original_weighting(points, means, inv_cov, norm_term, proposal_coeffs)
orig_duration = (time.time() - start_time) 
# Estimate total if we only ran 1
if N_REPEAT > 1:
    print(f"(Estimating Original time based on 1 run: {orig_duration:.4f}s)")
    orig_total = orig_duration * N_REPEAT
else:
    orig_total = orig_duration

# Optimized
start_time = time.time()
for _ in range(N_REPEAT):
    _ = optimized_weighting(points, means, inv_cov, norm_term, proposal_coeffs)
opt_duration = time.time() - start_time
opt_avg = opt_duration / N_REPEAT

print(f"\nResults:")
print(f"Original (Est. Avg): {orig_duration:.4f} s")
print(f"Optimized (Avg):     {opt_avg:.4f} s")
print(f"Speedup:             {orig_duration / opt_avg:.2f}x")
