import torch.distributed as dist
import torch.multiprocessing as mp
import os
import torch

def init_process(rank, size, fn, backend="gloo"):
    """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)

def train(rank: int, size: int):
    prediction = torch.tensor(
        [[1.,2.,3.,4.], [4.,7.,5.,8.]],
    requires_grad=True)
    label = torch.tensor([[0.,0.,0.,0.], [0.,0.,0.,0.]])
    loss = (prediction[rank] - label[rank]) ** 2
    loss.mean().backward()
    dist.all_reduce(prediction.grad.data, op=dist.ReduceOp.SUM)
    prediction.grad.data /= size
    print("gradient with data parallelism (multi-gpu) \n", prediction.grad)
    

if __name__ == "__main__":
    prediction = torch.tensor(
        [[1.,2.,3.,4.], [4.,7.,5.,8.]],
    requires_grad=True)
    label = torch.tensor([[0.,0.,0.,0.], [0.,0.,0.,0.]])

    loss = (prediction - label) ** 2
    loss.mean().backward()
    print("gradient with the whole batch\n", prediction.grad)

    prediction = torch.tensor(
        [[1.,2.,3.,4.], [4.,7.,5.,8.]],
    requires_grad=True)
    label = torch.tensor([[0.,0.,0.,0.], [0.,0.,0.,0.]])
    # do backward pass in two minibatches
    loss = (prediction[0] - label[0]) ** 2
    loss.mean().backward()
    loss = (prediction[1] - label[1]) ** 2
    loss.mean().backward()
    # divide the gradient by the size
    print("gradient accumulation \n", prediction.grad / 2)
    
    size = 2
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, train))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

optimizer.zero_grad()
start = 0
end = start + args.minibatch_size
mb_inds = b_inds[start:end]
fit_vloss(mb_inds)
print()
print(f"CASE 1: value function: forward and backward pass of the minibatch (size 256: i.e., data[0:256])")
print("agent.critic.weight.grad.sum() =", agent.critic.weight.grad.sum())

optimizer.zero_grad()
args.minibatch_size = 128
for start in [0, 128]:
    end = start + args.minibatch_size
    mb_inds = b_inds[start:end]
    fit_vloss(mb_inds)
print()
print(f"CASE 2: value function: forward and backward pass of 2 minibatches (size 128: i.e., data[0:128] and data[128:256])")
print("agent.critic.weight.grad.sum() / 2 =", agent.critic.weight.grad.sum() / 2)

optimizer.zero_grad()
start = 0
end = start + args.minibatch_size
mb_inds = b_inds[start:end]
fit_pgloss(mb_inds)
print()
print(f"CASE 3: policy function: forward and backward pass of the minibatch (size 256: i.e., data[0:256])")
print("agent.actor.weight.grad.sum() =", agent.actor.weight.grad.sum())

optimizer.zero_grad()
args.minibatch_size = 128
for start in [0, 128]:
    end = start + args.minibatch_size
    mb_inds = b_inds[start:end]
    fit_pgloss(mb_inds)
print()
print(f"CASE 4: policy function: forward and backward pass of 2 minibatches (size 128: i.e., data[0:128] and data[128:256])")
print("agent.actor.weight.grad.sum() / 2 =", agent.actor.weight.grad.sum() / 2)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical

class Agent(nn.Module):
    def __init__(self, action_n):
        super().__init__()
        self.network = nn.Sequential(
            nn.Conv2d(4, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 512),
            nn.ReLU(),
        )
        self.actor = nn.Linear(512, action_n)
        self.critic = nn.Linear(512, 1)

    def get_value(self, x):
        return self.critic(self.network(x / 255.0))

    def get_action_and_value(self, x, action=None):
        hidden = self.network(x / 255.0)
        logits = self.actor(hidden)
        probs = Categorical(logits=logits)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action), probs.entropy(), self.critic(hidden)

# setup
agent = Agent(4)
optimizer = optim.Adam(agent.parameters())
next_obs = torch.rand(8, 4, 84, 84)
action, newlogprob, entropy, newvalue = agent.get_action_and_value(next_obs)

optimizer.zero_grad()
_, newlogprob, _, newvalue = agent.get_action_and_value(next_obs, action)
newvalue.mean().backward()
print(f"`agent.critic.weight.grad.sum() = {agent.critic.weight.grad.sum()}` after fitting value loss using data[0:8]")

optimizer.zero_grad()
_, newlogprob, _, newvalue = agent.get_action_and_value(next_obs[0:4], action[0:4])
newvalue.mean().backward()
_, newlogprob, _, newvalue = agent.get_action_and_value(next_obs[4:8], action[4:8])
newvalue.mean().backward()
print(f"`agent.critic.weight.grad.sum() / 2 = {agent.critic.weight.grad.sum() / 2}` after fitting value loss using data[0:4] and data[4:8] respectively")

optimizer.zero_grad()
_, newlogprob, _, newvalue = agent.get_action_and_value(next_obs, action)
newlogprob.mean().backward()
print(f"`agent.actor.weight.grad.sum() = {agent.actor.weight.grad.sum()}` after fitting value loss using data[0:8]")

optimizer.zero_grad()
_, newlogprob, _, newvalue = agent.get_action_and_value(next_obs[0:4], action[0:4])
newlogprob.mean().backward()
_, newlogprob, _, newvalue = agent.get_action_and_value(next_obs[4:8], action[4:8])
newlogprob.mean().backward()
print(f"`agent.actor.weight.grad.sum() / 2 = {agent.actor.weight.grad.sum() / 2}` after fitting value loss using data[0:4] and data[4:8] respectively")