import torch
import torch.optim as optim
import torch.nn as nn
import numpy as np

class Worker:
    def __init__(self, args, logger):
        self.logger = logger
        self.lr = args.lr
        self.device = torch.device(args.device)
        self.args = args
        
    def train_net(self, worker_id, net, train_dataloader):
        net.to(self.device)
        net.train()
        optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=self.lr)
        
        if self.args.dataset == "synthetic": 
            criterion = nn.MSELoss().to(self.device)
        else:
            # this is not CONVEX!
            criterion = nn.CrossEntropyLoss().to(self.device)
        
        batch = np.random.choice(len(train_dataloader))
        single_batch = False if self.args.alg == "central" else True

        gradient = dict()
        for batch_idx, (x, target) in enumerate(train_dataloader):
            if single_batch and (batch_idx != batch):
                continue
            
            optimizer.zero_grad()
            
            x, target = x.to(self.device), target.to(self.device)
            
            out = net(x)
            loss = criterion(out, target)
            
            if self.args.lamb > 0:
                loss += self.args.lamb * sum([p.norm()**2 for p in net.parameters()])

            loss.backward()
            optimizer.step()
            
            # # accuracy
            # pred = out.argmax(dim=1, keepdim=True)
            # correct = pred.eq(target.view_as(pred)).sum().item()
            # print(f"Worker: {worker_id} Batch: {batch_idx} Accuracy: {correct/len(x)}")

            for k, p in net.named_parameters():
                if p.grad is None:
                    continue
                d_p = p.grad.data
                if k not in gradient:
                    gradient[k] = d_p.detach().clone()
                else:
                    gradient[k] += d_p.detach().clone()
            
            if single_batch:
                break
        
        # epoch_loss = sum(loss_list) / len(loss_list)
        epoch_loss = loss.item()
        self.logger.info('Worker: %d Loss: %f' % (worker_id, epoch_loss))
        return gradient

