import dgl.function as fn
import torch as th
import torch.nn as nn
from dgl.nn.functional import edge_softmax


class MLP(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.W = nn.Linear(in_dim, out_dim)

    def apply_edges(self, edges):
        h_e = edges.data["h"]
        h_u = edges.src["h"]
        h_v = edges.dst["h"]
        score = self.W(th.cat([h_e, h_u, h_v], -1))
        return {"score": score}

    def forward(self, g, e_feat, u_feat, v_feat):
        with g.local_scope():
            g.edges["forward"].data["h"] = e_feat
            g.nodes["u"].data["h"] = u_feat
            g.nodes["v"].data["h"] = v_feat
            g.apply_edges(self.apply_edges, etype="forward")
            return g.edges["forward"].data["score"]


class GASConv(nn.Module):
    """One layer of GAS."""

    def __init__(
        self,
        e_in_dim,
        u_in_dim,
        v_in_dim,
        e_out_dim,
        u_out_dim,
        v_out_dim,
        activation=None,
        dropout=0,
    ):
        super(GASConv, self).__init__()

        self.activation = activation
        self.dropout = nn.Dropout(dropout)

        self.e_linear = nn.Linear(e_in_dim, e_out_dim)
        self.u_linear = nn.Linear(u_in_dim, e_out_dim)
        self.v_linear = nn.Linear(v_in_dim, e_out_dim)

        self.W_ATTN_u = nn.Linear(u_in_dim, v_in_dim + e_in_dim)
        self.W_ATTN_v = nn.Linear(v_in_dim, u_in_dim + e_in_dim)

        # the proportion of h_u and h_Nu are specified as 1/2 in formula 8
        nu_dim = int(u_out_dim / 2)
        nv_dim = int(v_out_dim / 2)

        self.W_u = nn.Linear(v_in_dim + e_in_dim, nu_dim)
        self.W_v = nn.Linear(u_in_dim + e_in_dim, nv_dim)

        self.Vu = nn.Linear(u_in_dim, u_out_dim - nu_dim)
        self.Vv = nn.Linear(v_in_dim, v_out_dim - nv_dim)

    def forward(self, g, e_feat, u_feat, v_feat):
        with g.local_scope():
            g.nodes["u"].data["h"] = u_feat
            g.nodes["v"].data["h"] = v_feat
            g.edges["forward"].data["h"] = e_feat
            g.edges["backward"].data["h"] = e_feat

            # formula 3 and 4 (optimized implementation to save memory)
            g.nodes["u"].data.update({"he_u": self.u_linear(u_feat)})
            g.nodes["v"].data.update({"he_v": self.v_linear(v_feat)})
            g.edges["forward"].data.update({"he_e": self.e_linear(e_feat)})
            g.apply_edges(
                lambda edges: {
                    "he": edges.data["he_e"]
                    + edges.src["he_u"]
                    + edges.dst["he_v"]
                },
                etype="forward",
            )
            he = g.edges["forward"].data["he"]
            if self.activation is not None:
                he = self.activation(he)

            # formula 6
            g.apply_edges(
                lambda edges: {
                    "h_ve": th.cat([edges.src["h"], edges.data["h"]], -1)
                },
                etype="backward",
            )
            g.apply_edges(
                lambda edges: {
                    "h_ue": th.cat([edges.src["h"], edges.data["h"]], -1)
                },
                etype="forward",
            )

            # formula 7, self-attention
            g.nodes["u"].data["h_att_u"] = self.W_ATTN_u(u_feat)
            g.nodes["v"].data["h_att_v"] = self.W_ATTN_v(v_feat)

            # Step 1: dot product
            g.apply_edges(
                fn.e_dot_v("h_ve", "h_att_u", "edotv"), etype="backward"
            )
            g.apply_edges(
                fn.e_dot_v("h_ue", "h_att_v", "edotv"), etype="forward"
            )

            # Step 2. softmax
            g.edges["backward"].data["sfm"] = edge_softmax(
                g["backward"], g.edges["backward"].data["edotv"]
            )
            g.edges["forward"].data["sfm"] = edge_softmax(
                g["forward"], g.edges["forward"].data["edotv"]
            )

            # Step 3. Broadcast softmax value to each edge, and then attention is done
            g.apply_edges(
                lambda edges: {"attn": edges.data["h_ve"] * edges.data["sfm"]},
                etype="backward",
            )
            g.apply_edges(
                lambda edges: {"attn": edges.data["h_ue"] * edges.data["sfm"]},
                etype="forward",
            )

            # Step 4. Aggregate attention to dst,user nodes, so formula 7 is done
            g.update_all(
                fn.copy_e("attn", "m"), fn.sum("m", "agg_u"), etype="backward"
            )
            g.update_all(
                fn.copy_e("attn", "m"), fn.sum("m", "agg_v"), etype="forward"
            )

            # formula 5
            h_nu = self.W_u(g.nodes["u"].data["agg_u"])
            h_nv = self.W_v(g.nodes["v"].data["agg_v"])
            if self.activation is not None:
                h_nu = self.activation(h_nu)
                h_nv = self.activation(h_nv)

            # Dropout
            he = self.dropout(he)
            h_nu = self.dropout(h_nu)
            h_nv = self.dropout(h_nv)

            # formula 8
            hu = th.cat([self.Vu(u_feat), h_nu], -1)
            hv = th.cat([self.Vv(v_feat), h_nv], -1)

            return he, hu, hv


class GAS(nn.Module):
    def __init__(
        self,
        e_in_dim,
        u_in_dim,
        v_in_dim,
        e_hid_dim,
        u_hid_dim,
        v_hid_dim,
        out_dim,
        num_layers=2,
        dropout=0.0,
        activation=None,
    ):
        super(GAS, self).__init__()
        self.e_in_dim = e_in_dim
        self.u_in_dim = u_in_dim
        self.v_in_dim = v_in_dim
        self.e_hid_dim = e_hid_dim
        self.u_hid_dim = u_hid_dim
        self.v_hid_dim = v_hid_dim
        self.out_dim = out_dim
        self.num_layer = num_layers
        self.dropout = dropout
        self.activation = activation
        self.predictor = MLP(e_hid_dim + u_hid_dim + v_hid_dim, out_dim)
        self.layers = nn.ModuleList()

        # Input layer
        self.layers.append(
            GASConv(
                self.e_in_dim,
                self.u_in_dim,
                self.v_in_dim,
                self.e_hid_dim,
                self.u_hid_dim,
                self.v_hid_dim,
                activation=self.activation,
                dropout=self.dropout,
            )
        )

        # Hidden layers with n - 1 CompGraphConv layers
        for i in range(self.num_layer - 1):
            self.layers.append(
                GASConv(
                    self.e_hid_dim,
                    self.u_hid_dim,
                    self.v_hid_dim,
                    self.e_hid_dim,
                    self.u_hid_dim,
                    self.v_hid_dim,
                    activation=self.activation,
                    dropout=self.dropout,
                )
            )

    def forward(self, graph, e_feat, u_feat, v_feat):
        # For full graph training, directly use the graph
        # Forward of n layers of GAS
        for layer in self.layers:
            e_feat, u_feat, v_feat = layer(graph, e_feat, u_feat, v_feat)

        # return the result of final prediction layer
        return self.predictor(graph, e_feat, u_feat, v_feat)