# Tests for CDR3 detection
from igdiscover.species import find_cdr3
from igdiscover.utils import nt_to_aa


def split(s):
	for line in s.split('\n'):
		line = line.strip()
		if line:
			yield line.split()


def assert_cdr3_detection(chain, s):
	for amino_acids, sequence in split(s):
		for offset in range(3):
			target = sequence[offset:]
			match = find_cdr3(target, chain)
			assert match is not None
			assert nt_to_aa(target[match[0]:match[1]]) == amino_acids, (chain, amino_acids, offset)


def test_cdr3_detection_heavy():
	heavy = """
	ARRLHSGSYILFDY CAGGTGACCTTGAAGGAGTCTGGTCCTGCGCTGGTGAAACCCACACAGACCCTCACGCTGACCTGCACCTTCTCTGGGTTCTCACTCAGCACTAGTGGTATGGGTGTGGGCTGGATCCGTCAGCCCTCACGGAAGACCCTGGAGTGGCTTGCACACATTTATTGGAATGATGATAAATACTACAGCACATCGCTGAAGAGCAGGCTCACCATCTCCAAGGACACCTCCAAAAACCAGGTGGTTCTAACAATGACCAACATGGACCCTGTGGACACAGCCACATATTACTGTGCACGGAGACTTCATAGTGGGAGCTACATTCTCTTTGACTACTGGGGCCAGGGAGTCCTGGTCACCGTCTCCTCAGGGAGTGCATCCGCCCCAACCCTTTTCCCCCTCGTCTCCTGTGA

	ARIKWLRSPGYGYFDF CAGGTGACCTTGAAGGAGTCTGGTCCTGCGCTGGTGAGACCCACACAGACCCTCACTCTGACCTGCACCTTCTCTGGGTTCTCAATCAGCACCTCTGGAACAGGTGTGGGCTGGATCCGTCAGCCCCCAGGGAAGGCCCTGGAATGGCTTGCAAGCATTTATTGGACTGATGCTAAATACTATAGCACATCGCAGAAGAGCAGGCTCACCATCTCCAAGGACACCTCCAGAAACCAGGTGATTCTAACAATGACCAACATGGAGCCTGTGGACATAGCCACATATTTCTGTGCACGGATAAAGTGGCTGCGGTCCCCAGGCTATGGATACTTCGATTTCTGGGGCCCTGGCACCCCAATCACCATCTCCTCAGGGAGTGCATCCGCCCCAACCCTTTTCCCCCTCGTCTCCTGTGA

	ARHGIAAAGTHNWFDP TCAGCCGACAAGTCCATCAGCACCGCCTACCTGCAGTGGAGCAGCCTGAAGGCCTCGGACACCGCCATGGATTACTGTGCGAGACATGGGATAGCAGCAGCTGGTACCCACAACTGGTTCGACCCCTGGGGCCAGGGAACCCTGGTCACCGTCTCCTCAGGGAGTGCATCCGCCCCAACCCTTTTCCCCCTCGTCTCCTGTGAGAATTCCCCGTCGGCAGGTTGTT
	"""
	assert_cdr3_detection('VH', heavy)


def test_cdr3_detection_kappa():
	kappa = """
QQYDSSPRT TTCAGTGGCAGTGGAGCAGGGACAGATTTCACTCTCACCATCAGCAGTCTGGAACCTGAGGATGTCGCAACTTACTACTGTCAGCAGTATGATAGCAGCCCCCGGACGTTCGGCGCTGGGACCAAGCTGGAAATCAAACGGAGTGTGCAGAAGCCAACTATCTCCCTCTTCCCTCCATCATCTGAGGAGG

QQYSSYPYT GAGCTGGCCTCGGGAGTCCCAGCTCGCTTCAGTGGGAGTGGGTCAGGGACTTCTTTCTCTCTCACAATCAGCAACGTGGAGGCTGAAGATGTTGCAACCTATTACTGTCAGCAGTATAGCAGTTATCCGTACACGTTCGGCGCAGGGACCAAGCTGGAAATCAAACGGAGTGTGCAGAAGCCAACTATCTCCCTCTTCCCTCCATCATCTGAGGAGG

LQYDSSPYT ATTCTCACCATCAGCAGCCTGCAGCCTGAAGACTTTGCAACTTACTACTGTCTACAGTATGATAGTTCCCCGTACACGTTCGGCGCAGGGACCAAGCTGGAAATCAAACGGAGTGTGCAGAAGCCAACTATCTCCCTCTTCCCTCCATCATCTGAGGAGG

FQYYSGRLT ACAGACTTCACTCTCACCATCAGCAGCCTGCAGCCTGAGGACATTGCAGTTTATTACTGTTTCCAGTATTACAGCGGGAGACTCACGTTCGGAGGAGGGACCCGCTTGGAAATCAAACGGAGTGTGCAGAAGCCAACTATCTCCCTCTTCCCTCCATCATCTGAGGAGG
	"""
	assert_cdr3_detection('VK', kappa)


def test_cdr3_detection_lambda():
	lambda_ = """
LTYHGNSGTFV GGATCCAAAAACCCCTCAGCCAATGCAGGAATTTTGCTCATCTCTGAACTCCAGAATGAGGATGAGGCTGACTATTACTGTCTGACATATCATGGTAATAGTGGTACTTTTGTATTCGGTGGAGGAACCAAGCTGACCGTCCTAGGTCAGCCCAAGTCTGCCCCCACAGTCAGCCTGTTCC

QLWDANSTV ATGGCCACACTGACCATCACTGGCGCCCAGGGTGAGGACGAGGCCGACTATTGCTGTCAGTTGTGGGATGCTAACAGTACTGTGTTCGGTGGAGGAACCACGCTGACCGTCCTAGGTCAGCCCAAGTCTGCCCCCACAGTCAGCCTGTTCCCGCCCTCCTC

GVGYSGGYV GATCGCTACTTAACCATCTCCAACATCCAGCCTGAAGACGAGGCTGACTATTTCTGTGGTGTGGGTTATAGCGGTGGTTATGTATTCGGTGGAGGAACCAAGTTGACCGTCCTAGGTCAGCCCAAGTCTGCTCCCACAGTCAGCCTGTTCCCGCCCTCCTC
	"""
	assert_cdr3_detection('VL', lambda_)
