#!/usr/bin/env python3
"""
10_offtarget_flashfry.py -- THIRD independent off-target tool (local): FlashFry.
Runs discover + score against the cached hg38 SpCas9-NGG database and extracts, for each
candidate guide, the off-target count, CFD specificity, Hsu specificity and Doench on-target
-- the same community metrics CRISPOR/Cas-OFFinder report -- to cross-validate them.

Prereq (one-time, recorded in RERUN_NOTES): build ~/genomes/flashfry_hg38_cas9ngg via
  java -Xmx10g -jar FlashFry.jar index --tmpLocation <tmp> --database ~/genomes/flashfry_hg38_cas9ngg --reference ~/genomes/hg38.fa --enzyme SPCAS9NGG

Outputs: scripts/flashfry_summary.json (+ raw scored file for the record)
"""
import json, os, subprocess, sys, re, csv, collections
from Bio.Seq import Seq
csv.field_size_limit(sys.maxsize)

HERE = os.path.dirname(os.path.abspath(__file__))
JAR = os.path.expanduser("~/genomes/FlashFry.jar")
DB = os.path.expanduser("~/genomes/flashfry_hg38_cas9ngg")
if not (os.path.exists(DB) or os.path.exists(DB + ".bin") or os.path.exists(DB + "_2bit")):
    # FlashFry writes a single database file at the --database path; check generically
    if not any(os.path.exists(DB + s) for s in ("", ".db")) and not os.path.exists(DB):
        sys.exit(f"[flashfry] database not found at {DB}; build the index first (see RERUN_NOTES)")

guides = json.load(open(os.path.join(HERE, "guide_candidates.json")))["guides"]
gmap = {g["spacer"]: g for g in guides}

# input window = the same 300 bp gene-sense window submitted to CRISPOR
win = open(os.path.join(HERE, "crispor_submit_window.txt")).read().strip()
fa = os.path.join(HERE, "_ff_input.fa")
open(fa, "w").write(">TP53_Cterm_stop_window\n" + win + "\n")

sites = os.path.join(HERE, "_ff_sites.txt")
scored = os.path.join(HERE, "flashfry_scored.txt")
subprocess.run(["java", "-Xmx6g", "-jar", JAR, "discover", "--database", DB,
                "--fasta", fa, "--output", sites, "--maxMismatch", "4",
                "--maximumOffTargets", "5000"], check=True)
subprocess.run(["java", "-Xmx6g", "-jar", JAR, "score", "--input", sites, "--output", scored,
                "--database", DB, "--includeOTs",
                "--scoringMetrics", "hsu2013,doench2014ontarget,doench2016cfd,moreno2015,dangerous"],
               check=True)

# ---- parse ----
lines = [l.rstrip("\n") for l in open(scored) if l.strip()]
header = lines[0].split("\t")
print("[flashfry] scored columns:", header)
def colidx(substrs):
    for i, h in enumerate(header):
        hl = h.lower()
        if all(s.lower() in hl for s in substrs):
            return i
    return None
i_target = colidx(["target"])
i_otcount = colidx(["otcount"]) or colidx(["offtarget", "count"]) or colidx(["count"])
i_cfd = colidx(["cfd"])
i_hsu = colidx(["hsu"])
i_doench = colidx(["doench2014"]) or colidx(["doench", "ontarget"])
i_ots = None
for i, h in enumerate(header):
    if h.lower() in ("offtargets", "offtarget", "ots"):
        i_ots = i

def parse_ots(cell, spacer):
    """FlashFry offTargets entries look like '<23mer>_<count>_<mm>{Doench2016CFDScore=x}'.
    Mismatch count is derived by comparing the OT 20-nt protospacer to the guide spacer
    (robust). Returns (bins dict, n_off excl 0mm self, n at <=1mm, max OT CFD)."""
    bins = collections.Counter(); max_cfd = 0.0
    if cell and cell not in ("NA", "NONE", ""):
        for ent in cell.split(","):
            seq = ent.split("_")[0].upper()[:20]
            if len(seq) != 20:
                continue
            mm = sum(1 for a, b in zip(seq, spacer) if a != b)
            bins[mm] += 1
            m = re.search(r"Doench2016CFDScore=([0-9.eE+-]+)", ent)
            if m and mm > 0:
                max_cfd = max(max_cfd, float(m.group(1)))
    n_off = sum(v for k, v in bins.items() if k > 0)
    le1 = sum(v for k, v in bins.items() if 0 < k <= 1)
    return dict(sorted(bins.items())), n_off, le1, round(max_cfd, 3)

results = {}
for ln in lines[1:]:
    c = ln.split("\t")
    if i_target is None or i_target >= len(c):
        continue
    target23 = c[i_target].upper()
    spacer = target23[:20]
    # match to our guides (spacer is the 20 nt before NGG, either strand already resolved by FlashFry)
    g = gmap.get(spacer)
    if not g:
        # try reverse complement (FlashFry may report opposite-strand target)
        g = gmap.get(str(Seq(target23).reverse_complement())[3:23]) if len(target23) >= 23 else None
        if not g:
            continue
    def num(i):
        try: return float(c[i]) if (i is not None and i < len(c) and c[i] not in ("", "NA", "NONE")) else None
        except ValueError: return c[i]
    rec = {"flashfry_target": target23,
           "otCount": int(num(i_otcount)) if num(i_otcount) is not None else None,
           "cfd_specificityscore": num(i_cfd), "hsu_spec": num(i_hsu), "doench2014": num(i_doench)}
    if i_ots is not None and i_ots < len(c):
        bins, n_off, le1, max_cfd = parse_ots(c[i_ots], spacer)
        rec.update(ot_by_mm=bins, n_offtargets_seqmatched=n_off, n_le1mm=le1, max_cfd_offtarget=max_cfd)
    results[g["id"]] = rec

out = {"tool": "FlashFry (local; hg38 SpCas9-NGG database ~/genomes/flashfry_hg38_cas9ngg)",
       "note": "Third independent off-target tool. Same community metrics as CRISPOR/Cas-OFFinder "
               "(CFD, Hsu/MIT, Doench). Off-target count and specificity cross-checked vs both.",
       "scoring_metrics": "hsu2013,doench2014ontarget,doench2016cfd,moreno2015,dangerous",
       "guides": results}
json.dump(out, open(os.path.join(HERE, "flashfry_summary.json"), "w"), indent=2)
for gid in ["TP53ct_g3", "TP53ct_g2", "TP53ct_g4", "TP53ct_g1", "TP53ct_g5", "TP53ct_g6", "TP53ct_g7"]:
    r = results.get(gid)
    if r:
        print(f"  {gid}: otCount={r['otCount']} Hsuspec={r.get('hsu_spec')} "
              f"ot_by_mm={r.get('ot_by_mm')} <=1mm={r.get('n_le1mm')} maxCFD_OT={r.get('max_cfd_offtarget')}")
for f in (fa, sites):
    if os.path.exists(f): os.remove(f)
print("[ok] flashfry_summary.json written")
