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

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

if not str(torch.__version__).startswith("1.7.1"):
    raise RuntimeError(
        "MOLCLR requires torch==1.7.1+cpu. "
        "Found torch=%s. Recreate the MOLCLR .venv and run setup again." % torch.__version__
    )

try:
    from torch_geometric.data import Data, DataLoader
except Exception as exc:
    raise ImportError(
        "Could not import torch_geometric with the installed torch extension stack. "
        "Ensure torch==1.7.1+cpu and matching torch-scatter/torch-sparse wheels are installed."
    ) from exc


SCRIPT_DIR = Path(__file__).resolve().parent
REPO_DIR = SCRIPT_DIR / ".repo" / "MOLCLR"

if str(REPO_DIR) not in sys.path:
    sys.path.insert(0, str(REPO_DIR))

try:
    from models.gcn_molclr import GCN
    from models.ginet_molclr import GINet
except Exception as exc:
    raise ImportError(
        "Could not import MolCLR modules from the cloned repository. "
        "Ensure MOLCLR 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,
)

ATOM_LIST = list(range(1, 119))
CHIRALITY_LIST = [
    Chem.rdchem.ChiralType.CHI_UNSPECIFIED,
    Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
    Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
]
BOND_LIST = [BT.SINGLE, BT.DOUBLE, BT.TRIPLE, BT.AROMATIC]
BONDDIR_LIST = [
    Chem.rdchem.BondDir.NONE,
    Chem.rdchem.BondDir.ENDUPRIGHT,
    Chem.rdchem.BondDir.ENDDOWNRIGHT,
]


def _safe_index(value, choices, default_idx=0):
    try:
        return choices.index(value)
    except ValueError:
        return default_idx


def _resolve_device(device_arg):
    normalized = str(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("Unsupported device: %s" % 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 _normalize_variant(variant):
    normalized = str(variant).strip().lower()
    if normalized not in {"gin", "gcn"}:
        raise ValueError("Unknown MolCLR variant '%s'. Expected gin or gcn." % variant)
    return normalized


def mol_to_graph_data_obj_simple(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError("Invalid SMILES: %s" % smiles)

    type_idx = []
    chirality_idx = []
    for atom in mol.GetAtoms():
        atom_num = atom.GetAtomicNum()
        if atom_num not in ATOM_LIST:
            raise ValueError("Unsupported atomic number %s for MolCLR featurization." % atom_num)
        type_idx.append(ATOM_LIST.index(atom_num))
        chirality_idx.append(_safe_index(atom.GetChiralTag(), CHIRALITY_LIST, default_idx=0))

    x = torch.tensor(list(zip(type_idx, chirality_idx)), dtype=torch.long)

    row = []
    col = []
    edge_feat = []
    for bond in mol.GetBonds():
        start = bond.GetBeginAtomIdx()
        end = bond.GetEndAtomIdx()
        bond_type_idx = _safe_index(bond.GetBondType(), BOND_LIST, default_idx=0)
        bond_dir_idx = _safe_index(bond.GetBondDir(), BONDDIR_LIST, default_idx=0)

        row += [start, end]
        col += [end, start]
        edge_feat += [[bond_type_idx, bond_dir_idx], [bond_type_idx, bond_dir_idx]]

    if row:
        edge_index = torch.tensor([row, col], dtype=torch.long)
        edge_attr = torch.tensor(edge_feat, dtype=torch.long)
    else:
        edge_index = torch.empty((2, 0), dtype=torch.long)
        edge_attr = torch.empty((0, 2), dtype=torch.long)

    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)


def _build_model_and_checkpoint(variant, feat_dim):
    if variant == "gin":
        model = GINet(num_layer=5, emb_dim=300, feat_dim=feat_dim, drop_ratio=0, pool="mean")
        ckpt_subdir = "pretrained_gin"
    elif variant == "gcn":
        model = GCN(num_layer=5, emb_dim=300, feat_dim=feat_dim, drop_ratio=0, pool="mean")
        ckpt_subdir = "pretrained_gcn"
    else:
        raise ValueError("Unknown MolCLR variant '%s'." % variant)

    checkpoint_path = REPO_DIR / "ckpt" / ckpt_subdir / "checkpoints" / "model.pth"
    return model, checkpoint_path


def _load_checkpoint(model, checkpoint_path, device):
    if not checkpoint_path.exists():
        raise FileNotFoundError(
            "MolCLR checkpoint not found at %s. Ensure repository checkpoints are present."
            % checkpoint_path
        )

    state_obj = torch.load(str(checkpoint_path), map_location=device)
    if isinstance(state_obj, dict) and "state_dict" in state_obj:
        state_obj = state_obj["state_dict"]

    if isinstance(state_obj, dict):
        normalized = {}
        for key, value in state_obj.items():
            if key.startswith("module."):
                normalized[key[len("module.") :]] = value
            else:
                normalized[key] = value
        state_obj = normalized

    model.load_state_dict(state_obj)


def _encode_smiles(
    smiles_list,
    variant,
    feat_dim,
    embedding_type,
    batch_size,
    device,
    normalize,
):
    model, checkpoint_path = _build_model_and_checkpoint(variant=variant, feat_dim=feat_dim)
    model = model.to(device)
    _load_checkpoint(model, checkpoint_path, device)
    model.eval()

    graphs = []
    for smiles in smiles_list:
        graphs.append(mol_to_graph_data_obj_simple(smiles))

    loader = DataLoader(graphs, batch_size=batch_size, shuffle=False)

    chunks = []
    with torch.no_grad():
        for batch in loader:
            batch = batch.to(device)
            representation, projection = model(batch)
            if embedding_type == "representation":
                emb = representation
            else:
                emb = projection
            if normalize:
                emb = F.normalize(emb, dim=-1)
            chunks.append(emb.detach().cpu().numpy().astype(np.float32, copy=False))

    if not chunks:
        return np.empty((0, 0), dtype=np.float32)
    return np.concatenate(chunks, axis=0)


def main():
    parser = argparse.ArgumentParser(description="Encode molecules with MolCLR 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(
        "--variant",
        default="gcn",
        choices=["gin", "gcn"],
        help="MolCLR backbone checkpoint",
    )
    parser.add_argument(
        "--embedding",
        default="projected",
        choices=["projected", "representation"],
        help="Use projection head output or encoder representation",
    )
    parser.add_argument("--feat-dim", type=int, default=512)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument(
        "--device",
        default="auto",
        choices=["auto", "cpu", "cuda"],
        help="Inference device",
    )
    parser.add_argument(
        "--normalize",
        action="store_true",
        help="Apply L2 normalization to each output embedding",
    )
    args = parser.parse_args()

    variant = _normalize_variant(args.variant)
    device = _resolve_device(args.device)

    input_path = Path(args.file).resolve()
    output_dir = Path(args.output_dir).resolve()
    npy_path, map_path, failures_path = resolve_output_paths(output_dir, args.output_stem)

    lines = load_input_lines(input_path)
    records, failures = preprocess_lines(
        lines,
        input_format=args.input_format,
        model_input_type=args.model_input_type,
    )

    valid_records = []
    smiles_inputs = []
    for record in records:
        smiles_text = record["smiles"]
        try:
            _ = mol_to_graph_data_obj_simple(smiles_text)
        except Exception as exc:
            failures.append(
                {
                    "original_line": str(record.get("original_line", "")),
                    "input_text": record.get("input_text", ""),
                    "reason": "molclr_featurization_failed:%s" % exc,
                }
            )
            continue
        valid_records.append(record)
        smiles_inputs.append(smiles_text)

    if smiles_inputs:
        embeddings = _encode_smiles(
            smiles_inputs,
            variant=variant,
            feat_dim=args.feat_dim,
            embedding_type=args.embedding,
            batch_size=args.batch_size,
            device=device,
            normalize=args.normalize,
        )
    else:
        embeddings = np.empty((0, 0), dtype=np.float32)

    np.save(str(npy_path), embeddings)
    write_mapping_csv(map_path, valid_records)
    write_failures_csv(failures_path, failures)

    print("Saved embeddings to %s with shape %s." % (npy_path, embeddings.shape))
    print("Saved row mapping to %s." % map_path)
    if failures:
        print("Saved %d failed rows to %s." % (len(failures), failures_path))
    print("Model source: %s (variant=%s)" % (REPO_DIR, variant))


if __name__ == "__main__":
    main()
