#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from collections import Counter
from pathlib import Path

def read_fasta(file_path):
    file_path = Path(file_path)
    with open(file_path, 'r') as f:
        # Удаляем строки-заголовки (начинаются с '>') и объединяем всё в одну строку
        seq = ''.join(line.strip().upper() for line in f if not line.startswith('>'))
    return seq

def compute_base_frequencies(sequence):
    total_len = len(sequence)
    counts = Counter(sequence)
    return {base: counts.get(base, 0) / total_len for base in 'ACGT'}

def expected_occurrences(motif_prob, seq_len, motif_len):
    return (seq_len - motif_len + 1) * motif_prob

def z_score(observed, expected):
    if expected <= 0:
        return np.nan  
    return (observed - expected) / np.sqrt(expected)

def main():
    fasta_path = '/Users/alinadolina/Downloads/GCF_012849215.1_ASM1284921v1_genomic.fna'

    seq = read_fasta(fasta_path)
    L = len(seq)
    
    freq = compute_base_frequencies(seq)

    prob_fwd = freq['A']**2 * freq['G']**4
    prob_rev = freq['C']**4 * freq['T']**2
    
    motif_len = 6
    
    expected_fwd = expected_occurrences(prob_fwd, L, motif_len)
    expected_rev = expected_occurrences(prob_rev, L, motif_len)
 
    observed_fwd = 412
    observed_both = 792
    observed_rev = observed_both - observed_fwd  # 380
    
    # Z-score
    z_fwd = z_score(observed_fwd, expected_fwd)
    z_rev = z_score(observed_rev, expected_rev)
  
    print(f"Длина генома: {L} bp")
    print("Результаты для прямой цепи:")
    print(f"  Наблюдаемое: {observed_fwd}")
    print(f"  Ожидаемое:   {expected_fwd:.2f}")
    print(f"  Z-score:      {z_fwd:.4f}")
    
    print("\nРезультаты для комплементарной цепи:")
    print(f"  Наблюдаемое: {observed_rev}")
    print(f"  Ожидаемое:   {expected_rev:.2f}")
    print(f"  Z-score:      {z_rev:.4f}")