# Copyright 2012 by Wibowo Arindrarto.  All rights reserved.
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.

"""Tests for SearchIO writing."""

import os
import unittest

from Bio import SearchIO

from search_tests_common import compare_search_obj


class WriteCases(unittest.TestCase):
    def tearDown(self):
        if os.path.exists(self.out):
            os.remove(self.out)

    def parse_write_and_compare(
        self, source_file, source_format, out_file, out_format, **kwargs
    ):
        """Compare parsed QueryResults after they have been written to a file."""
        source_qresults = list(SearchIO.parse(source_file, source_format, **kwargs))
        SearchIO.write(source_qresults, out_file, out_format, **kwargs)
        out_qresults = list(SearchIO.parse(out_file, out_format, **kwargs))
        for source, out in zip(source_qresults, out_qresults):
            self.assertTrue(compare_search_obj(source, out))

    def read_write_and_compare(
        self, source_file, source_format, out_file, out_format, **kwargs
    ):
        """Compare read QueryResults after it has been written to a file."""
        source_qresult = SearchIO.read(source_file, source_format, **kwargs)
        SearchIO.write(source_qresult, out_file, out_format, **kwargs)
        out_qresult = SearchIO.read(out_file, out_format, **kwargs)
        self.assertTrue(compare_search_obj(source_qresult, out_qresult))


class BlastXmlWriteCases(WriteCases):

    fmt = "blast-xml"
    out = os.path.join("Blast", "test_write.xml")

    def test_write_single_from_blastxml(self):
        """Test blast-xml writing from blast-xml, BLAST 2.2.26+, single query (xml_2226_blastp_004.xml)."""
        source = os.path.join("Blast", "xml_2226_blastp_004.xml")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_multiple_from_blastxml(self):
        """Test blast-xml writing from blast-xml, BLAST 2.2.26+, multiple queries (xml_2226_blastp_001.xml)."""
        source = os.path.join("Blast", "xml_2226_blastp_001.xml")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)


class BlastTabWriteCases(WriteCases):

    fmt = "blast-tab"
    out = os.path.join("Blast", "test_write.txt")

    def test_write_single_from_blasttab(self):
        """Test blast-tab writing from blast-tab, BLAST 2.2.26+, single query (tab_2226_tblastn_004.txt)."""
        source = os.path.join("Blast", "tab_2226_tblastn_004.txt")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_multiple_from_blasttab(self):
        """Test blast-tab writing from blast-tab, BLAST 2.2.26+, multiple queries (tab_2226_tblastn_001.txt)."""
        source = os.path.join("Blast", "tab_2226_tblastn_001.txt")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_single_from_blasttabc(self):
        """Test blast-tabc writing from blast-tabc, BLAST 2.2.26+, single query (tab_2226_tblastn_008.txt)."""
        source = os.path.join("Blast", "tab_2226_tblastn_008.txt")
        self.parse_write_and_compare(
            source, self.fmt, self.out, self.fmt, comments=True
        )
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt, comments=True)

    def test_write_multiple_from_blasttabc(self):
        """Test blast-tabc writing from blast-tabc, BLAST 2.2.26+, multiple queries (tab_2226_tblastn_005.txt)."""
        source = os.path.join("Blast", "tab_2226_tblastn_005.txt")
        self.parse_write_and_compare(
            source, self.fmt, self.out, self.fmt, comments=True
        )

    def test_write_multiple_from_blasttabc_allfields(self):
        """Test blast-tabc writing from blast-tabc, BLAST 2.2.28+, multiple queries (tab_2228_tblastx_001.txt)."""
        source = os.path.join("Blast", "tab_2228_tblastx_001.txt")
        fields = [
            "qseqid",
            "qgi",
            "qacc",
            "qaccver",
            "qlen",
            "sseqid",
            "sallseqid",
            "sgi",
            "sallgi",
            "sacc",
            "saccver",
            "sallacc",
            "slen",
            "qstart",
            "qend",
            "sstart",
            "send",
            "qseq",
            "sseq",
            "evalue",
            "bitscore",
            "score",
            "length",
            "pident",
            "nident",
            "mismatch",
            "positive",
            "gapopen",
            "gaps",
            "ppos",
            "frames",
            "qframe",
            "sframe",
            "btop",
            "staxids",
            "sscinames",
            "scomnames",
            "sblastnames",
            "sskingdoms",
            "stitle",
            "salltitles",
            "sstrand",
            "qcovs",
            "qcovhsp",
        ]
        self.parse_write_and_compare(
            source, self.fmt, self.out, self.fmt, comments=True, fields=fields
        )


class HmmerTabWriteCases(WriteCases):

    fmt = "hmmer3-tab"
    out = os.path.join("Hmmer", "test_write.txt")

    def test_write_single_from_hmmertab(self):
        """Test hmmer3-tab writing from hmmer3-tab, HMMER 3.0, single query (tab_30_hmmscan_004.out)."""
        source = os.path.join("Hmmer", "tab_30_hmmscan_004.out")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_multiple_from_hmmertab(self):
        """Test hmmer3-tab writing from hmmer3-tab, HMMER 3.0, multiple queries (tab_30_hmmscan_001.out)."""
        source = os.path.join("Hmmer", "tab_30_hmmscan_001.out")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)


class HmmerDomtabWriteCases(WriteCases):

    out = os.path.join("Hmmer", "test_write.txt")

    def test_write_single_from_hmmscandomtab(self):
        """Test hmmscan-domtab writing from hmmscan-domtab, HMMER 3.0, single query (tab_30_hmmscan_004.out)."""
        source = os.path.join("Hmmer", "domtab_30_hmmscan_004.out")
        fmt = "hmmscan3-domtab"
        self.parse_write_and_compare(source, fmt, self.out, fmt)
        self.read_write_and_compare(source, fmt, self.out, fmt)

    def test_write_multiple_from_hmmscandomtab(self):
        """Test hmmscan-domtab writing from hmmscan-domtab, HMMER 3.0, multiple queries (tab_30_hmmscan_001.out)."""
        source = os.path.join("Hmmer", "domtab_30_hmmscan_001.out")
        fmt = "hmmscan3-domtab"
        self.parse_write_and_compare(source, fmt, self.out, fmt)

    def test_write_single_from_hmmsearchdomtab(self):
        """Test hmmsearch-domtab writing from hmmsearch-domtab, HMMER 3.0, single query (tab_30_hmmscan_004.out)."""
        source = os.path.join("Hmmer", "domtab_30_hmmsearch_001.out")
        fmt = "hmmsearch3-domtab"
        self.parse_write_and_compare(source, fmt, self.out, fmt)
        self.read_write_and_compare(source, fmt, self.out, fmt)


class BlatPslWriteCases(WriteCases):

    fmt = "blat-psl"
    out = os.path.join("Blat", "test_write.txt")

    def test_write_single_from_blatpsl(self):
        """Test blat-psl writing from blat-psl, single query (psl_34_004.psl)."""
        source = os.path.join("Blat", "psl_34_004.psl")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_single_from_blatpsl_protein_query(self):
        """Test blat-psl writing from blat-psl, single query (psl_35_002.psl)."""
        source = os.path.join("Blat", "psl_35_002.psl")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_multiple_from_blatpsl(self):
        """Test blat-psl writing from blat-psl, multiple queries (psl_34_001.psl)."""
        source = os.path.join("Blat", "psl_34_001.psl")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt)

    def test_write_single_from_blatpslx(self):
        """Test blat-pslx writing from blat-pslx, single query (pslx_34_004.pslx)."""
        source = os.path.join("Blat", "pslx_34_004.pslx")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt, pslx=True)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt, pslx=True)

    def test_write_single_from_blatpslx_protein_query(self):
        """Test blat-pslx writing from blat-pslx, single query (pslx_35_002.pslx)."""
        source = os.path.join("Blat", "pslx_35_002.pslx")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt, pslx=True)
        self.read_write_and_compare(source, self.fmt, self.out, self.fmt, pslx=True)

    def test_write_multiple_from_blatpslx(self):
        """Test blat-pslx writing from blat-pslx, multiple queries (pslx_34_001.pslx)."""
        source = os.path.join("Blat", "pslx_34_001.pslx")
        self.parse_write_and_compare(source, self.fmt, self.out, self.fmt, pslx=True)


if __name__ == "__main__":
    runner = unittest.TextTestRunner(verbosity=2)
    unittest.main(testRunner=runner)