import torch
from torch import nn
from tqdm import tqdm
import time

from algorithm.JOBCD.Parallel_updateX import Parallel_updateV
from UHstruct.fobj_val import UltraE_fval_obj

class JOBCD(nn.Module):
    def __init__(self, T, H, Q, P , config_yaml):
        super().__init__()
        self.device = config_yaml["device"]
        self.T = T.to(self.device)
        self.H = H.to(self.device)
        self.Q = nn.Parameter(data=Q.to(self.device))
        self.P = P.to(self.device)

        self.config_yaml = config_yaml
        self.maxiter = int(float(config_yaml["run"]["maxiter"]))
        self.Cmaxiter = int(float(config_yaml["run"]["Cmaxiter"]))
        self.d = config_yaml["datafeature"]["d"]
        self.p = config_yaml["datafeature"]["p"]
        self.stopt = config_yaml["run"]["stop_t"]
        self.stoplr = float(config_yaml["run"]["stop_lr"])
        self.Lconstant = float(config_yaml["JOBCD"]["Lconstant"])
        self.theta = torch.tensor(int(float(config_yaml["JOBCD"]["theta"])))
        self.stopindex = 0

        J = torch.eye(self.d)
        J[self.p:, self.p:] = -1 * torch.eye(self.d - self.p)
        self.J = J.to(torch.float32).to(self.device)
        self.Xgrad = None

    def Train(self):
        hist_Obj = torch.zeros(self.maxiter, 1).to(self.device)
        hist_t = torch.zeros(self.maxiter, 1).to(self.device)
        start_time = time.time()
        for iter in tqdm(range(self.maxiter), desc='Outiter'):
            with torch.no_grad():
                hist_Obj[iter] = UltraE_fval_obj(self.T, self.H, self.P, self.Q, self.config_yaml)

            self.innerhist_obj = torch.zeros(int(self.Cmaxiter/self.config_yaml['run']['inner_dispgap']))
            for Citer in range(self.Cmaxiter):
                self.innertrain(Citer)

            with torch.no_grad():
                hist_t[iter] = time.time() - start_time
                if hist_t[iter] > self.stopt:
                    self.stopindex = 1

                if self.stopindex == 1:
                    break


        hist_Obj = hist_Obj[hist_Obj != 0]
        hist_t = hist_t[:len(hist_Obj)]
        return hist_Obj, self.Q, hist_t

    def innertrain(self, Citer):
        # update all var
        self.Xgrad = self.fgrad()

        # generate Jocobbi matrix
        original_vector = torch.randperm(self.d)
        B = original_vector[torch.randperm(self.d)].view(-1, 2)

        gradX = self.Xgrad.to(torch.float32)
        Lconst = self.Lconstant
        theta = self.theta
        with torch.no_grad():
            self.Q = Parallel_updateV(self.Q, gradX, B, Lconst, theta, self.p)


    def fgrad(self):
        fval = UltraE_fval_obj(self.T, self.H, self.P, self.Q, self.config_yaml)
        fval.requires_grad_(True)
        fval.backward()

        if torch.isnan(self.Q.grad).any():
            self.Q.grad[torch.isnan(self.Q.grad)] = 0
        if torch.isinf(self.Q.grad).any():
            self.Q.grad[torch.isinf(self.Q.grad)] = 0

        grad_R = self.Q.grad
        self.Q.grad = None
        return grad_R


