#!/usr/bin/env python3
"""Train a target-specific PROTAC linker classifier.

For a given protac_module (e.g. ErlotinibVHL, BRD9VHL, CDK6CRBN):
  1. Load PROTAC-DB entries matching the target + E3 pair
  2. Extract linker SMILES from each (strip warhead + E3 substructures)
  3. Train LogisticRegression on Morgan FP (radius=2, 1024 bits) of
     target-specific linkers (positives) vs the existing generic negatives
  4. Save as protac_linker_clf_<module>.pkl

DeFoG generation can then load this classifier (when --protac_module matches)
and steer linker generation toward the chemistry distribution that's actually
represented in the target's PROTAC-DB active subset.

Usage:
    python train_target_specific_clf.py --protac-module BRD9VHL
    python train_target_specific_clf.py --protac-module CDK6CRBN
    python train_target_specific_clf.py --protac-module ErlotinibVHL
"""
import argparse
import importlib
import json
import pickle
import sys
from pathlib import Path

import numpy as np
from rdkit import Chem, RDLogger
from rdkit.Chem import AllChem, DataStructs
from sklearn.linear_model import LogisticRegression

RDLogger.DisableLog("rdApp.*")

# Resolve paths relative to the repo layout (scripts/ -> ../src/...).
_REPO_ROOT    = Path(__file__).resolve().parent.parent
CLF_DIR       = _REPO_ROOT / "src" / "classifier"
PROTACDB_JSON = _REPO_ROOT / "src" / "degrademaster" / "data" / "PROTAC" / "name.json"

_TARGET_FILTER = {
    "ErlotinibVHL": ("EGFR", "VHL"),
    "BRD9VHL":      ("BRD9", "VHL"),
    "BRD4VHL":      ("BRD4", "VHL"),
    "KRASVHL":      ("KRAS", "VHL"),
    "ARCRBN":       ("AR-V7", "CRBN"),
    "ARCRBNv2":     ("AR-V7", "CRBN"),
    "BRD4CRBN":     ("BRD4", "CRBN"),
    "CDK6CRBN":     ("CDK6", "CRBN"),
}

_TERMINI_FOR_NEGATIVES = {
    # Use the matching termini's existing negatives so the classifier learns
    # to distinguish target-specific linker chemistry from generic junk.
    "ErlotinibVHL": "protac_linker_clf.pkl",          # cooh
    "BRD9VHL":      "protac_linker_clf.pkl",          # cooh
    "BRD4VHL":      "protac_linker_clf_amino_acid.pkl",
    "KRASVHL":      "protac_linker_clf.pkl",
    "ARCRBN":       "protac_linker_clf_amine.pkl",
    "ARCRBNv2":     "protac_linker_clf.pkl",          # methyl uses generic
    "BRD4CRBN":     "protac_linker_clf_amine.pkl",
    "CDK6CRBN":     "protac_linker_clf.pkl",
}


def _morgan_fp(smi):
    mol = Chem.MolFromSmiles(smi)
    if mol is None:
        return None
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=1024)
    arr = np.zeros(1024, dtype=np.float32)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr


def _load_protac_module(module_name):
    sys.path.insert(0, str(CLF_DIR / "protac"))
    mod = importlib.import_module(module_name)
    # Prefer the loose SMARTS (WARHEAD_PAT / E3_PAT) which match the core scaffold
    # — full WARHEAD_MOL only matches one specific variant, missing PROTAC-DB
    # diversity (e.g., gefitinib vs erlotinib both = "EGFR" but fail strict match).
    wh = getattr(mod, "WARHEAD_PAT", None) or mod.WARHEAD_MOL
    e3 = getattr(mod, "E3_PAT", None) or mod.E3_MOL
    return wh, e3


def _extract_linker_smi(protac_mol, warhead_pat, e3_pat):
    """Get all atoms in PROTAC NOT belonging to either warhead-pattern or E3-pattern matches.
    Take the largest connected fragment after stripping. Returns None on ambiguity."""
    wh_atoms = set()
    for match in protac_mol.GetSubstructMatches(warhead_pat):
        wh_atoms.update(match)
    e3_atoms = set()
    for match in protac_mol.GetSubstructMatches(e3_pat):
        e3_atoms.update(match)
    if not wh_atoms or not e3_atoms:
        return None
    lk_atoms = set(range(protac_mol.GetNumAtoms())) - wh_atoms - e3_atoms
    if not lk_atoms:
        return None
    # Bonds at the linker boundary (exactly one endpoint in lk_atoms)
    bond_ids = [
        b.GetIdx() for b in protac_mol.GetBonds()
        if (b.GetBeginAtomIdx() in lk_atoms) != (b.GetEndAtomIdx() in lk_atoms)
    ]
    if not bond_ids:
        return None
    try:
        broken = Chem.FragmentOnBonds(protac_mol, bond_ids, addDummies=False)
        # Find the fragment containing the linker atoms
        frag_atom_lists = Chem.GetMolFrags(broken, asMols=False)
        for atom_list in frag_atom_lists:
            if set(atom_list).issubset(lk_atoms):
                frag_atoms = list(atom_list)
                if len(frag_atoms) >= 3:   # min linker size
                    sub_mol = Chem.PathToSubmol(broken, [b.GetIdx() for b in broken.GetBonds()
                                                          if b.GetBeginAtomIdx() in frag_atoms
                                                          and b.GetEndAtomIdx() in frag_atoms])
                    try:
                        Chem.SanitizeMol(sub_mol)
                        smi = Chem.MolToSmiles(sub_mol)
                        if smi and '.' not in smi:
                            return smi
                    except Exception:
                        pass
        # Fallback: take all fragments, find the one with exact linker atom count
        frags = Chem.GetMolFrags(broken, asMols=True)
        for frag in frags:
            if frag.GetNumHeavyAtoms() == len(lk_atoms):
                Chem.SanitizeMol(frag)
                return Chem.MolToSmiles(frag)
    except Exception:
        return None
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--protac-module", required=True, choices=list(_TARGET_FILTER.keys()))
    ap.add_argument("--out", default=None,
                    help="Output .pkl path (default: protac_linker_clf_<module>.pkl)")
    args = ap.parse_args()

    module = args.protac_module
    target, e3 = _TARGET_FILTER[module]
    out_path = Path(args.out) if args.out else CLF_DIR / f"protac_linker_clf_{module}.pkl"

    print(f"=== Training target-specific classifier for {module} ({target}+{e3}) ===")
    warhead_mol, e3_mol = _load_protac_module(module)

    # 1. Load PROTAC-DB target+E3 entries (positives)
    with open(PROTACDB_JSON) as f:
        name = json.load(f)
    target_entries = [n for n in name.values()
                      if target in n.get("tar_path", "")
                      and e3 in n.get("e3_ligase_path", "")
                      and "smiles" in n]
    print(f"  PROTAC-DB {target}+{e3} entries: {len(target_entries)}")

    # 2. Extract linkers
    linker_smis = set()
    for n in target_entries:
        m = Chem.MolFromSmiles(n["smiles"])
        if m is None:
            continue
        lk_smi = _extract_linker_smi(m, warhead_mol, e3_mol)
        if lk_smi:
            linker_smis.add(lk_smi)
    pos_smiles = sorted(linker_smis)
    print(f"  Extracted unique target-specific linkers (positives): {len(pos_smiles)}")

    if len(pos_smiles) < 10:
        print(f"  ⚠ Only {len(pos_smiles)} positives — classifier may overfit. Consider "
              f"augmenting with similar linkers or using termini-default classifier.")

    # 3. Load negatives from existing termini-matched classifier
    neg_pkl = CLF_DIR / _TERMINI_FOR_NEGATIVES[module]
    with open(neg_pkl, "rb") as f:
        existing = pickle.load(f)
    neg_smiles = list(existing["neg_smiles"])
    print(f"  Negatives (from {neg_pkl.name}): {len(neg_smiles)}")

    # 4. Featurize + train
    X_pos = [_morgan_fp(s) for s in pos_smiles]
    X_pos = [x for x in X_pos if x is not None]
    X_neg = [_morgan_fp(s) for s in neg_smiles]
    X_neg = [x for x in X_neg if x is not None]
    X = np.vstack(X_pos + X_neg)
    y = np.concatenate([np.ones(len(X_pos), dtype=int),
                        np.zeros(len(X_neg), dtype=int)])
    print(f"  Training set: {len(X_pos)} pos / {len(X_neg)} neg, {X.shape[1]}-dim FP")

    # Class-balanced LogReg (since pos is much smaller than neg)
    clf = LogisticRegression(class_weight="balanced", max_iter=2000)
    clf.fit(X, y)
    train_acc = clf.score(X, y)
    print(f"  Training accuracy: {train_acc*100:.1f}%")

    # Sanity: predict prob on a few positives + negatives
    sample_pos = clf.predict_proba(X_pos[:5])[:, 1]
    sample_neg = clf.predict_proba(X_neg[:5])[:, 1]
    print(f"  Sample pos probs: {[f'{p:.3f}' for p in sample_pos]}")
    print(f"  Sample neg probs: {[f'{p:.3f}' for p in sample_neg]}")

    # 5. Save
    with open(out_path, "wb") as f:
        pickle.dump({
            "clf":        clf,
            "pos_smiles": pos_smiles,
            "neg_smiles": neg_smiles,
            "module":     module,
            "target":     target,
            "e3":         e3,
            "n_pos":      len(X_pos),
            "n_neg":      len(X_neg),
        }, f)
    print(f"  Saved → {out_path}")


if __name__ == "__main__":
    main()