# Copyright (c) 2024 Bryan Briney
# Distributed under the terms of the MIT License.
# SPDX-License-Identifier: MIT


import pytest

from ..annotation.antibody import Antibody
from ..annotation.productivity import assess_productivity


@pytest.fixture
def productive_antibody():
    return Antibody(
        sequence="CAGGTTCAGCTGGTGCAGTCTGGAGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGTTACACCTTTAGCAGTTTTGGTATCAGCTGGGTGCGACAGGCCCCTGGGCAAGGGCTTGAGTGGCTGGGATGGAGCAGCACTGACAATGGTAACACAAACTATGCACAGAAGTTCCAGGGCAGAGTCACCATGACCACAGACACATCCACGACCACAGCCTACATGGAGCTGAGGAGCCTAAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGATTTAGGGCGGTGTACCAATACCGGGTGCTATCGCAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAG",
        sequence_aa="QVQLVQSGAEVKKPGASVKVSCKASGYTFSSFGISWVRQAPGQGLEWLGWSSTDNGNTNYAQKFQGRVTMTTDTSTTTAYMELRSLRSDDTAVYYCARDLGRCTNTGCYRNWFDPWGQGTLVTVSS",
        v_call="IGHV1-18*01",
        j_call="IGHJ5*01",
        junction_aa="CARDLGRCTNTGCYRNWFDPW",
    )


@pytest.fixture
def stop_codon_antibody():
    return Antibody(
        sequence="CAGGTTCAGCTGGTGCAGTCTGGAGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGTTACACCTTTAGCAGTTTTGGTATCAGCTGGGTGCGACAGGCCCCTGGGCAAGGGCTTGAGTGGCTGGGATGGAGCAGCACTGACAATGGTAACACAAACTATGCACAGAAGTTCCAGGGCAGAGTCACCATGACCACAGACACATCCACGACCACAGCCTACATGGAGCTGAGGAGCCTAAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGATTTAGGGCGGTGTACCAATACCGGGTGCTATCGCAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAG",
        sequence_aa="QVQLVQSGAEVKKPGASVKVSCKASGYTFSSFGISWVRQAPGQGLEWLGWSSTDNGNTNYAQKFQGRVT*TTDTSTTTAYMELRSLRSDDTAVYYCARDLGRCTNTGCYRNWFDPWGQGTLVTVSS",
        v_call="IGHV1-18*01",
        j_call="IGHJ5*01",
        junction_aa="CARDLGRCTNTGCYRNWFDPW",
    )


@pytest.fixture
def missing_conserved_cysteine_antibody():
    return Antibody(
        sequence="CAGGTTCAGCTGGTGCAGTCTGGAGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGTTACACCTTTAGCAGTTTTGGTATCAGCTGGGTGCGACAGGCCCCTGGGCAAGGGCTTGAGTGGCTGGGATGGAGCAGCACTGACAATGGTAACACAAACTATGCACAGAAGTTCCAGGGCAGAGTCACCATGACCACAGACACATCCACGACCACAGCCTACATGGAGCTGAGGAGCCTAAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGATTTAGGGCGGTGTACCAATACCGGGTGCTATCGCAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAG",
        sequence_aa="QVQLVQSGAEVKKPGASVKVSCKASGYTFSSFGISWVRQAPGQGLEWLGWSSTDNGNTNYAQKFQGRVTMTTDTSTTTAYMELRSLRSDDTAVYYAARDLGRCTNTGCYRNWFDPWGQGTLVTVSS",
        v_call="IGHV1-18*01",
        j_call="IGHJ5*01",
        junction_aa="AARDLGRCTNTGCYRNWFDPW",
    )


@pytest.fixture
def missing_conserved_tryptophan_antibody():
    return Antibody(
        sequence="CAGGTTCAGCTGGTGCAGTCTGGAGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGTTACACCTTTAGCAGTTTTGGTATCAGCTGGGTGCGACAGGCCCCTGGGCAAGGGCTTGAGTGGCTGGGATGGAGCAGCACTGACAATGGTAACACAAACTATGCACAGAAGTTCCAGGGCAGAGTCACCATGACCACAGACACATCCACGACCACAGCCTACATGGAGCTGAGGAGCCTAAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGATTTAGGGCGGTGTACCAATACCGGGTGCTATCGCAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAG",
        sequence_aa="QVQLVQSGAEVKKPGASVKVSCKASGYTFSSFGISWVRQAPGQGLEWLGWSSTDNGNTNYAQKFQGRVTMTTDTSTTTAYMELRSLRSDDTAVYYCARDLGRCTNTGCYRNWFDPAGQGTLVTVSS",
        v_call="IGHV1-18*01",
        j_call="IGHJ5*01",
        junction_aa="CARDLGRCTNTGCYRNWFDPA",
    )


@pytest.fixture
def locus_mismatch_antibody():
    return Antibody(
        sequence="CAGGTTCAGCTGGTGCAGTCTGGAGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGTTACACCTTTAGCAGTTTTGGTATCAGCTGGGTGCGACAGGCCCCTGGGCAAGGGCTTGAGTGGCTGGGATGGAGCAGCACTGACAATGGTAACACAAACTATGCACAGAAGTTCCAGGGCAGAGTCACCATGACCACAGACACATCCACGACCACAGCCTACATGGAGCTGAGGAGCCTAAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGATTTAGGGCGGTGTACCAATACCGGGTGCTATCGCAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAG",
        sequence_aa="QVQLVQSGAEVKKPGASVKVSCKASGYTFSSFGISWVRQAPGQGLEWLGWSSTDNGNTNYAQKFQGRVTMTTDTSTTTAYMELRSLRSDDTAVYYCARDLGRCTNTGCYRNWFDPWGQGTLVTVSS",
        v_call="IGHV1-18*01",
        j_call="IGKJ5*01",
        junction_aa="CARDLGRCTNTGCYRNWFDPW",
    )


@pytest.fixture
def ambiguous_nucleotide_antibody():
    return Antibody(
        sequence="CAGGTTCAGCTGGTGCAGTCTGGAGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGTTACACCTTTAGCAGTTTTGGTATCAGCTGGGTGCGACAGGCCCCTGGGCAAGGGCNTGAGTGGCTGGGATGGAGCAGCACTGACAATGGTAACACAAACTATGCACAGAAGTTCCAGGGCAGAGTCACCATGACCACAGACACATCCACGACCACAGCCTACATGGAGCTGAGGAGCCTAAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGATTTAGGGCGGTGTACCAATACCGGGTGCTATCGCAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAG",
        sequence_aa="QVQLVQSGAEVKKPGASVKVSCKASGYTFSSFGISWVRQAPGQGLEWLGWSSTDNGNTNYAQKFQGRVTMTTDTSTTTAYMELRSLRSDDTAVYYCARDLGRCTNTGCYRNWFDPWGQGTLVTVSS",
        v_call="IGHV1-18*01",
        j_call="IGHJ5*01",
        junction_aa="CARDLGRCTNTGCYRNWFDPW",
    )


def test_productive_antibody(productive_antibody):
    ab = assess_productivity(productive_antibody)
    assert ab.productive
    assert not ab.productivity_issues


def test_stop_codon_antibody(stop_codon_antibody):
    ab = assess_productivity(stop_codon_antibody)
    assert not ab.productive
    assert "stop codon" in ab.productivity_issues


def test_missing_conserved_cysteine_antibody(missing_conserved_cysteine_antibody):
    ab = assess_productivity(missing_conserved_cysteine_antibody)
    assert not ab.productive
    assert "conserved C" in ab.productivity_issues


def test_missing_conserved_tryptophan_antibody(missing_conserved_tryptophan_antibody):
    ab = assess_productivity(missing_conserved_tryptophan_antibody)
    assert not ab.productive
    assert "conserved W/F" in ab.productivity_issues


def test_locus_mismatch_antibody(locus_mismatch_antibody):
    ab = assess_productivity(locus_mismatch_antibody)
    assert not ab.productive
    assert "locus mismatch" in ab.productivity_issues


def test_ambiguous_nucleotide_antibody(ambiguous_nucleotide_antibody):
    ab = assess_productivity(ambiguous_nucleotide_antibody)
    assert not ab.productive
    assert "ambiguous nucleotide" in ab.productivity_issues
