"""
Mobiu-Q Client - Soft Algebra Optimizer
========================================
Cloud-connected optimizer for quantum, RL, and LLM applications.

Version: 3.2.1 - AUTO Mode with Maximize Support

NEW in v3.2:
- method='auto': Automatically selects best mode (boost/dampen/off) per problem
- Runs 3 virtual optimizers in parallel, picks winner every 20 steps
- DecayBoost option for safer optimization (less improvement, fewer hurts)

Method names:
- method='standard' (was 'vqe'): For smooth landscapes, chemistry, physics
- method='deep' (was 'qaoa'): For deep circuits, noisy hardware, complex optimization
- method='adaptive' (was 'rl'): For RL, LLM fine-tuning, high-variance problems
- method='auto': Automatic mode selection (NEW!)
- method='auto-safe': Auto with DecayBoost for fewer hurts (NEW!)

Usage:
    from mobiu_q import MobiuQCore
    
    # Auto-select best mode
    opt = MobiuQCore(license_key="your-key", method="auto")
    
    for step in range(100):
        params = opt.step(params, energy_fn)
    
    opt.end()
"""

import numpy as np
import requests
from typing import Optional, Tuple, List, Union
import os
import json
import warnings
import time
from collections import deque

# ═══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════════

API_ENDPOINT = os.environ.get(
    "MOBIU_Q_API_ENDPOINT",
    "https://us-central1-mobiu-q.cloudfunctions.net/mobiu_q_step"
)

LICENSE_KEY_FILE = os.path.expanduser("~/.mobiu_q_license")

# Available optimizers (must match server's AVAILABLE_OPTIMIZERS)
AVAILABLE_OPTIMIZERS = ["Adam", "AdamW", "NAdam", "AMSGrad", "SGD", "Momentum", "LAMB"]
DEFAULT_OPTIMIZER = "Adam"

# Method name mapping (new names + legacy support)
METHOD_ALIASES = {
    # New names (v2.5+)
    "standard": "standard",
    "deep": "deep", 
    "adaptive": "adaptive",
    # Auto mode (v3.2+)
    "auto": "auto",
    "auto-safe": "auto-safe",
    # Legacy names (backward compatibility)
    "vqe": "standard",
    "qaoa": "deep",
    "rl": "adaptive",
}

VALID_METHODS = list(METHOD_ALIASES.keys())


def get_license_key() -> Optional[str]:
    """Get license key from environment or file."""
    key = os.environ.get("MOBIU_Q_LICENSE_KEY")
    if key:
        return key
    
    if os.path.exists(LICENSE_KEY_FILE):
        with open(LICENSE_KEY_FILE, "r") as f:
            return f.read().strip()
    
    return None


def save_license_key(key: str):
    """Save license key to file."""
    with open(LICENSE_KEY_FILE, "w") as f:
        f.write(key)
    print(f"✅ License key saved to {LICENSE_KEY_FILE}")


# ═══════════════════════════════════════════════════════════════════════════════
# DEFAULT LEARNING RATE LOGIC
# ═══════════════════════════════════════════════════════════════════════════════

def get_default_lr(method: str, mode: str) -> float:
    """
    Get default learning rate based on method and mode.
    """
    method = METHOD_ALIASES.get(method, method)
    
    if method == 'adaptive':
        return 0.0003
    elif method == 'deep':
        return 0.1
    elif method in ('auto', 'auto-safe'):
        return 0.01  # Auto uses standard default
    elif mode == 'hardware':
        return 0.02
    else:  # standard + simulation
        return 0.01


# ═══════════════════════════════════════════════════════════════════════════════
# AUTO-SA ENGINE (NEW in v3.2!)
# ═══════════════════════════════════════════════════════════════════════════════

class AutoSAEngine:
    """
    Automatic Soft Algebra Mode Selection.
    
    Runs 3 virtual optimizers in parallel:
    - boost: Higher LR when making progress (for biased gradients)
    - dampen: Lower LR when gradients are noisy
    - off: Standard Adam baseline
    
    Every `period` steps, evaluates all modes and picks the best.
    
    Args:
        n: Number of parameters
        lr: Base learning rate
        period: Steps between mode evaluations (default: 20)
        decay_boost: If True, reduce boost strength as energy decreases (safer)
        maximize: If True, maximize metric instead of minimize (for RL/trading)
    """
    
    def __init__(self, n: int, lr: float = 0.01, period: int = 20, 
                 decay_boost: bool = False, maximize: bool = False):
        self.n = n
        self.lr = lr
        self.period = period
        self.decay_boost = decay_boost
        self.maximize = maximize
        self.t = 0
        
        # Three virtual optimizers
        self.modes = ['boost', 'dampen', 'off']
        self.vx = {m: None for m in self.modes}  # Virtual params
        self.vm = {m: np.zeros(n) for m in self.modes}  # Momentum
        self.vv = {m: np.zeros(n) for m in self.modes}  # Velocity
        self.venergy = {m: [] for m in self.modes}  # Energy history per mode
        
        # Soft Algebra state
        self.grad_hist = []
        self.energy_hist = []
        self.soft_r = 0.0  # Improvement EMA
        self.soft_s = 0.0  # Curvature EMA
        
        # Selection
        self.selected = 'boost'  # Start with boost
        self.initial_energy = None
        
    def step(self, x: np.ndarray, grad: np.ndarray, energy: float, 
             loss_fn: callable = None) -> np.ndarray:
        """
        Perform one AUTO-SA step.
        
        Args:
            x: Current parameters
            grad: Gradient
            energy: Current loss/energy
            loss_fn: Function to evaluate loss (needed for mode selection)
            
        Returns:
            Updated parameters
        """
        self.t += 1
        self.grad_hist.append(grad.copy())
        self.energy_hist.append(energy)
        
        if self.initial_energy is None:
            self.initial_energy = energy
        
        # Initialize virtual params on first step
        if self.vx['boost'] is None:
            for m in self.modes:
                self.vx[m] = x.copy()
        
        # Compute LRs for each mode
        lrs = self._compute_lrs(energy, grad)
        
        # Update ALL virtual optimizers
        for m in self.modes:
            # Adam update
            self.vm[m] = 0.9 * self.vm[m] + 0.1 * grad
            self.vv[m] = 0.999 * self.vv[m] + 0.001 * (grad ** 2)
            m_hat = self.vm[m] / (1 - 0.9 ** self.t)
            v_hat = self.vv[m] / (1 - 0.999 ** self.t)
            self.vx[m] = self.vx[m] - lrs[m] * m_hat / (np.sqrt(v_hat) + 1e-8)
        
        # Periodically evaluate and select best mode
        if loss_fn is not None and self.t % self.period == 0:
            for m in self.modes:
                self.venergy[m].append(loss_fn(self.vx[m]))
            
            if len(self.venergy['boost']) >= 1:
                recent = {m: self.venergy[m][-1] for m in self.modes}
                # maximize: pick highest, minimize: pick lowest
                if self.maximize:
                    self.selected = max(recent, key=recent.get)
                else:
                    self.selected = min(recent, key=recent.get)
        
        return self.vx[self.selected].copy()
    
    def _compute_lrs(self, energy: float, grad: np.ndarray) -> dict:
        """Compute learning rates for each mode."""
        
        # === Soft Algebra: a_t (curvature) ===
        if len(self.energy_hist) >= 3:
            E = self.energy_hist
            curv = abs(E[-1] - 2*E[-2] + E[-3])
            mean_E = np.mean(np.abs(E[-3:]))
            a = curv / (curv + mean_E + 1e-9)
        else:
            a = 0.1
        
        # === Soft Algebra: b_t (improvement) ===
        if len(self.energy_hist) >= 2:
            if self.maximize:
                # For maximize: improvement = current - previous (higher is better)
                imp = max(0, energy - self.energy_hist[-2])
            else:
                # For minimize: improvement = previous - current (lower is better)
                imp = max(0, self.energy_hist[-2] - energy)
            b = imp / (abs(energy) + 1e-9)
        else:
            b = 0
        
        # === Trust computation ===
        self.soft_r = 0.9 * self.soft_r + 0.1 * b
        self.soft_s = 0.9 * self.soft_s + 0.1 * a
        trust = abs(self.soft_r) / (abs(self.soft_r) + abs(self.soft_s) + 1e-9)
        
        # === Boost LR ===
        if self.decay_boost and self.initial_energy is not None:
            if self.maximize:
                # For maximize: decay as we get higher (closer to max)
                if self.initial_energy != 0:
                    ratio = self.initial_energy / (energy + 1e-9)
                    decay = min(1.0, max(0.1, ratio))
                else:
                    decay = 1.0
            else:
                # For minimize: decay as we get lower (closer to min)
                if self.initial_energy > 0:
                    decay = min(1.0, energy / self.initial_energy)
                    decay = max(0.1, decay)
                else:
                    decay = 1.0
        else:
            decay = 1.0
        
        lr_boost = self.lr * (1.0 + 0.5 * trust * decay)
        
        # === Dampen LR (based on gradient consistency) ===
        if len(self.grad_hist) >= 2:
            g1, g2 = self.grad_hist[-1], self.grad_hist[-2]
            cs = np.dot(g1, g2) / (np.linalg.norm(g1) * np.linalg.norm(g2) + 1e-9)
            noise = (1 - cs) / 2  # 0 = consistent, 1 = random
        else:
            noise = 0
        
        lr_dampen = self.lr * (1.0 - 0.5 * noise)
        
        return {
            'boost': lr_boost,
            'dampen': lr_dampen,
            'off': self.lr
        }
    
    def reset(self):
        """Reset for new optimization run."""
        self.t = 0
        self.vx = {m: None for m in self.modes}
        self.vm = {m: np.zeros(self.n) for m in self.modes}
        self.vv = {m: np.zeros(self.n) for m in self.modes}
        self.venergy = {m: [] for m in self.modes}
        self.grad_hist = []
        self.energy_hist = []
        self.soft_r = 0.0
        self.soft_s = 0.0
        self.selected = 'boost'
        self.initial_energy = None
    
    @property
    def current_mode(self) -> str:
        """Currently selected mode."""
        return self.selected


# ═══════════════════════════════════════════════════════════════════════════════
# UNIVERSAL FRUSTRATION ENGINE (CLIENT-SIDE)
# ═══════════════════════════════════════════════════════════════════════════════

class UniversalFrustrationEngine:
    """
    Logic Injection Engine that detects stagnation and boosts Learning Rate.
    Works entirely client-side for zero latency.
    """
    def __init__(self, base_lr: float, sensitivity: float = 0.05):
        self.base_lr = base_lr
        self.history = deque(maxlen=50)
        self.cooldown = 0
        self.best_metric = -float('inf')
        self.stagnation_counter = 0
        self.sensitivity = sensitivity 

    def get_lr_factor(self, current_metric: float) -> float:
        """
        Returns a multiplier factor (e.g. 1.0, 2.0, 3.0) for the LR.
        Assumes 'current_metric' is something we want to MAXIMIZE.
        """
        self.history.append(current_metric)
        
        if current_metric > self.best_metric:
            self.best_metric = current_metric
            self.stagnation_counter = 0
        else:
            self.stagnation_counter += 1

        if self.cooldown > 0:
            self.cooldown -= 1
            decay = 0.95 ** (30 - self.cooldown)
            return 1.0 + (2.0 * decay)

        if len(self.history) >= 20:
            recent_avg = np.mean(list(self.history)[-10:])
            old_avg = np.mean(list(self.history)[:10])
            is_stuck = (recent_avg < old_avg + abs(old_avg) * self.sensitivity)
            
            if is_stuck and self.stagnation_counter > 20:
                self.cooldown = 30
                self.history.clear()
                self.stagnation_counter = 0
                return 3.0

        return 1.0
    
    def reset(self):
        """Reset engine state for new run."""
        self.history.clear()
        self.cooldown = 0
        self.best_metric = -float('inf')
        self.stagnation_counter = 0


# ═══════════════════════════════════════════════════════════════════════════════
# MOBIU OPTIMIZER - UNIVERSAL WRAPPER
# ═══════════════════════════════════════════════════════════════════════════════

class MobiuOptimizer:
    """
    Universal Mobiu-Q Optimizer - wraps any optimizer with Soft Algebra intelligence.
    
    Auto-detects PyTorch optimizers and uses hybrid mode:
    - Cloud computes adaptive learning rate (Soft Algebra + Super-Equation)
    - Local PyTorch handles weight updates (fast, precise, GPU-accelerated)
    
    For non-PyTorch usage (quantum), delegates to MobiuQCore.
    """
    
    def __init__(
        self,
        optimizer_or_params,
        license_key: Optional[str] = None,
        method: str = "adaptive",
        mode: str = "simulation",
        use_soft_algebra: bool = True,
        sync_interval: Optional[int] = None,
        maximize: bool = False,
        verbose: bool = True,
        problem: Optional[str] = None,
        **kwargs
    ):
        self.license_key = license_key or get_license_key()
        if not self.license_key:
            raise ValueError(
                "License key required. Set MOBIU_Q_LICENSE_KEY environment variable, "
                "or pass license_key parameter, or run: mobiu-q activate YOUR_KEY"
            )
        
        # Handle deprecated 'problem' parameter
        if problem is not None:
            warnings.warn(
                "Parameter 'problem' is deprecated, use 'method' instead",
                DeprecationWarning,
                stacklevel=2
            )
            if method == "adaptive":
                method = problem
        
        # Validate method
        if method not in VALID_METHODS:
            raise ValueError(f"method must be one of {VALID_METHODS}, got '{method}'")
        
        self.method = METHOD_ALIASES.get(method, method)
        self._original_method = method
        self.verbose = verbose
        self.use_soft_algebra = use_soft_algebra
        
        # Auto-detect: Is this a PyTorch optimizer?
        self._is_pytorch = (
            hasattr(optimizer_or_params, 'step') and 
            hasattr(optimizer_or_params, 'param_groups') and
            hasattr(optimizer_or_params, 'zero_grad')
        )
        
        if self._is_pytorch:
            if sync_interval is None:
                sync_interval = 50
    
            self._backend = _MobiuPyTorchBackend(
                optimizer_or_params, 
                self.license_key, 
                self.method,
                use_soft_algebra=use_soft_algebra,
                sync_interval=sync_interval,
                maximize=maximize,
                verbose=verbose
            )
        else:
            self._backend = MobiuQCore(
                license_key=self.license_key,
                method=method,
                mode=mode,
                use_soft_algebra=use_soft_algebra,
                verbose=verbose,
                **kwargs
            )
    
    @property
    def problem(self) -> str:
        warnings.warn(
            "Attribute 'problem' is deprecated, use 'method' instead",
            DeprecationWarning,
            stacklevel=2
        )
        return self.method
    
    def step(self, *args, **kwargs):
        return self._backend.step(*args, **kwargs)
    
    def zero_grad(self):
        if hasattr(self._backend, 'zero_grad'):
            self._backend.zero_grad()
    
    def set_metric(self, metric: float):
        if hasattr(self._backend, 'set_metric'):
            self._backend.set_metric(metric)

    def end(self):
        self._backend.end()
    
    def new_run(self):
        if hasattr(self._backend, 'new_run'):
            self._backend.new_run()
    
    def reset(self):
        warnings.warn(
            "reset() is deprecated. Use new_run() for multi-seed experiments.",
            DeprecationWarning,
            stacklevel=2
        )
        if hasattr(self._backend, 'reset'):
            self._backend.reset()
        else:
            self.end()
            self._backend._start_session()
    
    def check_usage(self) -> dict:
        if hasattr(self._backend, 'check_usage'):
            return self._backend.check_usage()
        return {}
    
    def get_server_info(self) -> dict:
        if hasattr(self._backend, 'get_server_info'):
            return self._backend.get_server_info()
        return {
            "available_optimizers": AVAILABLE_OPTIMIZERS,
            "default_optimizer": DEFAULT_OPTIMIZER,
            "methods": ["standard", "deep", "adaptive", "auto", "auto-safe"],
        }
    
    @property
    def is_pytorch_mode(self) -> bool:
        return self._is_pytorch
    
    @property
    def energy_history(self) -> List[float]:
        return self._backend.energy_history
    
    @property
    def lr_history(self) -> List[float]:
        return self._backend.lr_history

    @property
    def warp_history(self) -> List[float]:
        return getattr(self._backend, 'warp_history', [])

    @property
    def sync_interval(self) -> Optional[int]:
        if hasattr(self._backend, 'sync_interval'):
            return self._backend.sync_interval
        return None

    @sync_interval.setter
    def sync_interval(self, value: int):
        if hasattr(self._backend, 'sync_interval'):
            self._backend.sync_interval = value
    
    @property
    def remaining_runs(self) -> Optional[int]:
        if hasattr(self._backend, 'remaining_runs'):
            return self._backend.remaining_runs
        return None
    
    @property
    def available_optimizers(self) -> List[str]:
        if hasattr(self._backend, 'available_optimizers'):
            return self._backend.available_optimizers
        return AVAILABLE_OPTIMIZERS
    
    @property
    def param_groups(self):
        if hasattr(self._backend, 'optimizer'):
            return self._backend.optimizer.param_groups
        return []
    
    @property
    def state(self):
        if hasattr(self._backend, 'optimizer'):
            return self._backend.optimizer.state
        return {}

    @property
    def current_mode(self) -> Optional[str]:
        """Get current AUTO-SA mode (only for auto/auto-safe methods)."""
        if hasattr(self._backend, 'current_mode'):
            return self._backend.current_mode
        return None

    def __del__(self):
        try:
            self.end()
        except:
            pass


# ═══════════════════════════════════════════════════════════════════════════════
# PYTORCH BACKEND (INTERNAL)
# ═══════════════════════════════════════════════════════════════════════════════

class _MobiuPyTorchBackend:
    """Internal backend for PyTorch hybrid mode."""
    
    def __init__(self, optimizer, license_key: str, method: str, 
                use_soft_algebra: bool = True, sync_interval: int = 50,
                maximize: bool = False, verbose: bool = True):
        self.optimizer = optimizer
        self.license_key = license_key
        self.method = method
        self.use_soft_algebra = use_soft_algebra
        self.verbose = verbose
        self.session_id = None
        self.api_endpoint = API_ENDPOINT
        
        self.base_lr = optimizer.param_groups[0]['lr']
        
        self.maximize = maximize
        self.frustration_engine = UniversalFrustrationEngine(base_lr=self.base_lr) if use_soft_algebra else None
        
        self.energy_history = []
        self.lr_history = []
        self.warp_history = []
        self._last_energy = None
        self._usage_info = None
        self._available_optimizers = AVAILABLE_OPTIMIZERS
        self.sync_interval = sync_interval
        self._local_step_count = 0
        self._accumulated_metric = 0.0
        self._metric_count = 0
        self._stored_metric = None
        
        self._start_session()
    
    def _start_session(self):
        try:
            r = requests.post(self.api_endpoint, json={
                'action': 'start',
                'license_key': self.license_key,
                'method': self.method,
                'mode': 'simulation',
                'base_lr': self.base_lr,
                'base_optimizer': 'Adam',
                'use_soft_algebra': self.use_soft_algebra,
                'maximize': self.maximize
            }, timeout=10)
            
            data = r.json()
            
            if data.get('success'):
                self.session_id = data['session_id']
                self._usage_info = data.get('usage', {})
                self._available_optimizers = data.get('available_optimizers', AVAILABLE_OPTIMIZERS)
                
                if self.verbose:
                    remaining = self._usage_info.get('remaining', 'unknown')
                    mode_str = f"method={self.method}, base_lr={self.base_lr}"
                    if self.sync_interval > 1:
                        mode_str += f", sync={self.sync_interval}"
                    if not self.use_soft_algebra:
                        mode_str += ", SA=off"
                    
                    if remaining == 'unlimited':
                        print(f"🚀 Mobiu-Q Hybrid session started (Pro tier) [{mode_str}]")
                    elif isinstance(remaining, int):
                        if remaining <= 2:
                            print(f"⚠️  Mobiu-Q Hybrid session started - LOW QUOTA: {remaining} runs remaining!")
                        else:
                            print(f"🚀 Mobiu-Q Hybrid session started ({remaining} runs remaining) [{mode_str}]")
                    else:
                        print(f"🚀 Mobiu-Q Hybrid session started [{mode_str}]")
            else:
                if self.verbose:
                    error = data.get('error', 'Unknown error')
                    print(f"⚠️  API start failed: {error}. Using constant LR.")
                    
        except Exception as e:
            if self.verbose:
                print(f"⚠️  Cannot connect to Mobiu-Q: {e}. Using constant LR.")
    
    def set_metric(self, metric: float):
        self._stored_metric = metric

    def step(self, metric: float = None):
        if metric is None:
            metric = self._stored_metric
    
        self._local_step_count += 1
        
        if self.frustration_engine and metric is not None:
            score = metric if self.maximize else -metric
            factor = self.frustration_engine.get_lr_factor(score)
            
            if factor > 1.0:
                new_lr = self.base_lr * factor
                for pg in self.optimizer.param_groups:
                    pg['lr'] = new_lr
                self.lr_history.append(new_lr)

        if metric is not None:
            self._accumulated_metric += metric
            self._metric_count += 1

        should_sync = (
            self.use_soft_algebra and
            self.session_id and self._metric_count > 0 and
            (self._local_step_count % self.sync_interval == 0)
        )

        if should_sync:
            avg_metric = self._accumulated_metric / self._metric_count
            energy_to_send = avg_metric
            
            try:
                r = requests.post(self.api_endpoint, json={
                    'action': 'step',
                    'license_key': self.license_key,
                    'session_id': self.session_id,
                    'params': [0.0], 
                    'gradient': [0.0],
                    'energy': energy_to_send
                }, timeout=1.0)
                
                data = r.json()
                if data.get('success'):
                    if 'adaptive_lr' in data:
                        self.base_lr = data['adaptive_lr']
                    
                    warp_factor = data.get('warp_factor', 1.0)
                    self.warp_history.append(warp_factor)
                    
                    if warp_factor != 1.0:
                        for pg in self.optimizer.param_groups:
                            for param in pg['params']:
                                if param.grad is not None:
                                    param.grad.data.mul_(warp_factor)
                    
            except: pass
            
            self._accumulated_metric = 0.0
            self._metric_count = 0

        self.optimizer.step()
    
    def zero_grad(self):
        self.optimizer.zero_grad()
    
    def new_run(self):
        self.energy_history.clear()
        self.lr_history.clear()
        self.warp_history.clear()
        self._last_energy = None
        self._local_step_count = 0
        self._accumulated_metric = 0.0
        self._metric_count = 0
        
        if self.frustration_engine:
            self.frustration_engine.reset()
    
        self.optimizer.state.clear()
        
        for group in self.optimizer.param_groups:
            group['lr'] = self.base_lr
        
        if self.session_id:
            try:
                requests.post(self.api_endpoint, json={
                    'action': 'reset',
                    'license_key': self.license_key,
                    'session_id': self.session_id
                }, timeout=5)
            except:
                pass
    
    def check_usage(self) -> dict:
        try:
            response = requests.post(
                self.api_endpoint,
                json={
                    "license_key": self.license_key,
                    "action": "usage"
                },
                timeout=10
            )
            data = response.json()
            if data.get("success"):
                self._usage_info = data.get("usage", {})
                return self._usage_info
        except:
            pass
        return {}
    
    def get_server_info(self) -> dict:
        try:
            response = requests.post(
                self.api_endpoint,
                json={
                    "license_key": self.license_key,
                    "action": "info"
                },
                timeout=10
            )
            data = response.json()
            if data.get("success"):
                return data
        except:
            pass
        return {
            "available_optimizers": AVAILABLE_OPTIMIZERS,
            "default_optimizer": DEFAULT_OPTIMIZER,
            "methods": ["standard", "deep", "adaptive", "auto", "auto-safe"],
        }
    
    @property
    def remaining_runs(self) -> Optional[int]:
        if self._usage_info:
            remaining = self._usage_info.get('remaining')
            if remaining == 'unlimited':
                return None
            return remaining
        return None
    
    @property
    def available_optimizers(self) -> List[str]:
        return self._available_optimizers
    
    def end(self):
        if self.session_id:
            try:
                response = requests.post(self.api_endpoint, json={
                    'action': 'end',
                    'license_key': self.license_key,
                    'session_id': self.session_id
                }, timeout=5)
                
                data = response.json()
                self._usage_info = data.get('usage', {})
                
                if self.verbose:
                    remaining = self._usage_info.get('remaining', 'unknown')
                    
                    if remaining == 'unlimited':
                        print(f"✅ Session ended (Pro tier)")
                    elif remaining == 0:
                        print(f"✅ Session ended")
                        print(f"❌ Quota exhausted! Upgrade at: https://app.mobiu.ai")
                    elif isinstance(remaining, int) and remaining <= 2:
                        print(f"✅ Session ended")
                        print(f"⚠️  Low quota warning: {remaining} runs remaining")
                    else:
                        print(f"✅ Session ended ({remaining} runs remaining)")
                    
            except Exception as e:
                if self.verbose:
                    print(f"⚠️  Could not cleanly end session: {e}")
            
            self.session_id = None


# ═══════════════════════════════════════════════════════════════════════════════
# MOBIU-Q CORE (Cloud Client) - WITH AUTO MODE SUPPORT
# ═══════════════════════════════════════════════════════════════════════════════

class MobiuQCore:
    """
    Mobiu-Q Optimizer - Cloud-connected version for quantum optimization.
    
    NEW in v3.2: method='auto' and method='auto-safe' for automatic mode selection.
    """
    
    def __init__(
        self,
        license_key: Optional[str] = None,
        method: str = "standard",
        mode: str = "simulation",
        base_lr: Optional[float] = None,
        base_optimizer: str = DEFAULT_OPTIMIZER,
        use_soft_algebra: bool = True,
        maximize: bool = False,
        offline_fallback: bool = True,
        verbose: bool = True,
        problem: Optional[str] = None,
    ):
        self.license_key = license_key or get_license_key()
        if not self.license_key:
            raise ValueError(
                "License key required. Set MOBIU_Q_LICENSE_KEY environment variable, "
                "or pass license_key parameter, or run: mobiu-q activate YOUR_KEY"
            )
        
        if problem is not None:
            warnings.warn(
                "Parameter 'problem' is deprecated, use 'method' instead",
                DeprecationWarning,
                stacklevel=2
            )
            if method == "standard":
                method = problem
        
        if mode == "standard":
            mode = "simulation"
        elif mode == "noisy":
            mode = "hardware"
        
        if method not in VALID_METHODS:
            raise ValueError(f"method must be one of {VALID_METHODS}, got '{method}'")
        
        internal_method = METHOD_ALIASES.get(method, method)
        
        if mode not in ("simulation", "hardware"):
            raise ValueError(f"mode must be 'simulation' or 'hardware', got '{mode}'")
        
        if base_optimizer not in AVAILABLE_OPTIMIZERS:
            raise ValueError(
                f"base_optimizer must be one of {AVAILABLE_OPTIMIZERS}, got '{base_optimizer}'"
            )
        
        self.method = internal_method
        self._original_method = method
        self.mode = mode
        self.base_lr = base_lr if base_lr is not None else get_default_lr(internal_method, mode)
        self.base_optimizer = base_optimizer
        self.use_soft_algebra = use_soft_algebra
        self.offline_fallback = offline_fallback
        self.verbose = verbose
        self.session_id = None
        self.api_endpoint = API_ENDPOINT
        self.maximize = maximize

        # AUTO-SA Engine (NEW!)
        self._auto_engine = None
        self._is_auto_mode = internal_method in ('auto', 'auto-safe')
        self._decay_boost = (internal_method == 'auto-safe')

        # Frustration Engine (for non-auto modes)
        self.frustration_engine = UniversalFrustrationEngine(base_lr=self.base_lr) if use_soft_algebra and not self._is_auto_mode else None
        self._current_lr = self.base_lr
        
        # Local state (for offline fallback)
        self._offline_mode = False
        self._local_m = None
        self._local_v = None
        self._local_t = 0
        
        # History
        self.energy_history = []
        self.lr_history = []
        
        self._run_count = 0
        self._usage_info = None
        self._available_optimizers = AVAILABLE_OPTIMIZERS
        
        # Start session (skip for auto mode - runs locally)
        if not self._is_auto_mode:
            self._start_session()
        else:
            if self.verbose:
                mode_name = "auto-safe (DecayBoost)" if self._decay_boost else "auto"
                print(f"🤖 Mobiu-Q AUTO mode started [{mode_name}, lr={self.base_lr}]")
    
    @property
    def problem(self) -> str:
        warnings.warn(
            "Attribute 'problem' is deprecated, use 'method' instead",
            DeprecationWarning,
            stacklevel=2
        )
        return self.method
    
    @property
    def current_mode(self) -> Optional[str]:
        """Get current AUTO-SA mode (boost/dampen/off)."""
        if self._auto_engine:
            return self._auto_engine.current_mode
        return None
    
    def _start_session(self):
        """Initialize optimization session with server."""
        try:
            response = requests.post(
                self.api_endpoint,
                json={
                    "license_key": self.license_key,
                    "action": "start",
                    "method": self.method,
                    "mode": self.mode,
                    "base_lr": self.base_lr,
                    "base_optimizer": self.base_optimizer,
                    "use_soft_algebra": self.use_soft_algebra,
                    "maximize": self.maximize
                },
                timeout=10
            )
            
            data = response.json()
            
            if not data.get("success"):
                error = data.get("error", "Unknown error")
                if "limit" in error.lower() or "quota" in error.lower():
                    print(f"❌ {error}")
                    print("   Upgrade at: https://app.mobiu.ai")
                raise RuntimeError(f"Failed to start session: {error}")
            
            self.session_id = data["session_id"]
            self._usage_info = data.get("usage", {})
            
            server_method = data.get("method", self.method)
            server_mode = data.get("mode", self.mode)
            server_lr = data.get("base_lr", self.base_lr)
            server_optimizer = data.get("base_optimizer", self.base_optimizer)
            self._available_optimizers = data.get("available_optimizers", AVAILABLE_OPTIMIZERS)
            
            if self.verbose:
                remaining = self._usage_info.get('remaining', 'unknown')
                
                mode_str = f"method={server_method}, mode={server_mode}, lr={server_lr}"
                if server_optimizer != DEFAULT_OPTIMIZER:
                    mode_str += f", optimizer={server_optimizer}"
                if not self.use_soft_algebra:
                    mode_str += ", SA=off"
                if self.maximize:
                    mode_str += ", maximize=True"
                
                if remaining == 'unlimited':
                    print(f"🚀 Mobiu-Q session started (Pro tier) [{mode_str}]")
                elif isinstance(remaining, int):
                    if remaining <= 2:
                        print(f"⚠️  Mobiu-Q session started - LOW QUOTA: {remaining} runs remaining!")
                    else:
                        print(f"🚀 Mobiu-Q session started ({remaining} runs remaining) [{mode_str}]")
                else:
                    print(f"🚀 Mobiu-Q session started [{mode_str}]")
                
        except requests.exceptions.RequestException as e:
            if self.offline_fallback:
                if self.verbose:
                    print(f"⚠️  Cannot connect to Mobiu-Q API: {e}")
                    print("   Running in offline fallback mode (plain Adam)")
                self._offline_mode = True
            else:
                raise RuntimeError(f"Cannot connect to Mobiu-Q API: {e}")
    
    def new_run(self):
        """Start a new optimization run within the same session."""
        self._run_count += 1
        
        self.energy_history.clear()
        self.lr_history.clear()
        self._local_m = None
        self._local_v = None
        self._local_t = 0
        
        # Reset AUTO-SA engine
        if self._auto_engine:
            self._auto_engine.reset()
        
        if self._offline_mode or not self.session_id:
            return
        
        if not self._is_auto_mode:
            try:
                response = requests.post(
                    self.api_endpoint,
                    json={
                        "license_key": self.license_key,
                        "session_id": self.session_id,
                        "action": "reset"
                    },
                    timeout=10
                )
                
                data = response.json()
                if not data.get("success"):
                    if self.verbose:
                        print(f"⚠️  Could not reset server state: {data.get('error')}")
                        
            except requests.exceptions.RequestException as e:
                if self.verbose:
                    print(f"⚠️  Could not reset server state: {e}")
    
    def step(
        self, 
        params: np.ndarray, 
        gradient_or_fn, 
        energy: float = None
    ) -> np.ndarray:
        """
        Perform one optimization step.
        
        For AUTO mode: runs locally with automatic mode selection.
        For other modes: uses cloud API.
        """
        # Auto-compute gradient if function provided
        if callable(gradient_or_fn):
            energy_fn = gradient_or_fn
    
            if self.mode == "hardware":
                gradient, energy = Demeasurement.spsa(energy_fn, params)
            else:
                gradient = Demeasurement.finite_difference(energy_fn, params)
                energy = energy_fn(params)
        else:
            gradient = gradient_or_fn
            if energy is None:
                raise ValueError("energy is required when providing gradient array")
    
        self.energy_history.append(energy)
        
        # === AUTO MODE: Run locally with AutoSAEngine ===
        if self._is_auto_mode:
            # Initialize engine on first call
            if self._auto_engine is None:
                self._auto_engine = AutoSAEngine(
                    n=len(params),
                    lr=self.base_lr,
                    period=20,
                    decay_boost=self._decay_boost,
                    maximize=self.maximize
                )
            
            # Get loss_fn for mode evaluation
            if callable(gradient_or_fn):
                loss_fn = gradient_or_fn
            else:
                loss_fn = None  # Can't evaluate modes without loss function
            
            new_params = self._auto_engine.step(params, gradient, energy, loss_fn)
            self.lr_history.append(self.base_lr)  # Track base LR
            return new_params

        # === Non-AUTO modes: Use existing logic ===
        
        # Frustration Engine
        if self.frustration_engine:
            score = energy if self.maximize else -energy
            factor = self.frustration_engine.get_lr_factor(score)
    
            if factor > 1.0:
                self._current_lr = self.base_lr * factor
                self.lr_history.append(self._current_lr)
            else:
                self._current_lr = self.base_lr
        
        if self._offline_mode:
            return self._offline_step(params, gradient)
        
        try:
            energy_to_send = energy
            for attempt in range(3):
                response = requests.post(
                    self.api_endpoint,
                    json={
                        "license_key": self.license_key,
                        "session_id": self.session_id,
                        "action": "step",
                        "params": params.tolist(),
                        "gradient": gradient.tolist(),
                        "energy": float(energy_to_send)
                    },
                    timeout=30
                )
                
                if response.status_code == 429:
                    if attempt < 2:
                        time.sleep(1)
                        continue
                    else:
                        raise RuntimeError("Rate limit exceeded.")
                break
            
            data = response.json()
            
            if not data.get("success"):
                error = data.get("error", "Unknown error")
                if self.offline_fallback:
                    if self.verbose:
                        print(f"⚠️  API error: {error}. Switching to offline mode.")
                    self._offline_mode = True
                    return self._offline_step(params, gradient)
                raise RuntimeError(f"Optimization step failed: {error}")
            
            new_params = np.array(data["new_params"])
            
            if "adaptive_lr" in data:
                self.lr_history.append(data["adaptive_lr"])
            
            return new_params
            
        except requests.exceptions.RequestException as e:
            if self.offline_fallback:
                if self.verbose:
                    print(f"⚠️  API connection lost: {e}. Switching to offline mode.")
                self._offline_mode = True
                return self._offline_step(params, gradient)
            raise
    
    def _offline_step(self, params: np.ndarray, gradient: np.ndarray) -> np.ndarray:
        """Fallback: plain Adam optimizer."""
        self._local_t += 1
        
        if self._local_m is None:
            self._local_m = np.zeros_like(gradient)
            self._local_v = np.zeros_like(gradient)
        
        lr = self._current_lr
        beta1, beta2, eps = 0.9, 0.999, 1e-8
        
        self._local_m = beta1 * self._local_m + (1 - beta1) * gradient
        self._local_v = beta2 * self._local_v + (1 - beta2) * (gradient ** 2)
        
        m_hat = self._local_m / (1 - beta1 ** self._local_t)
        v_hat = self._local_v / (1 - beta2 ** self._local_t)
        
        update = lr * m_hat / (np.sqrt(v_hat) + eps)
        return params - update
    
    def end(self):
        """End the optimization session."""
        if self._is_auto_mode:
            if self.verbose:
                if self._auto_engine:
                    print(f"✅ AUTO session ended (final mode: {self._auto_engine.current_mode})")
                else:
                    print(f"✅ AUTO session ended")
            return
        
        if self._offline_mode or not self.session_id:
            return
        
        try:
            response = requests.post(
                self.api_endpoint,
                json={
                    "license_key": self.license_key,
                    "session_id": self.session_id,
                    "action": "end"
                },
                timeout=10
            )
            
            data = response.json()
            self._usage_info = data.get("usage", {})
            
            if self.verbose:
                remaining = self._usage_info.get('remaining', 'unknown')
                
                if remaining == 'unlimited':
                    print(f"✅ Session ended (Pro tier)")
                elif remaining == 0:
                    print(f"✅ Session ended")
                    print(f"❌ Quota exhausted! Upgrade at: https://app.mobiu.ai")
                elif isinstance(remaining, int) and remaining <= 2:
                    print(f"✅ Session ended")
                    print(f"⚠️  Low quota warning: {remaining} runs remaining")
                else:
                    print(f"✅ Session ended ({remaining} runs remaining)")
                
        except Exception as e:
            if self.verbose:
                print(f"⚠️  Could not cleanly end session: {e}")
        
        self.session_id = None
    
    def reset(self):
        warnings.warn(
            "reset() is deprecated. Use new_run() for multi-seed experiments.",
            DeprecationWarning,
            stacklevel=2
        )
        self.end()
        self.energy_history.clear()
        self.lr_history.clear()
        if not self._is_auto_mode:
            self._start_session()
    
    def check_usage(self) -> dict:
        try:
            response = requests.post(
                self.api_endpoint,
                json={
                    "license_key": self.license_key,
                    "action": "usage"
                },
                timeout=10
            )
            data = response.json()
            if data.get("success"):
                self._usage_info = data.get("usage", {})
                return self._usage_info
        except:
            pass
        return {}
    
    def get_server_info(self) -> dict:
        try:
            response = requests.post(
                self.api_endpoint,
                json={
                    "license_key": self.license_key,
                    "action": "info"
                },
                timeout=10
            )
            data = response.json()
            if data.get("success"):
                return data
        except:
            pass
        return {
            "available_optimizers": AVAILABLE_OPTIMIZERS,
            "default_optimizer": DEFAULT_OPTIMIZER,
            "methods": ["standard", "deep", "adaptive", "auto", "auto-safe"],
        }
    
    @property
    def available_optimizers(self) -> List[str]:
        return self._available_optimizers
    
    @property
    def remaining_runs(self) -> Optional[int]:
        if self._usage_info:
            remaining = self._usage_info.get('remaining')
            if remaining == 'unlimited':
                return None
            return remaining
        return None

    def __del__(self):
        try:
            self.end()
        except:
            pass


# ═══════════════════════════════════════════════════════════════════════════════
# DEMEASUREMENT (Gradient Estimation) - Runs Locally
# ═══════════════════════════════════════════════════════════════════════════════

class Demeasurement:
    """Gradient estimation methods for quantum circuits."""
    
    @staticmethod
    def parameter_shift(
        circuit_fn, 
        params: np.ndarray, 
        shift: float = np.pi/2
    ) -> np.ndarray:
        """Parameter-shift rule gradient estimation."""
        grad = np.zeros_like(params)
        for i in range(len(params)):
            params_plus = params.copy()
            params_minus = params.copy()
            params_plus[i] += shift
            params_minus[i] -= shift
            grad[i] = (circuit_fn(params_plus) - circuit_fn(params_minus)) / 2.0
        return grad
    
    @staticmethod
    def finite_difference(
        circuit_fn, 
        params: np.ndarray,
        epsilon: float = 1e-3
    ) -> np.ndarray:
        """Finite difference gradient estimation."""
        grad = np.zeros_like(params)
        base_energy = circuit_fn(params)
        for i in range(len(params)):
            params_plus = params.copy()
            params_plus[i] += epsilon
            grad[i] = (circuit_fn(params_plus) - base_energy) / epsilon
        return grad
    
    @staticmethod
    def spsa(
        circuit_fn, 
        params: np.ndarray,
        c_shift: float = 0.1
    ) -> Tuple[np.ndarray, float]:
        """SPSA gradient estimation (2 evaluations regardless of param count)."""
        delta = np.random.choice([-1, 1], size=params.shape)
        
        params_plus = params + c_shift * delta
        params_minus = params - c_shift * delta
        
        energy_plus = circuit_fn(params_plus)
        energy_minus = circuit_fn(params_minus)
        
        grad = (energy_plus - energy_minus) / (2 * c_shift) * delta
        avg_energy = (energy_plus + energy_minus) / 2.0
        
        return grad, avg_energy


# ═══════════════════════════════════════════════════════════════════════════════
# CLI TOOLS
# ═══════════════════════════════════════════════════════════════════════════════

def activate_license(key: str):
    """Activate and save license key."""
    save_license_key(key)
    
    try:
        opt = MobiuQCore(license_key=key, verbose=False)
        opt.end()
        print("✅ License activated successfully!")
    except Exception as e:
        print(f"❌ License activation failed: {e}")


def check_status():
    """Check license status and remaining runs."""
    key = get_license_key()
    if not key:
        print("❌ No license key found")
        print("   Run: mobiu-q activate YOUR_KEY")
        return
    
    try:
        opt = MobiuQCore(license_key=key, verbose=False)
        usage = opt.check_usage()
        info = opt.get_server_info()
        opt.end()
        
        print("✅ License is active")
        if usage:
            print(f"   Tier: {usage.get('tier', 'unknown')}")
            print(f"   Used this month: {usage.get('used', 'unknown')}")
            print(f"   Remaining: {usage.get('remaining', 'unknown')}")
        if info:
            print(f"   Server version: {info.get('version', 'unknown')}")
            print(f"   Methods: {', '.join(info.get('methods', []))}")
    except Exception as e:
        print(f"❌ License check failed: {e}")


# ═══════════════════════════════════════════════════════════════════════════════
# EXPORTS
# ═══════════════════════════════════════════════════════════════════════════════

__version__ = "3.2.1"
__all__ = [
    # Universal optimizer
    "MobiuOptimizer",
    # AUTO-SA Engine (NEW!)
    "AutoSAEngine",
    # Frustration Engine
    "UniversalFrustrationEngine",
    # Core
    "MobiuQCore",
    "Demeasurement",
    # Utilities
    "activate_license",
    "check_status",
    "get_default_lr",
    "get_license_key",
    "save_license_key",
    # Constants
    "AVAILABLE_OPTIMIZERS",
    "DEFAULT_OPTIMIZER",
    "METHOD_ALIASES",
    "VALID_METHODS",
    "API_ENDPOINT",
]