"""
Test whether a SAM/BAM is correct.

Synopsis: $0 ref.fa file.bam

This script will work on sequences that are mapped to the opposite strand since in SAMv1.pdf, there is:

> For segments that have been mapped to the reverse strand, the recorded SEQ is reverse complemented from the original
> unmapped sequence and CIGAR, QUAL, and strand-sensitive optional fields are reversed and thus recorded
> consistently with the sequence bases as represented.
"""

import argparse

import pysam
import tqdm


class CigarOps:
    BAM_CMATCH = 0
    BAM_CMATCH_STR = "M"
    BAM_CINS = 1
    BAM_CINS_STR = "I"
    BAM_CDEL = 2
    BAM_CDEL_STR = "D"
    BAM_CREF_SKIP = 3
    BAM_CREF_SKIP_STR = "N"
    BAM_CSOFT_CLIP = 4
    BAM_CSOFT_CLIP_STR = "S"
    BAM_CHARD_CLIP = 5
    BAM_CHARD_CLIP_STR = "H"
    BAM_CPAD = 6
    BAM_CPAD_STR = "P"
    BAM_CEQUAL = 7
    BAM_CEQUAL_STR = "="
    BAM_CDIFF = 8
    BAM_CDIFF_STR = "X"

    INT_TO_STR = [
        BAM_CMATCH_STR,
        BAM_CINS_STR,
        BAM_CDEL_STR,
        BAM_CREF_SKIP_STR,
        BAM_CSOFT_CLIP_STR,
        BAM_CHARD_CLIP_STR,
        BAM_CPAD_STR,
        BAM_CEQUAL_STR,
        BAM_CDIFF_STR,
    ]
    STR_TO_INT = {s: i for i, s in enumerate(INT_TO_STR)}
    CONSUMES_QUERY = (
        True,  # BAM_CMATCH_STR
        True,  # BAM_CINS_STR
        False,  # BAM_CDEL_STR
        False,  # BAM_CREF_SKIP_STR
        True,  # BAM_CSOFT_CLIP_STR
        False,  # BAM_CHARD_CLIP_STR
        False,  # BAM_CPAD_STR
        True,  # BAM_CEQUAL_STR
        True,  # BAM_CDIFF_STR
    )
    CONSUMES_REFERENCE = (
        True,  # BAM_CMATCH_STR
        False,  # BAM_CINS_STR
        True,  # BAM_CDEL_STR
        True,  # BAM_CREF_SKIP_STR
        False,  # BAM_CSOFT_CLIP_STR
        False,  # BAM_CHARD_CLIP_STR
        False,  # BAM_CPAD_STR
        True,  # BAM_CEQUAL_STR
        True,  # BAM_CDIFF_STR
    )


def main():
    parser = argparse.ArgumentParser(description="Test whether a SAM/BAM is correct.")
    parser.add_argument("ref", help="Reference FASTA file")
    parser.add_argument("alignment", help="Alignment BAM/SAM file")
    args = parser.parse_args()
    flags = {
        "UNALIGNED": 0,
        "POS": 0,
        "NEG": 0,
    }
    read_id = 0
    with pysam.FastaFile(args.ref) as ref_file, pysam.AlignmentFile(
        args.alignment, "r" if args.alignment.endswith(".sam") else "rb", check_sq=False
    ) as alignment_file:
        for aln in tqdm.tqdm(alignment_file):
            read_id += 1
            if aln.is_unmapped:
                flags["UNALIGNED"] += 1
                continue
            if aln.is_reverse:
                flags["NEG"] += 1
            else:
                flags["POS"] += 1
            ref_seq = ref_file.fetch(aln.reference_name, aln.reference_start, aln.reference_end).upper()
            query_seq = aln.query_sequence.upper()
            ref_ptr = 0
            query_ptr = 0
            exon_no = 0
            genomic_ptr = aln.reference_start

            def where_we_are():
                return f"{aln.query_name}:Q:{query_ptr}/R:{ref_ptr}/G:{genomic_ptr}/A:{aln.reference_name}:{aln.reference_start}-{aln.reference_end}:{'-' if aln.is_reverse else '+'} (cigar:{cigar_id}/{len(aln.cigartuples)}:{cigar_len}{CigarOps.INT_TO_STR[cigar_op]}) (exon:{exon_no})\n"

            for cigar_id, (cigar_op, cigar_len) in enumerate(aln.cigartuples):
                if cigar_len == 0:
                    raise ValueError(f"Cigar length is zero {where_we_are()}")
                if cigar_op == CigarOps.BAM_CEQUAL:
                    ref_seg = ref_seq[ref_ptr : ref_ptr + cigar_len]
                    query_seg = query_seq[query_ptr : query_ptr + cigar_len]
                    assert ref_seg == query_seg, f"{query_seg} != {ref_seg} {where_we_are()}"
                elif cigar_op == CigarOps.BAM_CDIFF:
                    ref_seg = ref_seq[ref_ptr : ref_ptr + cigar_len]
                    query_seg = query_seq[query_ptr : query_ptr + cigar_len]
                    assert ref_seg != query_seg, f"{query_seg} == {ref_seg} {where_we_are()}"
                if CigarOps.CONSUMES_QUERY[cigar_op]:
                    query_ptr += cigar_len
                if CigarOps.CONSUMES_REFERENCE[cigar_op]:
                    ref_ptr += cigar_len
                    genomic_ptr += cigar_len
                if cigar_op == CigarOps.BAM_CREF_SKIP:
                    exon_no += 1
            assert (
                genomic_ptr == aln.reference_end
            ), f"genomic_ptr ({genomic_ptr}) != aln.reference_end ({aln.reference_end}) {where_we_are()}"
            assert query_ptr == len(
                query_seq
            ), f"query_ptr ({query_ptr}) != query_seq({len(query_seq)}) {where_we_are()}"
            assert ref_ptr == len(ref_seq), f"ref_ptr ({ref_ptr}) != ref_seq ({len(ref_seq)}) {where_we_are()}"
    print(flags)
    if read_id == 0:
        raise ValueError("No alignments found")


if __name__ == "__main__":
    main()
