#!/usr/bin/env python3
"""
02_design_guides.py -- Enumerate SpCas9 (NGG) guides whose blunt cut falls near the
TP53 C-terminal insertion point (first base of the native stop codon).

Coordinate convention: a phosphodiester bond is named by the index of the base
immediately 3' of it on the gene-sense (top) strand. Insertion bond = ins_gs.
SpCas9 makes a blunt cut 3 bp 5' of the PAM.
  +strand guide: PAM=top[j:j+3] (NGG), spacer=top[j-20:j], cut bond = j-3
  -strand guide: PAM=top[j:j+3] (CCN), spacer=revcomp(top[j+3:j+23]), cut bond = j+6
Distance to insertion (bp) = cut_bond - ins_gs.

Outputs: scripts/guide_candidates.json, scripts/guides_spacers.txt
"""
import json, os
from Bio.Seq import Seq

HERE = os.path.dirname(os.path.abspath(__file__))
locus = json.load(open(os.path.join(HERE, "locus.json")))
top = "".join(l.strip() for l in open(os.path.join(HERE, "tp53_genesense.fa")) if not l.startswith(">")).upper()
ins = locus["insertion_point_genesense_idx"]
L = len(top)
WIN_START = locus["window_chr_0based_halfopen"][0]

def gs_to_chr1(i):  # 1-based chr17(+) coordinate of gene-sense base i
    return WIN_START + (L - 1 - i) + 1

def gc(s):
    return round(100*(s.count("G")+s.count("C"))/len(s), 1)

MAX_CUT_DIST = 35      # cut must be within +/-35 bp of insertion point (HDR-permissive)
candidates = []
# scan a generous window so all qualifying guides are found
lo, hi = ins - 80, ins + 80
for j in range(max(20, lo), min(L-23, hi)):
    # +strand: PAM top[j:j+3]==N GG
    if top[j+1:j+3] == "GG":
        spacer = top[j-20:j]; pam = top[j:j+3]; cut = j-3
        if abs(cut-ins) <= MAX_CUT_DIST and set(spacer) <= set("ACGT"):
            ctx30 = top[j-24:j+6]
            candidates.append(dict(strand="+", spacer=spacer, pam=pam, cut_bond_gs=cut,
                                   dist_to_insert=cut-ins, ctx30=ctx30, pam_j=j))
    # -strand: PAM top[j:j+3]==CC N  (bottom-strand NGG)
    if top[j:j+2] == "CC":
        proto_top = top[j+3:j+23]; spacer = str(Seq(proto_top).reverse_complement())
        pam = str(Seq(top[j:j+3]).reverse_complement()); cut = j+6
        if abs(cut-ins) <= MAX_CUT_DIST and set(proto_top) <= set("ACGT"):
            ctx30 = str(Seq(top[j-3:j+27]).reverse_complement())
            candidates.append(dict(strand="-", spacer=spacer, pam=pam, cut_bond_gs=cut,
                                   dist_to_insert=cut-ins, ctx30=ctx30, pam_j=j))

# annotate
for k, c in enumerate(candidates):
    sp = c["spacer"]
    c["gc"] = gc(sp)
    c["len"] = len(sp)
    c["polyT"] = "TTTT" in sp                       # U6 pol-III terminator risk
    c["starts_with_G"] = sp[0] == "G"               # U6 transcription start pref
    c["spacer_for_U6"] = sp if sp[0]=="G" else "g"+sp  # lowercase g = appended G if needed
    c["target23_gs"] = (c["spacer"]+c["pam"]) if c["strand"]=="+" else None
    c["cut_chr17_1based"] = gs_to_chr1(c["cut_bond_gs"])
    # genomic 23-mer in gene-sense orientation for reference
    if c["strand"]=="+":
        c["target23_genesense"] = top[c["pam_j"]-20:c["pam_j"]+3]
    else:
        c["target23_genesense"] = top[c["pam_j"]:c["pam_j"]+23]
    c["id"] = f"TP53ct_g{k+1}"

# sort by absolute distance to insertion, then GC closeness to 50
candidates.sort(key=lambda c: (abs(c["dist_to_insert"]), abs(c["gc"]-50)))
for k, c in enumerate(candidates):
    c["rank_by_distance"] = k+1

json.dump({"insertion_point_genesense_idx": ins,
           "insertion_chr17_note": "insert placed before native stop codon "
               f"{locus['chrom']}:{locus['native_stop_chr17_1based'][0]}-{locus['native_stop_chr17_1based'][1]} (minus strand)",
           "n_candidates": len(candidates), "guides": candidates},
          open(os.path.join(HERE, "guide_candidates.json"), "w"), indent=2)

with open(os.path.join(HERE, "guides_spacers.txt"), "w") as fh:
    for c in candidates:
        fh.write(f"{c['id']}\t{c['spacer']}\t{c['pam']}\t{c['strand']}\tdist={c['dist_to_insert']}\tGC={c['gc']}\n")

print(f"{len(candidates)} candidate guides (cut within +/-{MAX_CUT_DIST} bp of insertion point)\n")
print(f"{'id':10} {'strand':6} {'spacer (5->3)':22} {'PAM':4} {'dist':>5} {'GC':>5} {'polyT':5} {'5pG':4} cut@chr17")
for c in candidates:
    print(f"{c['id']:10} {c['strand']:^6} {c['spacer']:22} {c['pam']:4} {c['dist_to_insert']:>5} "
          f"{c['gc']:>5} {str(c['polyT']):5} {str(c['starts_with_G']):4} {c['cut_chr17_1based']}")
