import argparse

# Создаём парсер и добавляем аргументы
parser = argparse.ArgumentParser(description='Поиск контигов с аномально высоким или низким покрытием')
parser.add_argument('fasta', help='Путь к файлу contigs.fa')
parser.add_argument('mean_coverage', type=float, help='Среднее покрытие трёх самых длинных контигов')
parser.add_argument('--factor', type=float, default=5.0, help='Коэффициент аномальности (по умолчанию: 5.0)')
args = parser.parse_args()

# Списки для хранения данных
high_contigs = []
low_contigs = []

# Чтение файла и извлечение данных из заголовков
with open(args.fasta, 'r') as f:
    for line in f:
        if line.startswith('>'):
            # Разбиваем заголовок по подчёркиваниям
            parts = line.strip().split('_')
            length = 0
            coverage = 0.0
            
            # Ищем length_ и cov_
            for i, part in enumerate(parts):
                if part == 'length' and i+1 < len(parts):
                    try:
                        length = int(parts[i+1])
                    except ValueError:
                        pass
                if part == 'cov' and i+1 < len(parts):
                    try:
                        coverage = float(parts[i+1])
                    except ValueError:
                        pass
            
            if length > 0 and coverage > 0:
                # Проверяем на аномальность
                if coverage > args.mean_coverage * args.factor:
                    high_contigs.append((length, coverage, line.strip()))
                elif coverage < args.mean_coverage / args.factor:
                    low_contigs.append((length, coverage, line.strip()))

# Сортировка по покрытию
high_contigs.sort(key=lambda x: x[1], reverse=True)
low_contigs.sort(key=lambda x: x[1])

# Вывод результатов
print(f'Анализ аномалий покрытия (среднее покрытие = {args.mean_coverage}, коэффициент = {args.factor})')
print()

if high_contigs:
    print(f'Контиги с аномально ВЫСОКИМ покрытием (> {args.mean_coverage * args.factor:.1f}):')
    print(f'Всего найдено: {len(high_contigs)}')
    for length, coverage, header in high_contigs:
        print(f'{header}')
        print(f'   Длина: {length}, Покрытие: {coverage:.2f}')
    print()
else:
    print('Контигов с аномально высоким покрытием не найдено')
    print()

if low_contigs:
    print(f'Контиги с аномально НИЗКИМ покрытием (< {args.mean_coverage / args.factor:.1f}):')
    print(f'Всего найдено: {len(low_contigs)}')
    for length, coverage, header in low_contigs:
        print(f'{header}')
        print(f'   Длина: {length}, Покрытие: {coverage:.2f}')
    print()
else:
    print('Контигов с аномально низким покрытием не найдено')
    print()

# Дополнительная статистика
total_contigs = len(high_contigs) + len(low_contigs)
print(f'Всего найдено аномальных контигов: {total_contigs}')