#!/usr/bin/env python3
"""
01_fetch_locus.py -- Fetch TP53 (NM_000546.6) and pin the stop-codon position on GRCh38.

Strategy
--------
1. efetch NM_000546.6 (RefSeq mRNA) -> mRNA sequence, CDS location, protein.
2. Extract the TP53 genomic window from the local hg38.2bit with twoBitToFa
   (plus strand), then reverse-complement to gene-sense (TP53 is on the MINUS strand).
3. Locate the stop codon in the gene-sense genomic sequence by exact-matching a
   60 nt anchor that straddles the stop (last 30 nt of CDS incl. stop + first 30 nt
   of 3'UTR). This region lies entirely within the last exon (exon 11), so it is
   contiguous in the genome -- the match gives an exact, build-independent anchor.
4. Define the C-terminal insertion point = first base of the native stop codon.
   Design reuses the endogenous stop + 3'UTR (insert is placed BEFORE the stop).

Outputs: scripts/locus.json, scripts/tp53_genesense.fa
"""
import json, os, subprocess, sys, time
from Bio import Entrez, SeqIO
from Bio.Seq import Seq

Entrez.email = "bioinformatics@labterna.com"

HERE = os.path.dirname(os.path.abspath(__file__))
TWOBIT = os.path.expanduser("~/genomes/hg38.2bit")
TWOBITTOFA = os.path.expanduser("~/genomes/twoBitToFa")

# Generous window around TP53 (hg38, minus strand). 0-based half-open (UCSC convention).
CHROM = "chr17"
WIN_START = 7_665_000   # 0-based
WIN_END   = 7_690_000   # exclusive
ACC = "NM_000546.6"

def fetch_mrna():
    for attempt in range(3):
        try:
            h = Entrez.efetch(db="nucleotide", id=ACC, rettype="gb", retmode="text")
            rec = SeqIO.read(h, "genbank"); h.close()
            return rec
        except Exception as e:
            sys.stderr.write(f"efetch attempt {attempt+1} failed: {e}\n"); time.sleep(3)
    raise RuntimeError("could not fetch " + ACC)

def main():
    rec = fetch_mrna()
    mrna = str(rec.seq).upper()
    cds = next(f for f in rec.features if f.type == "CDS")
    # GenBank CDS location is 0-based start, end-exclusive in Biopython; RefSeq CDS INCLUDES stop.
    cds_start = int(cds.location.start)
    cds_end   = int(cds.location.end)   # index just past the stop codon
    cds_nt = mrna[cds_start:cds_end]
    protein = str(Seq(cds_nt).translate(to_stop=False))
    stop_codon = cds_nt[-3:]
    assert protein.endswith("*"), f"CDS does not end in stop: ...{protein[-5:]}"
    assert stop_codon in ("TAA", "TAG", "TGA"), f"unexpected stop {stop_codon}"
    aa = protein[:-1]
    print(f"[mRNA] {ACC} len={len(mrna)}  CDS={cds_start+1}..{cds_end} (1-based, incl stop)")
    print(f"[prot] {len(aa)} aa  ...{aa[-12:]}  stop={stop_codon}")

    # 60 nt anchor straddling the stop codon: last 30 of CDS (incl stop) + first 30 of 3'UTR
    up = mrna[cds_end-30:cds_end]      # ...last 9 sense codons + stop
    dn = mrna[cds_end:cds_end+30]      # first 30 nt of 3'UTR
    anchor = up + dn
    assert len(anchor) == 60 and up[-3:] == stop_codon

    # Genomic window (plus strand) from local 2bit
    out_fa = os.path.join(HERE, "_win_plus.fa")
    subprocess.run([TWOBITTOFA, f"-seq={CHROM}", f"-start={WIN_START}", f"-end={WIN_END}",
                    TWOBIT, out_fa], check=True)
    plus = "".join(l.strip() for l in open(out_fa) if not l.startswith(">")).upper()
    assert len(plus) == WIN_END - WIN_START, f"window len {len(plus)} != {WIN_END-WIN_START}"
    # gene-sense = reverse complement (TP53 minus strand)
    genesense = str(Seq(plus).reverse_complement())

    # locate anchor in gene-sense
    idx = genesense.find(anchor)
    if idx < 0 or genesense.find(anchor, idx+1) != -1:
        raise RuntimeError(f"anchor match failed (idx={idx}, count={genesense.count(anchor)})")
    # stop codon occupies genesense[idx+27 : idx+30]; insertion point = first base of stop
    ins_gs = idx + 27
    assert genesense[ins_gs:ins_gs+3] == stop_codon
    # last sense codon ends at ins_gs (codon boundary, CDS in frame)
    last_codon = genesense[ins_gs-3:ins_gs]
    print(f"[locus] insertion point (gene-sense idx)={ins_gs}  last sense codon={last_codon} "
          f"({Seq(last_codon).translate()})  stop={genesense[ins_gs:ins_gs+3]}")

    # map gene-sense index -> chr17 (+) coordinate.
    # genesense[i] corresponds to plus[ (len-1) - i ]; chr coord (0-based) = WIN_START + (len-1 - i)
    L = len(genesense)
    def gs_to_chr0(i):  # 0-based chr coordinate of gene-sense base i
        return WIN_START + (L - 1 - i)
    stop_chr0_hi = gs_to_chr0(ins_gs)        # first base of stop in gene-sense = highest chr coord of the 3 (minus strand)
    stop_chr0_lo = gs_to_chr0(ins_gs+2)
    # 1-based inclusive chr range of the stop codon
    stop_chr_1based = (stop_chr0_lo + 1, stop_chr0_hi + 1)
    print(f"[locus] native stop codon at {CHROM}:{stop_chr_1based[0]}-{stop_chr_1based[1]} (minus strand)")

    locus = {
        "accession": ACC,
        "genome_build": "GRCh38/hg38",
        "twobit": "~/genomes/hg38.2bit",
        "chrom": CHROM, "strand": "-",
        "window_chr_0based_halfopen": [WIN_START, WIN_END],
        "mrna_len": len(mrna),
        "cds_1based_incl_stop": [cds_start+1, cds_end],
        "protein_len_aa": len(aa),
        "protein_cterm_12aa": aa[-12:],
        "stop_codon": stop_codon,
        "last_sense_codon": last_codon,
        "last_sense_aa": str(Seq(last_codon).translate()),
        "genesense_len": L,
        "insertion_point_genesense_idx": ins_gs,   # insert placed here, BEFORE native stop
        "native_stop_chr17_1based": list(stop_chr_1based),
        "genesense_anchor_60nt": anchor,
    }
    json.dump(locus, open(os.path.join(HERE, "locus.json"), "w"), indent=2)
    with open(os.path.join(HERE, "tp53_genesense.fa"), "w") as fh:
        fh.write(f">TP53_genesense {CHROM}:{WIN_START+1}-{WIN_END}_revcomp gene-sense(minus-strand)\n")
        for i in range(0, L, 70):
            fh.write(genesense[i:i+70] + "\n")
    os.remove(out_fa)
    print(f"[ok] wrote locus.json + tp53_genesense.fa ({L} bp gene-sense)")
    print(f"[check] C-term protein tail: ...{aa[-6:]} | (insert: linker-mCherry) | stop {stop_codon} + 3'UTR")

if __name__ == "__main__":
    main()
