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 GroverFeaturizer, GroverModel  # type: ignore
except Exception as exc:
    raise ImportError(
        "Could not import GROVER classes from huggingmolecules. "
        "Ensure GROVER 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_pretrained_name(variant: str, explicit_name: str) -> str:
    if explicit_name:
        return explicit_name
    normalized = variant.strip().lower()
    if normalized not in {"base", "large"}:
        raise ValueError(f"Unsupported GROVER variant: {variant}")
    return f"grover_{normalized}"


def _build_model(pretrained_name: str, device: torch.device):
    featurizer = GroverFeaturizer.from_pretrained(pretrained_name)
    model = GroverModel.from_pretrained(pretrained_name).to(device).eval()
    return featurizer, model


@torch.inference_mode()
def _encode_smiles_batches(
    smiles_inputs: List[str],
    model,
    featurizer,
    device: torch.device,
    batch_size: int,
    use_atoms: bool,
    use_bonds: bool,
    normalize: bool,
) -> np.ndarray:
    if not use_atoms and not use_bonds:
        raise ValueError("At least one GROVER view must be enabled: atoms and/or bonds.")

    captured = {"atoms": None, "bonds": None}

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

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

    atom_head = getattr(model, "mol_atom_from_atom_ffn", None)
    bond_head = getattr(model, "mol_atom_from_bond_ffn", None)
    if atom_head is None or bond_head is None:
        raise RuntimeError("GROVER model is missing expected FFN heads for atom/bond outputs.")

    hook_atoms = atom_head.register_forward_hook(_hook_atoms)
    hook_bonds = bond_head.register_forward_hook(_hook_bonds)

    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["atoms"] = None
            captured["bonds"] = None
            _ = model(batch)

            parts = []
            if use_atoms:
                if captured["atoms"] is None:
                    raise RuntimeError("Failed to capture GROVER atom embeddings.")
                parts.append(captured["atoms"])
            if use_bonds:
                if captured["bonds"] is None:
                    raise RuntimeError("Failed to capture GROVER bond embeddings.")
                parts.append(captured["bonds"])

            emb = parts[0] if len(parts) == 1 else torch.cat(parts, dim=1)
            if normalize:
                emb = F.normalize(emb, dim=-1)
            chunks.append(emb.cpu().numpy().astype(np.float32, copy=False))
    finally:
        hook_atoms.remove()
        hook_bonds.remove()

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


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode molecules with GROVER 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(
        "--variant",
        default="base",
        choices=["base", "large"],
        help="GROVER checkpoint variant",
    )
    parser.add_argument(
        "--pretrained-name",
        default="",
        help="Override checkpoint key (for example, grover_base or grover_large)",
    )
    parser.add_argument("--use-atoms", action="store_true", help="Include atom-view embeddings")
    parser.add_argument("--use-bonds", action="store_true", help="Include bond-view embeddings")
    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()

    # Conservative default keeps behavior stable if no view flag is passed.
    use_atoms = args.use_atoms or not args.use_bonds
    use_bonds = args.use_bonds

    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)

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

    pretrained_name = _resolve_pretrained_name(args.variant, args.pretrained_name)
    device = _resolve_device(args.device)
    featurizer, model = _build_model(pretrained_name, device)
    embeddings = _encode_smiles_batches(
        smiles_inputs=smiles_inputs,
        model=model,
        featurizer=featurizer,
        device=device,
        batch_size=args.batch_size,
        use_atoms=use_atoms,
        use_bonds=use_bonds,
        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}.")
    print(f"GROVER options: pretrained={pretrained_name}, atoms={use_atoms}, bonds={use_bonds}")
    if failures:
        print(f"Filtered {len(failures)} invalid rows. See {fail_path}.")


if __name__ == "__main__":
    main()
