"""Worked example for the codon optimization guide.

Builds an E. coli codon adaptiveness table from the K-12 MG1655 ribosomal
protein genes (Sharp and Li, 1987), takes the mature human interleukin-2
coding sequence, runs three codon optimization methods with DNA Chisel
(the engine behind the ProteinIQ codon optimization tool), and reports
CAI, GC content, rare codons, repeats and restriction sites before and after.

Requires: biopython>=1.84, dnachisel, python_codon_tables, numpy.
Network access is needed once, to download three NCBI records. Pass a shell
command that reads one RNA sequence on stdin and prints RNAfold output (for
example ViennaRNA's RNAfold --noPS) to add the 5' folding energies:

    python codon_optimization_example.py [output_dir] ["RNAfold --noPS"]
"""
import csv
import json
import math
import random
import re
import subprocess
import sys
from pathlib import Path

import numpy as np
import python_codon_tables as pct
from Bio import Entrez, SeqIO
from Bio.Seq import Seq
from Bio.SeqUtils import CodonAdaptationIndex
from dnachisel import CodonOptimize, DnaOptimizationProblem, EnforceTranslation

Entrez.email = "support@proteiniq.io"
OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).parent
SEED = 7
RNAFOLD = sys.argv[2] if len(sys.argv) > 2 else None
UPSTREAM, DOWNSTREAM = 30, 40  # fold window: -30 to +40 around the A of AUG

T7_ACCESSION = "NC_001604.1"  # phage T7; gene 10 (major head protein) leader
IL2_ACCESSION = "NM_000586.4"  # Homo sapiens interleukin 2 (IL2), mRNA
ECOLI_ACCESSION = "NC_000913.3"  # E. coli str. K-12 substr. MG1655
RIBOSOMAL = re.compile(r"^(rps|rpl|rpm)[A-Z]$")
RARE = ("AGA", "AGG", "CTA", "ATA", "CCC", "GGA")
SITES = {
    "NdeI": "CATATG", "NcoI": "CCATGG", "BamHI": "GGATCC", "EcoRI": "GAATTC",
    "HindIII": "AAGCTT", "XhoI": "CTCGAG", "XbaI": "TCTAGA", "NotI": "GCGGCCGC",
}


def fetch(accession):
    handle = Entrez.efetch(db="nuccore", id=accession, rettype="gbwithparts", retmode="text")
    return SeqIO.read(handle, "gb")


def codons(seq):
    return [seq[i : i + 3] for i in range(0, len(seq), 3)]


def gc(seq):
    return 100 * (seq.count("G") + seq.count("C")) / len(seq)


def longest_run(seq):
    return max(len(m.group(0)) for m in re.finditer(r"(.)\1*", seq))


def repeated_kmers(seq, k=12):
    seen, repeated = set(), set()
    for i in range(len(seq) - k + 1):
        kmer = seq[i : i + k]
        (repeated if kmer in seen else seen).add(kmer)
    return len(repeated)


def site_hits(seq):
    both = lambda s: s + "|" + str(Seq(s).reverse_complement())
    return {name: len(re.findall(f"(?=({both(site)}))", seq)) for name, site in SITES.items()}


def t7_leader():
    """Transcript start of the T7 phi10 promoter to the AUG of gene 10A."""
    t7 = fetch(T7_ACCESSION)
    gene = next(f for f in t7.features if f.type == "CDS" and f.qualifiers.get("product", [""])[0] == "major head protein")
    sequence = str(t7.seq)
    begin = int(gene.location.start)
    start = sequence.rindex("TAATACGACTCACTATAG", 0, begin) + len("TAATACGACTCACTATA")
    return sequence[start:begin]


def fold_energy(dna):
    rna = dna.replace("T", "U")
    out = subprocess.run(RNAFOLD, input=rna + "\n", capture_output=True, text=True, shell=True, check=True).stdout
    return float(out.strip().splitlines()[-1].rsplit("(", 1)[1].rstrip(") \n"))


def main():
    # Reference set: E. coli K-12 ribosomal protein genes
    genome = fetch(ECOLI_ACCESSION)
    reference = []
    for feature in genome.features:
        if feature.type == "CDS" and RIBOSOMAL.match(feature.qualifiers.get("gene", [""])[0]):
            cds = str(feature.extract(genome.seq))
            if len(cds) % 3 == 0:
                reference.append(SeqIO.SeqRecord(Seq(cds), id=feature.qualifiers["gene"][0]))
    index = CodonAdaptationIndex([r.seq for r in reference])

    # Genome-wide E. coli usage (the table the optimizer targets), as relative adaptiveness
    genome_table = pct.get_codons_table("e_coli")
    genome_w = {c: w for aa, group in genome_table.items() if aa != "*" for c, w in [(c, f / max(group.values())) for c, f in group.items()]}
    def genome_cai(seq):
        used = [c for c in codons(seq) if c in genome_w and c not in ("ATG", "TGG")]
        return math.exp(sum(math.log(genome_w[c]) for c in used) / len(used))

    with open(OUT / "ecoli_ribosomal_weights.csv", "w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["codon", "amino_acid", "w"])
        for codon, w in sorted(index.items()):
            writer.writerow([codon, str(Seq(codon).translate()), round(w, 4)])

    # Mature human IL-2: residues 21-153 of the precursor, with ATG start and TAA stop
    il2 = fetch(IL2_ACCESSION)
    cds_feature = next(f for f in il2.features if f.type == "CDS")
    cds = str(cds_feature.extract(il2.seq))
    native = "ATG" + cds[60:-3] + "TAA"  # drop the 20-residue signal peptide
    protein = str(Seq(native).translate())

    results = {"protein": protein, "reference_genes": len(reference),
               "reference_codons": sum(len(r.seq) // 3 for r in reference)}
    variants = {"native": native}
    def optimize(method, seed):
        random.seed(seed)
        np.random.seed(seed)
        problem = DnaOptimizationProblem(
            sequence=native,
            constraints=[EnforceTranslation()],
            objectives=[CodonOptimize(species="e_coli", method=method)],
            logger=None,
        )
        problem.resolve_constraints()
        problem.optimize()
        assert str(Seq(problem.sequence).translate()) == protein
        return problem.sequence

    for method in ("use_best_codon", "match_codon_usage"):
        variants[method] = optimize(method, SEED)

    leader = t7_leader() if RNAFOLD else None
    results["t7_leader_nt"] = len(leader) if leader else None
    for name, seq in variants.items():
        cs = codons(seq)
        results[name] = {
            "length_nt": len(seq),
            "CAI_ribosomal_reference": round(index.calculate(seq), 3),
            "CAI_genome_wide_table": round(genome_cai(seq), 3),
            "GC_percent": round(gc(seq), 1),
            "rare_codons": {c: cs.count(c) for c in RARE},
            "rare_codons_total": sum(cs.count(c) for c in RARE),
            "codons_changed_vs_native": sum(a != b for a, b in zip(cs, codons(native))),
            "longest_single_base_run": longest_run(seq),
            "repeated_12mers": repeated_kmers(seq),
            "restriction_sites": {k: v for k, v in site_hits(seq).items() if v},
            "first_48_nt": seq[:48],
        }
        if leader:
            window = leader[-UPSTREAM:] + seq[:DOWNSTREAM]
            results[name]["fold_window_dG_kcal_mol"] = fold_energy(window)

    # match_codon_usage is stochastic: repeat it over 20 seeds
    sweep = [optimize("match_codon_usage", seed) for seed in range(20)]
    results["match_codon_usage_20_seeds"] = {
        "distinct_sequences": len(set(sweep)),
        "CAI_ribosomal_reference_range": [round(min(index.calculate(x) for x in sweep), 3), round(max(index.calculate(x) for x in sweep), 3)],
        "GC_percent_range": [round(min(gc(x) for x in sweep), 1), round(max(gc(x) for x in sweep), 1)],
        "runs_with_a_listed_restriction_site": sum(any(site_hits(x).values()) for x in sweep),
    }
    if leader:
        energies = [fold_energy(leader[-UPSTREAM:] + x[:DOWNSTREAM]) for x in sweep]
        results["match_codon_usage_20_seeds"]["fold_window_dG_range"] = [min(energies), max(energies)]

    (OUT / "il2_variants.fasta").write_text("".join(f">{n}\n{s}\n" for n, s in variants.items()))
    (OUT / "results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()