import argparse
import sys
from pathlib import Path

import numpy as np
import torch

import clamp

# Make shared helpers at project root importable when running from models/<MODEL>/.
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 _call_encoder(obj, molecules):
    if hasattr(obj, "forward"):
        return obj.forward(molecules)
    return obj(molecules)


def get_molecule_embeddings(model, smiles_list):
    """Return (N, d) torch.Tensor of CLAMP molecule embeddings."""
    candidates = [
        "encode_molecules",
        "encode_compounds",
        "forward_molecules",
        "molecule_encoder",
        "compound_encoder",
        "mol_encoder",
    ]
    with torch.no_grad():
        for name in candidates:
            if hasattr(model, name):
                enc = getattr(model, name)
                try:
                    out = _call_encoder(enc, smiles_list)
                except TypeError:
                    continue
                if isinstance(out, torch.Tensor) and out.dim() == 2:
                    return out
    avail = [n for n in dir(model) if ("encode" in n.lower() or "encoder" in n.lower())]
    raise AttributeError(
        f"Could not find a molecule encoder on CLAMP model. "
        f"Candidates tried: {candidates}. Available: {avail}"
    )


def main():
    parser = argparse.ArgumentParser(description="Encode molecules with CLAMP 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("--device", default="cpu", choices=["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)

    smiles_in_order = [record["smiles"] for record in records]
    if not smiles_in_order:
        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)

    model = clamp.CLAMP(device=args.device)
    model.eval()

    with torch.no_grad():
        emb = get_molecule_embeddings(model, smiles_in_order)
        if args.normalize:
            emb = torch.nn.functional.normalize(emb, dim=-1)
        emb_np = emb.detach().cpu().numpy()

    np.save(npy_path, emb_np)

    write_mapping_csv(map_path, records)

    if failures:
        write_failures_csv(fail_path, failures)

    print(f"Saved embeddings to {npy_path} with shape {emb_np.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()
