#!/usr/bin/env python3
"""
11_variant_screen.py -- Variant screening under the guide binding sites (AGENTS.md
"Variant Screening"): a common SNP in a spacer seed or the PAM can cause allele-specific
cutting failure / dropout, which matters here because the edit is done in HEK293T's own genome.

Queries Ensembl Variation (GRCh38) for variants overlapping the g3 and g2 protospacer+PAM
windows, fetches MAF per variant, and flags any COMMON (MAF >= 1%) variant in a seed/PAM.

Outputs: scripts/variant_screen.json
"""
import json, os, sys, time, urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
SERVER = "https://rest.ensembl.org"
# protospacer(20)+PAM(3) genomic windows on chr17 (GRCh38), from the verified on-target hits:
#   g3  + strand  chr17:7,669,601-7,669,623  (spacer 7,669,601-620 ; PAM AGG 7,669,621-623)
#   g2  - strand  chr17:7,669,624-7,669,646  (cut on the CDS side; near the stop)
G3 = dict(id="TP53ct_g3", chrom="17", start=7669601, end=7669623, strand="+",
          spacer=(7669601, 7669620), pam=(7669621, 7669623))
G2 = dict(id="TP53ct_g2", chrom="17", start=7669624, end=7669646, strand="-",
          spacer=(7669627, 7669646), pam=(7669624, 7669626))  # - strand: PAM at low coords
WIN = (7669595, 7669652)

def get(path):
    for attempt in range(3):
        try:
            req = urllib.request.Request(SERVER + path,
                                         headers={"Content-Type": "application/json",
                                                  "User-Agent": "labterna-bioinformatics"})
            with urllib.request.urlopen(req, timeout=30) as r:
                return json.loads(r.read().decode())
        except Exception as e:
            sys.stderr.write(f"[ensembl] {path} attempt {attempt+1}: {e}\n"); time.sleep(2)
    return None

def post_maf(ids):
    """One bulk POST to /variation/homo_sapiens -> {id: minor_allele_freq} (max 200 ids)."""
    out = {}
    for i in range(0, len(ids), 190):
        chunk = ids[i:i+190]
        body = json.dumps({"ids": chunk}).encode()
        for attempt in range(3):
            try:
                req = urllib.request.Request(SERVER + "/variation/homo_sapiens",
                    data=body, headers={"Content-Type": "application/json",
                                        "Accept": "application/json",
                                        "User-Agent": "labterna-bioinformatics"})
                with urllib.request.urlopen(req, timeout=60) as r:
                    out.update(json.loads(r.read().decode())); break
            except Exception as e:
                sys.stderr.write(f"[ensembl POST] attempt {attempt+1}: {e}\n"); time.sleep(2)
    return out

def seed_or_pam(pos, g):
    """Is genomic pos in g's PAM, seed (PAM-proximal 12 nt of the 20-nt spacer), or distal?"""
    sp_lo, sp_hi = g["spacer"]; pam_lo, pam_hi = g["pam"]
    if pam_lo <= pos <= pam_hi:
        return "PAM"
    if sp_lo <= pos <= sp_hi:
        # seed = 12 nt adjacent to the PAM
        if g["strand"] == "+":
            return "seed" if pos >= sp_hi - 11 else "spacer-distal"
        else:
            return "seed" if pos <= sp_lo + 11 else "spacer-distal"
    return "outside"

def main():
    region = f"/overlap/region/human/{G3['chrom']}:{WIN[0]}-{WIN[1]}?feature=variation"
    variants = get(region) or []
    print(f"[ensembl] {len(variants)} variant(s) in chr17:{WIN[0]}-{WIN[1]}")
    ids = [v.get("id") for v in variants if v.get("id")]
    maf_map = post_maf(ids)
    records = []
    for v in variants:
        vid = v.get("id"); start = v.get("start"); end = v.get("end")
        det = maf_map.get(vid, {})
        maf = det.get("minor_allele_freq")
        loc3 = seed_or_pam(start, G3) if G3["start"] <= start <= G3["end"] else "outside"
        loc2 = seed_or_pam(start, G2) if G2["start"] <= start <= G2["end"] else "outside"
        rec = dict(id=vid, pos=start, end=end, alleles=v.get("alleles"),
                   consequence=v.get("consequence_type"), maf=maf,
                   minor_allele=det.get("minor_allele"),
                   g3_region=loc3, g2_region=loc2,
                   common=(maf is not None and maf >= 0.01))
        records.append(rec)
    common = [r for r in records if r["common"]]
    print(f"[maf] {len(common)} common (MAF>=1%) of {len(records)}:")
    for r in (common or []):
        print(f"  {r['id']} @chr17:{r['pos']} MAF={r['maf']} {r['consequence']} g3={r['g3_region']} g2={r['g2_region']}")
    if not common:
        print("  (none -- no common variant under either guide's protospacer/PAM)")

    def verdict(g):
        bad = [r for r in records if r["common"] and r[g+"_region"] in ("seed", "PAM")]
        bad_any = [r for r in records if r["common"] and r[g+"_region"] not in ("outside",)]
        return dict(common_in_seed_or_pam=[r["id"] for r in bad],
                    common_anywhere_in_protospacer=[r["id"] for r in bad_any],
                    verdict="PASS" if not bad_any else "REVIEW")
    out = dict(
        method="Ensembl Variation REST (GRCh38); common = global MAF >= 1%",
        window=f"chr17:{WIN[0]}-{WIN[1]}",
        retrieved="2026-06-30",
        n_variants=len(records), variants=records,
        g3=verdict("g3"), g2=verdict("g2"))
    json.dump(out, open(os.path.join(HERE, "variant_screen.json"), "w"), indent=2)
    print(f"\n[verdict] g3: {out['g3']['verdict']}  g2: {out['g2']['verdict']}")
    print("[ok] variant_screen.json written")

if __name__ == "__main__":
    main()
