
"""
Scientific validity tests for resequencing.
"""


from unittest import SkipTest
import os.path as op

from pbcore.io import GffReader, FastqReader, ContigSet
from pbcommand.models import FileTypes

from pysiv2.custom.base import TestStatisticsBase
from pysiv2.custom.utils import check_gff_order, calculate_qv


# XXX not to be confused with TestVariantsReport in test_resequencing.py
class TestVariantCallerOutput(TestStatisticsBase):
    """
    Test for variant calls made by GenomicConsensus in resequencing jobs.
    Unlike :module:`test_resequencing`, this does not rely on the variants
    report; instead it calculates metrics itself based on the output FASTQ
    file combined with the input ReferenceSet.  (Requires `blasr`)
    """
    TEST_ID = "variants"
    METRIC_IDS = [
        "consensus_qv",
        "number_of_variants",
        "number_of_fastx_records",
        "mean_consensus_quality",
        "coeff_of_variation_consensus_quality",
    ]

    @classmethod
    def getMetrics(cls):
        """
        Read in the variants GFF and consensus FASTQ and collect some basic
        statistics.
        """
        # FIXME this is much messier than I'd like because of the way the
        # current pbsmrtpipe datastore works
        cls.vcf_file = None
        for vcf_file in cls.datastore.files_by_type_and_source(
                file_type_id=FileTypes.VCF.file_type_id,
                task_name="gff2vcf"):
            assert cls.vcf_file is None
            cls.vcf_file = vcf_file
        cls.bed_file = None
        for bed_file in cls.datastore.files_by_type_and_source(
                file_type_id=FileTypes.BED.file_type_id,
                task_name="gff2bed"):
            assert cls.bed_file is None
            cls.bed_file = bed_file
        cls.variants_gff_file = None
        for gff_file in cls.datastore.files_by_type_and_source(
                file_type_id=FileTypes.GFF.file_type_id,
                task_name="variantcaller"):
            with GffReader(gff_file) as f:
                for header in f.headers:
                    h = header.split()
                    if h[0] == "##source" and h[1] == "GenomicConsensus":
                        cls.variants_gff_file = gff_file
                        n_rec = len([ r for r in f ])
                        cls.metric_dict["number_of_variants"] = n_rec
                        break
                else:
                    continue
                break
        fastq_file = cls.datastore._get_task_file(
            task_id="genomic_consensus.tasks.variantcaller",
            file_type_id=FileTypes.FASTQ.file_type_id)
        with FastqReader(fastq_file) as f:
            import scipy.stats
            n_rec = 0
            for fastq_record in f:
                n_rec += 1
                if n_rec > 1:
                    continue
                mean_quality = scipy.mean(fastq_record.quality)
                # XXX is this being done correctly?
                coeff_of_variation_quality = scipy.stats.variation(
                    fastq_record.quality)
                cls.metric_dict["mean_consensus_quality"] = mean_quality
                cls.metric_dict["coeff_of_variation_consensus_quality"] = \
                    coeff_of_variation_quality
            cls.metric_dict["number_of_fastq_records"] = n_rec
        ref_root, ref_seq = op.split(
            op.split(cls.entrypoints.reference_path)[0])
        reference_seqs = op.join(
            ref_root, ref_seq, 'sequence', '{s}.fasta'.format(s=ref_seq))
        try:
            qv = calculate_qv(fastq_file, reference_seqs, cls.job_dir)
            cls.metric_dict["consensus_qv"] = qv
        except ValueError as e:
            # FIXME this doesn't really work on multi-contig data
            cls.SKIP_TESTS["consensus_qv"] = \
                "Error attempting to calculate QV: {e}".format(e=str(e))
        # FIXME this breaks on unchunked jobs because file type is FASTA
        contig_set = cls.datastore._get_task_file(
            task_id="genomic_consensus.tasks.variantcaller",
            file_type_id=FileTypes.DS_CONTIG.file_type_id)
        with ContigSet(contig_set) as f:
            cls.metric_dict["number_of_fastx_records"] = len([r for r in f])

    def test_number_of_fastx_records_equal(self):
        """
        Check that the number of records in the FASTQ output is the same as the
        number of output contigs (verification for FASTQ gather bug 28405).
        """
        n_fastq = self.metric_dict["number_of_fastq_records"]
        n_fasta = self.metric_dict["number_of_fastx_records"]
        self.assertEqual(n_fastq, n_fasta)

    def test_variantcaller_gff_headers(self):
        """Check that GFF headers are preserved (even after chunking)"""
        self.assertFalse(self.variants_gff_file is None,
                         "No variants GFF file found - check file headers.")
        with GffReader(self.variants_gff_file) as f:
            for header in f.headers:
                fields = header.split()
                if fields[0] == "##sequence-region":
                    break
            else:
                self.fail("Can't find sequence-region header")

    def test_gff_to_vcf(self):
        """
        Test that the number of records in the VCF output file is the same as
        the number of variants reported.
        """
        if self.vcf_file is None:
            self.fail("Can't find VCF file")
        n_records = 0
        with open(self.vcf_file) as vcf:
            for line in vcf.readlines():
                if not line.startswith("#"):
                    n_records += 1
        self.assertEqual(n_records, self.metric_dict["number_of_variants"])

    def test_gff_to_bed(self):
        """
        Test that the number of records in the BED output file is the same as
        the number of variants reported.
        """
        if self.bed_file is None:
            self.fail("Can't find BED file")
        n_records = 0
        with open(self.bed_file) as bed:
            for line in bed.readlines():
                if not line.startswith("track name="):
                    n_records += 1
        self.assertEqual(n_records, self.metric_dict["number_of_variants"])
