import argparse
import os
import sys
from pathlib import Path
from typing import List

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

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_HM_CACHE = SCRIPT_DIR / ".cache" / "huggingmolecules"
# huggingmolecules snapshots cache path at import time, so set it before imports.
os.environ.setdefault("HUGGINGMOLECULES_CACHE", str(DEFAULT_HM_CACHE))

try:
    from huggingmolecules import MatFeaturizer, MatModel  # type: ignore
except Exception as exc:
    raise ImportError(
        "Could not import MAT classes from huggingmolecules. "
        "Ensure MAT 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 _build_model(pretrained_name: str, device: torch.device):
    featurizer = MatFeaturizer.from_pretrained(pretrained_name)
    model = MatModel.from_pretrained(pretrained_name).to(device).eval()
    return featurizer, model


def _pool_if_contextual(emb: torch.Tensor, batch) -> torch.Tensor:
    if emb.dim() != 3:
        return emb
    mask = getattr(batch, "mask", None)
    if mask is None:
        raise RuntimeError("MAT returned contextual embeddings but batch has no mask.")
    if mask.dim() == 3:
        mask = mask.squeeze(-1)
    m = mask.to(emb.dtype).unsqueeze(-1)
    return (emb * m).sum(1) / m.sum(1).clamp(min=1e-6)


@torch.inference_mode()
def _encode_smiles_batches(
    smiles_inputs: List[str],
    model,
    featurizer,
    device: torch.device,
    batch_size: int,
    normalize: bool,
) -> np.ndarray:
    pred_head = getattr(model, "generator", None)
    if pred_head is None or not hasattr(pred_head, "proj"):
        raise RuntimeError("Could not locate MAT generator.proj for embedding hook.")

    captured = {"emb": None}

    def _proj_hook(_module, inputs, _output):
        captured["emb"] = inputs[0].detach()

    handle = pred_head.proj.register_forward_hook(_proj_hook)
    chunks = []
    try:
        for start in range(0, len(smiles_inputs), batch_size):
            batch_smiles = smiles_inputs[start : start + batch_size]
            batch = featurizer(batch_smiles)
            if hasattr(batch, "to"):
                batch = batch.to(device)

            captured["emb"] = None
            _ = model(batch)

            emb = captured["emb"]
            if emb is None:
                raise RuntimeError("Embedding hook did not capture MAT representations.")

            emb = _pool_if_contextual(emb, batch)
            if normalize:
                emb = F.normalize(emb, dim=-1)

            chunks.append(emb.cpu().numpy().astype(np.float32, copy=False))
    finally:
        handle.remove()

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


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode molecules with MAT 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="selfies",
        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="mat_masking_20M",
        help="HuggingMolecules pretrained key for MAT weights",
    )
    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)

    # MAT featurizer consumes SMILES; preprocess_lines already provides canonical SMILES.
    smiles_inputs = [record["smiles"] for record in records]

    device = _resolve_device(args.device)
    featurizer, model = _build_model(args.pretrained_name, device)
    embeddings = _encode_smiles_batches(
        smiles_inputs=smiles_inputs,
        model=model,
        featurizer=featurizer,
        device=device,
        batch_size=args.batch_size,
        normalize=args.normalize,
    )

    np.save(npy_path, embeddings)
    write_mapping_csv(map_path, 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}.")
    if failures:
        print(f"Filtered {len(failures)} invalid rows. See {fail_path}.")


if __name__ == "__main__":
    main()
