import numpy as np


def polar_newton_schulz(a, iterations=5):
    """Muon-style semi-orthogonalization via normalized Newton-Schulz."""
    x = np.asarray(a, dtype=np.float64).copy()
    norm = np.linalg.norm(x, 2)
    if norm == 0:
        return np.zeros_like(x)
    x /= norm
    eye = np.eye(x.shape[1])
    for _ in range(iterations):
        x = 0.5 * x @ (3.0 * eye - x.T @ x)
    return x


class BiMaxwellMuon:
    """Two-relaxation-mode Muon state for one 2-D weight matrix."""
    def __init__(self, shape, lr=0.025, beta_fast=0.9, beta_slow=0.99,
                 weight_fast=0.5, ns_iterations=5):
        if not (0 <= weight_fast <= 1 and 0 <= beta_fast < beta_slow < 1):
            raise ValueError("require 0 <= beta_fast < beta_slow < 1 and valid weight")
        self.lr = lr
        self.beta_fast, self.beta_slow = beta_fast, beta_slow
        self.weight_fast = weight_fast
        self.ns_iterations = ns_iterations
        self.m_fast = np.zeros(shape, dtype=np.float64)
        self.m_slow = np.zeros(shape, dtype=np.float64)

    def step(self, weight, gradient):
        self.m_fast = self.beta_fast * self.m_fast + (1-self.beta_fast) * gradient
        self.m_slow = self.beta_slow * self.m_slow + (1-self.beta_slow) * gradient
        mixed = self.weight_fast*self.m_fast + (1-self.weight_fast)*self.m_slow
        update = polar_newton_schulz(mixed, self.ns_iterations)
        weight -= self.lr * update
        return weight, update