import torch
import torch.nn as nn
import dgl
from dgllife.model import GCN, GAT

class MolEnc(nn.Module):

    def __init__(self,
                 args,
                 in_dim,):
        super().__init__()

        # Whether to return node-level embeddings (sum_nodes, D) instead of a
        # pooled graph embedding (B, D).
        #
        # Backward compatible defaults:
        # - Historically, this was inferred from `args.model` for FLARE-style models.
        # - New: allow explicit override via `args.mol_return_emb` (preferred) or
        #   `args.return_mol_emb` (legacy-friendly alias).
        explicit = None
        if hasattr(args, "mol_return_emb"):
            explicit = getattr(args, "mol_return_emb")
        elif hasattr(args, "return_mol_emb"):
            explicit = getattr(args, "return_mol_emb")

        if explicit is not None:
            self.return_emb = bool(explicit)
        else:
            self.return_emb = False
            if args.model in ("filipContrastive", "crossAttenContrastive", "filipGlobalContrastive"):
                self.return_emb = True

        dropout = [args.gnn_dropout for _ in range(len(args.gnn_channels))]
        batchnorm = [True for _ in range(len(args.gnn_channels))]
        gnn_map = {
            "gcn": GCN(in_dim, args.gnn_channels, batchnorm = batchnorm, dropout = dropout),
        }
        self.GNN = gnn_map[args.gnn_type]
        pooling_type = getattr(args, "mol_pooling", "max")
        pool_map = {
            "max": dgl.nn.pytorch.glob.MaxPooling,
            "mean": dgl.nn.pytorch.glob.AvgPooling,
            "avg": dgl.nn.pytorch.glob.AvgPooling,
            "sum": dgl.nn.pytorch.glob.SumPooling,
        }
        if pooling_type not in pool_map:
            raise ValueError(f"Unsupported mol_pooling: {pooling_type}")
        self.pool = pool_map[pooling_type]()
        
        if not self.return_emb:
            self.fc1_graph = nn.Linear(args.gnn_channels[len(args.gnn_channels) - 1], args.gnn_hidden_dim * 2)
            self.fc2_graph = nn.Linear(args.gnn_hidden_dim * 2, args.final_embedding_dim)

            self.dropout = nn.Dropout(args.fc_dropout)
            self.relu = nn.ReLU()

    def forward(self, g) -> torch.Tensor:
        g1 = g
        f1 = g.ndata['h']

        f = self.GNN(g1, f1)
        if self.return_emb:
            return f
        h = self.pool(g1, f)
        h1 = self.relu(self.fc1_graph(h))
        h1 = self.dropout(h1)
        h1 = self.fc2_graph(h1)
        h1 = self.dropout(h1)

        return h1