"""
ScaledOptum — Pure-integer optimizer for BigInt-correlation ternary training.

For each ternary module:
  1. Calls module.update_corr() — pure-integer correlation accumulation
     score = Σ (grad_sign × T) per group  (int16)
     corr_accum += score  (int64, BigInt, never resets or clips)
  
  2. The CARRY step happens via the S computation in forward:
     S = 2^E × (1 + corr_accum / (step × gs))
     The corr_accum / (step × gs) is the continuous adjustment.
     
  3. E is never manually updated — the corr_accum BigInt provides
     the continuous gradient-driven adjustment to S.
     If desired, E can be slowly tracked toward the corr-derived S
     for better initialization at inference.
"""
import torch
from torch.optim import Optimizer


class ScaledOptum(Optimizer):
    """
    Pure-integer optimizer for ternary training with BigInt correlation.

    Calls update_corr() on each ternary module — no float state.
    """

    def __init__(self, params, lr=0.3, default_group_size=32):
        defaults = dict(lr=lr, default_group_size=default_group_size)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            if 'ternary_modules' in group:
                for mod in group['ternary_modules']:
                    mod.update_corr()

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.is_sparse:
                    grad = grad.to_dense()
                p.add_(-group['lr'] * grad.sign())

        return loss

    def add_ternary_modules(self, modules):
        if not self.param_groups:
            self.param_groups.append({'params': [], 'ternary_modules': [],
                                      'lr': 0.3, 'default_group_size': 32})
        for group in self.param_groups:
            if 'ternary_modules' not in group:
                group['ternary_modules'] = []
            for mod in modules:
                if mod not in group['ternary_modules']:
                    group['ternary_modules'].append(mod)