#!/usr/bin/env python
import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import time

from math import ceil
from random import Random
from torch.multiprocessing import Process
from torch.autograd import Variable
from torchvision import datasets, transforms


class Partition(object):
    """ Dataset-like object, but only access a subset of it. """

    def __init__(self, data, index):
        self.data = data
        self.index = index

    def __len__(self):
        return len(self.index)

    def __getitem__(self, index):
        data_idx = self.index[index]
        return self.data[data_idx]


class DataPartitioner(object):
    """ Partitions a dataset into different chuncks. """

    def __init__(self, data, sizes=[0.7, 0.2, 0.1], seed=1234):
        self.data = data
        self.partitions = []
        rng = Random()
        rng.seed(seed)
        data_len = len(data)
        indexes = [x for x in range(0, data_len)]
        rng.shuffle(indexes)

        for frac in sizes:
            part_len = int(frac * data_len)
            self.partitions.append(indexes[0:part_len])
            indexes = indexes[part_len:]

    def use(self, partition):
        return Partition(self.data, self.partitions[partition])


class Net(nn.Module):
    """ Network architecture. """

    def __init__(self):
        super(Net, self).__init__()
        '''第一层卷积层，卷积核为3*3，通道数为96，步距为1，原始图像大小为32*32，有R、G、B三个通道'''

        '''这样经过第一层卷积层之后，得到的feature map的大小为(32-3)/1+1=30,所以feature map的维度为96*30*30'''

        self.conv1 = nn.Conv2d(1, 96, kernel_size=3, stride=1)

        '''经过一次批归一化，将数据拉回到正态分布'''

        self.bn1 = nn.BatchNorm2d(96)

        '''第一层池化层，卷积核为3*3，步距为2，前一层的feature map的大小为30*30，通道数为96个'''

        '''这样经过第一层池化层之后，得到的feature map的大小为(30-3)/2+1=14,所以feature map的维度为96*14*14'''

        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2)

        '''第二层卷积层，卷积核为3*3，通道数为256，步距为1，前一层的feature map的大小为14*14，通道数为96个'''

        '''这样经过第一层卷积层之后，得到的feature map的大小为(14-3)/1+1=12,所以feature map的维度为256*12*12'''

        self.conv2 = nn.Conv2d(96, 256, kernel_size=3, stride=1)

        '''经过一次批归一化，将数据拉回到正态分布'''

        self.bn2 = nn.BatchNorm2d(256)

        '''第二层池化层，卷积核为3*3，步距为2，前一层的feature map的大小为12*12，通道数为256个'''

        '''这样经过第二层池化层之后，得到的feature map的大小为(12-3)/2+1=5,所以feature map的维度为256*5*5'''

        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2)

        '''第三层卷积层，卷积核为3*3，通道数为384，步距为1，前一层的feature map的大小为5*5，通道数为256个'''

        '''这样经过第一层卷积层之后，得到的feature map的大小为(5-3+2*1)/1+1=5,所以feature map的维度为384*5*5'''

        self.conv3 = nn.Conv2d(256, 384, kernel_size=3, padding=1, stride=1)

        '''第四层卷积层，卷积核为3*3，通道数为384，步距为1，前一层的feature map的大小为5*5，通道数为384个'''

        '''这样经过第一层卷积层之后，得到的feature map的大小为(5-3+2*1)/1+1=5,所以feature map的维度为384*5*5'''

        self.conv4 = nn.Conv2d(384, 384, kernel_size=3, padding=1, stride=1)

        '''第五层卷积层，卷积核为3*3，通道数为384，步距为1，前一层的feature map的大小为5*5，通道数为384个'''

        '''这样经过第一层卷积层之后，得到的feature map的大小为(5-3+2*1)/1+1=5,所以feature map的维度为256*5*5'''

        self.conv5 = nn.Conv2d(384, 256, kernel_size=3, padding=1, stride=1)

        '''第三层池化层，卷积核为3*3，步距为2，前一层的feature map的大小为5*5，通道数为256个'''

        '''这样经过第三层池化层之后，得到的feature map的大小为(5-3)/2+1=2,所以feature map的维度为256*2*2'''

        self.pool3 = nn.MaxPool2d(kernel_size=3, stride=2)

        '''经过第一层全连接层'''

        self.linear1 = nn.Linear(256, 512)

        '''经过第一次DropOut层'''

        self.dropout1 = nn.Dropout(0.5)

        '''经过第二层全连接层'''

        self.linear2 = nn.Linear(512, 1024)

        '''经过第二层DropOut层'''

        self.dropout2 = nn.Dropout(0.5)

        '''经过第三层全连接层，得到输出结果'''

        self.linear3 = nn.Linear(1024, 10)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = F.relu(out)
        out = self.pool1(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = F.relu(out)
        out = self.pool2(out)

        out = F.relu(self.conv3(out))

        out = F.relu(self.conv4(out))

        out = F.relu(self.conv5(out))

        out = self.pool3(out)

        out = out.reshape(-1, 256)

        out = F.relu(self.linear1(out))

        out = self.dropout1(out)

        out = F.relu(self.linear2(out))

        out = self.dropout2(out)

        out = self.linear3(out)
        out = F.log_softmax(out, dim=1)
        return out
    # #和非分布式一致的网络
    # def __init__(self):
    #     super(Net, self).__init__()
    #     self.conv1 = nn.Conv2d(1, 32, 3, 1)
    #     self.conv2 = nn.Conv2d(32, 64, 3, 1)
    #     self.dropout1 = nn.Dropout2d(0.25)
    #     self.dropout2 = nn.Dropout2d(0.5)
    #     self.fc1 = nn.Linear(9216, 128)
    #     self.fc2 = nn.Linear(128, 10)
    #
    # def forward(self, x):
    #     x = self.conv1(x)
    #     x = F.relu(x)
    #     x = self.conv2(x)
    #     x = F.relu(x)
    #     x = F.max_pool2d(x, 2)
    #     x = self.dropout1(x)
    #     x = torch.flatten(x, 1)
    #     x = self.fc1(x)
    #     x = F.relu(x)
    #     x = self.dropout2(x)
    #     x = self.fc2(x)
    #     output = F.log_softmax(x, dim=1)
    #     return output


def partition_dataset():
    """ Partitioning MNIST """
    dataset = datasets.MNIST(
        './data',
        train=True,
        download=True,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ]))
    size = dist.get_world_size()
    # bsz = 128 / float(size)
    bsz = 2
    print("进程的个数",size)
    print("每一批每个进程分得的训练集个数", bsz)
    partition_sizes = [1.0 / size for _ in range(size)]
    partition = DataPartitioner(dataset, partition_sizes)
    partition = partition.use(dist.get_rank())
    train_set = torch.utils.data.DataLoader(
        partition, batch_size=int(bsz), shuffle=True)
    return train_set, bsz


def average_gradients(model,group,mylist):
    """ Gradient averaging. """
    size = len(mylist)
    if dist.get_rank() in mylist:
        for param in model.parameters():
            # dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM, group=0)
            dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM,group = group)
            param.grad.data /= size


def run(rank, size):
    """ Distributed Synchronous SGD Example """
    torch.manual_seed(1234)
    train_set, bsz = partition_dataset()
    model = Net()
    model = model
#    model = model.cuda(rank)
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
    num_batches = ceil(len(train_set.dataset) / float(bsz))
    totalbegin = time.clock()
    totalbegin2 = time.time()
    mylist1 = list(range(64))
    mylist2 = list(range(64,128))
    mylist3 = list(range(128))
    group1 = dist.new_group(mylist1)
    group2 = dist.new_group(mylist2)
    group3 = dist.new_group(mylist3)
    for epoch in range(10):
        epoch_loss = 0.0
        begin = time.clock()
        for data, target in train_set:
            data, target = Variable(data), Variable(target)
#            data, target = Variable(data.cuda(rank)), Variable(target.cuda(rank))
            optimizer.zero_grad()
            output = model(data)
            loss = F.nll_loss(output, target)
            epoch_loss += loss.item()
            loss.backward()
            average_gradients(model,group1,mylist1)
            average_gradients(model,group2,mylist2)
            optimizer.step()
        end = time.clock()
        spendtime = (end-begin)/60
        print('Rank ',
              dist.get_rank(), ', epoch ', epoch, ': ',
              epoch_loss / num_batches,', spend time: ',spendtime)
        average_gradients(model, group3, mylist3)
    totalend =time.clock()
    totalend2 = time.time()
    print("total processrun time: ",(totalend-totalbegin)/60)
    print("total realrun time: ", (totalend2 - totalbegin2) / 60)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST(
            './data',
            train=False,
            transform=transforms.Compose(
                [
                    transforms.ToTensor(),
                    transforms.Normalize(
                        (0.1307,
                         ),
                        (0.3081,
                         ))])),
        batch_size=32,
        shuffle=True,
    )
    get_accuracy(test_loader, model)

def get_accuracy(test_loader, model):
    model.eval()
    correct_sum = 0
    # Use GPU to evaluate if possible
    device = torch.device("cpu")
    with torch.no_grad():
        for i, (data, target) in enumerate(test_loader):
            out = model(data)
            pred = out.argmax(dim=1, keepdim=True)
            pred, target = pred.to(device), target.to(device)
            correct = pred.eq(target.view_as(pred)).sum().item()
            correct_sum = correct_sum + correct
    print("Accuracy:", correct_sum / len(test_loader.dataset))


def init_processes(rank, size, fn, backend='mpi'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size)


if __name__ == "__main__":
    print("程序开始执行")
    print("all process BEGINtime is : ", time.strftime('%Y-%m-%d %H:%M:%S'))
    init_processes(0, 0, run, backend='mpi')
    print("all process ENDtime is : ", time.strftime('%Y-%m-%d %H:%M:%S'))
