# 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 unittest import skipIf
from .utils import (
    run,
    files_equal,
    datapath,
    cutpath,
    redirect_stderr,
    temporary_path,
    no_internet,
    no_import,
)


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',
    )


@skipIf(
    no_internet("https://ncbi.nlm.nih.gov") or no_import('srastream'),
    "No internet connection or srastream not importable",
)
def test_sra():
    run(
        '-b CTGGAGTTCAGACGTGTGCTCT --max-reads 100',
        'SRR2040662_trimmed.fq',
        sra_accn='SRR2040662',
    )
