import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import random
from scipy.stats import entropy as scipy_entropy
import time

# --- ПАРАМЕТРЫ ---
seqlen = 60
min_run, max_run = 1, 2
ANGLE_MAP = {'A': 60.0, 'C': 180.0, 'G': -60.0, 'T': -180.0, 'N': 0.0}
bases = ['A', 'C', 'G', 'T']
population_size = 10  # размер популяции "организмов"
survival_rate = 0.5  # процент выживших для следующего поколения

# --- ФУНКЦИИ ---

def bio_mutate(seq):
    r = random.random()
    if r < 0.70:
        idx = random.randint(0, len(seq)-1)
        orig = seq[idx]
        prob = random.random()
        if orig in 'AG':
            newbase = 'C' if prob < 0.65 else random.choice(['T', 'C'])
        elif orig in 'CT':
            newbase = 'G' if prob < 0.65 else random.choice(['A', 'G'])
        else:
            newbase = random.choice([b for b in bases if b != orig])
        seq = seq[:idx] + newbase + seq[idx+1:]
    elif r < 0.80:
        idx = random.randint(0, len(seq)-1)
        ins = ''.join(random.choices(bases, k=random.randint(1, 3)))
        seq = seq[:idx] + ins + seq[idx:]
        if len(seq) > seqlen:
            seq = seq[:seqlen]
    elif r < 0.90:
        if len(seq) > 4:
            idx = random.randint(0, len(seq)-2)
            dell = random.randint(1, min(3, len(seq)-idx))
            seq = seq[:idx] + seq[idx+dell:]
    else:
        if len(seq) > 10:
            start = random.randint(0, len(seq)-6)
            end = start + random.randint(3,6)
            subseq = seq[start:end][::-1]
            seq = seq[:start] + subseq + seq[end:]
    while len(seq) < seqlen:
        seq += random.choice(bases)
    return seq[:seqlen]

def compute_autocorr(profile):
    profile = profile - np.mean(profile)
    result = np.correlate(profile, profile, mode='full')
    result = result[result.size // 2:]
    norm = np.max(result) if np.max(result) != 0 else 1
    return result[:10]/norm

def compute_entropy(profile):
    vals, counts = np.unique(profile, return_counts=True)
    p = counts / counts.sum()
    return scipy_entropy(p, base=2)

def genetic_algorithm(population):
    """Эволюционный алгоритм для отбора и мутации."""
    # Отбор лучших организмов
    population.sort(key=lambda x: x[1])  # сортируем по фитнесу (энтропия)
    survivors = population[:int(population_size * survival_rate)]
    
    # Кроссовер: создаем новых организмов на основе выживших
    offspring = []
    for i in range(len(survivors) // 2):
        parent1, parent2 = survivors[i], survivors[-i-1]
        crossover_point = random.randint(0, seqlen)
        child1 = parent1[0][:crossover_point] + parent2[0][crossover_point:]
        child2 = parent2[0][:crossover_point] + parent1[0][crossover_point:]
        offspring.append((bio_mutate(child1), 0))
        offspring.append((bio_mutate(child2), 0))
    
    # Возвращаем новое поколение
    return survivors + offspring

# --- UI ---
st.title("🔴 Живой эфир мутаций ДНК")
start = st.button("▶️ Старт эфира")
stop = st.checkbox("⏹️ Остановить")

plot_placeholder = st.empty()

if start:
    # Начальная популяция
    population = [(random.choices(bases, k=seqlen), 0) for _ in range(population_size)]
    stat_bist_counts = []
    stat_entropy = []

    step = 0
    while True:
        if stop:
            st.warning("⏹️ Эфир остановлен пользователем.")
            break

        # Мутация и оценка каждого организма в популяции
        for i in range(population_size):
            seq, _ = population[i]
            torsion_profile = np.array([ANGLE_MAP.get(nt, 0.0) for nt in seq])
            ent = compute_entropy(torsion_profile)
            population[i] = (seq, ent)
        
        # Применяем эволюционный алгоритм
        population = genetic_algorithm(population)

        # Статистика для отображения
        stat_bist_counts.append(len(population))
        ent = np.mean([ind[1] for ind in population])  # средняя энтропия
        stat_entropy.append(ent)
        acorr = compute_autocorr(np.array([ANGLE_MAP.get(nt, 0.0) for nt in population[0][0]]))

        fig, axs = plt.subplots(3, 1, figsize=(10, 8))
        plt.subplots_adjust(hspace=0.45)

        axs[0].plot([ANGLE_MAP.get(nt, 0.0) for nt in population[0][0]], color='royalblue')
        axs[0].set_ylim(-200, 200)
        axs[0].set_title(f"Шаг {step}: {population[0][0]}")
        axs[0].set_ylabel("Торсионный угол")

        axs[1].plot(stat_bist_counts, '-o', color='crimson', markersize=4)
        axs[1].set_ylabel("Биомашины")
        axs[1].set_title("Количество машин")

        axs[2].bar(np.arange(6), acorr[:6], color='teal')
        axs[2].set_title(f"Автокорреляция / Энтропия: {ent:.2f}")
        axs[2].set_xlabel("Лаг")

        plot_placeholder.pyplot(fig)
        plt.close(fig)

        step += 1
        time.sleep(0.3)