import torch
import torch.optim as optim
import numpy as np
import logging

# Configure logging for loss monitoring
logging.basicConfig (level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger (__name__)


class Azure (optim.Optimizer):
    def __init__(self, params, lr=0.0007518383921113902, T0=2.2723218904585964, sigma=0.17181058166567398,
                 betas=(0.9, 0.999), eps=1e-8, sa_steps=5, sa_momentum=0.6612913488540948, clip_grad_norm=1.0):
        """
        Azure Sky Optimizer: A hybrid optimizer combining Simulated Annealing (SA) and Adam.

        Args:
            params (iterable): Iterable of parameters or dicts defining parameter groups.
            lr (float): Learning rate for Adam phase (default: 0.0007518383921113902).
            T0 (float): Initial temperature for SA (default: 2.2723218904585964).
            sigma (float): Perturbation strength for SA (default: 0.17181058166567398).
            betas (tuple): Adam's exponential decay rates (default: (0.9, 0.999)).
            eps (float): Adam's epsilon for numerical stability (default: 1e-8).
            sa_steps (int): Number of steps for SA phase (default: 5).
            sa_momentum (float): Momentum for SA updates (default: 0.6612913488540948).
            clip_grad_norm (float): Max norm for gradient clipping (default: 1.0).
        """
        # Process params to handle various input formats
        if isinstance (params, (list, tuple)) and isinstance (params [0], dict):
            # Handle parameter groups (e.g., [{'params': ..., 'lr': ...}, ...])
            param_groups = []
            for group in params:
                group_dict = group.copy ()
                if 'params' not in group_dict:
                    raise ValueError ("Each parameter group must contain a 'params' key")
                # Convert named_parameters() to a list of parameters if necessary
                if isinstance (group_dict ['params'], (list, tuple)) and isinstance (group_dict ['params'] [0], tuple):
                    group_dict ['params'] = [p for _, p in group_dict ['params']]
                param_groups.append (group_dict)
            params = param_groups
        else:
            # Handle direct parameter lists or named_parameters()
            if isinstance (params, (list, tuple)) and isinstance (params [0], tuple):
                params = [p for _, p in params]  # Convert named_parameters() to parameter list
            params = [{'params': params}]

        # Set defaults for each parameter group
        defaults = dict (lr=lr, T0=T0, sigma=sigma, betas=betas, eps=eps, sa_steps=sa_steps,
                         sa_momentum=sa_momentum, clip_grad_norm=clip_grad_norm)
        super ().__init__ (params, defaults)
        self.step_count = 0
        self.sa_active = True
        self.losses = []
        self.loss_window = 5
        self.loss_spike_threshold = 10.0

    def step(self, closure=None):
        """Performs a single optimization step."""
        loss = None
        if closure is not None:
            with torch.enable_grad ():
                loss = closure ()

        # Loss spike monitoring
        if loss is not None:
            self._monitor_loss (loss.item ())

        for group in self.param_groups:
            # Gradient clipping
            if group ['clip_grad_norm'] is not None:
                torch.nn.utils.clip_grad_norm_ (group ['params'], group ['clip_grad_norm'])

            for p in group ['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data

                # Dynamic Temperature Scaling
                T = self._compute_temperature (group)
                # Exploration-Exploitation Fusion
                alpha = self._compute_alpha (group)

                if self.sa_active:
                    noise = torch.randn_like (p.data) * group ['sigma'] * T
                    sa_update = noise
                else:
                    sa_update = torch.zeros_like (p.data)

                # Adam update
                state = self.state [p]
                if 'm' not in state:
                    state ['m'] = torch.zeros_like (p.data)
                    state ['v'] = torch.zeros_like (p.data)
                    state ['step'] = 0
                m, v = state ['m'], state ['v']
                beta1, beta2 = group ['betas']
                state ['step'] += 1
                m.mul_ (beta1).add_ (grad, alpha=1 - beta1)
                v.mul_ (beta2).addcmul_ (grad, grad, value=1 - beta2)
                m_hat = m / (1 - beta1 ** state ['step'])
                v_hat = v / (1 - beta2 ** state ['step'])
                # Use group-specific learning rate if provided
                lr = group.get ('lr', self.defaults ['lr'])
                adam_update = -lr * m_hat / (v_hat.sqrt () + group ['eps'])

                # Combined update
                update = alpha * adam_update + (1 - alpha) * sa_update
                p.data.add_ (update)

        self.step_count += 1
        if self.step_count >= self.param_groups [0] ['sa_steps']:
            self.sa_active = False
        return loss

    def _compute_temperature(self, group):
        """Dynamic Temperature Scaling based on step progress."""
        epoch_decay = 0.05  # Adjustable decay rate
        return group ['T0'] * (1.0 / (1.0 + epoch_decay * self.step_count))

    def _compute_alpha(self, group):
        """Exploration-Exploitation Fusion Schedule using sigmoid."""
        midpoint = group ['sa_steps'] / 2
        return 1 / (1 + np.exp (-(self.step_count - midpoint) / (midpoint / 5)))

    def _monitor_loss(self, loss):
        """Monitors for loss spikes and logs warnings."""
        self.losses.append (loss)
        if len (self.losses) > self.loss_window:
            self.losses.pop (0)
            avg_loss = sum (self.losses [:-1]) / (len (self.losses) - 1)
            current_loss = self.losses [-1]
            if current_loss > avg_loss * self.loss_spike_threshold:
                logger.warning (
                    f"Loss spike detected: {current_loss:.4f} > {avg_loss:.4f} * {self.loss_spike_threshold}")