'''Calculate pairwise identity matrix for alignment.

Reference: BioPython 1.83+, numpy 1.26+ | Verify API if version differs

Vectorized via numpy broadcasting; O(N^2 * L) but with numpy-internal loops.
For thousands of sequences this is dramatically faster than the pure-Python
double loop. Above ~10k sequences, switch to k-mer-based distance estimation
(e.g. mash) -- exact pairwise identity becomes prohibitive.
'''
from Bio import AlignIO
import numpy as np

def pairwise_identity(seq1, seq2):
    matches = sum(a == b and a != '-' for a, b in zip(seq1, seq2))
    aligned_positions = sum(a != '-' or b != '-' for a, b in zip(seq1, seq2))
    return matches / aligned_positions if aligned_positions > 0 else 0

def identity_matrix_vectorized(alignment):
    seq_array = np.array([list(str(r.seq)) for r in alignment])
    n = len(alignment)
    matrix = np.zeros((n, n))
    for i in range(n):
        equal = seq_array[i:i + 1] == seq_array
        valid = (seq_array[i:i + 1] != '-') & (seq_array != '-')
        matches = (equal & valid).sum(axis=1)
        denom = ((seq_array[i:i + 1] != '-') | (seq_array != '-')).sum(axis=1)
        matrix[i] = np.where(denom > 0, matches / denom, 0)
    return matrix

if __name__ == '__main__':
    alignment = AlignIO.read('alignment.fasta', 'fasta')
    n = len(alignment)
    seq_ids = [r.id for r in alignment]

    matrix = identity_matrix_vectorized(alignment)

    print('Pairwise Identity Matrix (%):')
    print(f'{"":>12}', ' '.join(f'{s[:8]:>8}' for s in seq_ids))
    for i, row in enumerate(matrix):
        print(f'{seq_ids[i][:12]:>12}', ' '.join(f'{v*100:>7.1f}%' for v in row))

    avg_identity = (matrix.sum() - n) / (n * (n - 1)) if n > 1 else 0
    print(f'\nAverage pairwise identity: {avg_identity*100:.1f}%')