def read_fasta(path):
    with open(path) as f:
        lines = f.readlines()
    # Skip lines starting with ">" (headers)
    seq = ''.join(line.strip() for line in lines if not line.startswith('>'))
    return seq.upper()

dna = read_fasta('h19.fasta')
print(f"DNA length: {len(dna)} nt")
print(f"First 100 nt: {dna[:100]}")
print(f"Unique bases: {set(dna)}")

rna = dna.replace('T', 'U')
print(f"RNA length: {len(rna)} nt")
print(f"First 100 nt: {rna[:100]}")
# Save for the next notebooks
with open('h19_rna.txt', 'w') as f:
    f.write(rna)
print("\n✓ Saved to h19_rna.txt")

from nupack import Strand, Model, pairs
import numpy as np

# Load the sequence we saved in step 1
with open('h19_rna.txt') as f:
    h19 = f.read().strip()
print(f"H19 loaded: {len(h19)} nt")

# Thermodynamic model: RNA at 37 °C (cell-free temperature)
model = Model(material='rna', celsius=37)

# Create the NUPACK Strand object
h19_strand = Strand(h19, name='h19')
print("✓ Ready to compute pairing probabilities")

from nupack import pairs

matrix = pairs(strands=[h19], model=model)
P_matrix = matrix.to_array()

print(f"Matrix shape: {P_matrix.shape}")

N = len(h19)

# In NUPACK 4, P[i,i] = probability that base i is UNPAIRED
prob_unpaired = np.diag(P_matrix)

print(f"prob_unpaired statistics:")
print(f"  Min:    {prob_unpaired.min():.3f}")
print(f"  Max:    {prob_unpaired.max():.3f}")
print(f"  Mean:   {prob_unpaired.mean():.3f}")

# Sanity check: each row should sum to ~1
row_sum = P_matrix[100, :].sum()

window = 30

# For each start position, average prob_unpaired over the next 30 nt
window_score = np.array([
    prob_unpaired[i:i+window].mean()
    for i in range(N - window + 1)
])

print(f"Total windows evaluated: {len(window_score)}")
print(f"Max score:  {window_score.max():.3f}")
print(f"Mean score: {window_score.mean():.3f}")

# Keep the top 30 candidates
top_indices = np.argsort(window_score)[::-1][:30]

print(f"\n{'Pos':>5} | {'Sequence (30 nt)':<32} | {'Score':>5}")
print("-"*55)
for idx in top_indices:
    fragment = h19[idx:idx+window]
    score = window_score[idx]
    print(f"{idx:>5} | {fragment} | {score:.3f}")

with open('candidates.txt', 'w') as f:
    for idx in top_indices:
        fragment = h19[idx:idx+window]
        score = window_score[idx]
        f.write(f"{idx}\t{fragment}\t{score:.4f}\n")
print("✓ Saved 30 candidates to candidates.txt")

def gc_content(seq):
    """Fraction of G+C in the sequence"""
    return (seq.count('G') + seq.count('C')) / len(seq)

def has_homopolymer(seq, length=5):
    """True if there are 5 or more identical bases in a row"""
    for base in 'AUGC':
        if base * length in seq:
            return True
    return False

def has_stop_in_critical_zone(trigger):
    """
    When the switch opens, the 9 nt after the AUG are trigger[9:18]
    read 5'→3'. Those 9 nt are translated as 3 consecutive codons.
    None of them can be a stop: UAA, UAG, UGA.
    """
    stops = {'UAA', 'UAG', 'UGA'}
    region = trigger[9:18]
    codons = [region[0:3], region[3:6], region[6:9]]
    return any(c in stops for c in codons), codons

def has_internal_aug(seq):
    """An internal AUG could initiate spurious translation"""
    return 'AUG' in seq

candidates = []
with open('candidates.txt') as f:
    for line in f:
        pos, seq, score = line.strip().split('\t')
        candidates.append((int(pos), seq, float(score)))

print(f"{'Position':>5} | {'Sequence':<32} | {'Score':>5} | {'GC':>4} | {'Homopolymers':>12} | {'Stop':>5} | {'AUG':>4} | OK")
print("=" * 95)

passed = []
for pos, seq, score in candidates:
    gc = gc_content(seq)
    homo = has_homopolymer(seq)
    stop_flag, codons = has_stop_in_critical_zone(seq)
    aug = has_internal_aug(seq)

    ok_gc = 0.40 <= gc <= 0.60
    passes = ok_gc and not homo and not stop_flag and not aug

    mark = "✓" if passes else "✗"
    homo_str = "yes" if homo else "no"
    stop_str = "yes" if stop_flag else "no"
    aug_str = "yes" if aug else "no"

    print(f"{pos:>5} | {seq} | {score:.3f} | {gc:.2f} | {homo_str:>12} | {stop_str:>5} | {aug_str:>4} | {mark}")

    if passes:
        passed.append((pos, seq, score))

print(f"\n→ {len(passed)} candidates pass ALL strict filters.")

final_list = passed if passed else passed_relaxed

if not final_list:
    print("No candidate passes. We'll need to adjust further.")
else:
    # The winner is the one with the highest score
    pos, seq, score = final_list[0]

    print("="*60)
    print("CHOSEN TRIGGER")
    print("="*60)
    print(f"Position in H19:      {pos}")
    print(f"Sequence (30 nt):     {seq}")
    print(f"Accessibility score:  {score:.3f}")
    print(f"GC content:           {gc_content(seq):.2f}")

    with open('chosen_trigger.txt', 'w') as f:
        f.write(seq)

    print("\nFunctional breakdown:")
    print(f"  trigger[0:18]  (pairs with stem):    {seq[0:18]}")
    print(f"  trigger[18:30] (pairs with toehold): {seq[18:30]}")
    print(f"  trigger[9:18]  (stop-free zone):     {seq[9:18]}")

    print("\n✓ Saved to chosen_trigger.txt")

def revcomp(seq):
    """Reverse complement of RNA"""
    comp = {'A':'U', 'U':'A', 'G':'C', 'C':'G'}
    return ''.join(comp[b] for b in reversed(seq))

with open('chosen_trigger.txt') as f:
    trigger = f.read().strip()

assert len(trigger) == 30, f"Trigger must be 30 nt, got {len(trigger)}"
print(f"Trigger loaded: {trigger} ({len(trigger)} nt)")

# Toehold a (12 nt): RC of the last 12 nt of the trigger
a_toehold = revcomp(trigger[18:30])

# Ascending stem b (18 nt): RC of the first 18 nt of the trigger
b_asc = revcomp(trigger[0:18])

# Loop (11 nt) with strong Shine-Dalgarno RBS (AGGAG embedded)
# From Green 2014 switch #2
loop = "AACAGAGGAGA"
assert len(loop) == 11

# Desc_top (6 nt before AUG): pairs with the last 6 nt of b_asc
desc_top = revcomp(b_asc[-6:])

# AUG bulge (3 nt unpaired)
aug = "AUG"

# Desc_bot (9 nt after AUG): pairs with the first 9 nt of b_asc
desc_bot = revcomp(b_asc[:9])

# Check that desc_bot has NO stop codons (read 5'→3', in frame with AUG)
desc_codons = [desc_bot[0:3], desc_bot[3:6], desc_bot[6:9]]
stops = {'UAA', 'UAG', 'UGA'}
assert not any(c in stops for c in desc_codons), \
    f"STOP in desc_bot! Codons: {desc_codons}. Need to pick another candidate."
print(f"✓ desc_bot stop-free. Codons: {desc_codons}")

# Green 2014: encodes N-L-A-A-A-Q-K
linker = "AACCUGGCGGCAGCGCAAAAG"
assert len(linker) == 21

# Check that linker has no stops (in frame with AUG)
linker_codons = [linker[i:i+3] for i in range(0, 21, 3)]
assert not any(c in stops for c in linker_codons), f"STOP in linker: {linker_codons}"
print(f"✓ linker stop-free. Codons: {linker_codons}")

print("\n" + "="*60)
print("SWITCH PARTS")
print("="*60)
print(f"  a_toehold  (12 nt):  {a_toehold}")
print(f"  b_asc      (18 nt):  {b_asc}")
print(f"  loop+RBS   (11 nt):  {loop}")
print(f"  desc_top    (6 nt):  {desc_top}")
print(f"  AUG         (3 nt):  {aug}")
print(f"  desc_bot    (9 nt):  {desc_bot}")
print(f"  linker     (21 nt):  {linker}")

# Regulatory part of the switch
regulatory_switch = a_toehold + b_asc + loop + desc_top + aug + desc_bot + linker
print(f"Regulatory switch: {len(regulatory_switch)} nt")
print(f"Expected: 12 + 18 + 11 + 6 + 3 + 9 + 21 = {12+18+11+6+3+9+21}")
print(f"\nFull sequence (regulatory):\n{regulatory_switch}")

with open('regulatory_switch.txt', 'w') as f:
    f.write(regulatory_switch)
print(f"\n✓ Saved to regulatory_switch.txt")

initial_orf = aug + desc_bot + linker
print(f"Initial ORF (without reporter, {len(initial_orf)} nt):")
print(f"  {initial_orf}")
print(f"\nTranslated codons:")

codons = [initial_orf[i:i+3] for i in range(0, len(initial_orf), 3)]

codon_table = {
    'AUG':'M', 'AAA':'K', 'AAG':'K', 'AAU':'N', 'AAC':'N',
    'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACU':'T',
    'AGA':'R', 'AGG':'R', 'AGU':'S', 'AGC':'S',
    'AUA':'I', 'AUC':'I', 'AUU':'I',
    'CAA':'Q', 'CAG':'Q', 'CAU':'H', 'CAC':'H',
    'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCU':'P',
    'CGA':'R', 'CGG':'R', 'CGU':'R', 'CGC':'R',
    'CUA':'L', 'CUG':'L', 'CUU':'L', 'CUC':'L',
    'GAA':'E', 'GAG':'E', 'GAU':'D', 'GAC':'D',
    'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCU':'A',
    'GGA':'G', 'GGG':'G', 'GGU':'G', 'GGC':'G',
    'GUA':'V', 'GUG':'V', 'GUU':'V', 'GUC':'V',
    'UAA':'*', 'UAG':'*', 'UGA':'*',
    'UAU':'Y', 'UAC':'Y',
    'UCA':'S', 'UCC':'S', 'UCG':'S', 'UCU':'S',
    'UGG':'W', 'UGU':'C', 'UGC':'C',
    'UUA':'L', 'UUG':'L', 'UUU':'F', 'UUC':'F',
}

aas = []
for c in codons:
    aa = codon_table.get(c, '?')
    aas.append(aa)
    print(f"  {c} → {aa}")

peptide = ''.join(aas)
print(f"\nN-terminal added to sfGFP: {peptide}")

if '*' in peptide:
    print("STOP CODON PRESENT! This will truncate translation.")
else:
    print("No stops. The ribosome will read through to sfGFP without issue."

from nupack import Strand, Tube, Model, mfe, complex_concentrations, SetSpec
import numpy as np

with open('regulatory_switch.txt') as f:
    switch_seq = f.read().strip()
with open('chosen_trigger.txt') as f:
    trigger_seq = f.read().strip()

print(f"Switch:  {switch_seq}")
print(f"         ({len(switch_seq)} nt)")
print(f"Trigger: {trigger_seq}")
print(f"         ({len(trigger_seq)} nt)")

model = Model(material='rna', celsius=37)
print("\n✓ Model configured (RNA, 37 °C)")

print("="*70)
print("TEST 1: MFE structure of the switch ALONE (no trigger)")
print("="*70)

switch_result = mfe(strands=[switch_seq], model=model)
switch_structure = str(switch_result[0].structure)
switch_energy = switch_result[0].energy

print(f"\nSequence:  {switch_seq}")
print(f"Structure: {switch_structure}")
print(f"ΔG = {switch_energy:.2f} kcal/mol")

# Annotate functional regions under the structure
# The switch has: G(1) + toehold(12) + asc_stem(18) + loop(11) + desc_top(6) + AUG(3) + desc_bot(9) + linker(21)
G_extra = 1
N_a = 12
N_b = 18
N_loop = 11
N_desc_top = 6
N_AUG = 3
N_desc_bot = 9
N_linker = 21

regions = []
regions.append(('G', 0, G_extra))
regions.append(('toehold', G_extra, G_extra+N_a))
regions.append(('asc_stem', G_extra+N_a, G_extra+N_a+N_b))
regions.append(('loop+RBS', G_extra+N_a+N_b, G_extra+N_a+N_b+N_loop))
regions.append(('desc_top', G_extra+N_a+N_b+N_loop, G_extra+N_a+N_b+N_loop+N_desc_top))
regions.append(('AUG', G_extra+N_a+N_b+N_loop+N_desc_top, G_extra+N_a+N_b+N_loop+N_desc_top+N_AUG))
regions.append(('desc_bot', G_extra+N_a+N_b+N_loop+N_desc_top+N_AUG, G_extra+N_a+N_b+N_loop+N_desc_top+N_AUG+N_desc_bot))
regions.append(('linker', G_extra+N_a+N_b+N_loop+N_desc_top+N_AUG+N_desc_bot, len(switch_seq)))

print("\nRegion breakdown:")
for name, start, end in regions:
    region_struct = switch_structure[start:end]
    region_seq = switch_seq[start:end]
    n_paired = sum(1 for c in region_struct if c in '()')
    n_unpaired = region_struct.count('.')
    print(f"  {name:>10} ({start:>2}-{end-1:>2}): {region_seq}")
    print(f"  {'':>10}  {'':>2}     {region_struct}  → {n_paired} paired, {n_unpaired} unpaired")

print("="*70)
print("TEST 2: MFE structure of the SWITCH + TRIGGER complex")
print("="*70)

duplex_result = mfe(strands=[switch_seq, trigger_seq], model=model)
duplex_structure = str(duplex_result[0].structure)
duplex_energy = duplex_result[0].energy

print(f"\nDuplex structure:")
print(f"  {duplex_structure}")
print(f"\nΔG duplex = {duplex_energy:.2f} kcal/mol")
print(f"ΔG switch alone = {switch_energy:.2f} kcal/mol")
print(f"ΔΔG (duplex - switch alone) = {duplex_energy - switch_energy:.2f} kcal/mol")
print(f"  (the more negative, the stronger the drive to open the switch)")

print("="*70)
print("TEST 3: Equilibrium concentrations (100 nM each)")
print("="*70)

A = Strand(switch_seq, name='switch')
B = Strand(trigger_seq, name='trigger')

t1 = Tube(
    strands={A: 100e-9, B: 100e-9},  # 100 nM each
    complexes=SetSpec(max_size=2),
    name='reaction'
)

from nupack import tube_analysis
analysis = tube_analysis(tubes=[t1], model=model, compute=['mfe'])
print(analysis)