from sqt.dna import (reverse_complement, n_intervals, intervals_complement,
	amino_acid_regex, GENETIC_CODE, nt_to_aa)

def test_complement_string():
	rc = reverse_complement
	assert rc('') == ''
	assert rc('A') == 'T'
	assert rc('C') == 'G'
	assert rc('TG') == 'CA'
	assert rc('N') == 'N'
	assert rc('a') == 't'

	assert rc('ACGTUMRWSYKVHDBN') == 'NVHDBMRSWYKAACGT'
	assert rc('acgtumrwsykvhdbn') == 'nvhdbmrswykaacgt'
	assert rc('ACGTUMRWSYKVHDBNacgtumrwsykvhdbn') == 'nvhdbmrswykaacgtNVHDBMRSWYKAACGT'

#ACGTUMRWSYKVHDBN
#TGCAAKYWSRMBDHVN


def test_complement_bytes():
	rc = reverse_complement
	assert rc(b'') == b''
	assert rc(b'A') == b'T'
	assert rc(b'C') == b'G'
	assert rc(b'TG') == b'CA'
	assert rc(b'N') == b'N'
	assert rc(b'a') == b't'

	assert rc(b'ACGTUMRWSYKVHDBN') == b'NVHDBMRSWYKAACGT'
	assert rc(b'acgtumrwsykvhdbn') == b'nvhdbmrswykaacgt'
	assert rc(b'ACGTUMRWSYKVHDBNacgtumrwsykvhdbn') == b'nvhdbmrswykaacgtNVHDBMRSWYKAACGT'


def test_n_intervals():
	assert list(n_intervals(b'', N=ord(b'N'))) == []
	assert list(n_intervals(b'N', N=ord(b'N'))) == [(0, 1)]
	assert list(n_intervals(b'n', N=ord(b'N'))) == [(0, 1)]
	assert list(n_intervals(b'an', N=ord(b'N'))) == [(1, 2)]
	assert list(n_intervals(b'ACGTNNAC', N=ord(b'N'))) == [(4, 6)]
	assert list(n_intervals(b'NCGTNNACN', N=ord(b'N'))) == [(0, 1), (4, 6), (8, 9)]


def test_n_intervals():
	assert list(n_intervals('')) == []
	assert list(n_intervals('N')) == [(0, 1)]
	assert list(n_intervals('n')) == [(0, 1)]
	assert list(n_intervals('an')) == [(1, 2)]
	assert list(n_intervals('ACGTNNAC')) == [(4, 6)]
	assert list(n_intervals('NCGTNNACN')) == [(0, 1), (4, 6), (8, 9)]


def test_intervals_complement():
	assert list(intervals_complement([], length=10)) == [(0, 10)]
	assert list(intervals_complement([(0, 10)], length=10)) == []
	assert list(intervals_complement([(0, 2), (4, 6)], length=10)) == [(2, 4), (6,10)]
	assert list(intervals_complement([(1, 2), (4, 6)], length=10)) == [(0,1), (2, 4), (6,10)]
	assert list(intervals_complement([(0, 1), (3, 10)], length=5)) == [(1, 3)]
	assert list(intervals_complement([(2, 10)], length=5)) == [(0, 2)]
	assert list(intervals_complement([(0, 10)], length=5)) == []


def test_amino_acid_regex():
	for codon1, aa1 in GENETIC_CODE.items():
		r = amino_acid_regex(aa1, compile=True)
		for codon2, aa2 in GENETIC_CODE.items():
			m = r.match(codon2)
			assert bool(m) == (aa1 == aa2)
			if m:
				assert m.group(0) == codon2


def test_nt_to_aa():
	assert nt_to_aa('') == ''
	assert nt_to_aa('A') == '*'
	assert nt_to_aa('AC') == '*'
	assert nt_to_aa('AAA') == 'K'
	assert nt_to_aa('AAAG') == 'K*'
	assert nt_to_aa('AAAGG') == 'K*'
	assert nt_to_aa('AAAATG') == 'KM'
	assert nt_to_aa('TGA') == '*'
	assert nt_to_aa('TGAT') == '**'
	assert nt_to_aa('TGATT') == '**'
	assert nt_to_aa('TAGTTT') == '*F'
	assert nt_to_aa('TAATTTC') == '*F*'
	assert nt_to_aa('taatttatgc') == '*FM*'

