import sys

def read_fasta(filename):
    """Читает FASTA файл и возвращает словарь {ID: последовательность}."""
    seqs = {}
    current_id = ""
    current_seq = []
    
    with open(filename, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith('>'):
                if current_id:
                    seqs[current_id] = "".join(current_seq)
                # Берем первое слово из заголовка как ID (чтобы игнорировать лишние пробелы/описания)
                current_id = line[1:].split()[0]
                current_seq = []
            else:
                current_seq.append(line.upper())
                
    if current_id:
        seqs[current_id] = "".join(current_seq)
        
    return seqs

def get_column_signatures(seq_dict, seq_ids, length):
    """
    Преобразует каждую колонку в кортеж биологических координат.
    Например, колонка может выглядеть так: (5, -1, 4), что значит:
    5-я аминокислота белка 1, гэп (-1), 4-я аминокислота белка 3.
    """
    signatures = []
    # Храним текущую биологическую координату для каждой последовательности
    current_pos = {seq_id: 0 for seq_id in seq_ids}
    
    for col_idx in range(length):
        col_sig = []
        is_all_gaps = True
        
        for seq_id in seq_ids:
            char = seq_dict[seq_id][col_idx]
            if char != '-':
                col_sig.append(current_pos[seq_id])
                current_pos[seq_id] += 1
                is_all_gaps = False
            else:
                col_sig.append(-1) # -1 означает гэп
                
        if is_all_gaps:
            signatures.append(None) # Игнорируем колонки, состоящие только из гэпов
        else:
            signatures.append(tuple(col_sig))
            
    return signatures

def compare(file1, file2, out=None):
    s1 = read_fasta(file1)
    s2 = read_fasta(file2)
    
    # 1. Проверяем, что наборы последовательностей одинаковые
    seq_ids1 = set(s1.keys())
    seq_ids2 = set(s2.keys())
    
    if seq_ids1 != seq_ids2:
        print("Ошибка: В файлах разные наборы последовательностей (разные заголовки >)!")
        return
        
    sorted_ids = sorted(list(seq_ids1))
    
    # 2. Проверяем, что сами белки (без учета гэпов) абсолютно одинаковые
    for seq_id in sorted_ids:
        ungapped1 = s1[seq_id].replace('-', '')
        ungapped2 = s2[seq_id].replace('-', '')
        if ungapped1 != ungapped2:
            print(f"Ошибка: Исходная последовательность {seq_id} различается в двух файлах!")
            return

    # Длины выравниваний
    m1 = len(list(s1.values())[0])
    m2 = len(list(s2.values())[0])
    
    # 3. Получаем биологические "сигнатуры" (координаты) каждой колонки
    sig1 = get_column_signatures(s1, sorted_ids, m1)
    sig2 = get_column_signatures(s2, sorted_ids, m2)
    
    # Создаем словарь для быстрого поиска колонок второго выравнивания: {координаты -> индекс_колонки}
    sig2_dict = {sig: j for j, sig in enumerate(sig2) if sig is not None}
    
    all_pairs = []
    
    # 4. Ищем одинаковые колонки
    for i, sig in enumerate(sig1):
        if sig is not None and sig in sig2_dict:
            # Сохраняем индексы в формате 1-based (для человека)
            all_pairs.append((i + 1, sig2_dict[sig] + 1))
            
    # 5. Группируем в блоки
    blocks = []
    single_pairs = []
    
    if all_pairs:
        block_start1, block_start2 = all_pairs[0]
        prev1, prev2 = all_pairs[0]
        
        for k in range(1, len(all_pairs)):
            curr1, curr2 = all_pairs[k]
            
            # Если колонки идут строго подряд в обоих выравниваниях
            if curr1 == prev1 + 1 and curr2 == prev2 + 1:
                prev1, prev2 = curr1, curr2
            else:
                block_len = prev1 - block_start1 + 1
                if block_len >= 2:
                    blocks.append((block_start1, prev1, block_start2, prev2))
                else:
                    single_pairs.append((block_start1, block_start2))
                    
                block_start1, block_start2 = curr1, curr2
                prev1, prev2 = curr1, curr2
                
        # Не забываем про последний блок
        block_len = prev1 - block_start1 + 1
        if block_len >= 2:
            blocks.append((block_start1, prev1, block_start2, prev2))
        else:
            single_pairs.append((block_start1, block_start2))

    # 6. Запись и вывод
    if out:
        with open(out, 'w') as f:
            f.write("Блоки одинаково выровненных колонок:\n")
            for b in blocks:
                f.write(f"({b[0]},{b[1]})=({b[2]},{b[3]})\n")
            if single_pairs:
                f.write("\nОдиночные одинаковые колонки:\n")
                for p in single_pairs:
                    f.write(f"{p[0]}\t{p[1]}\n")
        print(f"Результат записан в: {out}")
        
    total_in_blocks = sum(b[1] - b[0] + 1 for b in blocks)
    print(f"Длина 1 выравнивания: {m1}")
    print(f"Длина 2 выравнивания: {m2}")
    print(f"Найдено блоков: {len(blocks)}")
    print(f"Колонок в блоках: {total_in_blocks}")
    print(f"Одиночных колонок: {len(single_pairs)}")
    print(f"Всего одинаковых колонок: {total_in_blocks + len(single_pairs)}")
    print(f"% в блоках от 1 выравнивания: {round(100 * total_in_blocks / m1, 1)}")
    print(f"% в блоках от 2 выравнивания: {round(100 * total_in_blocks / m2, 1)}")

if __name__ == "__main__":
    if len(sys.argv) == 2 and sys.argv[1] == '-h':
        print('Использование: python compare.py файл1 файл2 [выход]')
        print('Файлы: выравнивание в формате FASTA')
        print('Выход: блоки (s1,f1)=(s2,f2) длиной >= 2')
        print('       и список одиночных одинаковых колонок')
    elif len(sys.argv) >= 3:
        out_file = sys.argv[3] if len(sys.argv) > 3 else None
        compare(sys.argv[1], sys.argv[2], out_file)
    else:
        print('Укажите два файла. Используйте -h для справки')