import json
import os
import pathlib
import shutil
import zipfile

from os.path import basename, join
from tempfile import TemporaryDirectory
from unittest import TestCase, main
from unittest.mock import Mock, patch

import numpy

from numpy.testing import assert_allclose

from cogent3 import DNA
from cogent3.app import align as align_app
from cogent3.app import io as io_app
from cogent3.app.composable import NotCompleted
from cogent3.app.data_store import WritableZippedDataStore
from cogent3.app.io import write_db
from cogent3.core.alignment import ArrayAlignment, SequenceCollection
from cogent3.core.profile import PSSM, MotifCountsArray, MotifFreqsArray
from cogent3.evolve.fast_distance import DistanceMatrix
from cogent3.maths.util import safe_log
from cogent3.util.table import Table


__author__ = "Gavin Huttley"
__copyright__ = "Copyright 2007-2020, The Cogent Project"
__credits__ = ["Gavin Huttley"]
__license__ = "BSD-3"
__version__ = "2020.12.21a"
__maintainer__ = "Gavin Huttley"
__email__ = "Gavin.Huttley@anu.edu.au"
__status__ = "Alpha"


class TestIo(TestCase):
    basedir = "data"

    def tearDown(self) -> None:
        path = pathlib.Path("delme.tinydb")
        if path.exists():
            path.unlink()

    def test_findall(self):
        """find all files recursively"""
        found = list(io_app.findall(self.basedir, suffix=".fasta"))
        self.assertTrue(len(found) > 1)
        found = list(io_app.findall(self.basedir, suffix=".fasta", limit=2))
        self.assertTrue(len(found) == 2)

        # and with a suffix
        found = list(io_app.findall(self.basedir, suffix=".fasta*"))
        self.assertTrue(len(found) > 2)

    def test_findall_zip(self):
        """find all files recursively in a zip archive"""
        with TemporaryDirectory(dir=".") as dirname:
            zip_path = join(dirname, "new")
            shutil.make_archive(zip_path, "zip", self.basedir)
            zip_path = zip_path + ".zip"  # because shutil adds the suffix
            found = list(io_app.findall(zip_path, suffix="fasta"))
            self.assertTrue(len(found) > 1)
            found = list(io_app.findall(zip_path, suffix="fasta", limit=2))
            self.assertTrue(len(found) == 2)

            # and with a suffix
            found = list(io_app.findall(zip_path, suffix=".fasta*"))
            self.assertTrue(len(found) > 2)

    def test_define_data_store(self):
        """returns an iterable data store"""
        found = io_app.get_data_store(self.basedir, suffix=".fasta")
        self.assertTrue(len(found) > 1)
        found = io_app.get_data_store(self.basedir, suffix=".fasta", limit=2)
        self.assertTrue(len(found) == 2)

        # and with a suffix
        found = list(io_app.get_data_store(self.basedir, suffix=".fasta*"))
        self.assertTrue(len(found) > 2)

        # with a wild-card suffix
        found = list(io_app.get_data_store(self.basedir, suffix="*"))
        self.assertEqual(len(os.listdir(self.basedir)), len(found))

        # raises ValueError if suffix not provided or invalid
        with self.assertRaises(ValueError):
            _ = io_app.get_data_store(self.basedir)

        with self.assertRaises(ValueError):
            _ = io_app.get_data_store(self.basedir, 1)

    def test_load_aligned(self):
        """correctly loads aligned seqs"""

        def validate(paths, loader):
            loaded = list(map(loader, paths))
            for i, aln in enumerate(loaded):
                self.assertTrue(len(aln) > 10)
                self.assertIsInstance(aln, ArrayAlignment)
                self.assertEqual(aln.info.source, paths[i])

        fasta_paths = io_app.get_data_store(self.basedir, suffix=".fasta", limit=2)
        fasta_loader = io_app.load_aligned(format="fasta")
        validate(fasta_paths, fasta_loader)

    def test_load_aligned_nexus(self):
        """should handle nexus too"""
        nexus_paths = io_app.get_data_store(self.basedir, suffix="nex")
        loader = io_app.load_aligned(format="nexus")
        results = [loader(m) for m in nexus_paths]
        for result in results:
            self.assertIsInstance(result, ArrayAlignment)

    def test_load_aligned_from_zip(self):
        """correctly loads aligned seqs from a zip archive"""

        def validate(paths, loader):
            loaded = list(map(loader, paths))
            for i, aln in enumerate(loaded):
                self.assertTrue(len(aln) > 10)
                self.assertIsInstance(aln, ArrayAlignment)
                # paths is only the basename when workjing with zip archives
                # whereas the inpath will have full path of zip archive
                self.assertEqual(aln.info.source, paths[i])
                self.assertEqual(aln.info.source, paths[i])

        with TemporaryDirectory(dir=".") as dirname:
            zip_path = join(dirname, self.basedir.replace(".zip", ""))
            shutil.make_archive(
                base_name=zip_path, root_dir=".", format="zip", base_dir=self.basedir
            )
            zip_path = zip_path + ".zip"  # because shutil adds the suffix
            fasta_paths = list(io_app.findall(zip_path, suffix=".fasta", limit=2))
            fasta_loader = io_app.load_aligned(format="fasta")
            validate(fasta_paths, fasta_loader)

    def test_load_unaligned(self):
        """load_unaligned returns degapped sequence collections"""
        fasta_paths = io_app.get_data_store(self.basedir, suffix=".fasta", limit=2)
        fasta_loader = io_app.load_unaligned(format="fasta")
        for i, seqs in enumerate(map(fasta_loader, fasta_paths)):
            self.assertIsInstance(seqs, SequenceCollection)
            self.assertTrue("-" not in "".join(seqs.to_dict().values()))
            self.assertEqual(seqs.info.source, fasta_paths[i])

        # returns NotCompleted when it's given an alignment/sequence
        # collection
        got = fasta_loader(seqs)
        self.assertIsInstance(got, NotCompleted)

    def test_write_seqs(self):
        """correctly writes sequences out"""
        fasta_paths = list(io_app.findall(self.basedir, suffix=".fasta", limit=2))
        fasta_loader = io_app.load_aligned(format="fasta")
        alns = list(map(fasta_loader, fasta_paths))
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_seqs(dirname, if_exists="ignore")
            wrote = list(map(writer, alns))
            written = list(io_app.findall(dirname, suffix="fasta"))
            for i, wrote in enumerate(written):
                self.assertEqual(alns[i].info.stored, join(dirname, wrote))

    def test_load_json(self):
        """correctly loads an object from json"""
        from cogent3.app.data_store import make_record_for_json

        data = make_record_for_json("delme", DNA, True)
        data = json.dumps(data)
        # straight directory
        with TemporaryDirectory(dir=".") as dirname:
            outpath = join(dirname, "delme.json")
            with open(outpath, "w") as outfile:
                outfile.write(data)
            reader = io_app.load_json()
            got = reader(outpath)
            self.assertIsInstance(got, DNA.__class__)
            self.assertEqual(got, DNA)

        # zipped directory
        with TemporaryDirectory(dir=".") as dirname:
            zip_path = join(dirname, "delme.zip")
            outpath = "delme/delme.json"
            with zipfile.ZipFile(zip_path, "a") as out:
                out.writestr(outpath, data)

            dstore = io_app.get_data_store(zip_path, suffix="json")
            member = dstore.get_member("delme.json")
            reader = io_app.load_json()
            got = reader(member)
            self.assertIsInstance(got, DNA.__class__)
            self.assertEqual(got, DNA)

    def test_write_db_load_db(self):
        """correctly write/load from tinydb"""
        # straight directory
        with TemporaryDirectory(dir=".") as dirname:
            outpath = join(dirname, "delme")
            writer = write_db(outpath, create=True, if_exists="ignore")
            mock = patch("data.source", autospec=True)
            mock.to_json = DNA.to_json
            mock.source = join("blah", "delme.json")
            got = writer(mock)
            writer.data_store.db.close()
            dstore = io_app.get_data_store(f"{outpath}.tinydb", suffix="json")
            reader = io_app.load_db()
            got = reader(dstore[0])
            dstore.close()
            self.assertIsInstance(got, DNA.__class__)
            self.assertEqual(got, DNA)

    def test_write_db_load_db2(self):
        """correctly write/load built-in python from tinydb"""
        with TemporaryDirectory(dir=".") as dirname:
            outpath = join(dirname, "delme")
            writer = write_db(outpath, create=True, if_exists="ignore")
            data = dict(a=[1, 2], b="string")
            m = writer(data, identifier=join("blah", "delme.json"))
            writer.data_store.db.close()
            dstore = io_app.get_data_store(f"{outpath}.tinydb", suffix="json")
            reader = io_app.load_db()
            got = reader(dstore[0])
            dstore.close()
            self.assertEqual(got, data)

    def test_load_db_failure_json_file(self):
        """informative load_db error message when given a json file path"""
        # todo this test has a trapped exception about being unable to delete
        # a file
        with TemporaryDirectory(dir=".") as dirname:
            outpath = join(dirname, "delme")
            writer = write_db(outpath, create=True, if_exists="ignore")
            mock = patch("data.source", autospec=True)
            mock.to_json = DNA.to_json
            mock.source = join("blah", "delme.json")
            got = writer(mock)
            writer.data_store.db.close()
            dstore = io_app.get_data_store(f"{outpath}.tinydb", suffix="json")
            reader = io_app.load_db()
            outpath = join(dirname, "dummy.json")
            with open(outpath, mode="w") as outfile:
                outfile.write("\n\n")

            got = reader(outpath)
            self.assertIsInstance(got, NotCompleted)
            self.assertTrue("json" in got.message)

    def test_load_tabular(self):
        """correctly loads tabular data"""
        rows = [[1, 2], [3, 4], [5, 6.5]]
        table = Table(["A", "B"], data=rows)
        load_table = io_app.load_tabular(sep="\t", with_header=True)
        with TemporaryDirectory(dir=".") as dirname:
            outpath = join(dirname, "delme.tsv")
            table.write(outpath)
            new = load_table(outpath)
            self.assertEqual(new.title, "")
            self.assertEqual(type(new[0, "B"]), type(table[0, "B"]))
            self.assertEqual(type(new[0, "A"]), type(table[0, "A"]))
            outpath = join(dirname, "delme2.tsv")
            with open(outpath, "w") as out:
                out.write("\t".join(table.header[:1]) + "\n")
                for row in table.tolist():
                    row = "\t".join(map(str, row))
                    out.write(row + "\n")
            result = load_table(outpath)
            self.assertIsInstance(result, NotCompleted)

        with TemporaryDirectory(dir=".") as dirname:
            outpath = join(dirname, "delme.zip")
            dstore = WritableZippedDataStore(outpath, suffix="tsv", create=True)
            dstore.write("sample1.tsv", table.to_string("tsv"))
            new = load_table(dstore[0])
            self.assertEqual(type(new[0, "B"]), type(table[0, "B"]))
            self.assertEqual(type(new[0, "A"]), type(table[0, "A"]))

    def test_write_tabular_motif_counts_array(self):
        """correctly writes tabular data for MotifCountsArray"""

        data = [[2, 4], [3, 5], [4, 8]]
        mca = MotifCountsArray(data, "AB")
        loader = io_app.load_tabular(sep="\t")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(mca, identifier=outpath)
            new = loader(outpath)
            # when written to file in tabular form
            # the loaded table will have dim-1 dim-2 as column labels
            # and the key-values pairs listed below; in dict form...
            expected = {
                0: {"dim-1": 0, "dim-2": "A", "value": 2},
                1: {"dim-1": 0, "dim-2": "B", "value": 4},
                2: {"dim-1": 1, "dim-2": "A", "value": 3},
                3: {"dim-1": 1, "dim-2": "B", "value": 5},
                4: {"dim-1": 2, "dim-2": "A", "value": 4},
                5: {"dim-1": 2, "dim-2": "B", "value": 8},
            }
            self.assertEqual(expected, new.to_dict())

    def test_write_tabular_motif_freqs_array(self):
        """correctly writes tabular data for MotifFreqsArray"""

        data = [[0.3333, 0.6667], [0.3750, 0.625], [0.3333, 0.6667]]
        mfa = MotifFreqsArray(data, "AB")
        loader = io_app.load_tabular(sep="\t")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(mfa, identifier=outpath)
            new = loader(outpath)
            # when written to file in tabular form
            # the loaded table will have dim-1 dim-2 as column labels
            # and the key-values pairs listed below; in dict form...
            expected = {
                0: {"dim-1": 0, "dim-2": "A", "value": 0.3333},
                1: {"dim-1": 0, "dim-2": "B", "value": 0.6667},
                2: {"dim-1": 1, "dim-2": "A", "value": 0.3750},
                3: {"dim-1": 1, "dim-2": "B", "value": 0.6250},
                4: {"dim-1": 2, "dim-2": "A", "value": 0.3333},
                5: {"dim-1": 2, "dim-2": "B", "value": 0.6667},
            }
            self.assertEqual(expected, new.to_dict())

    def test_write_tabular_pssm(self):
        """correctly writes tabular data for PSSM"""

        # data from test_profile
        data = numpy.array(
            [
                [0.1, 0.3, 0.5, 0.1],
                [0.25, 0.25, 0.25, 0.25],
                [0.05, 0.8, 0.05, 0.1],
                [0.7, 0.1, 0.1, 0.1],
                [0.6, 0.15, 0.05, 0.2],
            ]
        )
        pssm = PSSM(data, "ACTG")
        loader = io_app.load_tabular(sep="\t")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(pssm, identifier=outpath)
            new = loader(outpath)
            expected = safe_log(data) - safe_log(numpy.array([0.25, 0.25, 0.25, 0.25]))
            for i in range(len(expected)):
                j = i // 4
                self.assertTrue(
                    numpy.isclose(new.array[i][2], expected[j][i - j], atol=0.0001)
                )

    def test_write_tabular_distance_matrix(self):
        """correctly writes tabular data for DistanceMatrix"""
        data = {(0, 0): 0, (0, 1): 4, (1, 0): 4, (1, 1): 0}
        matrix = DistanceMatrix(data)
        loader = io_app.load_tabular(sep="\t")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(matrix, identifier=outpath)
            new = loader(outpath)
            # when written to file in tabular form
            # the loaded table will have dim-1 dim-2 as column labels
            # and the key-values pairs listed below; in dict form...
            expected = {
                0: {"dim-1": 0, "dim-2": 1, "value": 4},
                1: {"dim-1": 1, "dim-2": 0, "value": 4},
            }
            self.assertEqual(expected, new.to_dict())

    def test_write_tabular_table(self):
        """correctly writes tabular data"""
        rows = [[1, 2], [3, 4], [5, 6.5]]
        table = Table(["A", "B"], data=rows)
        loader = io_app.load_tabular(sep="\t")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(table, identifier=outpath)
            new = loader(outpath)
            self.assertEqual(table.to_dict(), new.to_dict())

    def test_load_tabular_motif_counts_array(self):
        """correctly loads tabular data for MotifCountsArray"""

        data = [[2, 4], [3, 5], [4, 8]]
        mca = MotifCountsArray(data, "AB")
        loader = io_app.load_tabular(sep="\t", as_type="motif_counts")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(mca, identifier=outpath)
            new = loader(outpath)
            self.assertEqual(mca.to_dict(), new.to_dict())

    def test_load_tabular_motif_freqs_array(self):
        """correctly loads tabular data for MotifFreqsArray"""

        data = [[0.3333, 0.6667], [0.3750, 0.625], [0.3333, 0.6667]]
        mfa = MotifFreqsArray(data, "AB")
        loader = io_app.load_tabular(sep="\t", as_type="motif_freqs")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(mfa, identifier=outpath)
            new = loader(outpath)
            self.assertEqual(mfa.to_dict(), new.to_dict())

    def test_load_tabular_pssm(self):
        """correctly loads tabular data for PSSM"""

        # data from test_profile
        data = [
            [0.1, 0.3, 0.5, 0.1],
            [0.25, 0.25, 0.25, 0.25],
            [0.05, 0.8, 0.05, 0.1],
            [0.7, 0.1, 0.1, 0.1],
            [0.6, 0.15, 0.05, 0.2],
        ]
        pssm = PSSM(data, "ACTG")
        loader = io_app.load_tabular(sep="\t", as_type="pssm")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(pssm, identifier=outpath)
            new = loader(outpath)
            assert_allclose(pssm.array, new.array, atol=0.0001)

    def test_load_tabular_distance_matrix(self):
        """correctly loads tabular data for DistanceMatrix"""
        data = {(0, 0): 0, (0, 1): 4, (1, 0): 4, (1, 1): 0}
        matrix = DistanceMatrix(data)
        loader = io_app.load_tabular(sep="\t", as_type="distances")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(matrix, identifier=outpath)
            new = loader(outpath)
            self.assertEqual(matrix.to_dict(), new.to_dict())

    def test_load_tabular_table(self):
        """correctly loads tabular data"""
        rows = [[1, 2], [3, 4], [5, 6.5]]
        table = Table(["A", "B"], data=rows)
        loader = io_app.load_tabular(sep="\t", as_type="table")
        with TemporaryDirectory(dir=".") as dirname:
            writer = io_app.write_tabular(data_path=dirname, format="tsv")
            outpath = join(dirname, "delme.tsv")
            writer.write(table, identifier=outpath)
            new = loader(outpath)
            self.assertEqual(table.to_dict(), new.to_dict())

    def test_write_json_with_info(self):
        """correctly writes an object with info attribute from json"""
        # create a mock object that pretends like it's been derived from
        # something
        from cogent3.util.union_dict import UnionDict

        with TemporaryDirectory(dir=".") as dirname:
            outdir = join(dirname, "delme")
            mock = Mock()
            mock.to_rich_dict = DNA.to_rich_dict
            mock.info = UnionDict(source=join("blah", "delme.json"))
            writer = io_app.write_json(outdir, create=True)
            _ = writer(mock)
            reader = io_app.load_json()
            got = reader(join(outdir, "delme.json"))
            self.assertEqual(got, DNA)

        # now with a zipped archive
        with TemporaryDirectory(dir=".") as dirname:
            outdir = join(dirname, "delme.zip")
            mock = Mock()
            mock.to_rich_dict = DNA.to_rich_dict
            mock.info = UnionDict(source=join("blah", "delme.json"))
            writer = io_app.write_json(outdir, create=True)
            identifier = writer(mock)
            reader = io_app.load_json()
            got = reader(writer.data_store[0])
            self.assertEqual(got, DNA)
            expect = join(outdir.replace(".zip", ""), "delme.json")
            if expect.startswith("." + os.sep):
                expect = expect[2:]
            self.assertEqual(identifier, expect)

    def test_write_json_no_info(self):
        """correctly writes an object with out an info attribute from json"""
        # create a mock object that pretends like it's been derived from
        # something
        with TemporaryDirectory(dir=".") as dirname:
            outdir = join(dirname, "delme")
            mock = patch("data.source", autospec=True)
            mock.to_rich_dict = DNA.to_rich_dict
            mock.source = join("blah", "delme.json")
            writer = io_app.write_json(outdir, create=True)
            _ = writer(mock)
            reader = io_app.load_json()
            got = reader(writer.data_store[0])
            self.assertEqual(got, DNA)

        # now with a zipped archive
        with TemporaryDirectory(dir=".") as dirname:
            outdir = join(dirname, "delme.zip")
            mock = patch("data.source", autospec=True)
            mock.to_rich_dict = DNA.to_rich_dict
            mock.source = join("blah", "delme.json")
            writer = io_app.write_json(outdir, create=True)
            identifier = writer(mock)
            reader = io_app.load_json()
            # checking loadable from a data store member too
            got = reader(writer.data_store[0])
            self.assertEqual(got, DNA)
            expect = join(outdir.replace(".zip", ""), "delme.json")
            if expect.startswith("." + os.sep):
                expect = expect[2:]
            self.assertEqual(identifier, expect)

    def test_restricted_usage_of_tinydb_suffix(self):
        """can only use tinydb in a load_db, write_db context"""
        with TemporaryDirectory(dir=".") as dirname:
            outdir = join(dirname, "delme.tinydb")
            for writer_class in (
                io_app.write_seqs,
                io_app.write_json,
                io_app.write_tabular,
            ):
                with self.assertRaises(ValueError):
                    writer_class(outdir, create=True, if_exists="skip")
            # but OK for write_db
            w = io_app.write_db(outdir, create=True, if_exists="skip")
            w.data_store.close()

    def test_write_db_parallel(self):
        """writing with overwrite in parallel should reset db"""
        dstore = io_app.get_data_store(self.basedir, suffix="fasta")
        members = dstore.filtered(callback=lambda x: "brca1.fasta" not in x.split("/"))
        reader = io_app.load_unaligned()
        aligner = align_app.align_to_ref()
        writer = write_db("delme.tinydb", create=True, if_exists="overwrite")
        process = reader + aligner + writer

        r = process.apply_to(members, logger=False, show_progress=False, parallel=True)

        expect = [str(m) for m in process.data_store]
        process.data_store.close()

        # now get read only and check what's in there
        result = io_app.get_data_store("delme.tinydb")
        got = [str(m) for m in result]
        result.close()

        self.assertEqual(got, expect)


if __name__ == "__main__":
    main()
