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


class MLP(nn.Module):

    def __init__(self, in_dim: int, out_dim: int, num_layers: int,
                 hidden: int, drop_rate: float, is_output_layer: bool = False):
        super().__init__()
        self.num_layers = num_layers
        self.is_output_layer = is_output_layer

        self.layers = nn.ModuleList()
        self.layers.append(nn.Linear(in_dim, hidden))
        for _ in range(1, num_layers - 1):
            self.layers.append(nn.Linear(hidden, hidden))
        self.layers.append(nn.Linear(hidden, out_dim))

        self.batch_norms = nn.ModuleList()
        for _ in range(num_layers - 1):
            self.batch_norms.append(nn.BatchNorm1d(hidden))
        if not is_output_layer:
            self.batch_norms.append(nn.BatchNorm1d(out_dim))

        self.dropout = nn.Dropout(drop_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for i in range(self.num_layers):
            x = self.layers[i](x)
            if i != self.num_layers - 1 or not self.is_output_layer:
                shape = x.shape
                x = x.view(-1, shape[-1])
                x = self.batch_norms[i](x)
                x = x.view(shape)
                x = self.dropout(F.relu(x))
        return x


class GCNConv(nn.Module):

    def __init__(self, in_dim: int, out_dim: int):
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.fc = nn.Linear(in_dim, out_dim)
        self.AA = None

    def set_adjacency(self, AA: torch.Tensor):
        self.AA = AA

    def forward(self, R: torch.Tensor) -> torch.Tensor:
        n, batch_size, in_dim = R.shape

        if in_dim > self.out_dim:
            R = self.fc(R)
            R = R.view(n, batch_size * self.out_dim)
            R = self.AA @ R
            R = R.view(n, batch_size, self.out_dim)
        else:
            R = R.view(n, batch_size * in_dim)
            R = self.AA @ R
            R = R.view(n, batch_size, in_dim)
            R = self.fc(R)

        return R


class MPNNConv(nn.Module):

    def __init__(self, node_dim: int, out_dim: int):
        super().__init__()
        self.node_dim = node_dim
        self.out_dim = out_dim
        self.message_fn = nn.Sequential(
            nn.Linear(2 * node_dim + 1, out_dim),
            nn.ReLU(),
            nn.Linear(out_dim, out_dim),
        )
        self.edge_index = None
        self.edge_values = None
        self.n = None

    def set_edges(self, edge_index: torch.Tensor, edge_values: torch.Tensor, n: int):
        self.edge_index = edge_index
        self.edge_values = edge_values
        self.n = n

    def forward(self, R: torch.Tensor) -> torch.Tensor:
        n, batch_size, node_dim = R.shape
        src, dst = self.edge_index

        h_src = R[src]
        h_dst = R[dst]
        e = self.edge_values.view(-1, 1, 1).expand(-1, batch_size, 1)

        msg_input = torch.cat([h_src, h_dst, e], dim=-1)
        messages = self.message_fn(msg_input)

        out = torch.zeros(n, batch_size, self.out_dim, dtype=R.dtype, device=R.device)
        idx = dst.view(-1, 1, 1).expand(-1, batch_size, self.out_dim)
        out.scatter_add_(0, idx, messages)

        return out


class ContextAggregator(nn.Module):

    def __init__(self, context_dim: int, output_dim: int, hidden: int):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(context_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, output_dim)
        )
        self.attention = nn.Sequential(
            nn.Linear(output_dim, 1),
            nn.Softmax(dim=-2)
        )

    def forward(self, context: torch.Tensor) -> torch.Tensor:
        encoded = self.encoder(context)
        weights = self.attention(encoded)
        aggregated = (encoded * weights).sum(dim=-2)
        return aggregated


class ContextResGCN(nn.Module):

    def __init__(self, num_layers: int, embed: int, hidden: int,
                 drop_rate: float, num_context_pairs: int,
                 scale_input: bool = True, dtype: torch.dtype = torch.float32):
        super().__init__()
        self.dtype = dtype
        self.num_layers = num_layers
        self.embed = embed
        self.hidden = hidden
        self.scale_input = scale_input
        self.num_context_pairs = num_context_pairs

        context_feature_dim = 2 * num_context_pairs
        context_embed_dim = embed // 2

        self.context_aggregator = ContextAggregator(
            context_dim=2,
            output_dim=context_embed_dim,
            hidden=hidden
        )

        input_dim = 1 + 1 + context_embed_dim

        self.mlp_initial = MLP(input_dim, embed, 4, hidden, drop_rate)
        self.mlp_final = MLP(embed, 1, 4, hidden, drop_rate, is_output_layer=True)

        self.gconv = nn.ModuleList()
        self.skip = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        for _ in range(num_layers):
            self.gconv.append(GCNConv(embed, embed))
            self.skip.append(nn.Linear(embed, embed))
            self.batch_norms.append(nn.BatchNorm1d(embed))

        self.dropout = nn.Dropout(drop_rate)
        self.AA = None

    def set_matrix(self, A: torch.Tensor):
        self.AA = self._scale_by_spectral_radius(A).to(self.dtype)
        for layer in self.gconv:
            layer.set_adjacency(self.AA)

    def _scale_by_spectral_radius(self, A: torch.Tensor) -> torch.Tensor:
        if A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo().coalesce()
        else:
            A_coo = A.coalesce() if A.is_sparse else A

        indices = A_coo.indices()
        values = A_coo.values().abs()
        n = A.shape[0]

        row_sums = torch.zeros(n, dtype=torch.float64, device=A.device)
        row_sums.scatter_add_(0, indices[0], values)

        col_sums = torch.zeros(n, dtype=torch.float64, device=A.device)
        col_sums.scatter_add_(0, indices[1], values)

        gamma = min(row_sums.max().item(), col_sums.max().item())

        scaled_values = A_coo.values() / gamma
        scaled_A = torch.sparse_coo_tensor(indices, scaled_values, A.shape)

        return scaled_A.coalesce().to_sparse_csc()

    def _extract_diagonal(self, A: torch.Tensor) -> torch.Tensor:
        if A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo()
        else:
            A_coo = A

        indices = A_coo.indices()
        values = A_coo.values()
        n = A.shape[0]

        diag_mask = indices[0] == indices[1]
        diag_values = torch.zeros(n, dtype=torch.float64, device=A.device)
        diag_values[indices[0, diag_mask]] = values[diag_mask]

        return diag_values

    def forward(self, b: torch.Tensor, diag_A: torch.Tensor,
                context_pairs: torch.Tensor) -> torch.Tensor:
        n, batch_size = b.shape

        if self.scale_input:
            scaling = torch.linalg.vector_norm(b, dim=0) / np.sqrt(n)
            scaling = scaling.clamp(min=1e-8)
            b = b / scaling

        context_aggregated = self.context_aggregator(context_pairs)

        b_expanded = b.unsqueeze(-1)
        diag_expanded = diag_A.view(n, 1, 1).expand(n, batch_size, 1)
        context_expanded = context_aggregated.view(n, 1, -1).expand(n, batch_size, -1)

        x = torch.cat([b_expanded, diag_expanded, context_expanded], dim=-1)
        x = self.mlp_initial(x)

        for i in range(self.num_layers):
            x = self.gconv[i](x) + self.skip[i](x)
            x = x.view(n * batch_size, self.embed)
            x = self.batch_norms[i](x)
            x = x.view(n, batch_size, self.embed)
            x = self.dropout(F.relu(x))

        z = self.mlp_final(x)
        z = z.view(n, batch_size)

        if self.scale_input:
            z = z * scaling

        return z

    def _extract_diagonal(self, A: torch.Tensor) -> torch.Tensor:
        if A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo()
        else:
            A_coo = A

        indices = A_coo.indices()
        values = A_coo.values()
        n = A.shape[0]

        diag_mask = indices[0] == indices[1]
        diag_values = torch.zeros(n, dtype=torch.float64, device=A.device)
        diag_values[indices[0, diag_mask]] = values[diag_mask]

        return diag_values


class BaselineResGCN(nn.Module):

    def __init__(self, num_layers: int, embed: int, hidden: int,
                 drop_rate: float, scale_input: bool = True,
                 dtype: torch.dtype = torch.float32):
        super().__init__()
        self.dtype = dtype
        self.num_layers = num_layers
        self.embed = embed
        self.hidden = hidden
        self.scale_input = scale_input

        self.mlp_initial = MLP(1, embed, 4, hidden, drop_rate)
        self.mlp_final = MLP(embed, 1, 4, hidden, drop_rate, is_output_layer=True)

        self.gconv = nn.ModuleList()
        self.skip = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        for _ in range(num_layers):
            self.gconv.append(GCNConv(embed, embed))
            self.skip.append(nn.Linear(embed, embed))
            self.batch_norms.append(nn.BatchNorm1d(embed))

        self.dropout = nn.Dropout(drop_rate)
        self.AA = None

    def set_matrix(self, A: torch.Tensor):
        self.AA = self._scale_by_spectral_radius(A).to(self.dtype)
        for layer in self.gconv:
            layer.set_adjacency(self.AA)

    def _scale_by_spectral_radius(self, A: torch.Tensor) -> torch.Tensor:
        if A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo().coalesce()
        else:
            A_coo = A.coalesce() if A.is_sparse else A

        indices = A_coo.indices()
        values = A_coo.values().abs()
        n = A.shape[0]

        row_sums = torch.zeros(n, dtype=torch.float64, device=A.device)
        row_sums.scatter_add_(0, indices[0], values)

        col_sums = torch.zeros(n, dtype=torch.float64, device=A.device)
        col_sums.scatter_add_(0, indices[1], values)

        gamma = min(row_sums.max().item(), col_sums.max().item())

        scaled_values = A_coo.values() / gamma
        scaled_A = torch.sparse_coo_tensor(indices, scaled_values, A.shape)

        return scaled_A.coalesce().to_sparse_csc()

    def forward(self, b: torch.Tensor) -> torch.Tensor:
        n, batch_size = b.shape

        if self.scale_input:
            scaling = torch.linalg.vector_norm(b, dim=0) / np.sqrt(n)
            scaling = scaling.clamp(min=1e-8)
            b = b / scaling

        x = b.unsqueeze(-1)
        x = self.mlp_initial(x)

        for i in range(self.num_layers):
            x = self.gconv[i](x) + self.skip[i](x)
            x = x.view(n * batch_size, self.embed)
            x = self.batch_norms[i](x)
            x = x.view(n, batch_size, self.embed)
            x = self.dropout(F.relu(x))

        z = self.mlp_final(x)
        z = z.view(n, batch_size)

        if self.scale_input:
            z = z * scaling

        return z


class ContextResMPNN(nn.Module):

    def __init__(self, num_layers: int, embed: int, hidden: int,
                 drop_rate: float, num_context_pairs: int,
                 scale_input: bool = True, dtype: torch.dtype = torch.float32):
        super().__init__()
        self.dtype = dtype
        self.num_layers = num_layers
        self.embed = embed
        self.hidden = hidden
        self.scale_input = scale_input
        self.num_context_pairs = num_context_pairs

        context_embed_dim = embed // 2

        self.context_aggregator = ContextAggregator(
            context_dim=2,
            output_dim=context_embed_dim,
            hidden=hidden
        )

        input_dim = 1 + 1 + context_embed_dim

        self.mlp_initial = MLP(input_dim, embed, 4, hidden, drop_rate)
        self.mlp_final = MLP(embed, 1, 4, hidden, drop_rate, is_output_layer=True)

        self.mpnn = nn.ModuleList()
        self.skip = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        for _ in range(num_layers):
            self.mpnn.append(MPNNConv(embed, embed))
            self.skip.append(nn.Linear(embed, embed))
            self.batch_norms.append(nn.BatchNorm1d(embed))

        self.dropout = nn.Dropout(drop_rate)

    def set_matrix(self, A: torch.Tensor):
        if A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo().coalesce()
        else:
            A_coo = A.coalesce() if A.is_sparse else A

        indices = A_coo.indices()
        values = A_coo.values()
        n = A.shape[0]

        row_abs = torch.zeros(n, dtype=torch.float64, device=A.device)
        row_abs.scatter_add_(0, indices[0], values.abs())
        col_abs = torch.zeros(n, dtype=torch.float64, device=A.device)
        col_abs.scatter_add_(0, indices[1], values.abs())
        gamma = min(row_abs.max().item(), col_abs.max().item())

        scaled_values = (values / gamma).to(self.dtype)

        for layer in self.mpnn:
            layer.set_edges(indices, scaled_values, n)

    def forward(self, b: torch.Tensor, diag_A: torch.Tensor,
                context_pairs: torch.Tensor) -> torch.Tensor:
        n, batch_size = b.shape

        if self.scale_input:
            scaling = torch.linalg.vector_norm(b, dim=0) / np.sqrt(n)
            scaling = scaling.clamp(min=1e-8)
            b = b / scaling

        context_aggregated = self.context_aggregator(context_pairs)

        b_expanded = b.unsqueeze(-1)
        diag_expanded = diag_A.view(n, 1, 1).expand(n, batch_size, 1)
        context_expanded = context_aggregated.view(n, 1, -1).expand(n, batch_size, -1)

        x = torch.cat([b_expanded, diag_expanded, context_expanded], dim=-1)
        x = self.mlp_initial(x)

        for i in range(self.num_layers):
            x = self.mpnn[i](x) + self.skip[i](x)
            x = x.view(n * batch_size, self.embed)
            x = self.batch_norms[i](x)
            x = x.view(n, batch_size, self.embed)
            x = self.dropout(F.relu(x))

        z = self.mlp_final(x)
        z = z.view(n, batch_size)

        if self.scale_input:
            z = z * scaling

        return z
