"""
telepathy_enhancer.py
ML function that iteratively enhances human telepathic energy (inter-brain coherence)
using reinforcement learning / Bayesian optimization.

Based on Q-CCT Telepathy Framework: shared collapse → maximized coherence.
"""

import numpy as np
from typing import Callable, Tuple, Optional
from scipy.stats import pearsonr
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, Matern, WhiteKernel
from scipy.special import erf

class TelepathicEnergyOptimizer:
    """
    Machine Learning optimizer that learns the best sensory feedback to maximize
    inter-brain coherence (a proxy for telepathic energy).
    
    The human (or pair of humans) receives feedback parameters x (e.g., frequency,
    intensity, delay, visual pattern). The environment returns a telepathy score y.
    The optimizer iterates to find x* that maximizes y.
    
    Use case: real-time EEG from two subjects, compute coherence in alpha/gamma bands.
    """

    def __init__(self,
                 param_bounds: np.ndarray,
                 n_initial_random: int = 8,
                 acquisition_func: str = 'ei',
                 noise_level: float = 0.05):
        """
        Args:
            param_bounds: shape (n_params, 2) – lower and upper bound for each parameter.
            n_initial_random: number of random exploration trials before GP model.
            acquisition_func: 'ei' (expected improvement) or 'ucb' (upper confidence bound).
            noise_level: assumed observation noise.
        """
        self.bounds = param_bounds
        self.n_params = param_bounds.shape[0]
        self.n_init = n_initial_random
        self.acq = acquisition_func
        self.noise = noise_level

        # History of tried parameters and observed telepathy scores
        self.X = []          # list of parameter vectors
        self.y = []          # list of telepathy scores (higher = better)

        # Gaussian Process surrogate model
        kernel = Matern(nu=2.5) * RBF(length_scale=np.ones(self.n_params)) + WhiteKernel(noise_level=noise_level)
        self.gp = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=True, n_restarts_optimizer=5)

    def _suggest_parameter(self) -> np.ndarray:
        """Use acquisition function to propose next best parameter set."""
        if len(self.X) < self.n_init:
            # Random exploration
            return np.random.uniform(self.bounds[:, 0], self.bounds[:, 1])
        # Update GP with all data
        X_arr = np.array(self.X)
        y_arr = np.array(self.y).reshape(-1, 1)
        self.gp.fit(X_arr, y_arr.ravel())

        # Grid search over parameter space
        n_grid = 30  # per dimension (simple grid, can be replaced with LHS)
        grids = [np.linspace(self.bounds[i,0], self.bounds[i,1], n_grid) for i in range(self.n_params)]
        if self.n_params == 1:
            candidates = grids[0].reshape(-1,1)
        else:
            # meshgrid
            meshes = np.meshgrid(*grids, indexing='ij')
            candidates = np.stack([m.ravel() for m in meshes], axis=1)

        # Acquisition values
        mu, sigma = self.gp.predict(candidates, return_std=True)
        if self.acq == 'ei':
            # Expected Improvement
            best_y = np.max(self.y)
            imp = mu - best_y
            Z = imp / (sigma + 1e-9)
            ei = imp * self._cdf(Z) + sigma * self._pdf(Z)
            acq_val = ei
        else:  # UCB
            beta = 2.0
            acq_val = mu + beta * sigma

        best_idx = np.argmax(acq_val)
        return candidates[best_idx]

    @staticmethod
    def _cdf(x):
        return 0.5 * (1 + erf(x / np.sqrt(2)))
    @staticmethod
    def _pdf(x):
        return np.exp(-0.5 * x**2) / np.sqrt(2 * np.pi)

    def iterate(self,
                telepathy_measurement_fn: Callable[[np.ndarray], float],
                max_iterations: int = 30,
                verbose: bool = True) -> Tuple[np.ndarray, list]:
        """
        Main iterative ML loop.

        Args:
            telepathy_measurement_fn: function that takes a parameter vector x
                (e.g., [frequency, intensity, visual_delay]) and returns a telepathy score.
                This function must run one trial (e.g., 60 sec of EEG recording) and
                compute inter-brain coherence averaged over the trial.
            max_iterations: number of optimization steps.
            verbose: print progress.

        Returns:
            best_parameters (np.ndarray), history of scores (list).
        """
        for i in range(max_iterations):
            # Suggest next parameter set
            x_next = self._suggest_parameter()
            # Perform one trial and measure telepathy energy
            y_next = telepathy_measurement_fn(x_next)

            # Append to history
            self.X.append(x_next.tolist())
            self.y.append(y_next)

            if verbose:
                print(f"Iter {i+1}/{max_iterations} | x = {np.round(x_next,3)} | telepathy score = {y_next:.4f} | best so far = {max(self.y):.4f}")

        best_idx = np.argmax(self.y)
        best_params = np.array(self.X[best_idx])
        return best_params, self.y

# ----------------------------------------------------------------------
# Example of a real telepathy measurement function (mock / substitute)
# In a real setup, replace with EEG coherence between two subjects.
# ----------------------------------------------------------------------
def mock_telepathy_measurement(params: np.ndarray) -> float:
    """
    Mock function: simulates that certain audio‑visual frequencies enhance coherence.
    True optimum: params = [7.83 Hz (Schumann), 0.7 intensity, 0.05 sec delay].
    """
    freq, intensity, delay = params
    # Gaussian‑like peak at optimum
    ideal_freq = 7.83
    ideal_intensity = 0.7
    ideal_delay = 0.05
    score = 0.5 * np.exp(-((freq - ideal_freq)**2) / 8.0) \
            + 0.3 * np.exp(-((intensity - ideal_intensity)**2) / 0.1) \
            + 0.2 * np.exp(-((delay - ideal_delay)**2) / 0.002)
    # add noise
    score += np.random.normal(0, 0.03)
    return np.clip(score, 0.0, 1.0)

# ----------------------------------------------------------------------
# Real EEG‑based telepathy function (pseudo‑code – integrate with hardware)
# ----------------------------------------------------------------------
def real_telepathy_measurement(params: np.ndarray) -> float:
    """
    Use two EEG headsets (e.g., Muse 2, OpenBCI). Present auditory/visual stimulus
    defined by params. Compute inter‑brain coherence in alpha (8‑12 Hz) and gamma (30‑45 Hz).
    Return a combined score: coherence + phase locking value.
    """
    # 1. Apply stimulus (sound frequency = params[0], intensity = params[1], delay = params[2])
    # 2. Record 60 seconds of EEG from two subjects
    # 3. Compute coherence between homologous channels (e.g., Fz, Pz) using Welch's method
    # 4. Telepathy score = mean(coherence_alpha) + 0.5*mean(coherence_gamma)
    # Placeholder – replace with actual hardware integration.
    raise NotImplementedError("Integrate with real EEG device and stimulus generator.")

# ----------------------------------------------------------------------
# Run iterative optimization
# ----------------------------------------------------------------------
if __name__ == "__main__":
    # Parameter space: [frequency (Hz), intensity (0..1), delay (sec)]
    bounds = np.array([[4.0, 12.0],   # Theta/Alpha range
                       [0.2, 1.0],    # stimulus intensity
                       [0.01, 0.2]])  # delay between left/right ear, or visual flicker offset

    optimizer = TelepathicEnergyOptimizer(param_bounds=bounds,
                                          n_initial_random=6,
                                          acquisition_func='ei',
                                          noise_level=0.02)

    # Use mock for demonstration (replace with real_telepathy_measurement in lab)
    best_params, history = optimizer.iterate(telepathy_measurement_fn=mock_telepathy_measurement,
                                             max_iterations=25,
                                             verbose=True)

    print("\n" + "="*50)
    print(f"OPTIMAL TELEPATHIC PARAMETERS found:")
    print(f"  Frequency    : {best_params[0]:.2f} Hz")
    print(f"  Intensity    : {best_params[1]:.2f}")
    print(f"  Delay        : {best_params[2]:.3f} sec")
    print(f"  Max telepathy score : {max(history):.4f}")
    print("="*50)
