# coding: utf-8
# TODO
# test with the --output option
# test reading from standard input
from io import StringIO
import os
from pytest import raises
import sys
from atropos.commands import execute_cli, get_command
from .utils import (
    run, files_equal, datapath, cutpath, redirect_stderr, temporary_path)

def test_example():
    run('-N -b ADAPTER', 'example.fa', 'example.fa')

def test_small():
    run('-b TTAGACATATCTCCGTCG', 'small.fastq', 'small.fastq')

def test_empty():
    '''empty input'''
    run('-a TTAGACATATCTCCGTCG', 'empty.fastq', 'empty.fastq')

def test_newlines():
    '''DOS/Windows newlines'''
    run('-e 0.12 -b TTAGACATATCTCCGTCG', 'dos.fastq', 'dos.fastq')

def test_lowercase():
    '''lowercase adapter'''
    run('-b ttagacatatctccgtcg', 'lowercase.fastq', 'small.fastq')


def test_rest():
    '''-r/--rest-file'''
    with temporary_path('rest.tmp') as rest_tmp:
        run(['-b', 'ADAPTER', '-N', '-r', rest_tmp], "rest.fa", "rest.fa")
        assert files_equal(datapath('rest.txt'), rest_tmp)


def test_restfront():
    with temporary_path("rest.txt") as path:
        run(['-g', 'ADAPTER', '-N', '-r', path], "restfront.fa", "rest.fa")
        assert files_equal(datapath('restfront.txt'), path)


def test_discard():
    '''--discard'''
    run("-b TTAGACATATCTCCGTCG --discard", "discard.fastq", "small.fastq")


def test_discard_untrimmed():
    '''--discard-untrimmed'''
    run('-b CAAGAT --discard-untrimmed', 'discard-untrimmed.fastq', 'small.fastq')


def test_plus():
    '''test if sequence name after the "+" is retained'''
    run("-e 0.12 -b TTAGACATATCTCCGTCG", "plus.fastq", "plus.fastq")


def test_extensiontxtgz():
    '''automatic recognition of "_sequence.txt.gz" extension'''
    run("-b TTAGACATATCTCCGTCG", "s_1_sequence.txt", "s_1_sequence.txt.gz")


def test_format():
    '''the -f/--format parameter'''
    run("-f fastq -b TTAGACATATCTCCGTCG", "small.fastq", "small.myownextension")


def test_minimum_length():
    '''-m/--minimum-length'''
    run("-c -m 5 -a 330201030313112312", "minlen.fa", "lengths.fa")


def test_too_short():
    '''--too-short-output'''
    run("-c -m 5 -a 330201030313112312 --too-short-output tooshort.tmp.fa", "minlen.fa", "lengths.fa")
    assert files_equal(datapath('tooshort.fa'), "tooshort.tmp.fa")
    os.remove('tooshort.tmp.fa')


def test_too_short_no_primer():
    '''--too-short-output and --trim-primer'''
    run("-c -m 5 -a 330201030313112312 --trim-primer --too-short-output tooshort.tmp.fa", "minlen.noprimer.fa", "lengths.fa")
    assert files_equal(datapath('tooshort.noprimer.fa'), "tooshort.tmp.fa")
    os.remove('tooshort.tmp.fa')


def test_maximum_length():
    '''-M/--maximum-length'''
    run("-c -M 5 -a 330201030313112312", "maxlen.fa", "lengths.fa")


def test_too_long():
    '''--too-long-output'''
    run("-c -M 5 --too-long-output toolong.tmp.fa -a 330201030313112312", "maxlen.fa", "lengths.fa")
    assert files_equal(datapath('toolong.fa'), "toolong.tmp.fa")
    os.remove('toolong.tmp.fa')


def test_length_tag():
    '''454 data; -n and --length-tag'''
    run("-n 3 -e 0.1 --length-tag length= " \
        "-b TGAGACACGCAACAGGGGAAAGGCAAGGCACACAGGGGATAGG "\
        "-b TCCATCTCATCCCTGCGTGTCCCATCTGTTCCCTCCCTGTCTCA", '454.fa', '454.fa')

def test_overlap_a():
    '''-O/--overlap with -a (-c omitted on purpose)'''
    run("-O 10 -a 330201030313112312 -e 0.0 -N", "overlapa.fa", "overlapa.fa")

def test_overlap_b():
    '''-O/--overlap with -b'''
    run("-O 10 -b TTAGACATATCTCCGTCG -N", "overlapb.fa", "overlapb.fa")

def test_qualtrim():
    '''-q with low qualities'''
    run("-q 10 -a XXXXXX", "lowqual.fastq", "lowqual.fastq")

def test_qualbase():
    '''-q with low qualities, using ascii(quality+64) encoding'''
    run("-q 10 --quality-base 64 -a XXXXXX", "illumina64.fastq", "illumina64.fastq")

def test_quality_trim_only():
    '''only trim qualities, do not remove adapters'''
    run("-q 10 --quality-base 64", "illumina64.fastq", "illumina64.fastq")

def test_twoadapters():
    '''two adapters'''
    run("-a AATTTCAGGAATT -a GTTCTCTAGTTCT", "twoadapters.fasta", "twoadapters.fasta")

def test_polya():
    '''poly-A tails'''
    run("-m 24 -O 10 -a AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", "polya.fasta", "polya.fasta")

def test_polya_brace_notation():
    '''poly-A tails'''
    run("-m 24 -O 10 -a A{35}", "polya.fasta", "polya.fasta")

def test_mask_adapter():
    '''mask adapter with N (reads maintain the same length)'''
    run("-b CAAG -n 3 --mask-adapter", "anywhere_repeat.fastq", "anywhere_repeat.fastq")

def test_gz_multiblock():
    '''compressed gz file with multiple blocks (created by concatenating two .gz files)'''
    run("-b TTAGACATATCTCCGTCG", "small.fastq", "multiblock.fastq.gz")

def test_suffix():
    '''-y/--suffix parameter, combined with _F3'''
    run("-c -e 0.12 -a 1=330201030313112312 -y _my_suffix_{name} --strip-f3", "suffix.fastq", "solid.csfasta", qualfile='solid.qual')

def test_read_wildcard():
    '''test wildcards in reads'''
    run("--match-read-wildcards -b ACGTACGT", "wildcard.fa", "wildcard.fa")

def test_adapter_wildcard():
    '''wildcards in adapter'''
    for adapter_type, expected in (
            ("-a", "wildcard_adapter.fa"),
            ("-b", "wildcard_adapter_anywhere.fa")):
        with temporary_path("wildcardtmp.txt") as wildcardtmp:
            run("--wildcard-file {0} {1} ACGTNNNACGT".format(wildcardtmp, adapter_type),
                expected, "wildcard_adapter.fa")
            with open(wildcardtmp) as wct:
                lines = wct.readlines()
            lines = [ line.strip() for line in lines ]
            assert lines == ['AAA 1', 'GGG 2', 'CCC 3b', 'TTT 4b']

def test_wildcard_N():
    '''test 'N' wildcard matching with no allowed errors'''
    run("-e 0 -a GGGGGGG --match-read-wildcards", "wildcardN.fa", "wildcardN.fa")

def test_illumina_adapter_wildcard():
    run("-a VCCGAMCYUCKHRKDCUBBCNUWNSGHCGU", "illumina.fastq", "illumina.fastq.gz")

def test_adapter_front():
    '''test adapter in front'''
    run("--front ADAPTER -N", "examplefront.fa", "example.fa")

def test_literal_N():
    '''test matching literal 'N's'''
    run("-N -e 0.2 -a NNNNNNNNNNNNNN", "trimN3.fasta", "trimN3.fasta")

def test_literal_N2():
    run("-N -O 1 -g NNNNNNNNNNNNNN", "trimN5.fasta", "trimN5.fasta")

def test_literal_N_brace_notation():
    '''test matching literal 'N's'''
    run("-N -e 0.2 -a N{14}", "trimN3.fasta", "trimN3.fasta")

def test_literal_N2_brace_notation():
    run("-N -O 1 -g N{14}", "trimN5.fasta", "trimN5.fasta")

def test_anchored_front():
    run("-g ^FRONTADAPT -N", "anchored.fasta", "anchored.fasta")

def test_anchored_front_ellipsis_notation():
    run("-a FRONTADAPT... -N", "anchored.fasta", "anchored.fasta")

def test_anchored_back():
    run("-a BACKADAPTER$ -N", "anchored-back.fasta", "anchored-back.fasta")

def test_anchored_back_no_indels():
    run("-a BACKADAPTER$ -N --no-indels", "anchored-back.fasta", "anchored-back.fasta")


def test_no_indels():
    run('-a TTAGACATAT -g GAGATTGCCA --no-indels', 'no_indels.fasta', 'no_indels.fasta')


def test_issue_46():
    '''issue 46 - IndexError with --wildcard-file'''
    with temporary_path("wildcardtmp.txt") as wildcardtmp:
        run("--anywhere=AACGTN --wildcard-file={0}".format(wildcardtmp), "issue46.fasta", "issue46.fasta")

def test_strip_suffix():
    run("--strip-suffix _sequence -a XXXXXXX", "stripped.fasta", "simple.fasta")


def test_info_file():
    # The true adapter sequence in the illumina.fastq.gz data set is
    # GCCTAACTTCTTAGACTGCCTTAAGGACGT (fourth base is different)
    #
    with temporary_path("infotmp.txt") as infotmp:
        run(["--info-file", infotmp, '-a', 'adapt=GCCGAACTTCTTAGACTGCCTTAAGGACGT'], "illumina.fastq", "illumina.fastq.gz")
        assert files_equal(cutpath('illumina.info.txt'), infotmp)


def test_info_file_times():
    with temporary_path("infotmp.txt") as infotmp:
        run(["--info-file", infotmp, '--times', '2', '-a', 'adapt=GCCGAACTTCTTA', '-a', 'adapt2=GACTGCCTTAAGGACGT'], "illumina5.fastq", "illumina5.fastq")
        assert files_equal(cutpath('illumina5.info.txt'), infotmp)


def test_info_file_fasta():
    with temporary_path("infotmp.txt") as infotmp:
        # Just make sure that it runs
        run(['--info-file', infotmp, '-a', 'TTAGACATAT', '-g', 'GAGATTGCCA', '--no-indels'], 'no_indels.fasta', 'no_indels.fasta')


def test_named_adapter():
    run("-a MY_ADAPTER=GCCGAACTTCTTAGACTGCCTTAAGGACGT", "illumina.fastq", "illumina.fastq.gz")


def test_adapter_with_U():
    run("-a GCCGAACUUCUUAGACUGCCUUAAGGACGU", "illumina.fastq", "illumina.fastq.gz")


def test_no_trim():
    ''' --no-trim '''
    run("--no-trim --discard-untrimmed -a CCCTAGTTAAAC", 'no-trim.fastq', 'small.fastq')


def test_bzip2():
    '''test bzip2 support'''
    run('-b TTAGACATATCTCCGTCG', 'small.fastq', 'small.fastq.bz2')


try:
    import lzma

    def test_xz():
        '''test xz support'''
        run('-b TTAGACATATCTCCGTCG', 'small.fastq', 'small.fastq.xz')
except ImportError:
    pass


def test_qualfile_only():
    with raises(SystemExit), redirect_stderr():
        execute_cli(['-sq', datapath('E3M.qual')])


def test_no_args():
    with redirect_stderr():
        assert execute_cli() != 0


def test_two_fastqs():
    with raises(SystemExit), redirect_stderr():
        execute_cli(['-pe1', datapath('paired.1.fastq'), '-pe2', datapath('paired.2.fastq')])


def test_anchored_no_indels():
    '''anchored 5' adapter, mismatches only (no indels)'''
    run('-g ^TTAGACATAT --no-indels -e 0.1', 'anchored_no_indels.fasta', 'anchored_no_indels.fasta')


def test_anchored_no_indels_wildcard_read():
    '''anchored 5' adapter, mismatches only (no indels), but wildcards in the read count as matches'''
    run('-g ^TTAGACATAT --match-read-wildcards --no-indels -e 0.1', 'anchored_no_indels_wildcard.fasta', 'anchored_no_indels.fasta')


def test_anchored_no_indels_wildcard_adapt():
    '''anchored 5' adapter, mismatches only (no indels), but wildcards in the adapter count as matches'''
    run('-g ^TTAGACANAT --no-indels -e 0.1', 'anchored_no_indels.fasta', 'anchored_no_indels.fasta')


def test_unconditional_cut_front():
    run('-u 5', 'unconditional-front.fastq', 'small.fastq')


def test_unconditional_cut_back():
    run('-u -5', 'unconditional-back.fastq', 'small.fastq')


def test_unconditional_cut_both():
    run('-u -5 -u 5', 'unconditional-both.fastq', 'small.fastq')


def test_untrimmed_output():
    with temporary_path('untrimmed.tmp.fastq') as tmp:
        run(['-a', 'TTAGACATATCTCCGTCG', '--untrimmed-output', tmp], 'small.trimmed.fastq', 'small.fastq')
        assert files_equal(cutpath('small.untrimmed.fastq'), tmp)


def test_adapter_file():
    run('-a file:' + datapath('adapter.fasta'), 'illumina.fastq', 'illumina.fastq.gz')

def test_adapter_file_5p_anchored():
    run('-N -g file:' + datapath('prefix-adapter.fasta'), 'anchored.fasta', 'anchored.fasta')

def test_adapter_file_3p_anchored():
    run('-N -a file:' + datapath('suffix-adapter.fasta'), 'anchored-back.fasta', 'anchored-back.fasta')


def test_adapter_file_5p_anchored_no_indels():
    run('-N --no-indels -g file:' + datapath('prefix-adapter.fasta'), 'anchored.fasta', 'anchored.fasta')


def test_adapter_file_3p_anchored_no_indels():
    run('-N --no-indels -a file:' + datapath('suffix-adapter.fasta'), 'anchored-back.fasta', 'anchored-back.fasta')


def test_demultiplex():
    multiout = os.path.join(os.path.dirname(__file__), 'data', 'tmp-demulti.{name}.fasta')
    params = ['-a', 'first=AATTTCAGGAATT', '-a', 'second=GTTCTCTAGTTCT', '-o', multiout, '-se', datapath('twoadapters.fasta')]
    command = get_command('trim')
    result = command.execute(params)
    assert isinstance(result, tuple)
    assert len(result) == 2
    assert result[0] == 0
    assert files_equal(cutpath('twoadapters.first.fasta'), multiout.format(name='first'))
    assert files_equal(cutpath('twoadapters.second.fasta'), multiout.format(name='second'))
    assert files_equal(cutpath('twoadapters.unknown.fasta'), multiout.format(name='unknown'))
    os.remove(multiout.format(name='first'))
    os.remove(multiout.format(name='second'))
    os.remove(multiout.format(name='unknown'))


def test_max_n():
    run('--max-n 0', 'maxn0.fasta', 'maxn.fasta')
    run('--max-n 1', 'maxn1.fasta', 'maxn.fasta')
    run('--max-n 2', 'maxn2.fasta', 'maxn.fasta')
    run('--max-n 0.2', 'maxn0.2.fasta', 'maxn.fasta')
    run('--max-n 0.4', 'maxn0.4.fasta', 'maxn.fasta')


def test_quiet_is_quiet():
    captured_standard_output = StringIO()
    captured_standard_error = StringIO()
    try:
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        sys.stdout = captured_standard_output
        sys.stderr = captured_standard_error
        execute_cli(['-o', '/dev/null', '--quiet', '-a', 'XXXX', '-se', datapath('illumina.fastq.gz')])
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr
    print(captured_standard_output.getvalue())
    assert captured_standard_output.getvalue() == ''
    assert captured_standard_error.getvalue() == ''


def test_nextseq():
    run('--nextseq-trim 22', 'nextseq.fastq', 'nextseq.fastq')


def test_linked():
    run('-a AAAAAAAAAA...TTTTTTTTTT', 'linked.fasta', 'linked.fasta')


def test_fasta():
    run('-a TTAGACATATCTCCGTCG', 'small.fasta', 'small.fastq')

def test_custom_bisulfite_1():
    run('-b TTAGACATATCTCCGTCG -q 0,0 --bisulfite 2,2,1,1', 'small.fastq', 'small.fastq')
def test_custom_bisulfite_2():
    run('-b TTAGACATATCTCCGTCG -q 0,0 --bisulfite 15,15,1,1', 'small_mincut1.fastq', 'small.fastq')
def test_custom_bisulfite_3():
    run('-b TTAGACATATCTCCGTCG -q 0,0 --bisulfite 2,2,1,0', 'small_mincut2.fastq', 'small.fastq')
def test_custom_bisulfite_4():
    run('-b TTAGACATATCTCCGTCG -q 0,0 --bisulfite 2,2,0,0', 'small_mincut3.fastq', 'small.fastq')
