import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import random
from scipy.stats import entropy as scipy_entropy

# --- НАСТРОЙКИ ---
seqlen = 60
steps = 120
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']

def find_local_min_runs(profile, min_run=1, max_run=2):
    result = []
    N = len(profile)
    i = 0
    while i < N:
        run_val = profile[i]
        run_length = 1
        while i + run_length < N and profile[i + run_length] == run_val:
            run_length += 1
        if min_run <= run_length <= max_run:
            result.append((i, i + run_length - 1, run_val))
        i += run_length
    return result

# --- Более биологичные мутации ---
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]
            subseq = subseq[::-1]
            seq = seq[:start] + subseq + seq[end:]
    while len(seq) < seqlen:
        seq += random.choice(bases)
    if len(seq) > seqlen:
        seq = seq[:seqlen]
    return seq

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  # только лаги 0..9

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

# --- Дополнительный анализ стромбистов ---
def analyze_strombists(runs, seqlen):
    counts = len(runs)
    lengths = [end - start + 1 for start, end, _ in runs]
    angle_freq = {}
    heatmap_row = np.zeros(seqlen)
    for start, end, val in runs:
        for pos in range(start, end + 1):
            heatmap_row[pos] = 1
        angle_freq[val] = angle_freq.get(val, 0) + 1
    return counts, lengths, angle_freq, heatmap_row

# --- Начальная цепь ---
seq = ''.join(random.choices(bases, k=seqlen))
stat_bist_counts = []
stat_entropy = []
stat_autocorr = []
stat_strombists = []

fig, axs = plt.subplots(4, 1, figsize=(10, 10))
plt.subplots_adjust(hspace=0.45)
lags_shown = 6

def draw_world(seq, axs, step, cnt_hist, ent_hist, ac_hist, st_hist):
    torsion_profile = np.array([ANGLE_MAP.get(nt, 0.0) for nt in seq])
    runs = find_local_min_runs(torsion_profile, min_run, max_run)
    st_count, st_lengths, st_angle_freq, st_heatmap_row = analyze_strombists(runs, seqlen)
    
    axs[0].cla()
    axs[1].cla()
    axs[2].cla()
    axs[3].cla()

    axs[0].plot(torsion_profile, color='royalblue', label="Торсионный угол")
    for start, end, val in runs:
        axs[0].axvspan(start, end, color="red", alpha=0.3)
        axs[0].plot(range(start, end+1), torsion_profile[start:end+1], 'ro', markersize=5)
    axs[0].set_ylim(-200, 200)
    axs[0].set_xlabel("Позиция")
    axs[0].set_ylabel("Торсионный угол (град.)")
    axs[0].set_title(f"Шаг {step}: {seq}\nЧисло машин: {st_count}, энтропия: {ent_hist[-1]:.2f}")
    axs[0].legend()

    # История динамики "машин"
    axs[1].plot(cnt_hist, '-o', color='crimson', markersize=4)
    axs[1].set_xlabel("Шаг")
    axs[1].set_ylabel("Число машин")
    axs[1].set_ylim(0, max(10, max(cnt_hist)+1))
    axs[1].set_title("Динамика: число 'биомашин'")

    # Автокорреляция для текущего шага
    axs[2].bar(np.arange(lags_shown), ac_hist[-1][:lags_shown], color='teal', alpha=0.7)
    axs[2].set_xlabel("Лаг")
    axs[2].set_ylabel("Автокорреляция")
    axs[2].set_title("Автокорреляция углового профиля (структурность) и энтропия")
    axs[2].text(0.70, 0.70, f"Энтропия: {ent_hist[-1]:.2f}", transform=axs[2].transAxes)

    # Карта стромбистов
    axs[3].plot(st_heatmap_row, color='orange', label="Карта стромбистов", linewidth=2)
    axs[3].set_ylim(0, 1)
    axs[3].set_xlabel("Позиция")
    axs[3].set_ylabel("Стромбист (1 - стабильность)")
    axs[3].set_title(f"Карты стромбистов на шаге {step}")
    axs[3].legend()

def animate(i):
    global seq, stat_bist_counts, stat_entropy, stat_autocorr, stat_strombists
    if i == 0:
        stat_bist_counts.clear()
        stat_entropy.clear()
        stat_autocorr.clear()
        stat_strombists.clear()
    else:
        seq = bio_mutate(seq)
    torsion_profile = np.array([ANGLE_MAP.get(nt, 0.0) for nt in seq])
    runs = find_local_min_runs(torsion_profile, min_run, max_run)
    stat_bist_counts.append(len(runs))
    ent = compute_entropy(torsion_profile)
    stat_entropy.append(ent)
    acorr = compute_autocorr(torsion_profile)
    stat_autocorr.append(acorr)
    st_count, st_lengths, st_angle_freq, st_heatmap_row = analyze_strombists(runs, seqlen)
    stat_strombists.append((st_count, st_lengths, st_angle_freq))
    draw_world(seq, axs, i, stat_bist_counts, stat_entropy, stat_autocorr, stat_strombists)
    return axs

anim = FuncAnimation(
    fig, animate, frames=steps, interval=600, repeat=False, blit=False
)

plt.show()