"""Reproduce an external, species-level codon-usage comparison for native TasA.

This creates a separate candidate CDS. It does not edit Kernel or the v2 plasmid.
Requires the pinned packages in RootGuard_codon_comparison_requirements.txt.
"""
from pathlib import Path
from collections import Counter
import csv
import hashlib
import importlib.metadata
import json
import random
import re
import numpy as np
import python_codon_tables as pct
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.SeqFeature import SeqFeature, FeatureLocation
from dnachisel import DnaOptimizationProblem, CodonOptimize, EnforceTranslation, AvoidChanges


ROOT = Path(__file__).resolve().parent
SEED = 20261001
random.seed(SEED)
np.random.seed(SEED)
native = str(next(SeqIO.parse(ROOT / 'RootGuard_TasA_CDS_786bp.fasta', 'fasta')).seq).upper()
protein_reference = str(next(SeqIO.parse(ROOT / 'RootGuard_TasA_protein_261aa.faa', 'fasta')).seq)
v2 = next(SeqIO.parse(ROOT / 'RootGuard_pHT01_Pveg_TasA_v2_Kernel_export_2026-10-01.gb', 'genbank'))
assert len(native) == 786 and native[:3] == 'ATG' and native[-3:] == 'TAA'
assert str(v2.seq)[2796:3582].upper() == native
assert str(Seq(native).translate(table=11)) == protein_reference + '*'
assert len(protein_reference) == 261

table = pct.get_codons_table('b_subtilis_1423')
location = (0, 783)
objective = CodonOptimize(codon_usage_table=table, method='match_codon_usage', location=location)
problem = DnaOptimizationProblem(
    sequence=native,
    constraints=[
        EnforceTranslation(genetic_table='Bacterial', translation=protein_reference, location=location),
        AvoidChanges(location=(0, 3)),
        AvoidChanges(location=(783, 786)),
    ],
    objectives=[objective],
    logger=None,
)
score_before = float(problem.objectives[0].evaluate(problem).score)
problem.resolve_constraints()
problem.optimize()
candidate = problem.sequence
score_after = float(problem.objectives[0].evaluate(problem).score)
assert problem.all_constraints_pass()
assert score_after >= score_before

# Independently check the output with Biopython, outside the optimizer checks.
translated = str(Seq(candidate).translate(table=11))
checks = {
    'length_786_bp': len(candidate) == 786,
    'only_acgt': set(candidate) <= set('ACGT'),
    'start_atg_preserved': candidate[:3] == 'ATG',
    'terminal_taa_preserved': candidate[-3:] == 'TAA',
    'protein_identical_261_aa': translated == protein_reference + '*',
    'no_internal_in_frame_stops': '*' not in translated[:-1],
    'native_matches_verified_v2': str(v2.seq)[2796:3582].upper() == native,
    'optimizer_constraints_pass': bool(problem.all_constraints_pass()),
}
assert all(checks.values())

def gc(seq):
    return 100.0 * sum(b in 'GC' for b in seq) / len(seq)

def homopolymer(seq):
    return max(len(m.group()) for m in re.finditer(r'(A+|C+|G+|T+)', seq))

rows = []
for i in range(0, 783, 3):
    a, b = native[i:i+3], candidate[i:i+3]
    if a != b:
        rows.append({'codon_number_1_based': i // 3 + 1, 'cds_nt_start_1_based': i + 1,
                     'native_codon': a, 'candidate_codon': b,
                     'amino_acid_unchanged': protein_reference[i // 3]})

prefix = 'RootGuard_TasA_Bsub_DNAChisel_candidate_v1_2026-10-01'
record = SeqRecord(Seq(candidate), id='RootGuard_TasA_Bsub_CO_candidate_v1',
                   name='TasA_Bsub_CO_v1',
                   description='786bp external DNAChisel match_codon_usage candidate; B_subtilis_1423 species-level table; native 261aa protein preserved; not a Kernel run or validated expression improvement')
SeqIO.write(record, ROOT / (prefix + '.fasta'), 'fasta')
record.annotations.update({'molecule_type': 'DNA', 'topology': 'linear'})
record.features = [SeqFeature(FeatureLocation(0, 786, strand=1), type='CDS', qualifiers={
    'label': ['TasA_Bsub_CO_candidate_v1'], 'gene': ['tasA'], 'codon_start': ['1'],
    'transl_table': ['11'], 'translation': [protein_reference],
    'note': ['Separate external candidate; original protein WP_004398632.1; b_subtilis_1423 is a species-level codon table, not proof of the production strain; no added signal peptide or tag'],
})]
SeqIO.write(record, ROOT / (prefix + '.gb'), 'genbank')
(ROOT / (prefix + '_DNA.txt')).write_text(candidate + '\n')
with (ROOT / 'RootGuard_TasA_codon_changes_2026-10-01.csv').open('w', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=['codon_number_1_based', 'cds_nt_start_1_based', 'native_codon', 'candidate_codon', 'amino_acid_unchanged'])
    writer.writeheader()
    writer.writerows(rows)
(ROOT / 'RootGuard_Bsub_1423_codon_table_snapshot_2026-10-01.json').write_text(json.dumps(table, indent=2) + '\n')

report = {
    'analysis_date_WIB': '2026-10-01', 'execution': 'external local Python run, not Kernel',
    'software_versions': {x: importlib.metadata.version(x) for x in ['dnachisel', 'python_codon_tables', 'biopython', 'numpy']},
    'seed': SEED,
    'intended_host': 'Bacillus subtilis; production strain not selected',
    'table': 'b_subtilis_1423; package species-level Kazusa/NCBI-derived snapshot; not a curated highly expressed-gene set',
    'method': 'CodonOptimize(method=match_codon_usage) over 261 sense codons',
    'constraints': ['same 261-aa bacterial-code protein', 'native ATG retained', 'native TAA retained'],
    'not_constrained_or_predicted': ['GC target', 'restriction-site avoidance', 'mRNA folding', 'secretion efficiency', 'vendor synthesis acceptance'],
    'native': {'length_bp': len(native), 'gc_percent': gc(native), 'max_homopolymer_nt': homopolymer(native), 'sha256': hashlib.sha256(native.encode()).hexdigest()},
    'candidate': {'length_bp': len(candidate), 'gc_percent': gc(candidate), 'max_homopolymer_nt': homopolymer(candidate), 'sha256': hashlib.sha256(candidate.encode()).hexdigest()},
    'nucleotide_substitutions': sum(a != b for a, b in zip(native, candidate)),
    'synonymous_codons_changed': len(rows),
    'protein_length_aa_without_stop': len(protein_reference),
    'match_codon_usage_objective_score_before': score_before,
    'match_codon_usage_objective_score_after': score_after,
    'objective_interpretation': 'Closer to zero is better for this discrepancy objective; this is not CAI and not measured expression. Rounding/discrete codon counts can prevent a perfect profile match.',
    'CAI': 'Not computed/reported; avoid claiming that this codon-profile objective is a CAI result.',
    'checks': checks,
    'candidate_not_imported_or_screened_in_Kernel': True,
    'native_v2_unchanged': True,
    'assignment_acceptance': 'External alternative; instructor acceptance of replacing the requested Kernel optimization step has not been confirmed.',
    'sources': [
        'https://edinburgh-genome-foundry.github.io/DnaChisel/ref/builtin_specifications.html',
        'https://github.com/Edinburgh-Genome-Foundry/python_codon_tables',
        'https://www.kazusa.or.jp/codon/',
    ],
}
(ROOT / 'RootGuard_TasA_external_codon_comparison_QC_2026-10-01.json').write_text(json.dumps(report, indent=2) + '\n')
print(json.dumps(report, indent=2))