from riot_na import (
    AirrRearrangementEntryAA,
    AirrRearrangementEntryNT,
    ChainType,
    Locus,
    Scheme,
    create_riot_aa,
    create_riot_nt,
)

SEQUENCES_NT = {
    "H": "AACAACACATGTCCAATGTCCTCTCCACAGACACTGAACACACTGACTCTAACCATGGGAAGGAGCTGGATCTTTCTCTTCCTCCTGTCAGGAACTGCAGGTGTCCACTCTGAGGTCCAGCTGCAACAGTCTGGACCTGTGCTGGTGAAGCCTGGGGCTTCAGTGAAGATGTCCTGTAAGGCTTCTGGATACACATTCACTGACTACTATATGAACTGGGTGAAGCAGAGCCATGGAAAGAGACTTGAGTGGATTGGAGTTATTAATCCTTACAACGGTGGTACTAACTATAACCAGAAGTTCAAGGGCAAGGCCACATTGACTGTTGACAAGTCCTCCAGCACAGCCTACATGGAGCTCAACAGCCTGACATCTGAGGACTCTGCAGTCTATTACTGTGCAGATGGGATTATTACGAATTGGTATTTCGATGTCTGGGGCACAGGGACCACGGTCACCGTCTCCTCAGCCAAAACGACACCCCCATCTGTCTATCCACTGGCCCCTGGATCTGCTGCCCAAACTAACTCCATGGTGACCCTGGGATGCCTGGTCAAGGGCTATTTCCCTGAGCCAGTGACAGTGACCTGGAACTCTGGATCCCTGTCCAGCGGTGTGCACACCTTCCCAGCTGTCCTGCAGTCTGACCTCTACACTCTGAGCAGCTCAGTGACTGTCCCCTCCAGCACCTGGCCCAGCCAGACCGTCACCTGCAACGTTGCCCACCCGGCCAGCAGCACCAAGGTGGACAAGAAAATTGTGCCCAGGGATTGTGGTTGTAAGCCTTGCATATGTACAGTCCCAGAAGT",
    "L": "GCTGACCAATATTGAAAAGAATAGACCTGGTTTGTGAATTATGGCCTGGATTTCACTTATACTCTCTCTCCTGGCTCTCAGCTCAGGGGCCATTTCCCAGGCTGTTGTGACTCAGGAATCTGCACTCACCACATCACCTGGTGAAACAGTCACACTCACTTGTCGCTCAAGTACTGGGGCTGTTACAACTAGTAACTATGCCAACTGGGTCCAAGAAAAACCAGATCATTTATTCACTGGTCTAATAGGTGGTACCAACAACCGAGCTCCAGGTGTTCCTGCCAGATTCTCAGGCTCCCTGATTGGAGACAAGGCTGCCCTCACCATCACAGGGGCACAGACTGAGGATGAGGCAATATATTTCTGTGCTCTATGGTACAGCAACCATTTCCACAATGACATGTGTAGATGGGGAAGTAGAACAAGAACACTCTGGTACAGTCTCATAACT",
}


def test_e2e_nucleotides():
    riot_numbering_nt = create_riot_nt()

    for scheme in Scheme:
        airr_heavy: AirrRearrangementEntryNT = riot_numbering_nt.run_on_sequence("", SEQUENCES_NT["H"], scheme=scheme)
        assert airr_heavy.v_call
        assert airr_heavy.j_call
        assert ChainType.from_locus(Locus(airr_heavy.locus)) == ChainType.HEAVY

        airr_light: AirrRearrangementEntryNT = riot_numbering_nt.run_on_sequence("", SEQUENCES_NT["L"], scheme=scheme)
        assert airr_light.v_call
        assert airr_light.j_call
        assert ChainType.from_locus(Locus(airr_light.locus)) == ChainType.LIGHT


SEQUENCES_AA = {
    "H": "QVQLVQSGVEVKKPGASVKVSCKASGYTFTNYYMYWVRQAPGQGLEWMGGINPSNGGTNFNEKFKNRVTLTTDSSTTTAYMELKSLQFDDTAVYYCARRDYRFDMGFDYWGQGTTVTVSSASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKKVEPKSCDKTHHHHHH",
    "L": "EIVLTQSPATLSLSPGERATLSCRASKGVSTSGYSYLHWYQQKPGQAPRLLIYLASYLESGVPARFSGSGSGTDFTLTISSLEPEDFAVYYCQHSRDLPLTFGGGTKVEIKRTVAAPSVFIFPPSDEQLKSGTASVVCLLNNFYPREAKVQWKVDNALQSGNSQESVTEQDSKDSTYSLSSTLTLSKADYEKHKVYACEVTHQGLSSPVTKSFNRGEC",
}


def test_e2e_amino_acids():
    riot_numbering_aa = create_riot_aa()

    for scheme in Scheme:
        airr_heavy: AirrRearrangementEntryAA = riot_numbering_aa.run_on_sequence("", SEQUENCES_AA["H"], scheme=scheme)
        assert airr_heavy.v_call
        assert airr_heavy.j_call
        assert ChainType.from_locus(Locus(airr_heavy.locus)) == ChainType.HEAVY

        airr_light: AirrRearrangementEntryAA = riot_numbering_aa.run_on_sequence("", SEQUENCES_AA["L"], scheme=scheme)
        assert airr_light.v_call
        assert airr_light.j_call
        assert ChainType.from_locus(Locus(airr_light.locus)) == ChainType.LIGHT
