import subprocess
from pathlib import Path


def prepare_protein(input_pdb: Path, output_pdb: Path, remove_hetatoms: bool = True) -> None:
    """PDB 전처리: 물·이온·리간드 제거 + H 추가 + 부분 전하 보존.
    
    실전 옵션: pdb2pqr · PDBFixer · Schrödinger Protein Prep · UCSF ChimeraX.
    """
    from pdbfixer import PDBFixer
    from openmm.app import PDBFile
    
    fixer = PDBFixer(filename=str(input_pdb))
    fixer.findMissingResidues()
    fixer.findMissingAtoms()
    fixer.addMissingAtoms()
    fixer.addMissingHydrogens(pH=7.4)
    if remove_hetatoms:
        fixer.removeHeterogens(keepWater=False)
    
    with open(output_pdb, "w") as f:
        PDBFile.writeFile(fixer.topology, fixer.positions, f)


def pdb_to_pdbqt(pdb_path: Path, pdbqt_path: Path) -> None:
    """AutoDock Vina용 PDBQT 형식 변환 (OpenBabel).
    
    -xr: rigid receptor (side chain 유연성 배제, 도킹 속도↑).
    -xh: H 유지, -xn: N amide charge, -xr r: rigid.
    """
    subprocess.run(
        ["obabel", str(pdb_path), "-O", str(pdbqt_path), "-xr"],
        check=True, capture_output=True,
    )


def identify_binding_site(pdb_path: Path, known_ligand_pdb: Path | None = None) -> tuple[float, float, float]:
    """도킹 box center. 알려진 리간드 있으면 그 중심, 없으면 pocket 예측 도구 활용.
    
    Mpro 6LU7 시나리오: co-crystallized inhibitor 있으므로 그 중심 사용.
    """
    if known_ligand_pdb:
        # 리간드 원자 평균 좌표
        from Bio.PDB import PDBParser
        parser = PDBParser(QUIET=True)
        structure = parser.get_structure("lig", str(known_ligand_pdb))
        coords = [atom.get_coord() for atom in structure.get_atoms()]
        import numpy as np
        center = np.mean(coords, axis=0)
        return (float(center[0]), float(center[1]), float(center[2]))
    else:
        # fpocket · PocketFinder 등 활용 (여기서는 개념 stub)
        raise NotImplementedError("known ligand 없으면 fpocket 통합 필요")

from rdkit import Chem
from rdkit.Chem import AllChem


def smiles_to_sdf(smiles: str, sdf_path: Path, num_conformers: int = 3) -> Path:
    """SMILES → 여러 3D conformer → SDF.
    
    여러 conformer는 DiffDock에 다양성 힌트로 유용 (일부 접근은 단일 conformer만).
    """
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Invalid SMILES: {smiles}")
    mol = Chem.AddHs(mol)
    
    # 여러 conformer 생성 후 에너지 최소화
    conf_ids = AllChem.EmbedMultipleConfs(
        mol,
        numConfs=num_conformers,
        params=AllChem.ETKDGv3(),
    )
    for conf_id in conf_ids:
        try:
            AllChem.MMFFOptimizeMolecule(mol, confId=conf_id, maxIters=500)
        except Exception:
            pass
    
    writer = Chem.SDWriter(str(sdf_path))
    for conf_id in conf_ids:
        writer.write(mol, confId=conf_id)
    writer.close()
    return sdf_path

import subprocess


def run_diffdock(
    protein_pdb: Path,
    ligand_sdf: Path,
    output_dir: Path,
    num_poses: int = 20,
    num_inference_steps: int = 40,
    device: str = "cuda",
    use_pocket: bool = False,
    pocket_center: tuple[float, float, float] | None = None,
) -> list[Path]:
    """DiffDock 실행 → 여러 pose SDF 파일 반환.
    
    DiffDock 저장소의 inference.py CLI 호출 방식.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    
    cmd = [
        "python", "-m", "inference",
        "--protein_path", str(protein_pdb),
        "--ligand", str(ligand_sdf),
        "--out_dir", str(output_dir),
        "--samples_per_complex", str(num_poses),
        "--inference_steps", str(num_inference_steps),
        "--batch_size", "10",
    ]
    if use_pocket and pocket_center:
        # DiffDock-Pocket 확장 (좌표 힌트)
        cmd += ["--pocket_center", ",".join(f"{c:.2f}" for c in pocket_center)]
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"DiffDock 실패 STDERR: {result.stderr[:500]}")
        raise RuntimeError("DiffDock 실행 실패")
    
    pose_files = sorted(output_dir.glob("complex_0/rank*.sdf"))
    return pose_files

def validate_pose_with_posebusters(
    pose_sdf: Path,
    protein_pdb: Path,
) -> dict:
    """PoseBusters로 pose validity 검사.
    
    Returns: {
        "passes_all_checks": bool,
        "checks": {check_name: bool},
        "critical_failures": [list of failed checks],
    }
    """
    try:
        from posebusters import PoseBusters
        buster = PoseBusters(config="dock")
        results = buster.bust(
            mol_pred=[pose_sdf],
            mol_cond=protein_pdb,
        )
        # results는 pandas DataFrame
        checks = results.iloc[0].to_dict()
        critical_failures = [k for k, v in checks.items() if v is False and "clash" in k.lower() or "geometry" in k.lower()]
        passes_all = all(v for v in checks.values() if isinstance(v, bool))
        return {
            "passes_all_checks": passes_all,
            "checks": checks,
            "critical_failures": critical_failures,
        }
    except ImportError:
        # PoseBusters 미설치 시 간단 clash 검사로 fallback
        return {"passes_all_checks": True, "checks": {}, "critical_failures": []}


import numpy as np
from Bio.PDB import PDBParser, NeighborSearch


VDW_RADII = {
    "H": 1.20, "C": 1.70, "N": 1.55, "O": 1.52, "F": 1.47,
    "P": 1.80, "S": 1.80, "Cl": 1.75, "Br": 1.85, "I": 1.98,
}


def simple_clash_check(
    protein_pdb: Path,
    ligand_sdf: Path,
    clash_threshold: float = 0.7,
    max_clash_count: int = 3,
) -> bool:
    """clash 있으면 True (PoseBusters 대체 fallback)."""
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("protein", str(protein_pdb))
    protein_atoms = list(structure.get_atoms())
    ns = NeighborSearch(protein_atoms)
    
    mol = Chem.SDMolSupplier(str(ligand_sdf), removeHs=False)[0]
    if mol is None:
        return True
    conf = mol.GetConformer()
    
    clash_count = 0
    for i, atom in enumerate(mol.GetAtoms()):
        pos = conf.GetAtomPosition(i)
        lig_vdw = VDW_RADII.get(atom.GetSymbol(), 1.7)
        nearby = ns.search([pos.x, pos.y, pos.z], 5.0)
        for prot_atom in nearby:
            prot_vdw = VDW_RADII.get(prot_atom.element, 1.7)
            distance = np.linalg.norm(
                np.array([pos.x, pos.y, pos.z]) - prot_atom.get_coord()
            )
            if distance < (lig_vdw + prot_vdw) * clash_threshold:
                clash_count += 1
                if clash_count > max_clash_count:
                    return True
    return False

def vina_local_score(
    receptor_pdbqt: Path,
    ligand_pdbqt: Path,
    center: tuple[float, float, float],
    size: tuple[float, float, float] = (20, 20, 20),
    exhaustiveness: int = 1,      # local optimize라 낮게
    output_pdbqt: Path | None = None,
) -> float:
    """Vina local optimization + score. exhaustiveness=1로 local minimize.
    
    Returns: 결합 에너지 (kcal/mol, 음수일수록 강력).
    """
    cmd = [
        "vina",
        "--receptor", str(receptor_pdbqt),
        "--ligand", str(ligand_pdbqt),
        "--center_x", str(center[0]),
        "--center_y", str(center[1]),
        "--center_z", str(center[2]),
        "--size_x", str(size[0]),
        "--size_y", str(size[1]),
        "--size_z", str(size[2]),
        "--exhaustiveness", str(exhaustiveness),
        "--num_modes", "1",
        "--local_only",
    ]
    if output_pdbqt:
        cmd += ["--out", str(output_pdbqt)]
    
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    # Vina 출력에서 첫 mode의 affinity 파싱
    for line in result.stdout.splitlines():
        stripped = line.strip()
        if stripped.startswith("1 ") or stripped.startswith("1\t"):
            parts = stripped.split()
            try:
                return float(parts[1])  # kcal/mol
            except (IndexError, ValueError):
                pass
    return 0.0

from dataclasses import dataclass


@dataclass
class PoseResult:
    ligand_id: str
    pose_path: Path
    diffdock_confidence: float
    posebusters_passes: bool
    vina_score: float
    final_score: float
    metadata: dict


def ensemble_score(
    diffdock_conf: float,
    posebusters_passes: bool,
    vina_score: float,
    weight_dl: float = 0.3,
    weight_vina: float = 0.6,
    invalid_penalty: float = 5.0,
) -> float:
    """앙상블 스코어. 낮을수록 강력 (Vina 기준)."""
    # DL score → Vina 스케일로 대략 변환 (-5.0 근처가 uncertain, -8.0 이하 강력 결합)
    dl_component = -diffdock_conf * 3.0
    combined = weight_vina * vina_score + weight_dl * dl_component
    if not posebusters_passes:
        combined += invalid_penalty  # 물리학적으로 이상한 pose는 penalty
    return combined


def rank_ligand_poses(
    ligand_id: str,
    diffdock_poses: list[dict],       # [{pose_path, confidence}]
    receptor_pdb: Path,
    receptor_pdbqt: Path,
    binding_center: tuple[float, float, float],
) -> list[PoseResult]:
    """리간드 하나의 여러 pose를 Vina 재검증 + PoseBusters + 앙상블."""
    results = []
    for pose in diffdock_poses:
        pose_pdbqt = pose["pose_path"].with_suffix(".pdbqt")
        try:
            subprocess.run(
                ["obabel", str(pose["pose_path"]), "-O", str(pose_pdbqt)],
                check=True, capture_output=True,
            )
        except subprocess.CalledProcessError:
            continue
        
        vina_score = vina_local_score(receptor_pdbqt, pose_pdbqt, binding_center)
        pb_result = validate_pose_with_posebusters(pose["pose_path"], receptor_pdb)
        final = ensemble_score(
            pose["confidence"], pb_result["passes_all_checks"], vina_score,
        )
        results.append(PoseResult(
            ligand_id=ligand_id,
            pose_path=pose["pose_path"],
            diffdock_confidence=pose["confidence"],
            posebusters_passes=pb_result["passes_all_checks"],
            vina_score=vina_score,
            final_score=final,
            metadata={
                "posebusters_details": pb_result["checks"],
                "critical_failures": pb_result["critical_failures"],
            },
        ))
    return sorted(results, key=lambda r: r.final_score)

def render_top_poses(
    receptor_pdb: Path,
    top_poses: list[PoseResult],
    output_dir: Path,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    for i, pose in enumerate(top_poses):
        script = f"""
load {receptor_pdb}, receptor
load {pose.pose_path}, ligand
hide everything
show cartoon, receptor
show sticks, ligand
show sticks, receptor within 5 of ligand
color grey70, receptor
color yellow, ligand
zoom ligand, 5
bg_color white
ray 1200, 900
png {output_dir / f"{pose.ligand_id}_rank{i+1}.png"}, dpi=150
quit
"""
        script_path = output_dir / f"render_{pose.ligand_id}_{i}.pml"
        script_path.write_text(script)
        subprocess.run(["pymol", "-cq", str(script_path)], check=True, capture_output=True)

import pandas as pd


def hybrid_docking_pipeline(
    protein_pdb: Path,
    known_ligand_pdb: Path | None,
    ligand_smiles_list: list[tuple[str, str]],  # [(ligand_id, smiles)]
    work_dir: Path,
    top_k: int = 10,
    poses_per_ligand: int = 10,
    device: str = "cuda",
) -> pd.DataFrame:
    """전체 하이브리드 도킹 스크리닝."""
    work_dir.mkdir(parents=True, exist_ok=True)
    
    print("[1/5] 단백질 전처리")
    prepared_pdb = work_dir / "protein_prepared.pdb"
    prepare_protein(protein_pdb, prepared_pdb)
    receptor_pdbqt = work_dir / "protein_prepared.pdbqt"
    pdb_to_pdbqt(prepared_pdb, receptor_pdbqt)
    binding_center = identify_binding_site(prepared_pdb, known_ligand_pdb)
    print(f"      binding center: {binding_center}")
    
    all_results = []
    for ligand_id, smiles in ligand_smiles_list:
        print(f"[2/5] {ligand_id} 3D conformer + DiffDock")
        ligand_sdf = work_dir / "ligands" / f"{ligand_id}.sdf"
        try:
            smiles_to_sdf(smiles, ligand_sdf, num_conformers=3)
            pose_files = run_diffdock(
                prepared_pdb, ligand_sdf,
                work_dir / "diffdock" / ligand_id,
                num_poses=poses_per_ligand, device=device,
            )
        except Exception as e:
            print(f"      실패 ({ligand_id}): {e}")
            continue
        
        # DiffDock confidence는 rank 파일명 or 별도 metadata에서 파싱
        diffdock_poses = [
            {"pose_path": p, "confidence": max(0.0, 1.0 - 0.05 * i)}  # rank 낮을수록 높은 confidence
            for i, p in enumerate(pose_files)
        ]
        
        print(f"[3/5] {ligand_id} PoseBusters + Vina 재검증")
        ranked = rank_ligand_poses(
            ligand_id, diffdock_poses, prepared_pdb, receptor_pdbqt, binding_center,
        )
        all_results.extend(ranked[:3])  # 리간드당 top-3만
    
    # 4. 전체 랭킹
    all_results.sort(key=lambda r: r.final_score)
    df = pd.DataFrame([{
        "ligand_id": r.ligand_id,
        "diffdock_conf": r.diffdock_confidence,
        "posebusters_pass": r.posebusters_passes,
        "vina_score": r.vina_score,
        "final_score": r.final_score,
        "critical_failures": ";".join(r.metadata["critical_failures"]),
        "pose_path": str(r.pose_path),
    } for r in all_results])
    df.to_csv(work_dir / "docking_results.csv", index=False)
    
    print(f"[4/5] Top-{top_k} PyMOL 렌더링")
    render_top_poses(prepared_pdb, all_results[:top_k], work_dir / "renders")
    
    print("[5/5] 완료")
    return df


# 실행 예시 (SARS-CoV-2 Mpro 시나리오)
# result_df = hybrid_docking_pipeline(
#     protein_pdb=Path("./6LU7.pdb"),
#     known_ligand_pdb=Path("./6LU7_ligand.pdb"),  # co-crystallized N3 inhibitor
#     ligand_smiles_list=[
#         ("nirmatrelvir", "CC1(C)C2CC1C(NC(=O)C1CCCN1C(=O)C(NC(=O)OC(C)(C)C)C(C)(C)C)C(=O)NC(C#N)CC2=O"),
#         ("ensitrelvir", "..."),
#         # 편 07 GNN 상위 20 후보 추가
#     ],
#     work_dir=Path("./mpro_docking_output"),
# )