import torch
from torch import nn
import torch.nn.functional as F

import pandas as pd
import numpy as np
from torch.autograd import Variable

from mpi4py import MPI

import sys
from .utils import err_simulation
sys.path.insert(0, '../compress_gradient') # ~ adds folder to visible path for the next import
from compress_gradient import compress

# ~ the following lines need to be commented by default, they are only for reproducibility
# SEED_=428
# torch.manual_seed(SEED_)

# we use LeNet here for our simple case
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)
        self.ceriation = nn.CrossEntropyLoss()
    def forward(self, x):
        x = self.conv1(x)
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(x)
        x = x.view(-1, 4*4*50)
        x = self.fc1(x)
        x = self.fc2(x)
        #loss = self.ceriation(x, target)
        return x
    def name(self):
        return 'lenet'

class LeNetSplit(nn.Module):
    '''
    this is a module that we split the module and do backward process layer by layer
    please don't call this module for normal uses, this is a hack and run slower than
    the automatic chain rule version
    '''
    def __init__(self):
        super(LeNetSplit, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)
        
        self.maxpool2d = nn.MaxPool2d(2, stride=2)
        self.relu = nn.ReLU()

        self.full_modules = [self.conv1, self.conv2, self.fc1, self.fc2]
        self._init_channel_index = len(self.full_modules)*2

        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        self.output = []
        self.input = []
        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.conv1(x)
        self.output.append(x)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.maxpool2d(x)
        self.output.append(x)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.relu(x)
        self.output.append(x)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.conv2(x)
        self.output.append(x)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.maxpool2d(x)
        self.output.append(x)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.relu(x)
        self.output.append(x)

        x = x.view(-1, 4*4*50)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.fc1(x)
        self.output.append(x)

        x = Variable(x.data, requires_grad=True)
        self.input.append(x)
        x = self.fc2(x)
        self.output.append(x)
        return x

    @property
    def fetch_init_channel_index(self):
        return self._init_channel_index

    # this is not used ???
    # see fc_nn for similar code
    def backward_normal(self, g, communicator, req_send_check, cur_step, fail_workers, err_mode, compress_grad):
        mod_avail_index = len(self.full_modules)-1
        #channel_index = len(self.full_modules)*2-2
        channel_index = self._init_channel_index - 2
        mod_counters_ = [0]*len(self.full_modules)
        for i, output in reversed(list(enumerate(self.output))):
            req_send_check[-1].wait()
            if i == (len(self.output) - 1):
                # for last node, use g
                output.backward(g)
                # get gradient here after some sanity checks:
                tmp_grad = self.full_modules[mod_avail_index].weight.grad
                if not pd.isnull(tmp_grad):
                    grads = tmp_grad.data.numpy().astype(np.float64)
                    ############################### simulation here #########################################
                    if communicator.Get_rank() in fail_workers:
                        simulation_grad = err_simulation(grad=grads, mode=err_mode)
                        if compress_grad == 'compress':
                            _compressed_grad = compress(simulation_grad)
                            req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                        else:
                            req_isend = communicator.Isend([simulation_grad, MPI.DOUBLE], dest=0, tag=88+channel_index)
                    else:
                        if compress_grad == 'compress':
                            _compressed_grad = compress(grads)
                            req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                        else:                     
                            req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
                    #########################################################################################
                    req_send_check.append(req_isend)
                    # update counters
                    mod_avail_index-=1
                    channel_index-=1
                else:
                    continue
            else:
                output.backward(self.input[i+1].grad.data)
                tmp_grad_weight = self.full_modules[mod_avail_index].weight.grad
                tmp_grad_bias = self.full_modules[mod_avail_index].bias.grad
                if not pd.isnull(tmp_grad_weight) and not pd.isnull(tmp_grad_bias):
                    # we always send bias first
                    if mod_counters_[mod_avail_index] == 0:
                        grads = tmp_grad_bias.data.numpy().astype(np.float64)
                        ############################### simulation here #########################################
                        if communicator.Get_rank() in fail_workers:
                            simulation_grad = err_simulation(grad=grads, mode=err_mode)
                            if compress_grad == 'compress':
                                _compressed_grad = compress(simulation_grad)
                                req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                            else:
                                req_isend = communicator.Isend([simulation_grad, MPI.DOUBLE], dest=0, tag=88+channel_index)
                        else:
                            if compress_grad == 'compress':
                                _compressed_grad = compress(grads)
                                req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                            else:                     
                                req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
                        #########################################################################################
                        req_send_check.append(req_isend)
                        channel_index-=1
                        mod_counters_[mod_avail_index]+=1
                    elif mod_counters_[mod_avail_index] == 1:
                        grads = tmp_grad_weight.data.numpy().astype(np.float64)
                        ############################### simulation here #########################################
                        if communicator.Get_rank() in fail_workers:
                            simulation_grad = err_simulation(grad=grads, mode=err_mode)
                            if compress_grad == 'compress':
                                _compressed_grad = compress(simulation_grad)
                                req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                            else:
                                req_isend = communicator.Isend([simulation_grad, MPI.DOUBLE], dest=0, tag=88+channel_index)
                        else:
                            if compress_grad == 'compress':
                                _compressed_grad = compress(grads)
                                req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                            else:                     
                                req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
                        #########################################################################################
                        req_send_check.append(req_isend)
                        channel_index-=1
                        mod_counters_[mod_avail_index]+=1
                        # update counters
                        mod_avail_index-=1
                else:
                    continue
        if mod_counters_[0] == 1:
            req_send_check[-1].wait()
            grads = tmp_grad_weight.data.numpy().astype(np.float64)
            ############################### simulation here #########################################
            if communicator.Get_rank() in fail_workers:
                simulation_grad = err_simulation(grad=grads, mode=err_mode)
                if compress_grad == 'compress':
                    _compressed_grad = compress(simulation_grad)
                    req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                else:
                    req_isend = communicator.Isend([simulation_grad, MPI.DOUBLE], dest=0, tag=88+channel_index)
            else:
                if compress_grad == 'compress':
                    _compressed_grad = compress(grads)
                    req_isend = communicator.isend(_compressed_grad, dest=0, tag=88+channel_index)
                else:                     
                    req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
            #########################################################################################
            req_send_check.append(req_isend)
        return req_send_check

    # this is not used ???
    def backward_signal_kill(self, g, communicator, req_send_check, cur_step):
        '''
        This killer is triggered by signals bcasting from master, channel of 
        signal is kept checking by each worker to determine if they're the 
        straggler
        '''
        mod_avail_index = len(self.full_modules)-1
        channel_index = self._init_channel_index - 2
        mod_counters_ = [0]*len(self.full_modules)

        # should kill flag
        should_kill = False

        for i, output in reversed(list(enumerate(self.output))):
        ############################ killing process on workers #####################################
            for _ in range(10000):
                status = MPI.Status()
                communicator.Iprobe(0, 77, status)
                if status.Get_source() == 0:
                    print("Worker {}, Cur Step: {} I'm the straggler, killing myself!".format(communicator.Get_rank(), cur_step))
                    tmp = communicator.recv(source=0, tag=77)
                    should_kill = True
                    break
            if should_kill:
                break
        ############################################################################################

            if i == (len(self.output) - 1):
                # for last node, use g
                output.backward(g)
                # get gradient here after some sanity checks:
                tmp_grad = self.full_modules[mod_avail_index].weight.grad
                if not pd.isnull(tmp_grad):
                    grads = tmp_grad.data.numpy().astype(np.float64)
                    req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
                    req_send_check.append(req_isend)
                    # update counters
                    mod_avail_index-=1
                    channel_index-=1
                else:
                    continue
            else:
                output.backward(self.input[i+1].grad.data)
                tmp_grad_weight = self.full_modules[mod_avail_index].weight.grad
                tmp_grad_bias = self.full_modules[mod_avail_index].bias.grad
                if not pd.isnull(tmp_grad_weight) and not pd.isnull(tmp_grad_bias):
                    # we always send bias first
                    if mod_counters_[mod_avail_index] == 0:
                        grads = tmp_grad_bias.data.numpy().astype(np.float64)
                        req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
                        req_send_check.append(req_isend)
                        channel_index-=1
                        mod_counters_[mod_avail_index]+=1
                    elif mod_counters_[mod_avail_index] == 1:
                        grads = tmp_grad_weight.data.numpy().astype(np.float64)
                        req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
                        req_send_check.append(req_isend)
                        channel_index-=1
                        mod_counters_[mod_avail_index]+=1
                        # update counters
                        mod_avail_index-=1
                else:
                    continue
        if mod_counters_[0] == 1:
            grads = tmp_grad_weight.data.numpy().astype(np.float64)
            req_isend = communicator.Isend([grads, MPI.DOUBLE], dest=0, tag=88+channel_index)
            req_send_check.append(req_isend)
        return req_send_check

    def backward_timeout_kill(self, g, communicator, req_send_check):
        """do we even need this?"""
        pass

    # this is not used ???
    # see fc_nn for similar code
    def backward_coded(self, g, cur_step):
        grad_aggregate_list = []

        mod_avail_index = len(self.full_modules)-1
        #channel_index = len(self.full_modules)*2-2
        channel_index = self._init_channel_index - 2
        mod_counters_ = [0]*len(self.full_modules)
        for i, output in reversed(list(enumerate(self.output))):
            if i == (len(self.output) - 1):
                # for last node, use g
                output.backward(g)
            else:
                output.backward(self.input[i+1].grad.data)
                tmp_grad_weight = self.full_modules[mod_avail_index].weight.grad
                tmp_grad_bias = self.full_modules[mod_avail_index].bias.grad
                # specific for this fc nn setting
                if not pd.isnull(tmp_grad_weight) and not pd.isnull(tmp_grad_bias):
                    # we always send bias first
                    if mod_counters_[mod_avail_index] == 0:
                        grads = tmp_grad_bias.data.numpy().astype(np.float64)
                        grad_aggregate_list.append(grads)
                        channel_index-=1
                        mod_counters_[mod_avail_index]+=1
                    elif mod_counters_[mod_avail_index] == 1:
                        grads = tmp_grad_weight.data.numpy().astype(np.float64)
                        grad_aggregate_list.append(grads)
                        channel_index-=1
                        mod_counters_[mod_avail_index]+=1
                        # update counters
                        mod_avail_index-=1
                else:
                    continue
        if mod_counters_[0] == 1:
            grads = tmp_grad_weight.data.numpy().astype(np.float64)
            grad_aggregate_list.append(grads)
        return grad_aggregate_list