import torch
import torch.nn as nn
from torchvision import transforms
from torch.utils.data import DataLoader, TensorDataset, Dataset
from torch.utils.data.sampler import SubsetRandomSampler
from torch import optim
import pandas as pd
import sys
sys.path.append('./proto')
import trainer_pb2
import trainer_pb2_grpc
import time
from concurrent import futures
import logging
import grpc
import argparse
import base64
import io
import ipfshttpclient

torch.nn.Module.dump_patches = True

def fullmodel2base64(model):
    buffer = io.BytesIO()
    torch.save(model, buffer)
    bg = buffer.getvalue()
    return base64.b64encode(bg).decode()

def base642fullmodel(modbase64):
    inputrpc = bytes(modbase64.encode())
    inputrpc_ = base64.b64decode(inputrpc)
    loadmodel = torch.load(io.BytesIO(inputrpc_))
    return loadmodel

# class Model(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.hidden = nn.Linear(784, 20)
#         self.output = nn.Linear(20, 10)

#     def forward(self, x):
#         x = self.hidden(x)
#         x = torch.sigmoid(x)
#         x = self.output(x)
#         return x
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.cnn = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5),
                                     nn.ReLU(inplace=True),
                                     nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5),
                                     nn.ReLU(inplace=True),
                                     nn.MaxPool2d(kernel_size=2),
                                     nn.Dropout(0.25),
                                     nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
                                     nn.ReLU(inplace=True),
                                     nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3),
                                     nn.ReLU(inplace=True),
                                     nn.MaxPool2d(kernel_size=2, stride=2),
                                     nn.Dropout(0.25))
        
        self.classifier = nn.Sequential(nn.Linear(576, 256),
                                       nn.Dropout(0.5),
                                       nn.Linear(256, 47))

        
    def forward(self, x):
        x = self.cnn(x)
        x = x.view(x.size(0), -1) # flatten layer
        x = self.classifier(x)
        
        return x



class Trainer(trainer_pb2_grpc.TrainerServicer):
    def __init__(self, csvdata, device, batch):
        self.dloader = getdataloader(csvdata, batch=batch)
        self.device = device
        # self.batch = batch
        while True:
            try:
                self.client = ipfshttpclient.connect("/ip4/192.168.127.12/tcp/5001/http")
                break
            except:
                print("Waiting for ipfs services at : 192.168.127.12:5001")
                time.sleep(1)

    def Train(self, request, result):
        #print(request.BaseModel)
        print("Training...")
        result = trainOneEp(self.client.cat(request.BaseModel).decode(), self.dloader, self.device)
        hashresult = self.client.add_str(result)
        return trainer_pb2.TrainResult(Round=request.Round, Result=hashresult)


def serve(data, port, dev, bat):
    print("Read dataset : ",data)
    print("Using : ",dev)
    print("Port : ",port)
    time.sleep(2)
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    trainer_pb2_grpc.add_TrainerServicer_to_server(Trainer(data, dev, bat), server)

    server.add_insecure_port('0.0.0.0:'+port)
    server.start()
    server.wait_for_termination()

def trainOneEp(bmodel, dloader, device):
    #return bmodel
    model = Model()
    model = base642fullmodel(bmodel)
    #return fullmodel2base64(model)
    print(model)
    # loss_function = nn.CrossEntropyLoss()
    # optimizer = optim.SGD(model.parameters(), lr=0.005)
    optimizer = optim.RMSprop(model.parameters(), lr=0.001)
    loss_function = nn.CrossEntropyLoss()
    if (device=="GPU"):
        model.cuda()

    model.train()
    for data, target in dloader:
        if (device=="GPU"):
            data = data.cuda()
            target = target.cuda()

        optimizer.zero_grad()
        #data = data.view(data.size(0),-1)

        output = model(data.float())

        loss = loss_function(output, target)

        loss.backward()

        optimizer.step()

    #model.eval()
    #print(model)
    if (device=="GPU"):
        model.cpu()
    bmodel_ = fullmodel2base64(model)
    #print(bmodel_)
    return bmodel_


# def getdataloader(dset = '/home/tedbest/Documents/mnist_train_0.csv'):
#     #print(dset)
#     train = pd.read_csv(dset)

#     train_labels = train['label'].values
#     train_data = train.drop(labels = ['label'], axis = 1)
#     train_data = train_data.values.reshape(-1,28, 28)

#     train_images_tensor = torch.tensor(train_data)/255.0
#     train_labels_tensor = torch.tensor(train_labels)
#     mnist = TensorDataset(train_images_tensor, train_labels_tensor)

#     trainloader = DataLoader(mnist, batch_size=256, shuffle= True)

#     return trainloader
class MNISTDataset(Dataset):
    """MNIST dataset"""
    
    def __init__(self, feature, target, transform=None):
        
        self.X = []
        self.Y = target
            
        if transform is not None:
            for i in range(len(feature)):
                self.X.append(transform(feature[i]))
    
    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        
        return self.X[idx], self.Y[idx]

def getdataloader(dset = './mnist_test.csv', batch=256):
    #print(dset)
    train = pd.read_csv(dset)

    train_labels = train['label'].values
    train_data = train.drop(labels = ['label'], axis = 1)
    train_data = train_data.values.reshape(-1,28, 28)
    
    featuresTrain = torch.from_numpy(train_data)
    targetsTrain = torch.from_numpy(train_labels)
    
    data_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomAffine(degrees=45, translate=(0.1, 0.1), scale=(0.8, 1.2)),
        transforms.ToTensor()]
        )

    train_set = MNISTDataset(featuresTrain.float(), targetsTrain, transform=data_transform)
    
    trainloader = torch.utils.data.DataLoader(train_set, batch_size = batch, shuffle = True, num_workers=4)
    return trainloader



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default="/home/tedbest/Documents/mnist_train_0.csv")
    parser.add_argument('--port', type=str, default="63387")
    parser.add_argument('--device', type=str, default="CPU") # GPU/CPU
    parser.add_argument('--batch', type=int, default=256)
    parser.add_argument('-f')
    args = parser.parse_args()

    if (args.device=="GPU"):
        if torch.cuda.is_available():
            print("GPU found.")
        else:
            print("GPU not found.")
            exit(0)

    logging.basicConfig()
    serve(args.data, args.port, args.device, args.batch)
import torch
import torch.nn as nn
from torchvision import transforms
from torch.utils.data import DataLoader, TensorDataset, Dataset
from torch.utils.data.sampler import SubsetRandomSampler
from torch import optim
import pandas as pd
import sys
sys.path.append('./proto')
import trainer_pb2
import trainer_pb2_grpc
import time
from concurrent import futures
import logging
import grpc
import argparse
import base64
import io
import ipfshttpclient

torch.nn.Module.dump_patches = True

def fullmodel2base64(model):
    buffer = io.BytesIO()
    torch.save(model, buffer)
    bg = buffer.getvalue()
    return base64.b64encode(bg).decode()

def base642fullmodel(modbase64):
    inputrpc = bytes(modbase64.encode())
    inputrpc_ = base64.b64decode(inputrpc)
    loadmodel = torch.load(io.BytesIO(inputrpc_))
    return loadmodel

# class Model(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.hidden = nn.Linear(784, 20)
#         self.output = nn.Linear(20, 10)

#     def forward(self, x):
#         x = self.hidden(x)
#         x = torch.sigmoid(x)
#         x = self.output(x)
#         return x
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.cnn = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5),
                                     nn.ReLU(inplace=True),
                                     nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5),
                                     nn.ReLU(inplace=True),
                                     nn.MaxPool2d(kernel_size=2),
                                     nn.Dropout(0.25),
                                     nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
                                     nn.ReLU(inplace=True),
                                     nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3),
                                     nn.ReLU(inplace=True),
                                     nn.MaxPool2d(kernel_size=2, stride=2),
                                     nn.Dropout(0.25))
        
        self.classifier = nn.Sequential(nn.Linear(576, 256),
                                       nn.Dropout(0.5),
                                       nn.Linear(256, 47))

        
    def forward(self, x):
        x = self.cnn(x)
        x = x.view(x.size(0), -1) # flatten layer
        x = self.classifier(x)
        
        return x



class Trainer(trainer_pb2_grpc.TrainerServicer):
    def __init__(self, csvdata, device, batch):
        self.dloader = getdataloader(csvdata, batch=batch)
        self.device = device
        # self.batch = batch
        while True:
            try:
                self.client = ipfshttpclient.connect("/ip4/192.168.127.12/tcp/5001/http")
                break
            except:
                print("Waiting for ipfs services at : 192.168.127.12:5001")
                time.sleep(1)

    def Train(self, request, result):
        #print(request.BaseModel)
        print("Training...")
        result = trainOneEp(self.client.cat(request.BaseModel).decode(), self.dloader, self.device)
        hashresult = self.client.add_str(result)
        return trainer_pb2.TrainResult(Round=request.Round, Result=hashresult)


def serve(data, port, dev, bat):
    print("Read dataset : ",data)
    print("Using : ",dev)
    print("Port : ",port)
    time.sleep(2)
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    trainer_pb2_grpc.add_TrainerServicer_to_server(Trainer(data, dev, bat), server)

    server.add_insecure_port('0.0.0.0:'+port)
    server.start()
    server.wait_for_termination()

def trainOneEp(bmodel, dloader, device):
    #return bmodel
    model = Model()
    model = base642fullmodel(bmodel)
    #return fullmodel2base64(model)
    print(model)
    # loss_function = nn.CrossEntropyLoss()
    # optimizer = optim.SGD(model.parameters(), lr=0.005)
    optimizer = optim.RMSprop(model.parameters(), lr=0.001)
    loss_function = nn.CrossEntropyLoss()
    if (device=="GPU"):
        model.cuda()

    model.train()
    for data, target in dloader:
        if (device=="GPU"):
            data = data.cuda()
            target = target.cuda()

        optimizer.zero_grad()
        #data = data.view(data.size(0),-1)

        output = model(data.float())

        loss = loss_function(output, target)

        loss.backward()

        optimizer.step()

    #model.eval()
    #print(model)
    if (device=="GPU"):
        model.cpu()
    bmodel_ = fullmodel2base64(model)
    #print(bmodel_)
    return bmodel_


# def getdataloader(dset = '/home/tedbest/Documents/mnist_train_0.csv'):
#     #print(dset)
#     train = pd.read_csv(dset)

#     train_labels = train['label'].values
#     train_data = train.drop(labels = ['label'], axis = 1)
#     train_data = train_data.values.reshape(-1,28, 28)

#     train_images_tensor = torch.tensor(train_data)/255.0
#     train_labels_tensor = torch.tensor(train_labels)
#     mnist = TensorDataset(train_images_tensor, train_labels_tensor)

#     trainloader = DataLoader(mnist, batch_size=256, shuffle= True)

#     return trainloader
class MNISTDataset(Dataset):
    """MNIST dataset"""
    
    def __init__(self, feature, target, transform=None):
        
        self.X = []
        self.Y = target
            
        if transform is not None:
            for i in range(len(feature)):
                self.X.append(transform(feature[i]))
    
    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        
        return self.X[idx], self.Y[idx]

def getdataloader(dset = './mnist_test.csv', batch=256):
    #print(dset)
    train = pd.read_csv(dset)

    train_labels = train['label'].values
    train_data = train.drop(labels = ['label'], axis = 1)
    train_data = train_data.values.reshape(-1,28, 28)
    
    featuresTrain = torch.from_numpy(train_data)
    targetsTrain = torch.from_numpy(train_labels)
    
    data_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomAffine(degrees=45, translate=(0.1, 0.1), scale=(0.8, 1.2)),
        transforms.ToTensor()]
        )

    train_set = MNISTDataset(featuresTrain.float(), targetsTrain, transform=data_transform)
    
    trainloader = torch.utils.data.DataLoader(train_set, batch_size = batch, shuffle = True, num_workers=4)
    return trainloader



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default="/home/tedbest/Documents/mnist_train_0.csv")
    parser.add_argument('--port', type=str, default="63387")
    parser.add_argument('--device', type=str, default="CPU") # GPU/CPU
    parser.add_argument('--batch', type=int, default=256)
    parser.add_argument('-f')
    args = parser.parse_args()

    if (args.device=="GPU"):
        if torch.cuda.is_available():
            print("GPU found.")
        else:
            print("GPU not found.")
            exit(0)

    logging.basicConfig()
    serve(args.data, args.port, args.device, args.batch)
# class Model(nn.Module): #     def __init__(self): #         super().__init__() #         self.hidden = nn.Linear(784, 20) #         self.output = nn.Linear(20, 10) #     def forward(self, x): #         x = self.hidden(x) #         x = torch.sigmoid(x) #         x = self.output(x) #         return x # flatten layer # self.batch = batch #print(request.BaseModel) #return bmodel #return fullmodel2base64(model) # loss_function = nn.CrossEntropyLoss() # optimizer = optim.SGD(model.parameters(), lr=0.005) #data = data.view(data.size(0),-1) #model.eval() #print(model) #print(bmodel_) # def getdataloader(dset = '/home/tedbest/Documents/mnist_train_0.csv'): #     #print(dset) #     train = pd.read_csv(dset) #     train_labels = train['label'].values #     train_data = train.drop(labels = ['label'], axis = 1) #     train_data = train_data.values.reshape(-1,28, 28) #     train_images_tensor = torch.tensor(train_data)/255.0 #     train_labels_tensor = torch.tensor(train_labels) #     mnist = TensorDataset(train_images_tensor, train_labels_tensor) #     trainloader = DataLoader(mnist, batch_size=256, shuffle= True) #     return trainloader MNIST dataset #print(dset) # GPU/CPU