import os
import random
from threading import Lock
import atexit

import numpy as np
import torch
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc
import torch.nn as nn


WORLD_RANK = RANK = int(os.getenv('RANK'))
WORLD_SIZE = int(os.getenv('WORLD_SIZE'))

AUTOGRAD_LOCK = Lock()


fmt = '{name} ({tensor.__class__.__module__}.{tensor.__class__.__qualname__}): param={tensor.data}, grad={tensor.grad}, grad_fn={tensor.grad_fn}'.format


def get_model():
    model = nn.Linear(1, 1)
    nn.init.ones_(model.weight)
    nn.init.zeros_(model.bias)
    return model


def worker_init():
    random.seed(RANK)
    np.random.seed(RANK)
    torch.manual_seed(RANK)

    torch.backends.cudnn.enabled = False
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

    print(f'Worker init => Rank {RANK}')


def main():
    rpc.init_rpc(name=f'worker{RANK}', rank=RANK, world_size=WORLD_SIZE)
    atexit.register(rpc.shutdown, graceful=True)

    rpc.api._barrier(worker_names={f'worker{i}' for i in range(WORLD_SIZE)})
    worker_init()
    rpc.api._barrier(worker_names={f'worker{i}' for i in range(WORLD_SIZE)})

    if RANK == 0:
        model = get_model()
        train(model)
    else:
        pass  # listening for RPC calls


def compute_loss(model_rref, x):
    if isinstance(model_rref, nn.Module):
        model = model_rref
    else:
        model = model_rref.to_here()

    print()
    print('Before RPC forward:')
    for name, param in model.named_parameters():
        print(f'    worker{RANK} => {fmt(name=name, tensor=param)}')

    out = model(x)
    return out.mean()


def convert_param_type(model):
    for param in model.parameters():
        param.__class__ = torch.Tensor


def train(model):
    x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])

    print()
    print('Before forward:')
    for name, param in model.named_parameters():
        print(f'    {fmt(name=name, tensor=param)}')

    # convert_param_type(model)
    model_rref = rpc.RRef(model)
    with dist_autograd.context() as ctx:
        loss0 = rpc.rpc_sync('worker0', compute_loss, args=(model_rref, x[: len(x) // 2]))
        loss1 = rpc.rpc_sync('worker1', compute_loss, args=(model_rref, x[len(x) // 2 :]))
        loss = torch.mean(torch.stack([loss0, loss1]))
        print()
        print(f'Loss: {loss}')
        dist_autograd.backward(ctx, [loss])

        with AUTOGRAD_LOCK:
            all_local_grads = dist_autograd.get_gradients(ctx)
            for p, g in all_local_grads.items():
                if p.grad is not None:
                    p.grad = p.grad.add(g)
                else:
                    p.grad = g

    print()
    print('After backward:')
    for name, param in model.named_parameters():
        print(f'    {fmt(name=name, tensor=param)}')


if __name__ == '__main__':
    main()

weight.grad = tensor([[0.7500]])
bais.grad = tensor([[0.7500]])

    convert_param_type(model)  # change __class__ of all `nn.Parameter`s
    model_rref = rpc.RRef(model)
    with dist_autograd.context() as ctx:
        loss0 = rpc.rpc_sync('worker0', compute_loss, args=(model_rref, x[: len(x) // 2]))
        loss1 = rpc.rpc_sync('worker1', compute_loss, args=(model_rref, x[len(x) // 2 :]))
        loss = torch.mean(torch.stack([loss0, loss1]))
        print()
        print(f'Loss: {loss}')
        dist_autograd.backward(ctx, [loss])

        with AUTOGRAD_LOCK:
            all_local_grads = dist_autograd.get_gradients(ctx)
            for p, g in all_local_grads.items():
                if p.grad is not None:
                    p.grad = p.grad.add(g)
                else:
                    p.grad = g

weight.grad = tensor([[2.5000]])
bais.grad = tensor([[1.]])