#!/usr/bin/env python3
"""
03_offtarget_casoffinder.py -- Exhaustive genome-wide off-target enumeration for all
candidate guides using Cas-OFFinder (local, hg38.2bit), then annotate every hit against
RefSeq genes. Two PAM passes:
    NGG  up to 4 mismatches   (canonical, high-activity)
    NAG  up to 3 mismatches   (main residual-activity alternative PAM)
Seed = PAM-proximal 12 nt (positions 9-20, 1-based from 5'); mismatches here are the
functionally dangerous ones. This directly addresses the client's prior failure (a
pseudogene off-target missed by a single tool): every genic/exonic hit is flagged.

Outputs: scripts/offtarget_results.json, scripts/offtarget_summary.tsv
"""
import json, os, subprocess, gzip, urllib.request, collections

HERE = os.path.dirname(os.path.abspath(__file__))
GENOME = os.path.expanduser("~/genomes/hg38.2bit")
CASOFF = os.path.expanduser("~/genomes/cas-offinder")
REFGENE = os.path.expanduser("~/genomes/hg38_refGene.txt.gz")
guides = json.load(open(os.path.join(HERE, "guide_candidates.json")))["guides"]

# ---- ensure refGene table cached (gene/exon annotation) ----
if not os.path.exists(REFGENE):
    url = "https://hgdownload.soe.ucsc.edu/goldenPath/hg38/database/refGene.txt.gz"
    print(f"[dl] {url}")
    urllib.request.urlretrieve(url, REFGENE)
# build per-chrom interval list: (txStart, txEnd, name2, strand, exonStarts, exonEnds)
genes = collections.defaultdict(list)
with gzip.open(REFGENE, "rt") as fh:
    for line in fh:
        f = line.rstrip("\n").split("\t")
        # refGene: bin name chrom strand txStart txEnd cdsStart cdsEnd exonCount exonStarts exonEnds score name2 ...
        chrom, strand = f[2], f[3]
        if "_" in chrom:  # skip alt/random contigs for annotation clarity
            continue
        txS, txE = int(f[4]), int(f[5])
        es = [int(x) for x in f[9].rstrip(",").split(",") if x]
        ee = [int(x) for x in f[10].rstrip(",").split(",") if x]
        genes[chrom].append((txS, txE, f[12], strand, es, ee))
for c in genes: genes[c].sort()
print(f"[refGene] loaded genes on {len(genes)} chroms")

def annotate(chrom, pos0, length=23):
    """Return (gene_symbol or '', region) for an off-target spanning [pos0, pos0+length)."""
    s, e = pos0, pos0 + length
    hits = []
    for (txS, txE, name2, strand, es, ee) in genes.get(chrom, []):
        if txE < s: continue
        if txS > e: break
        if txS <= e and txE >= s:
            exonic = any(xs <= e and xe >= s for xs, xe in zip(es, ee))
            hits.append((name2, "exon" if exonic else "intron"))
    if not hits:
        return ("", "intergenic")
    # prefer an exonic annotation if any
    for n, r in hits:
        if r == "exon":
            return (n, "exon")
    return (hits[0][0], "intron")

def run_casoffinder(pam_suffix, max_mm, tag):
    """pam_suffix e.g. 'GG' (NGG) or 'AG' (NAG). Pattern = 21 N + suffix."""
    inp = os.path.join(HERE, f"_co_in_{tag}.txt")
    out = os.path.join(HERE, f"_co_out_{tag}.txt")
    with open(inp, "w") as fh:
        fh.write(GENOME + "\n")
        fh.write("N"*21 + pam_suffix + "\n")
        for g in guides:
            fh.write(f"{g['spacer']}NNN {max_mm}\n")
    print(f"[cas-offinder] {tag}: pattern N21{pam_suffix}, max_mm={max_mm} ...")
    subprocess.run([CASOFF, inp, "C", out], check=True,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    rows = []
    for line in open(out):
        p = line.rstrip("\n").split("\t")
        if len(p) < 6: continue
        q, chrom, pos, matched, strand, mm = p[0], p[1], int(p[2]), p[3], p[4], int(p[5])
        rows.append(dict(spacer=q.replace("N",""), chrom=chrom, pos0=pos, matched=matched,
                         strand=strand, mm=mm, pam_class=("NGG" if pam_suffix=="GG" else "NAG")))
    os.remove(inp); os.remove(out)
    return rows

allrows = run_casoffinder("GG", 4, "ngg") + run_casoffinder("AG", 3, "nag")

# index rows by spacer
by_spacer = collections.defaultdict(list)
for r in allrows:
    by_spacer[r["spacer"]].append(r)

def seed_mismatches(spacer, matched):
    """matched is 23-mer (20 spacer + 3 PAM), lowercase = mismatch in cas-offinder output.
    Seed = PAM-proximal 12 nt = spacer positions 9..20 (1-based)."""
    proto = matched[:20]
    mmpos = [i+1 for i in range(20) if proto[i].islower()]   # 1-based from 5'
    seed = [p for p in mmpos if p >= 9]
    return mmpos, seed

results = {}
summary_rows = []
for g in guides:
    rows = by_spacer.get(g["spacer"], [])
    # the on-target itself appears as a 0-mm NGG hit; identify & exclude from off-target tallies
    offs = []
    on_target = None
    for r in rows:
        mmpos, seed = seed_mismatches(g["spacer"], r["matched"])
        gene, region = annotate(r["chrom"], r["pos0"])
        rec = dict(chrom=r["chrom"], pos1=r["pos0"]+1, strand=r["strand"], mm=r["mm"],
                   pam_class=r["pam_class"], matched=r["matched"], mm_positions=mmpos,
                   seed_mm=seed, gene=gene, region=region)
        if r["mm"] == 0 and r["pam_class"] == "NGG" and r["chrom"] == "chr17":
            on_target = rec
        else:
            offs.append(rec)
    # tallies
    bins = collections.Counter((o["pam_class"], o["mm"]) for o in offs)
    genic = [o for o in offs if o["gene"]]
    exonic = [o for o in genic if o["region"] == "exon"]
    # most concerning off-targets: fewest mismatches, fewest seed mismatches, prefer genic/exonic
    offs_sorted = sorted(offs, key=lambda o: (o["mm"], len(o["seed_mm"]), o["region"]!="exon"))
    results[g["id"]] = dict(spacer=g["spacer"], pam=g["pam"], strand=g["strand"],
                            dist_to_insert=g["dist_to_insert"], on_target=on_target,
                            n_offtargets=len(offs),
                            bins={f"{k[0]}_{k[1]}mm": v for k, v in sorted(bins.items())},
                            n_genic=len(genic), n_exonic=len(exonic),
                            exonic_offtargets=exonic[:25],
                            top_offtargets=offs_sorted[:15])
    # worst genic hit by fewest mm
    worst_genic = sorted(genic, key=lambda o:(o["mm"], len(o["seed_mm"])))[:1]
    wg = worst_genic[0] if worst_genic else None
    summary_rows.append((g["id"], g["strand"], g["dist_to_insert"], g["gc"], len(offs),
                         len(genic), len(exonic),
                         f"{wg['gene']}({wg['region']},{wg['mm']}mm)" if wg else "-",
                         bins.get(("NGG",0),0)+bins.get(("NGG",1),0)))

with open(os.path.join(HERE, "offtarget_results.json"), "w") as fh:
    json.dump(results, fh, indent=2)

with open(os.path.join(HERE, "offtarget_summary.tsv"), "w") as fh:
    hdr = ["guide","strand","dist","GC","off_total","off_genic","off_exonic","worst_genic_hit","NGG_0-1mm"]
    fh.write("\t".join(hdr)+"\n")
    print("\n"+"\t".join(hdr))
    for row in summary_rows:
        line = "\t".join(str(x) for x in row)
        fh.write(line+"\n"); print(line)
print("\n[ok] offtarget_results.json + offtarget_summary.tsv")
print("Note: on-target 0-mm NGG self-hit excluded from off-target tallies.")
