#!/usr/bin/env python3
"""
Input
--------------
A PyG HeteroData graph containing:
  graph["protein"].x                       [num_proteins, 1280]
  graph["orthogroup"].node_id              [num_orthogroups_in_graph]
  graph[edge_type].edge_index              [2, E] for each edge type

Call
--------
# Inspect/load only:
python inference.py --weights model.safetensors --config config.json

# Score one existing protein-protein pair by node index:
python inference.py --weights model.safetensors --config config.json \
    --graph inference_graph.pt --edge 123 456

# Score several pairs:
python inference.py --weights model.safetensors --config config.json \
    --graph inference_graph.pt --edges query_edges.tsv
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import torch
import torch.nn as nn
from safetensors.torch import load_file
from torch_geometric.nn import HGTConv


def build_pair_features(z: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
    """Pair feature order used by the ProtLink MLP decoder."""
    zu = z[edges[:, 0]]
    zv = z[edges[:, 1]]
    return torch.cat([zu, zv, torch.abs(zu - zv), zu * zv], dim=-1)


class PairMLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 512, dropout: float = 0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.net(features).squeeze(-1)


class HGTEncoder(nn.Module):
    def __init__(
        self,
        protein_input_dim: int,
        hidden_dim: int,
        heads: int,
        layers: int,
        dropout: float,
        metadata,
        node_type_counts: dict[str, int],
    ):
        super().__init__()
        self.protein_proj = nn.Linear(protein_input_dim, hidden_dim)

        self.node_type_embeddings = nn.ModuleDict()
        for node_type, count in sorted(node_type_counts.items()):
            if node_type == "protein":
                continue
            self.node_type_embeddings[node_type] = nn.Embedding(max(1, count), hidden_dim)

        self.convs = nn.ModuleList(
            [HGTConv(hidden_dim, hidden_dim, metadata, heads=heads) for _ in range(layers)]
        )
        self.norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(layers)])
        self.dropout = nn.Dropout(dropout)
        self.output_dim = hidden_dim

        for module in self.node_type_embeddings.values():
            nn.init.normal_(module.weight, std=1.0 / math.sqrt(hidden_dim))

    def forward(self, data):
        x_dict = {"protein": self.protein_proj(data["protein"].x)}
        for node_type, embedding in self.node_type_embeddings.items():
            x_dict[node_type] = embedding(data[node_type].node_id)

        for conv, norm in zip(self.convs, self.norms):
            out_dict = conv(x_dict, data.edge_index_dict)
            next_x_dict = {}
            for node_type, out in out_dict.items():
                out = norm(out)
                out = torch.relu(out)
                out = self.dropout(out)
                if out.shape[-1] == x_dict[node_type].shape[-1]:
                    out = out + x_dict[node_type]
                next_x_dict[node_type] = out
            x_dict = next_x_dict

        return x_dict["protein"]


class PairMLPDecoder(nn.Module):
    def __init__(self, latent_dim: int, hidden_dim: int, dropout: float):
        super().__init__()
        self.mlp = PairMLP(latent_dim * 4, hidden_dim=hidden_dim, dropout=dropout)

    def forward(self, z: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
        return self.mlp(build_pair_features(z, edges))


class ProtLinkModel(nn.Module):
    def __init__(self, config: dict):
        super().__init__()
        a = config["architecture"]
        g = config["graph_schema"]

        node_types = list(g["node_types"])
        edge_types = [tuple(x) for x in g["edge_types"]]
        metadata = (node_types, edge_types)

        self.encoder = HGTEncoder(
            protein_input_dim=int(a["protein_input_dim"]),
            hidden_dim=int(a["hidden_dim"]),
            heads=int(a["hgt_heads"]),
            layers=int(a["hgt_layers"]),
            dropout=float(a["hgt_dropout"]),
            metadata=metadata,
            node_type_counts={"orthogroup": int(g["num_orthogroups"])},
        )
        self.decoder = PairMLPDecoder(
            latent_dim=int(a["hidden_dim"]),
            hidden_dim=int(a["decoder_hidden_dim"]),
            dropout=float(a["decoder_dropout"]),
        )

    def encode(self, graph):
        return self.encoder(graph)

    def forward(self, graph, query_edges: torch.Tensor):
        z = self.encode(graph)
        return self.decoder(z, query_edges)


def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_model(config_path: str, weights_path: str, device: torch.device) -> tuple[ProtLinkModel, dict]:
    config = load_config(config_path)
    model = ProtLinkModel(config)
    state = load_file(weights_path, device="cpu")
    model.load_state_dict(state, strict=True)
    model.to(device).eval()
    return model, config


def canonicalize_edges(edges: torch.Tensor) -> torch.Tensor:
    """Training PPIs were treated as undirected; keep a stable u<v ordering."""
    lo = torch.minimum(edges[:, 0], edges[:, 1])
    hi = torch.maximum(edges[:, 0], edges[:, 1])
    return torch.stack([lo, hi], dim=1)


def read_query_edges(path: str) -> torch.Tensor:
    rows = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.replace(",", "\t").split()
            if len(parts) < 2:
                raise ValueError(f"Expected two integer node indices: {line}")
            rows.append((int(parts[0]), int(parts[1])))
    if not rows:
        raise ValueError("No query edges found.")
    return torch.tensor(rows, dtype=torch.long)


def validate_graph(graph, config: dict):
    g = config["graph_schema"]
    a = config["architecture"]

    expected_node_types = set(g["node_types"])
    expected_edge_types = {tuple(x) for x in g["edge_types"]}

    if set(graph.node_types) != expected_node_types:
        raise ValueError(
            f"Node types mismatch: got {graph.node_types}, expected {sorted(expected_node_types)}"
        )
    if set(graph.edge_types) != expected_edge_types:
        raise ValueError(
            f"Edge types mismatch: got {graph.edge_types}, expected {sorted(expected_edge_types)}"
        )
    if graph["protein"].x.ndim != 2 or graph["protein"].x.shape[1] != int(a["protein_input_dim"]):
        raise ValueError(
            f'graph["protein"].x must have shape [N, {a["protein_input_dim"]}]'
        )
    if "node_id" not in graph["orthogroup"]:
        raise ValueError('graph["orthogroup"].node_id is required.')
    if int(graph["orthogroup"].node_id.max()) >= int(g["num_orthogroups"]):
        raise ValueError(
            "The graph contains an orthogroup node_id outside the checkpoint embedding table."
        )


@torch.inference_mode()
def score_edges(model, graph, edges: torch.Tensor, threshold: float):
    edges = canonicalize_edges(edges)
    logits = model(graph, edges)
    probs = torch.sigmoid(logits)
    preds = probs >= threshold
    return logits, probs, preds


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--config", default="config.json")
    p.add_argument("--weights", default="model.safetensors")
    p.add_argument("--graph", default=None, help="Trusted torch-saved PyG HeteroData.")
    p.add_argument("--edge", nargs=2, type=int, metavar=("U", "V"))
    p.add_argument("--edges", default=None, help="TSV/whitespace file with two protein indices per row.")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = p.parse_args()

    device = torch.device(args.device)
    model, config = load_model(args.config, args.weights, device)
    print(
        f"Loaded {args.weights} strictly: "
        f"HGT hidden={config['architecture']['hidden_dim']}, "
        f"layers={config['architecture']['hgt_layers']}, "
        f"heads={config['architecture']['hgt_heads']}."
    )

    if args.graph is None:
        print("No --graph supplied; model construction and strict weight loading succeeded.")
        return

    if args.edge is None and args.edges is None:
        raise SystemExit("Provide --edge U V or --edges query_edges.tsv when using --graph.")
    if args.edge is not None and args.edges is not None:
        raise SystemExit("Use only one of --edge or --edges.")

    # HeteroData is a Python object; only load graph files you trust.
    graph = torch.load(args.graph, map_location="cpu", weights_only=False)
    validate_graph(graph, config)
    graph = graph.to(device)

    if args.edge is not None:
        query_edges = torch.tensor([args.edge], dtype=torch.long, device=device)
    else:
        query_edges = read_query_edges(args.edges).to(device)

    num_proteins = int(graph["protein"].num_nodes)
    if query_edges.min() < 0 or query_edges.max() >= num_proteins:
        raise ValueError(f"Query edge index outside [0, {num_proteins - 1}].")

    threshold = float(config["checkpoint"]["decision_threshold"])
    logits, probs, preds = score_edges(model, graph, query_edges, threshold)

    for (u, v), logit, prob, pred in zip(
        canonicalize_edges(query_edges).cpu().tolist(),
        logits.cpu().tolist(),
        probs.cpu().tolist(),
        preds.cpu().tolist(),
    ):
        print(
            json.dumps(
                {
                    "protein_u_index": int(u),
                    "protein_v_index": int(v),
                    "logit": float(logit),
                    "probability": float(prob),
                    "threshold": threshold,
                    "predicted_interaction": bool(pred),
                }
            )
        )


if __name__ == "__main__":
    main()