#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (C) 2016 by Gaik Tamazian
# gaik (dot) tamazian (at) gmail (dot) com

from evolspace import pdist
from evolspace import exception
import h5py
import itertools
import logging
import numpy as np
import os
import pyfaidx
import tempfile
import unittest

path = os.path.dirname(__file__)
os.chdir(path)


class TestDistanceCalculator(unittest.TestCase):
    def setUp(self):
        self.__fasta = os.path.join(
            'data', 'seq.fa'
        )
        self.__block_num = 3
        self.__output = []
        for i in xrange(self.__block_num):
            self.__output.append(tempfile.NamedTemporaryFile().name)

        # determine the number of sequences in the FASTA file
        with pyfaidx.Fasta(self.__fasta) as fasta_file:
            self.__seq_num = len(fasta_file.keys())

    def test_distance_calculator_single_block(self):
        dc = pdist.DistanceCalculator(self.__fasta, 1, 1)
        dc.calculate_distances(self.__output[0])
        with h5py.File(self.__output[0]) as produced_file:
            self.assertEqual(len(produced_file["distances"]),
                             self.__seq_num)

    def test_distance_calculator_multiple_blocks(self):
        for i in xrange(self.__block_num):
            dc = pdist.DistanceCalculator(self.__fasta, cur_block=i+1,
                                          block_num=self.__block_num)
            dc.calculate_distances(self.__output[i])
        dist_seqs = set()
        for dist_filename in self.__output:
            with h5py.File(dist_filename) as dist_file:
                cur_seqs = set(dist_file["distances"])
                self.assertFalse(dist_seqs.intersection(cur_seqs))
                dist_seqs.update(cur_seqs)
        self.assertEqual(len(dist_seqs), self.__seq_num)

    def tearDown(self):
        for i in self.__output:
            if os.path.isfile(i):
                os.unlink(i)


class TestCheckPDistFile(unittest.TestCase):
    def setUp(self):
        self.fasta = os.path.join(
            "data", "seq.fa"
        )
        self.block_num = 3
        self.pdist_files = []
        for i in xrange(self.block_num):
            self.pdist_files.append(tempfile.NamedTemporaryFile().name)
        self.pdist = self.pdist_files[0]
        self.merged_pdist = tempfile.NamedTemporaryFile().name

        self.text_dist = tempfile.NamedTemporaryFile().name

        # determine the number of sequences in the FASTA file and
        # their names
        with pyfaidx.Fasta(self.fasta) as fasta_file:
            self.seq_names = fasta_file.keys()
            self.seq_num = len(self.seq_names)

        logging.disable(logging.ERROR)

    def prepare_pdist_file(self):
        dc = pdist.DistanceCalculator(self.fasta)
        dc.calculate_distances(self.pdist)

    def prepare_pdist_block_files(self):
        for i in xrange(self.block_num):
            dc = pdist.DistanceCalculator(self.fasta, cur_block=i+1,
                                          block_num=self.block_num)
            dc.calculate_distances(self.pdist_files[i])

    def test_check_correct_file(self):
        self.prepare_pdist_file()
        self.assertTrue(pdist.check_pdist_file(self.pdist))
        self.assertTrue(pdist.check_pdist_file(self.pdist,
                                               complete=True))

    def test_check_missing_attributes(self):
        # check against files with missing attributes
        for i in ("fasta", "distance", "seq_id_hash"):
            self.prepare_pdist_file()
            with h5py.File(self.pdist, "a") as pdist_file:
                del pdist_file.attrs[i]

            self.assertFalse(pdist.check_pdist_file(self.pdist))

    def test_check_missing_sec(self):
        # create a file with a missing section
        self.prepare_pdist_file()
        with h5py.File(self.pdist, "a") as pdist_file:
            del pdist_file["header/cur_ids"]

        self.assertFalse(pdist.check_pdist_file(self.pdist))

    def test_check_missing_seq(self):
        # create a file with an array of distances from a sequence
        # that is missing from the header
        self.prepare_pdist_file()
        with h5py.File(self.pdist, "a") as pdist_file:
            pdist_file.create_dataset("distances/extra_seq",
                                      (5, 1),
                                      data=np.zeros(5))
        self.assertFalse(pdist.check_pdist_file(self.pdist))

    def test_is_master_pdist_file(self):
        # check if against a master file
        self.prepare_pdist_file()
        self.assertTrue(pdist.is_master_pdist_file(self.pdist))
        # check against a secondary file
        with h5py.File(self.pdist, "a") as pdist_file:
            del pdist_file["header/seq"]
            del pdist_file["header/mask"]
        self.assertFalse(pdist.is_master_pdist_file(self.pdist))

    def test_is_complete_pdist_file(self):
        # check against a complete file
        self.prepare_pdist_file()
        self.assertTrue(pdist.is_complete_pdist_file(self.pdist))
        # check against an incomplete file
        with h5py.File(self.pdist, "a") as pdist_file:
            del pdist_file["distances/" + self.seq_names[-1]]
        self.assertFalse(pdist.is_complete_pdist_file(self.pdist))

    def test_check_incomplete_dist(self):
        # create a file which misses distance array for some of its
        # sequences
        self.prepare_pdist_file()
        with h5py.File(self.pdist, "a") as pdist_file:
            del pdist_file["distances/" + self.seq_names[0]]
        self.assertTrue(pdist.check_pdist_file(self.pdist))

    def test_check_get_bin_counts(self):
        self.prepare_pdist_file()
        pdist.get_bin_counts(self.pdist, 10)

    def test_check_incorrect_dist_array(self):
        # create a file with a distance array of an inappropriate
        # length
        self.prepare_pdist_file()
        with h5py.File(self.pdist, "a") as pdist_file:
            seq_id = self.seq_names[self.seq_num / 2]
            temp = pdist_file["distances/" + seq_id]
            del pdist_file["distances/" + seq_id]
            pdist_file["distances/" + seq_id] = np.repeat(temp, 2)

        self.assertTrue(pdist.check_pdist_file(self.pdist))
        self.assertFalse(pdist.check_pdist_file(self.pdist,
                                                complete=True))

    def test_check_compatibility(self):
        self.prepare_pdist_block_files()
        for i, j in itertools.combinations(self.pdist_files, 2):
            self.assertTrue(pdist.check_compatibility(i, j))

        # check against the case of different FASTA files
        with h5py.File(self.pdist_files[1]) as second_block_file:
            second_block_file.attrs["fasta"] = "another_file.fa"
        self.assertFalse(pdist.check_compatibility(
            self.pdist_files[0], self.pdist_files[1]))

        # check against the different hash value of sequence IDs
        self.prepare_pdist_block_files()
        with h5py.File(self.pdist_files[1]) as second_block_file:
            second_block_file.attrs["seq_id_hash"] = "another_value"
        self.assertFalse(pdist.check_compatibility(
            self.pdist_files[0], self.pdist_files[1]))

        # check against repetitive distance array
        self.prepare_pdist_block_files()
        with h5py.File(self.pdist_files[0]) as first_block_file:
            first_block_file.create_dataset("distances/" +
                                            self.seq_names[-1],
                                            data=np.zeros(self.seq_num))
        self.assertFalse(pdist.check_compatibility(
            self.pdist_files[0], self.pdist_files[-1]))

    def test_merge_dist_files(self):
        self.prepare_pdist_block_files()
        pdist.merge_dist_files(self.pdist_files, self.merged_pdist)
        # make sure that the merged file is complete
        self.assertTrue(pdist.is_complete_pdist_file(
            self.merged_pdist))

        with self.assertRaises(exception.DistFileMergeError):
            pdist.merge_dist_files(self.pdist_files[1:],
                                   self.merged_pdist)

        # make a secondary file incompatible with a primary one by
        # setting another FASTA file name
        with h5py.File(self.pdist_files[2]) as secondary_file:
            secondary_file.attrs["fasta"] = "another.fa"
        with self.assertRaises(exception.DistFileMergeError):
            pdist.merge_dist_files(self.pdist_files, self.merged_pdist)

        # do the same think but set another distance measure
        self.prepare_pdist_block_files()
        with h5py.File(self.pdist_files[2]) as secondary_file:
            secondary_file.attrs["distance"] = "another distance"
        with self.assertRaises(exception.DistFileMergeError):
            pdist.merge_dist_files(self.pdist_files, self.merged_pdist)

    def test_output_distances(self):
        self.prepare_pdist_file()
        for seq_ids in (None, self.seq_names[:self.seq_num/2]):
            n = len(seq_ids) if seq_ids is not None else self.seq_num
            pdist.output_distances(self.pdist, self.text_dist,
                                   sequence_ids=seq_ids)
            distances = np.genfromtxt(
                self.text_dist, skip_header=1,
                usecols=range(1, n + 1)
            )
            self.assertEqual(distances.shape, (n, n))
            self.assertTrue(np.all(distances.diagonal() == 0))

        with self.assertRaises(exception.MissingSeqError):
            pdist.output_distances(self.pdist, self.text_dist,
                                   ["missing_seq"])

        self.prepare_pdist_block_files()
        pdist.output_distances(self.pdist_files[0], self.text_dist)

    def test_show_pdist_info(self):
        self.prepare_pdist_file()
        pdist.show_pdist_info(self.pdist)

    def tearDown(self):
        for i in self.pdist_files + [self.merged_pdist]:
            if os.path.isfile(i):
                os.unlink(i)
