import pysam
from Bio import SeqIO
import pandas as pd
from pydantic import BaseModel

class ReferenceMetadata(BaseModel):
    ref_id : str
    genus : str
    species : str
    ref_len : int

if __name__ == '__main__':
    input_bamfile = snakemake.input[0]
    reference_genomes = snakemake.input[1]
    output_parquet = snakemake.output[0]

    # Get Reference Sequence ID -> Taxonomy mapping
    ref_tax_data = {}
    for ref in SeqIO.parse(reference_genomes,'fasta'):
        ref_tax_data[ref.id] = {
            'genus' : ref.description.split(' ')[1].lower(),
            'species' : '_'.join(ref.description.split(' ')[1:3]).lower(),
            'ref_length' : len(ref.seq)
        }

    # Open Bamfile for reading
    bamfile = pysam.AlignmentFile(input_bamfile,'rb')

    # Transcribe Taxonomy from Sequnce ID dict to bam header position list for fast lookup
    ref_data = []
    for ref_id in bamfile.header.references:
        ref_data.append(
            ReferenceMetadata(
                ref_id = ref_id,
                genus = ref_tax_data[ref_id]['genus'],
                species = ref_tax_data[ref_id]['species'],
                ref_len = ref_tax_data[ref_id]['ref_length']
            )
        )

    # Parse alignments
    data = []
    n_records = 0
    n_unmapped = 0
    n_secondary = 0
    n_supplementary = 0
    n_primary = 0
    for record in bamfile.fetch():
        n_records += 1
        if record.is_unmapped:
            n_unmapped += 1
            continue
        if record.is_secondary:
            n_secondary += 1
            continue
        if record.is_supplementary:
            n_supplementary += 1
            continue
        n_primary += 1
        nm_count = -1
        for tag in record.get_tags():
            if tag[0] == 'NM':
                nm_count = tag[1]
        assert nm_count != -1
        assert isinstance(nm_count,int)
        rlen = record.infer_read_length()
        assert rlen is not None
        qlen = record.infer_query_length()
        assert qlen is not None
        ref = ref_data[record.reference_id]
        data.append(
            {
                'read_id' : record.query_name,
                'ref_name' : ref.ref_id,
                'genus' : ref.genus,
                'species' : ref.species,
                'identity' : 1 - (nm_count / qlen),
                'aligned_len' : qlen,
                'query_len' : rlen,
            }
        )
    df = pd.DataFrame(data)
    df.to_parquet(output_parquet,index=False)
    print(f'Processed {input_bamfile}:')
    print(f'\t{n_records} alignment records')
    print(f'\t{n_primary} primary alignments')
    print(f'\t{n_unmapped} unmapped reads')
    print(f'\t{n_secondary+n_supplementary} non-primary alignments')
    print(f'\t{df.shape[0]} alignment records written to disk')
    print('Per-Genus counts:')
    print(df.groupby('genus').aggregate(n_reads=('read_id','nunique')).reset_index())