import os

import torch
import torch.optim as optim


def get_token_id(lang, tp, d):
    """
    Get token id faster by caching results in `d`
    """
    token_id = d.get(lang)
    if not token_id:
        token_id = tp.token_id(lang)
        d[lang] = token_id
    return token_id


def build_optimizer(model, learning_rate, warump_steps):
    return AdamInverseSqrtWithWarmup(model.parameters(), lr=learning_rate, betas=(0.9, 0.98),
                                     warmup_updates=warump_steps)


def init_distributed(options):
    if options.local_rank >= 0:
        os.environ['MASTER_ADDR'] = '127.0.0.1'
        os.environ['MASTER_PORT'] = '29500'
        os.environ['WORLD_SIZE'] = str(torch.cuda.device_count())
        torch.distributed.init_process_group(backend='nccl', world_size=torch.cuda.device_count(),
                                             rank=options.local_rank)


class AdamInverseSqrtWithWarmup(optim.Adam):
    """
    Decay the LR based on the inverse square root of the update number.
    We also support a warmup phase where we linearly increase the learning rate
    from some initial learning rate (`warmup-init-lr`) until the configured
    learning rate (`lr`). Thereafter we decay proportional to the number of
    updates, with a decay factor set to align with the configured learning rate.
    During warmup:
        lrs = torch.linspace(warmup_init_lr, lr, warmup_updates)
        lr = lrs[update_num]
    After warmup:
        lr = decay_factor / sqrt(update_num)
    where
        decay_factor = lr * sqrt(warmup_updates)
    """

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, warmup_updates=4000, warmup_init_lr=1e-7):
        super().__init__(
            params,
            lr=warmup_init_lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
        )
        self.warmup_updates = warmup_updates
        self.warmup_init_lr = warmup_init_lr
        # linearly warmup for the first warmup_updates
        warmup_end_lr = lr
        self.lr_step = (warmup_end_lr - warmup_init_lr) / warmup_updates
        # then, decay prop. to the inverse square root of the update number
        self.decay_factor = warmup_end_lr * warmup_updates ** 0.5
        for param_group in self.param_groups:
            param_group['num_updates'] = 0
        self.max_lr = lr

    def get_lr_for_step(self, num_updates):
        # update learning rate
        if num_updates < self.warmup_updates:
            return self.warmup_init_lr + num_updates * self.lr_step
        else:
            return max(self.warmup_init_lr, min(self.max_lr, self.decay_factor * (num_updates ** -0.5)))

    def step(self, closure=None):
        super().step(closure)
        for param_group in self.param_groups:
            param_group['num_updates'] += 1
            param_group['lr'] = self.get_lr_for_step(param_group['num_updates'])

    def reset(self):
        for param_group in self.param_groups:
            param_group['num_updates'] = 0
