from riot_na.alignment.aa_gene_alignments import create_vjc_aligner_aa
from riot_na.data.model import Locus, Organism


def test_produce_aa_alignments():
    # given
    aligner = create_vjc_aligner_aa(use_segment_aligner=True)

    query_sequence = "QVQLQQWGAGLLKPSETLSLTCAVFGGSFSGYYWSWIRQPPGKGLEWIGEINHRGNTNDNPSLKSRVTISVDTSKNQFALKLSSVTAADTAVYYCARERGYTYGNFDHWGQGTLVTVSSASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKKVEPKSCDKTHTCPPCPAPELLGGPSVFLFPPKPKDTLMISRTPEVTCVVVDVSHEDPEVKFNWYVDGVEVHNAKTKPREEQYNSTYRVVSVLTVLHQDWLNGKEYKCKVSNKALPAPIEKTISKAKGQPREPQVYTLPPSRDELTKNQVSLTCLVKGFYPSDIAVEWESNGQPENNYKTTPPVLDSDGSFFLYSKLTVDKSRWQQGNVFSCSVMHEALHNHYTQKSLSLSPGK"

    # when
    alignments = aligner.produce_aa_alignments(query_sequence)[0]

    # then
    assert alignments.v is not None
    assert alignments.v.target_id == "IGHV4-34*01"
    assert alignments.v.alignment_score == 192.68417758360758
    assert alignments.v.seq_identity == 0.9484536082474226
    assert alignments.v.e_value == 2.2632803302266517e-141
    assert alignments.v.q_start == 0
    assert alignments.v.q_end == 97
    assert alignments.v.t_start == 0
    assert alignments.v.t_end == 97
    assert alignments.v.cigar == "97M"
    assert alignments.v.species == Organism.HOMO_SAPIENS
    assert alignments.v.locus == Locus.IGH
    assert (
        alignments.v.q_seq
        == "QVQLQQWGAGLLKPSETLSLTCAVFGGSFSGYYWSWIRQPPGKGLEWIGEINHRGNTNDNPSLKSRVTISVDTSKNQFALKLSSVTAADTAVYYCARERGYTYGNFDHWGQGTLVTVSSASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKKVEPKSCDKTHTCPPCPAPELLGGPSVFLFPPKPKDTLMISRTPEVTCVVVDVSHEDPEVKFNWYVDGVEVHNAKTKPREEQYNSTYRVVSVLTVLHQDWLNGKEYKCKVSNKALPAPIEKTISKAKGQPREPQVYTLPPSRDELTKNQVSLTCLVKGFYPSDIAVEWESNGQPENNYKTTPPVLDSDGSFFLYSKLTVDKSRWQQGNVFSCSVMHEALHNHYTQKSLSLSPGK"
    )
    assert (
        alignments.v.t_seq
        == "QVQLQQWGAGLLKPSETLSLTCAVYGGSFSGYYWSWIRQPPGKGLEWIGEINHSGSTNYNPSLKSRVTISVDTSKNQFSLKLSSVTAADTAVYYCAR"
    )

    assert alignments.j is not None
    assert alignments.j.target_id == "IGHJ1*01"
    assert alignments.j.alignment_score == 32.01169295201648
    assert alignments.j.seq_identity == 0.9285714285714286
    assert alignments.j.e_value == 9.01920682236379e-18
    assert alignments.j.q_start == 8
    assert alignments.j.q_end == 22
    assert alignments.j.t_start == 3
    assert alignments.j.t_end == 17
    assert alignments.j.cigar == "14M"
    assert alignments.j.species == Organism.HOMO_SAPIENS
    assert alignments.j.locus == Locus.IGH
    assert (
        alignments.j.q_seq
        == "ERGYTYGNFDHWGQGTLVTVSSASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKKVEPKSCDKTHTCPPCPAPELLGGPSVFLFPPKPKDTLMISRTPEVTCVVVDVSHEDPEVKFNWYVDGVEVHNAKTKPREEQYNSTYRVVSVLTVLHQDWLNGKEYKCKVSNKALPAPIEKTISKAKGQPREPQVYTLPPSRDELTKNQVSLTCLVKGFYPSDIAVEWESNGQPENNYKTTPPVLDSDGSFFLYSKLTVDKSRWQQGNVFSCSVMHEALHNHYTQKSLSLSPGK"
    )
    assert alignments.j.t_seq == "AEYFQHWGQGTLVTVSS"

    assert alignments.c is not None
    assert alignments.c.target_id == "IGHG1"
    assert alignments.c.alignment_score == 674.7016314783808
    assert alignments.c.seq_identity == 1.0
    assert alignments.c.e_value == 0.0
    assert alignments.c.q_start == 0
    assert alignments.c.q_end == 328
    assert alignments.c.t_start == 0
    assert alignments.c.t_end == 328
    assert alignments.c.cigar == "328M"
    assert alignments.c.species == Organism.HOMO_SAPIENS
    assert alignments.c.locus == Locus.IGH
    assert (
        alignments.c.q_seq
        == "ASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKKVEPKSCDKTHTCPPCPAPELLGGPSVFLFPPKPKDTLMISRTPEVTCVVVDVSHEDPEVKFNWYVDGVEVHNAKTKPREEQYNSTYRVVSVLTVLHQDWLNGKEYKCKVSNKALPAPIEKTISKAKGQPREPQVYTLPPSRDELTKNQVSLTCLVKGFYPSDIAVEWESNGQPENNYKTTPPVLDSDGSFFLYSKLTVDKSRWQQGNVFSCSVMHEALHNHYTQKSLSLSPGK"
    )
    assert (
        alignments.c.t_seq
        == "ASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKKVEPKSCDKTHTCPPCPAPELLGGPSVFLFPPKPKDTLMISRTPEVTCVVVDVSHEDPEVKFNWYVDGVEVHNAKTKPREEQYNSTYRVVSVLTVLHQDWLNGKEYKCKVSNKALPAPIEKTISKAKGQPREPQVYTLPPSRDELTKNQVSLTCLVKGFYPSDIAVEWESNGQPENNYKTTPPVLDSDGSFFLYSKLTVDKSRWQQGNVFSCSVMHEALHNHYTQKSLSLSPELQLEESCAEAQDGELDGLWTTITIFITLFLLSVCYSATVTFFKVKWIFSSVVDLKQTIIPDYRNMIGQGA"
    )


def test_produce_aa_alignments_therapeutic():
    # given
    aligner = create_vjc_aligner_aa(use_segment_aligner=True)

    query_sequence = "EVQLVESGGGSVQPGGSLRLSCTASGFTISRSYWICWVRQAPGKGLEWVGCIYGDNDITPLYANWAKGRFTISRDTSKNTVYLQMNSLRAEDTATYYCARLGYADYAYDLWGQGTTVTVSS"

    # when
    alignments = aligner.produce_aa_alignments(query_sequence)[0]

    # then
    assert alignments.v is not None
    assert alignments.v.target_id == "IGHV3-66*01"


if __name__ == "__main__":
    test_produce_aa_alignments()
