
from unittest import SkipTest
import logging
import json
import os.path as op

from pbcommand.pb_io.report import load_report_from_json
from pbcommand.models import FileTypes
from pbcore.io import ContigSet

from pysiv2.custom.base import TestReportStatistics

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)

class TestIsoSeqCluster(TestReportStatistics):
    """
    Test for cluster stage, including the separately run ICE scripts.  The
    report should include information collected by the ice_quiver_postprocess
    task.
    """
    REPORT_ID = "isoseq_cluster"
    TEST_ID = REPORT_ID
    METRIC_IDS = [
        "num_consensus_isoforms",
        "avg_consensus_isoform_length",
        "num_polished_hq_isoforms",
        "num_polished_lq_isoforms",
    ]

    # TODO also check output files

    @classmethod
    def getMetrics(cls):
        super(TestIsoSeqCluster, cls).getMetrics()
        cls.hq_fasta_file = cls.lq_fasta_file = None
        for file_id, file_info in cls.datastore.get_file_dict().iteritems():
            if file_info.file_type_id == FileTypes.DS_CONTIG.file_type_id:
                file_name = op.basename(file_info.path)
                if file_name.startswith("hq_isoforms"):
                    cls.hq_fasta_file = file_info.path
                    with ContigSet(cls.hq_fasta_file) as ds:
                        n = len(ds)
                        cls.metric_dict["num_polished_hq_isoforms_fasta"] = n
                elif file_name.startswith("lq_isoforms"):
                    cls.lq_fasta_file = file_info.path
                    with ContigSet(cls.lq_fasta_file) as ds:
                        n = len(ds)
                        cls.metric_dict["num_polished_lq_isoforms_fasta"] = n

    def test_number_of_polished_hq_isoforms(self):
        """
        Check that the number of high-quality reads in the JSON report is
        consistent with the output ContigSet.
        """
        report_value = self._get_stat("num_polished_hq_isoforms")
        if report_value is None:
            raise SkipTest("Missing num_polished_hq_isoforms")
        self.assertEqual(self.metric_dict["num_polished_hq_isoforms_fasta"],
                         self._get_stat("num_polished_hq_isoforms"))

    def test_number_of_polished_lq_isoforms(self):
        """
        Check that the number of low-quality reads in the JSON report is
        consistent with the output ContigSet.
        """
        self.assertEqual(self.metric_dict["num_polished_lq_isoforms_fasta"],
                         self._get_stat("num_polished_lq_isoforms"))

    def test_total_number_of_consensus_isoforms(self):
        """
        Check that the numbers of high- and low-quality isoforms add up.
        """
        report_value = self._get_stat("num_polished_hq_isoforms")
        if report_value is None:
            raise SkipTest("Missing num_polished_hq_isoforms")
        total = self._get_stat("num_polished_lq_isoforms") + \
                self._get_stat("num_polished_hq_isoforms")
        self.assertTrue(self._get_stat("num_consensus_isoforms") >= total)

    def test_hq_lq_isoforms_avgqv_hist(self):
        """Check hq_lq_isoforms_avg_qv_hist.png exists."""
        images = [str(p.image) for g in self.getReport()._plotgroups for p in g.plots]
        self.assertTrue('hq_lq_isoforms_avgqv_hist.png' in images)

    def test_consensus_isoforms_readlength_hist(self):
        """Check consensus_isoforms_readlength_hist.png exists."""
        images = [str(p.image) for g in self.getReport()._plotgroups for p in g.plots]
        self.assertTrue('consensus_isoforms_readlength_hist.png' in images)
