"""Distributed GAT training, targeting XPU devices.
PVC has 2 tiles, each reports itself as a separate
device. DDP approach allows us to employ both tiles.

Additional requirements:
    IPEX (intel_extension_for_pytorch)
    oneCCL (oneccl_bindings_for_pytorch)

    We need to import both these modules, as they extend
    torch module with XPU/oneCCL related functionality.

Run with:
    mpirun -np 2 python distributed_sampling_xpu.py
"""

import copy
import os
import os.path as osp
from typing import Any, Tuple, Union

import intel_extension_for_pytorch  # noqa
import oneccl_bindings_for_pytorch  # noqa
import torch
import torch.distributed as dist
import torch.nn.functional as F
from ogb.nodeproppred import Evaluator, PygNodePropPredDataset
from torch import Tensor
from torch.nn import Linear as Lin
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm

from torch_geometric.loader import NeighborLoader
from torch_geometric.nn import GATConv
from torch_geometric.profile import get_stats_summary, profileit


class GAT(torch.nn.Module):
    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        out_channels: int,
        num_layers: int,
        heads: int,
    ):
        super().__init__()

        self.num_layers = num_layers

        self.convs = torch.nn.ModuleList()
        self.convs.append(GATConv(dataset.num_features, hidden_channels,
                                  heads))
        for _ in range(num_layers - 2):
            self.convs.append(
                GATConv(heads * hidden_channels, hidden_channels, heads))
        self.convs.append(
            GATConv(heads * hidden_channels, out_channels, heads,
                    concat=False))

        self.skips = torch.nn.ModuleList()
        self.skips.append(Lin(dataset.num_features, hidden_channels * heads))
        for _ in range(num_layers - 2):
            self.skips.append(
                Lin(hidden_channels * heads, hidden_channels * heads))
        self.skips.append(Lin(hidden_channels * heads, out_channels))

    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
        for i, (conv, skip) in enumerate(zip(self.convs, self.skips)):
            x = conv(x, edge_index) + skip(x)
            if i != self.num_layers - 1:
                x = F.elu(x)
                x = F.dropout(x, p=0.5, training=self.training)
        return x

    def inference(
        self,
        x_all: Tensor,
        device: Union[str, torch.device],
        subgraph_loader: NeighborLoader,
    ) -> Tensor:
        pbar = tqdm(total=x_all.size(0) * self.num_layers)
        pbar.set_description("Evaluating")

        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        for i in range(self.num_layers):
            xs = []
            for batch in subgraph_loader:
                x = x_all[batch.n_id].to(device)
                edge_index = batch.edge_index.to(device)
                x = self.convs[i](x, edge_index) + self.skips[i](x)
                x = x[:batch.batch_size]
                if i != self.num_layers - 1:
                    x = F.elu(x)
                xs.append(x.cpu())

                pbar.update(batch.batch_size)

            x_all = torch.cat(xs, dim=0)

        pbar.close()

        return x_all


@profileit('xpu')
def train_step(model: Any, optimizer: Any, x: Tensor, edge_index: Tensor,
               y: Tensor, bs: int) -> float:
    optimizer.zero_grad()
    out = model(x, edge_index)[:bs]
    loss = F.cross_entropy(out, y[:bs].squeeze())
    loss.backward()
    optimizer.step()
    return float(loss)


def run(rank: int, world_size: int, dataset: PygNodePropPredDataset):
    device = f"xpu:{rank}"

    split_idx = dataset.get_idx_split()
    split_idx["train"] = (split_idx["train"].split(
        split_idx["train"].size(0) // world_size, dim=0)[rank].clone())
    data = dataset[0].to(device, "x", "y")

    kwargs = dict(batch_size=1024, num_workers=0, pin_memory=True)
    train_loader = NeighborLoader(data, input_nodes=split_idx["train"],
                                  num_neighbors=[10, 10, 5], **kwargs)

    if rank == 0:
        subgraph_loader = NeighborLoader(copy.copy(data), num_neighbors=[-1],
                                         **kwargs)
        evaluator = Evaluator(name="ogbn-products")

    torch.manual_seed(12345)
    model = GAT(dataset.num_features, 128, dataset.num_classes, num_layers=3,
                heads=4).to(device)
    model = DDP(model, device_ids=[device])
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(1, 21):
        epoch_stats = []
        model.train()
        for batch in train_loader:
            batch = batch.to(device)
            loss, stats = train_step(model, optimizer, batch.x,
                                     batch.edge_index, batch.y,
                                     batch.batch_size)
            epoch_stats.append(stats)

        dist.barrier()

        if rank == 0:
            print(f"Epoch: {epoch:02d}, Loss: {loss:.4f}")

        print(f"Epoch: {epoch:02d}, Rank: {rank}, "
              f"Stats: {get_stats_summary(epoch_stats)}")

        if rank == 0 and epoch % 5 == 0:  # Evaluation on a single GPU
            model.eval()
            with torch.no_grad():
                out = model.module.inference(data.x, device, subgraph_loader)

            y_true = data.y.to(out.device)
            y_pred = out.argmax(dim=-1, keepdim=True)

            train_acc = evaluator.eval({
                "y_true": y_true[split_idx["train"]],
                "y_pred": y_pred[split_idx["train"]],
            })["acc"]
            val_acc = evaluator.eval({
                "y_true": y_true[split_idx["valid"]],
                "y_pred": y_pred[split_idx["valid"]],
            })["acc"]
            test_acc = evaluator.eval({
                "y_true": y_true[split_idx["test"]],
                "y_pred": y_pred[split_idx["test"]],
            })["acc"]

            print(f"Train: {train_acc:.4f}, Val: {val_acc:.4f}, "
                  f"Test: {test_acc:.4f}")

        dist.barrier()

    dist.destroy_process_group()


def get_dist_params() -> Tuple[int, int, str]:
    master_addr = "127.0.0.1"
    master_port = "29500"
    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port

    mpi_rank = int(os.environ.get("PMI_RANK", -1))
    mpi_world_size = int(os.environ.get("PMI_SIZE", -1))
    rank = mpi_rank if mpi_world_size > 0 else os.environ.get("RANK", 0)
    world_size = (mpi_world_size if mpi_world_size > 0 else os.environ.get(
        "WORLD_SIZE", 1))

    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)

    init_method = f"tcp://{master_addr}:{master_port}"

    return rank, world_size, init_method


if __name__ == "__main__":
    rank, world_size, init_method = get_dist_params()
    dist.init_process_group(backend="ccl", init_method=init_method,
                            world_size=world_size, rank=rank)

    path = osp.join(osp.dirname(osp.realpath(__file__)), "../../data",
                    "ogbn-products")
    dataset = PygNodePropPredDataset("ogbn-products", path)

    run(rank, world_size, dataset)