#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 26 15:12:03 2020

@author: garima
"""

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from filelock import FileLock
import numpy as np

import ray

#helper functions 
def get_data_loader():
    """download dataset"""
    
    mnist_transforms=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307, ), (0.3081, ))])
    
    #we add FileLock here because multiple workers will want to download data, and this may cause overwrite 
    with FileLock(os.path.expanduser("~/data.lock")):
        train_loader= torch.utils.data.DataLoader(datasets.MNIST(
            "~/data",
            train=True,
            download=True,
            transform=mnist_transforms),
            batch_size=128,shuffle=True)
        test_loader= torch.utils.data.DataLoader(
            datasets.MNIST("~/data", train=False, transform=mnist_transforms),
            batch_size=128, shuffle=True)
    return train_loader, test_loader

def evaluate(model, test_loader):
    """Evaluate the accuracy of the model on a validation set"""
    model.eval()
    correct=0
    total=0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            #to execute faster
            if batch_idx*len(data)>1024:
                break
            outputs=model(data)
            _, predicted = torch.max(outputs.data,1)
            total+=target.size(0)
            correct+=(predicted==target).sum().item()
    return 100.*correct/total


## Set up - define the Neural network
class ConvNet(nn.Module):
        """Small ConvNet for MNIST"""
        def __init__(self):
            super(ConvNet, self).__init__()
            self.conv1=nn.Conv2d(1,3, kernel_size=3)
            self.fc= nn.Linear(192,10)
        
        def forward(self,x):
            x = F.relu(F.max_pool2d(self.conv1(x),3))
            x = x.view(-1,192)
            x = self.fc(x)
            return F.log_softmax(x,dim=1)
        
        def get_weights(self):
            return{k:v.cpu() for k,v in self.state_dict().items()}
        
        def set_weights(self, weights):
            self.load_state_dict(weights)
            
        def get_gradients(self):
            grads = []
            for p in self.parameters():
                grad = None if p.grad is None else p.grad.data.cpu().numpy()
                grad.append(grad)
            return grads
        
        def set_gradients(self, gradients):
            for g, p in zip(gradients, self.parameters()):
                if g is not None:
                    p.grad= torch.from_numpy(g)
           

#defining Parameter Server
                    
@ray.remote
class ParameterServer(object):
        def __init__(self,lr):
            self.model=ConvNet()
            self.optimizer=torch.optim.SGD(self.model.parameters(), lr=lr)
            
        def apply_gradients(self, *gradients):
            summed_gradients = [
                np.stack(gradient_zip).sum(axis=0)
                for gradient_zip in zip(*gradients)
                ]
            self.optimizer.zero_grad()
            self.model.set_gradients(summed_gradients)
            self.optimizer.step()
            return self.model.get_weights()
        
        def get_weights(self):
            return self.model.get_weights()
        
# define the worker
        
@ray.remote
class DataWorker(object):
    def __init__(self):
        self.model=ConvNet()
        self.data_iterator =iter(get_data_loader()[0])
    
    def compute_gradients(self,weights):
        self.model.set_weights(weights)
        try:
            data,target = next(self.data_iterator)
        except StopIteration: # when the epoch ends, start a new eopch
            self.data_iterator=iter(get_data_loader()[0])
            data, target = next(self.data_iterator)
        self.model.zero_grad()
        output= self.model(data)
        loss=F.nll_loss(output,target)
        loss.backward()
        return self.model.get_gradients()
    
                  
        
iterations = 2
num_workers=2

ray.init(ignore_reinit_error=True)
ps=ParameterServer.remote(1e-2)
workers=[DataWorker.remote() for i in range(num_workers)]
 
model= ConvNet() 
test_loader= get_data_loader()[1]
        
print("Running synchronous parameter server training")  
current_weights= ps.get_weights.remote()    
for i in range(iterations):
    gradients=[
        worker.compute_gradients.remote(current_weights) for worker in workers
        ]
        
    
model.get_gradients()  
  
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        



