import argparse
import os
import sys
from collections import deque
from pathlib import Path
from typing import Dict, List, Optional, Tuple

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

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_HF_HOME = SCRIPT_DIR / ".cache" / "huggingface"
DEFAULT_TRANSFORMERS_CACHE = DEFAULT_HF_HOME / "transformers"
DEFAULT_LOCAL_MODEL_DIR = DEFAULT_HF_HOME / "graphormer-base-pcqm4mv2"
INF_SPATIAL_POS = 510

# Ensure huggingface cache is model-local by default.
os.environ.setdefault("HF_HOME", str(DEFAULT_HF_HOME))
os.environ.setdefault("TRANSFORMERS_CACHE", str(DEFAULT_TRANSFORMERS_CACHE))
os.environ.setdefault("HF_HUB_CACHE", str(DEFAULT_HF_HOME / "hub"))
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
os.environ.setdefault("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1")

try:
    from transformers import GraphormerForGraphClassification  # type: ignore
except Exception as exc:
    raise ImportError(
        "Could not import GraphormerForGraphClassification from transformers. "
        "Ensure GRAPHORMER setup completed successfully."
    ) from exc

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from embeddings_preprocess import (
    load_input_lines,
    preprocess_lines,
    resolve_output_paths,
    write_failures_csv,
    write_mapping_csv,
)


def _resolve_device(device_arg: str) -> torch.device:
    normalized = device_arg.strip().lower()
    if normalized == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if normalized not in {"cpu", "cuda"}:
        raise ValueError(f"Unsupported device: {device_arg}")
    if normalized == "cuda" and not torch.cuda.is_available():
        print("CUDA requested but unavailable; falling back to CPU.", file=sys.stderr)
        return torch.device("cpu")
    return torch.device(normalized)


def _resolve_model_source(model_path: str, pretrained_name: str) -> Tuple[str, bool]:
    explicit = model_path.strip()
    if explicit:
        candidate = Path(explicit).expanduser()
        if candidate.exists():
            return str(candidate.resolve()), True
        return explicit, False

    if DEFAULT_LOCAL_MODEL_DIR.exists():
        return str(DEFAULT_LOCAL_MODEL_DIR.resolve()), True

    return pretrained_name, False


def _build_model(model_source: str, local_files_only: bool, device: torch.device):
    # Checkpoint is GraphormerForGraphClassification; we use its encoder hidden states
    # as molecule embeddings and ignore classifier head mismatch if present.
    model = GraphormerForGraphClassification.from_pretrained(
        model_source,
        local_files_only=local_files_only,
        ignore_mismatched_sizes=True,
    ).to(device).eval()
    return model


def _convert_to_single_emb(x: np.ndarray, offset: int = 512) -> np.ndarray:
    feature_num = x.shape[1] if len(x.shape) > 1 else 1
    feature_offset = 1 + np.arange(0, feature_num * offset, offset, dtype=np.int64)
    return x + feature_offset


def _bfs_dist_prev(adj_list: List[List[int]], src: int) -> Tuple[List[int], List[int]]:
    n = len(adj_list)
    dist = [INF_SPATIAL_POS] * n
    prev = [-1] * n
    dist[src] = 0
    q: deque[int] = deque([src])

    while q:
        u = q.popleft()
        for v in adj_list[u]:
            if dist[v] == INF_SPATIAL_POS:
                dist[v] = dist[u] + 1
                prev[v] = u
                q.append(v)

    return dist, prev


def _reconstruct_path(prev: List[int], src: int, dst: int) -> Optional[List[int]]:
    if src == dst:
        return [src]
    if prev[dst] == -1:
        return None

    path = [dst]
    cur = dst
    while cur != src:
        cur = prev[cur]
        if cur == -1:
            return None
        path.append(cur)
    path.reverse()
    return path


def _smiles_to_processed_item(smiles: str) -> Dict[str, np.ndarray]:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Invalid SMILES for Graphormer: {smiles}")

    num_nodes = mol.GetNumAtoms()
    node_feat = np.array([[atom.GetAtomicNum()] for atom in mol.GetAtoms()], dtype=np.int64)

    edge_feat_size = 1
    edge_index_i: List[int] = []
    edge_index_j: List[int] = []
    edge_attr: List[List[int]] = []

    for bond in mol.GetBonds():
        i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        bond_type = int(bond.GetBondTypeAsDouble())
        edge_index_i += [i, j]
        edge_index_j += [j, i]
        edge_attr += [[bond_type], [bond_type]]

    if edge_attr:
        edge_attr_arr = np.asarray(edge_attr, dtype=np.int64)
    else:
        edge_attr_arr = np.empty((0, edge_feat_size), dtype=np.int64)

    attn_edge_type = np.zeros((num_nodes, num_nodes, edge_feat_size), dtype=np.int64)
    if edge_attr_arr.shape[0] > 0:
        edge_vals = _convert_to_single_emb(edge_attr_arr) + 1
        for idx in range(len(edge_index_i)):
            attn_edge_type[edge_index_i[idx], edge_index_j[idx]] = edge_vals[idx]

    adj_bool = np.zeros((num_nodes, num_nodes), dtype=bool)
    adj_list: List[List[int]] = [[] for _ in range(num_nodes)]
    for i, j in zip(edge_index_i, edge_index_j):
        adj_bool[i, j] = True
        adj_list[i].append(j)

    spatial_pos = np.full((num_nodes, num_nodes), INF_SPATIAL_POS, dtype=np.int64)
    prev_all: List[List[int]] = []
    max_dist = 0
    for src in range(num_nodes):
        dist, prev = _bfs_dist_prev(adj_list, src)
        prev_all.append(prev)
        for dst, d in enumerate(dist):
            spatial_pos[src, dst] = d
            if d != INF_SPATIAL_POS and d > max_dist:
                max_dist = d

    max_dist = max(1, int(max_dist))
    input_edges = np.zeros((num_nodes, num_nodes, max_dist, edge_feat_size), dtype=np.int64)

    for src in range(num_nodes):
        prev = prev_all[src]
        for dst in range(num_nodes):
            d = spatial_pos[src, dst]
            if d == INF_SPATIAL_POS or d <= 0:
                continue
            path = _reconstruct_path(prev, src, dst)
            if not path or len(path) < 2:
                continue
            for hop, (u, v) in enumerate(zip(path[:-1], path[1:])):
                if hop >= max_dist:
                    break
                input_edges[src, dst, hop, :] = attn_edge_type[u, v, :]

    input_nodes = _convert_to_single_emb(node_feat) + 1
    attn_bias = np.zeros((num_nodes + 1, num_nodes + 1), dtype=np.float32)
    in_degree = np.sum(adj_bool, axis=1).astype(np.int64)

    item = {
        "input_nodes": input_nodes + 1,
        "attn_bias": attn_bias,
        "attn_edge_type": attn_edge_type,
        "spatial_pos": spatial_pos.astype(np.int64) + 1,
        "in_degree": in_degree + 1,
        "out_degree": in_degree + 1,
        "input_edges": input_edges + 1,
    }
    return item


def _collate_graphormer(features: List[Dict[str, np.ndarray]], spatial_pos_max: int = 20) -> Dict[str, torch.Tensor]:
    max_node_num = max(f["input_nodes"].shape[0] for f in features)
    node_feat_size = features[0]["input_nodes"].shape[1]
    edge_feat_size = features[0]["attn_edge_type"].shape[2]
    max_dist = max(f["input_edges"].shape[2] for f in features)
    edge_input_size = features[0]["input_edges"].shape[3]
    batch_size = len(features)

    batch = {
        "attn_bias": torch.zeros(batch_size, max_node_num + 1, max_node_num + 1, dtype=torch.float),
        "attn_edge_type": torch.zeros(batch_size, max_node_num, max_node_num, edge_feat_size, dtype=torch.long),
        "spatial_pos": torch.zeros(batch_size, max_node_num, max_node_num, dtype=torch.long),
        "in_degree": torch.zeros(batch_size, max_node_num, dtype=torch.long),
        "input_nodes": torch.zeros(batch_size, max_node_num, node_feat_size, dtype=torch.long),
        "input_edges": torch.zeros(
            batch_size,
            max_node_num,
            max_node_num,
            max_dist,
            edge_input_size,
            dtype=torch.long,
        ),
    }

    for ix, f in enumerate(features):
        attn_bias = torch.tensor(f["attn_bias"], dtype=torch.float)
        attn_edge_type = torch.tensor(f["attn_edge_type"], dtype=torch.long)
        spatial_pos = torch.tensor(f["spatial_pos"], dtype=torch.long)
        in_degree = torch.tensor(f["in_degree"], dtype=torch.long)
        input_nodes = torch.tensor(f["input_nodes"], dtype=torch.long)
        input_edges = torch.tensor(f["input_edges"], dtype=torch.long)

        if torch.any(spatial_pos >= spatial_pos_max):
            attn_bias[1:, 1:][spatial_pos >= spatial_pos_max] = float("-inf")

        batch["attn_bias"][ix, : attn_bias.shape[0], : attn_bias.shape[1]] = attn_bias
        batch["attn_edge_type"][ix, : attn_edge_type.shape[0], : attn_edge_type.shape[1], :] = attn_edge_type
        batch["spatial_pos"][ix, : spatial_pos.shape[0], : spatial_pos.shape[1]] = spatial_pos
        batch["in_degree"][ix, : in_degree.shape[0]] = in_degree
        batch["input_nodes"][ix, : input_nodes.shape[0], :] = input_nodes
        batch["input_edges"][
            ix,
            : input_edges.shape[0],
            : input_edges.shape[1],
            : input_edges.shape[2],
            :,
        ] = input_edges

    batch["out_degree"] = batch["in_degree"]
    return batch


@torch.inference_mode()
def _encode_graphs(
    processed_items: List[Dict[str, np.ndarray]],
    model,
    device: torch.device,
    batch_size: int,
    normalize: bool,
) -> np.ndarray:
    chunks = []

    for start in range(0, len(processed_items), batch_size):
        batch_items = processed_items[start : start + batch_size]
        batch = _collate_graphormer(batch_items)
        batch = {k: v.to(device) for k, v in batch.items()}

        outputs = model.encoder(
            input_nodes=batch["input_nodes"],
            input_edges=batch["input_edges"],
            attn_bias=batch["attn_bias"],
            in_degree=batch["in_degree"],
            out_degree=batch["out_degree"],
            spatial_pos=batch["spatial_pos"],
            attn_edge_type=batch["attn_edge_type"],
        )

        emb = outputs.last_hidden_state[:, 0, :]
        if normalize:
            emb = F.normalize(emb, dim=-1)
        chunks.append(emb.cpu().numpy().astype(np.float32, copy=False))

    if not chunks:
        raise RuntimeError("No embeddings were produced by GRAPHORMER.")
    return np.concatenate(chunks, axis=0)


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode molecules with GRAPHORMER and save embeddings")
    parser.add_argument("--file", required=True, help="Path to input .txt file")
    parser.add_argument(
        "--input-format",
        default="selfies",
        choices=["selfies", "smiles"],
        help="Format of rows in input file",
    )
    parser.add_argument(
        "--model-input-type",
        default="smiles",
        choices=["selfies", "smiles"],
        help="Expected model input type from metadata",
    )
    parser.add_argument("--output-dir", default=".", help="Directory to write outputs")
    parser.add_argument("--output-stem", default="embeddings", help="Base output filename stem")
    parser.add_argument(
        "--pretrained-name",
        default="clefourrier/graphormer-base-pcqm4mv2",
        help="Hugging Face repo id or local model path",
    )
    parser.add_argument(
        "--model-path",
        default="",
        help="Optional explicit local model directory (overrides default cache lookup)",
    )
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument(
        "--device",
        default="auto",
        choices=["auto", "cpu", "cuda"],
        help="Inference device",
    )
    parser.add_argument(
        "--no-normalize",
        dest="normalize",
        action="store_false",
        help="Disable L2-normalization before saving",
    )
    parser.set_defaults(normalize=True)
    args = parser.parse_args()

    input_path = Path(args.file).resolve()
    if not input_path.exists():
        print(f"Input file not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    output_dir = Path(args.output_dir).resolve()
    npy_path, map_path, fail_path = resolve_output_paths(output_dir, args.output_stem)

    lines = load_input_lines(input_path)
    records, failures = preprocess_lines(lines, args.input_format, args.model_input_type)
    if not records:
        print("No valid molecules after circularity validation. Aborting.", file=sys.stderr)
        if failures:
            print(f"Failures: {len(failures)} (see {fail_path})", file=sys.stderr)
            write_failures_csv(fail_path, failures)
        sys.exit(1)

    processed_items: List[Dict[str, np.ndarray]] = []
    filtered_records = []
    for record in records:
        try:
            processed_items.append(_smiles_to_processed_item(record["model_input"]))
            filtered_records.append(record)
        except Exception:
            failures.append(
                {
                    "original_line": record["original_line"],
                    "input_text": record["input_text"],
                    "reason": "graph_preprocess_failed",
                }
            )

    if not processed_items:
        print("No valid graph inputs after Graphormer preprocessing. Aborting.", file=sys.stderr)
        if failures:
            write_failures_csv(fail_path, failures)
        sys.exit(1)

    model_source, local_files_only = _resolve_model_source(args.model_path, args.pretrained_name)
    device = _resolve_device(args.device)
    model = _build_model(
        model_source=model_source,
        local_files_only=local_files_only,
        device=device,
    )

    embeddings = _encode_graphs(
        processed_items=processed_items,
        model=model,
        device=device,
        batch_size=args.batch_size,
        normalize=args.normalize,
    )

    np.save(npy_path, embeddings)
    write_mapping_csv(map_path, filtered_records)
    if failures:
        write_failures_csv(fail_path, failures)

    print(f"Saved embeddings to {npy_path} with shape {embeddings.shape}.")
    print(f"Saved row mapping to {map_path}.")
    print(f"Model source: {model_source} (local_files_only={local_files_only})")
    if failures:
        print(f"Filtered {len(failures)} invalid rows. See {fail_path}.")


if __name__ == "__main__":
    main()