from __future__ import annotations

import inspect
import math

import torch


def compute_elf_lr(lr: float, blr: float, effective_batch_size: int) -> float:
    if lr is not None and lr > 0:
        return float(lr)
    return float(blr) * float(effective_batch_size) / 256.0


def create_learning_rate_fn(
    num_train_steps: int,
    num_warmup_steps: int,
    learning_rate: float,
    schedule: str = "constant",
    min_lr: float = 0.0,
):
    """ELF pytorch_elf linear warmup schedule, returning absolute LR."""
    alpha = (min_lr / learning_rate) if learning_rate > 0 else 0.0

    def fn(step: int) -> float:
        step = int(step)
        if num_warmup_steps > 0 and step < num_warmup_steps:
            return learning_rate * step / max(1, num_warmup_steps)
        if schedule == "cosine":
            progress = (step - num_warmup_steps) / max(1, num_train_steps - num_warmup_steps)
            progress = min(max(progress, 0.0), 1.0)
            cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
            return learning_rate * (alpha + (1.0 - alpha) * cosine)
        return learning_rate

    return fn


def build_optimizer(args, model: torch.nn.Module, world_size: int) -> torch.optim.Optimizer:
    effective_batch = int(args.batch_size) * int(world_size) * int(args.grad_accum)
    args.effective_batch_size = effective_batch
    args.lr = compute_elf_lr(float(args.lr), float(args.blr), effective_batch)
    if args.optimizer == "muon":
        from utils.muon_utils import muon_with_aux_adam

        return muon_with_aux_adam(model, lr=args.lr)
    params = [p for p in model.parameters() if p.requires_grad]
    fused = "fused" in inspect.signature(torch.optim.AdamW).parameters and torch.cuda.is_available()
    return torch.optim.AdamW(
        params,
        lr=args.lr,
        weight_decay=args.weight_decay,
        betas=(args.adam_beta1, args.adam_beta2),
        eps=args.adam_eps,
        fused=fused,
    )


def build_lr_scheduler(args, optimizer: torch.optim.Optimizer):
    lr_fn = create_learning_rate_fn(
        num_train_steps=args.steps,
        num_warmup_steps=args.warmup_steps,
        learning_rate=args.lr,
        schedule=args.lr_schedule,
        min_lr=args.min_lr,
    )
    base_lr = optimizer.param_groups[0]["lr"]
    return torch.optim.lr_scheduler.LambdaLR(
        optimizer,
        lr_lambda=lambda step: lr_fn(step) / max(base_lr, 1e-12),
    )