from tinygrad import Tensor,nn
import math

class CosineLR:
    def __init__(self,optm,totalSteps,maxlr,minlr):
        self.optm = optm
        self.maxlr = maxlr
        self.minlr = minlr
        self.totalSteps = totalSteps
        self.steps = 0

    def step(self):
        lr = self.minlr + 0.5 * (self.maxlr - self.minlr) * (1 + math.cos((self.steps / self.totalSteps) * math.pi))
        for o in self.optm:
            o.lr = lr
        self.optm.step()
        self.steps += 1

    def zero_grad(self):
        self.optm.zero_grad()


def llmOptimizer(params,steps,maxlr,minlr):
    muon_params = [p for p in params if len(p.shape) >= 2]
    adamw_params = [p for p in params if len(p.shape) < 2]

    o1 = nn.optim.Muon(muon_params, lr=maxlr)
    o2 = nn.optim.AdamW(adamw_params, lr=maxlr)
    optimizer = nn.optim.OptimizerGroup(o1,o2)
    return CosineLR(optimizer,steps,maxlr,minlr)