import time
import math
import torch
import torch.autograd as autograd

from .distributed import reduce_mean

from .cgd_utils import zero_grad, general_conjugate_gradient, Hvp_vec


class ACGD(object):
    def __init__(self, max_params, min_params,
                 max_reducer=None, min_reducer=None,
                 lr_max=1e-3, lr_min=1e-3,
                 backward_mode=False,
                 eps=1e-5, beta=0.99,
                 tol=1e-12, atol=1e-20,
                 solve_x=False, collect_info=False):
        self.backwardMode = backward_mode
        self.max_reducer = max_reducer
        self.min_reducer = min_reducer
        self.max_params = list(max_params)
        self.min_params = list(min_params)
        self.state = {'lr_max': lr_max, 'lr_min': lr_min,
                      'eps': eps, 'solve_x': solve_x,
                      'tol': tol, 'atol': atol,
                      'beta': beta, 'step': 0,
                      'old_max': None, 'old_min': None,  # start point of CG
                      'sq_exp_avg_max': None, 'sq_exp_avg_min': None}  # save last update
        self.info = {'grad_x': None, 'grad_y': None,
                     'hvp_x': None, 'hvp_y': None,
                     'cg_x': None, 'cg_y': None,
                     'time': 0, 'iter_num': 0}
        self.collect_info = collect_info

    def zero_grad(self):
        zero_grad(self.max_params)
        zero_grad(self.min_params)

    def get_info(self):
        if self.info['grad_x'] is None:
            print('Warning! No update information stored. Set collect_info=True before call this method')
        return self.info

    def state_dict(self):
        return self.state

    def load_state_dict(self, state_dict):
        self.state.update(state_dict)
        print('Load state: {}'.format(state_dict))

    def set_lr(self, lr_max, lr_min):
        self.state.update({'lr_max': lr_max, 'lr_min': lr_min})
        print('Maximizing side learning rate: {:.4f}\n '
              'Minimizing side learning rate: {:.4f}'.format(lr_max, lr_min))

    def step(self, loss, trigger):
        '''
        Update model weights using ACGD
        :param loss: objective function
        :param trigger: scalar, dummy loss to trigger DDP gradient reduction w.r.t. minimizer model parameters
        :return:
        '''
        lr_max = self.state['lr_max']
        lr_min = self.state['lr_min']
        beta = self.state['beta']
        eps = self.state['eps']
        tol = self.state['tol']
        atol = self.state['atol']
        time_step = self.state['step'] + 1
        self.state['step'] = time_step
        should_rebuild = time_step < 2

        grad_x = autograd.grad(loss, self.max_params, create_graph=True)
        grad_x_vec = torch.cat([g.contiguous().view(-1) for g in grad_x])
        grad_y = autograd.grad(loss, self.min_params, create_graph=True)
        grad_y_vec = torch.cat([g.contiguous().view(-1) for g in grad_y])

        grad_x_vec_d = grad_x_vec.clone().detach()
        grad_y_vec_d = grad_y_vec.clone().detach()
        # reduce leaf nodes across all devices
        grad_x_vec_d = reduce_mean(grad_x_vec_d)
        grad_y_vec_d = reduce_mean(grad_y_vec_d)

        sq_avg_x = self.state['sq_exp_avg_max']
        sq_avg_y = self.state['sq_exp_avg_min']
        sq_avg_x = torch.zeros_like(grad_x_vec_d, requires_grad=False) if sq_avg_x is None else sq_avg_x
        sq_avg_y = torch.zeros_like(grad_y_vec_d, requires_grad=False) if sq_avg_y is None else sq_avg_y
        sq_avg_x.mul_(beta).addcmul_(grad_x_vec_d, grad_x_vec_d, value=1-beta)
        sq_avg_y.mul_(beta).addcmul_(grad_y_vec_d, grad_y_vec_d, value=1-beta)

        bias_correction = 1 - beta ** time_step
        lr_max = math.sqrt(bias_correction) * lr_max / sq_avg_x.sqrt().add(eps)
        lr_min = math.sqrt(bias_correction) * lr_min / sq_avg_y.sqrt().add(eps)

        scaled_grad_x = torch.mul(lr_max, grad_x_vec_d)
        scaled_grad_y = torch.mul(lr_min, grad_y_vec_d)

        hvp_x_vec = Hvp_vec(grad_y_vec, self.max_params, scaled_grad_y,
                            backward=self.backwardMode,
                            retain_graph=True,
                            trigger=trigger,
                            rebuild=should_rebuild)  # h_xy * d_y
        hvp_y_vec = Hvp_vec(grad_x_vec, self.min_params, scaled_grad_x,
                            backward=self.backwardMode,
                            retain_graph=True,
                            trigger=trigger,
                            rebuild=should_rebuild)  # h_yx * d_x
        p_x = torch.add(grad_x_vec_d, - hvp_x_vec)
        p_y = torch.add(grad_y_vec_d, hvp_y_vec)
        if self.collect_info:
            norm_px = torch.norm(hvp_x_vec, p=2).item()
            norm_py = torch.norm(hvp_y_vec, p=2).item()
            timer = time.time()

        if self.state['solve_x']:
            p_y.mul_(lr_min.sqrt())
            cg_y, iter_num = general_conjugate_gradient(grad_x=grad_y_vec,
                                                        grad_y=grad_x_vec,
                                                        x_params=self.min_params,
                                                        y_params=self.max_params,
                                                        x_reducer=self.min_reducer,
                                                        y_reducer=self.max_reducer,
                                                        rebuild=should_rebuild,
                                                        trigger=trigger,
                                                        b=p_y, x=self.state['old_min'],
                                                        tol=tol, atol=atol,
                                                        lr_x=lr_min, lr_y=lr_max,
                                                        backward=self.backwardMode)
            old_min = cg_y.detach_()
            min_update = cg_y.mul(- lr_min.sqrt())
            hcg = Hvp_vec(grad_y_vec, self.max_params, min_update,
                          self.backwardMode,
                          trigger=trigger, reducer=self.max_reducer,
                          rebuild=should_rebuild).detach_()
            hcg.add_(grad_x_vec_d)
            max_update = hcg.mul(lr_max)
            old_max = hcg.mul(lr_max.sqrt())
        else:
            p_x.mul_(lr_max.sqrt())
            cg_x, iter_num = general_conjugate_gradient(grad_x=grad_x_vec,
                                                        grad_y=grad_y_vec,
                                                        x_params=self.max_params,
                                                        y_params=self.min_params,
                                                        x_reducer=self.max_reducer,
                                                        y_reducer=self.min_reducer,
                                                        rebuild=should_rebuild,
                                                        trigger=trigger,
                                                        b=p_x, x=self.state['old_max'],
                                                        tol=tol, atol=atol,
                                                        lr_x=lr_max, lr_y=lr_min,
                                                        backward=self.backwardMode)
            old_max = cg_x.detach_()
            max_update = cg_x.mul(lr_max.sqrt())
            hcg = Hvp_vec(grad_x_vec, self.min_params, max_update,
                          self.backwardMode,
                          trigger=trigger, reducer=self.min_reducer,
                          rebuild=should_rebuild).detach_()
            hcg.add_(grad_y_vec_d)
            min_update = hcg.mul(- lr_min)
            old_min = hcg.mul(lr_min.sqrt())
        self.state.update({'old_max': old_max, 'old_min': old_min,
                           'sq_exp_avg_max': sq_avg_x, 'sq_exp_avg_min': sq_avg_y})

        if self.collect_info:
            timer = time.time() - timer
            self.info.update({'time': timer, 'iter_num': iter_num,
                              'hvp_x': norm_px, 'hvp_y': norm_py})

        index = 0
        for p in self.max_params:
            p.data.add_(max_update[index: index + p.numel()].reshape(p.shape))
            index += p.numel()
        assert index == max_update.numel(), 'Maximizer CG size mismatch'

        index = 0
        for p in self.min_params:
            p.data.add_(min_update[index: index + p.numel()].reshape(p.shape))
            index += p.numel()
        assert index == min_update.numel(), 'Minimizer CG size mismatch'

        if self.collect_info:
            norm_gx = torch.norm(grad_x_vec, p=2).item()
            norm_gy = torch.norm(grad_y_vec, p=2).item()
            norm_cgx = torch.norm(max_update, p=2).item()
            norm_cgy = torch.norm(min_update, p=2).item()
            self.info.update({'grad_x': norm_gx, 'grad_y': norm_gy,
                              'cg_x': norm_cgx, 'cg_y': norm_cgy})
        self.state['solve_x'] = False if self.state['solve_x'] else True