import glob
import os
from pathlib import Path
import re

REPDIR = Path(__file__).resolve().parents[2] / "data" / "p53" / "pocket_reps"
Y220C_WALL = {219, 220, 221, 222, 223}


def site_drug(mi):
    out = REPDIR / f"micro{mi:03d}_out"
    info = out / f"micro{mi:03d}_info.txt"
    drug = {}
    if info.exists():
        cur = None
        for l in open(info):
            m = re.match(r"\s*Pocket (\d+)", l)
            if m:
                cur = int(m.group(1))
            if cur and "Druggability Score" in l:
                drug[cur] = float(l.split(":")[1])
    best = 0.0
    for pf in glob.glob(str(out / "pockets" / "pocket*_atm.pdb")):
        pid = int(re.findall(r"\d+", os.path.basename(pf))[0])
        res = set()
        for l in open(pf):
            if l[:6] in ("ATOM  ", "HETATM"):
                try:
                    res.add(int(l[22:26]))
                except ValueError:
                    pass
        if res & Y220C_WALL:
            best = max(best, drug.get(pid, 0.0))
    return best