import torch
from torch import Tensor
import torch.nn as nn
from torch.autograd import Variable
from torch.optim import Optimizer
import numpy as np


class SGHMC_annealing(Optimizer):
    def __init__(self, 
                params,
                lr=1e-3, 
                alpha=1.0,
                ):

        defaults = dict(lr=lr, alpha=alpha)
        super().__init__(params, defaults)
        
    def update_parameters(self, parameter_name, parameter_value):
        for group in self.param_groups:
            group[parameter_name] = parameter_value
            
    # def learning_rate(self, lr_init):
    #     lr = lr_init
    #     return lr

    @torch.no_grad()
    def step(self, current_step, pseudo_population, closure=None):
        ''' One sigle step of LKTD algorithm
            observation (Tensor):  
            measurement (Tensor):
        '''
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            alpha = group['alpha']
            w_sd = np.sqrt(lr / pseudo_population)
            for p in group['params']:
                if p.grad is None:
                    continue
                state = self.state[p]
                if current_step == 0:
                    state['momentum'] = torch.zeros_like(p)
                if 'momentum' not in state:
                    state['momentum'] = torch.zeros_like(p)
                v = state['momentum']
                v = (1 - alpha) * v + lr * p.grad + np.sqrt(2 * alpha) * w_sd * torch.randn_like(p, device=p.device)
                p.sub_(v)


if __name__ == "__main__":
    pass