import torch
from detectron2.config import LazyCall as L
from detectron2.solver.build import get_default_optimizer_params
# from ovdino.configs.common.optimizer import c_adamw, c_lion
# from ovdino.configs.common.optimizer import c_lion

SGD = L(torch.optim.SGD)(
    params=L(get_default_optimizer_params)(
        # params.model is meant to be set to the model object, before instantiating
        # the optimizer.
        weight_decay_norm=0.0
    ),
    lr=0.02,
    momentum=0.9,
    weight_decay=1e-4,
)


AdamW = L(torch.optim.AdamW)(
    params=L(get_default_optimizer_params)(
        # params.model is meant to be set to the model object, before instantiating
        # the optimizer.
        base_lr="${..lr}",
        weight_decay_norm=0.0,
    ),
    lr=1e-4,
    betas=(0.9, 0.999),
    weight_decay=0.1,
)


# # ###########################   CadamW      ######################################
# import torch
# from torch.optim import AdamW

# class CautiousAdamW(AdamW):
#     def step(self, closure=None):
#         loss = None
#         if closure is not None:
#             with torch.enable_grad():
#                 loss = closure()

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

#                 if len(state) == 0:
#                     # Lazy state initialization
#                     state['step'] = 0
#                     state['exp_avg'] = torch.zeros_like(p.data)
#                     state['exp_avg_sq'] = torch.zeros_like(p.data)

#                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
#                 beta1, beta2 = group['betas']
#                 state['step'] += 1

#                 # update exp_avg and exp_avg_sq
#                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
#                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

#                 denom = (exp_avg_sq.sqrt() / (1 - beta2 ** state['step'])).add_(group['eps'])
#                 step_size = group['lr'] / (1 - beta1 ** state['step'])

#                 update = exp_avg / denom

#                 # cautious update only if direction matches
#                 mask = (update * grad > 0).to(grad.dtype)
#                 norm_factor = mask.mean().clamp(min=1e-3)
#                 cautious_update = update * mask / norm_factor

#                 p.data.add_(cautious_update, alpha=-step_size)

#         return loss


# CAdamW = L(CautiousAdamW)(
#     params=L(get_default_optimizer_params)(
#         base_lr="${..lr}",
#         weight_decay_norm=0.0,
#     ),
#     lr=1e-4,
#     betas=(0.9, 0.999),
#     weight_decay=0.1,
# )

