"""
Data-driven tests for sequence annotation Features.
"""
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import ga4gh.datarepo as datarepo
import ga4gh.datamodel as datamodel
import ga4gh.datamodel.datasets as datasets
import ga4gh.datamodel.references as references
import ga4gh.datamodel.sequenceAnnotations as sequenceAnnotations
import ga4gh.protocol as protocol
import tests.datadriven as datadriven
import tests.paths as paths

_datasetName = "ds"

_discontinuousTestData = {
    "featureSetName": "discontinuous",
    "referenceName": "apidb|Pf3D7_13",
    "totalFeatures": 30,
    "sampleFeatureId": 4409955920,
    "sampleParentId": 4409956112,
    "sampleStart": 820942,
    "sampleEnd": 821379,
    "sampleStrand": protocol.POS_STRAND,
    "sampleSiblings": 2,
    "region": [0, 2**32],
    "ontologyRestriction": ["gene", ],
    "featuresWithOntology": 3
}

_gencodeV21Set1TestData = {
    "featureSetName": "gencodeV21Set1",
    "referenceName": "chr1",
    "totalFeatures": 543,
    "sampleFeatureId": 4397111632,
    "sampleParentId": 4397111312,
    "sampleStart": 804776,
    "sampleEnd": 804832,
    "sampleStrand": protocol.POS_STRAND,
    "sampleSiblings": 5,
    "region": [0, 2**32],
    "ontologyRestriction": ["gene", ],
    "featuresWithOntology": 58
}

_sacCerTestTestData = {
    "featureSetName": "sacCerTest",
    "referenceName": "chrI",
    "totalFeatures": 33,
    "sampleFeatureId": 4324354512,
    "sampleParentId": None,
    "sampleStart": 337,
    "sampleEnd": 801,
    "sampleStrand": protocol.NEG_STRAND,
    "sampleSiblings": 20,
    "region": [0, 2**32],
    "ontologyRestriction": ["gene", ],
    "featuresWithOntology": 11
}

_specialCasesTestTestData = {
    "featureSetName": "specialCasesTest",
    "referenceName": "2L",
    "totalFeatures": 4,
    "sampleFeatureId": 4554844304,
    "sampleParentId": None,
    "sampleStart": 22229583,
    "sampleEnd": 22229699,
    "sampleStrand": protocol.NEG_STRAND,
    "sampleSiblings": 2,
    "region": [0, 2**32],
    "ontologyRestriction": ["gene", ],
    "featuresWithOntology": 0
}

_testDataForFeatureSetName = {
    "discontinuous": _discontinuousTestData,
    "gencodeV21Set1": _gencodeV21Set1TestData,
    "sacCerTest": _sacCerTestTestData,
    "specialCasesTest": _specialCasesTestTestData
}


def _getFeatureCompoundId(dataSetName, featureSetName, featureId):
    splits = [dataSetName, featureSetName, str(featureId)]
    joined = datamodel.CompoundId.join(splits)
    obfuscated = datamodel.CompoundId.obfuscate(joined)
    return obfuscated


def testFeatureSets():
    testDataDir = "tests/data/datasets/dataset1/sequenceAnnotations"
    for test in datadriven.makeTests(testDataDir, FeatureSetTests, "*.db"):
        yield test


class FeatureSetTests(datadriven.DataDrivenTest):
    """
    Re-parses source GFF3 files, compares the results with the contents
    of corresponding sequenceAnnotations.Feature objects.
    """
    def __init__(self, featureSetLocalName, dataPath):
        """
        :param localId: Name of the GFF3 resource corresponding to a pair
        of files, .db and .gff3
        :param dataPath: string representing full path to the .db file
        :return:
        """
        self._dataset = datasets.Dataset(_datasetName)
        self._repo = datarepo.SqlDataRepository(paths.testDataRepo)
        self._repo.open(datarepo.MODE_READ)
        self._ontology = self._repo.getOntologyByName(paths.ontologyName)
        self._referenceSet = references.AbstractReferenceSet("test_rs")
        featureSetLocalName = featureSetLocalName[:-3]  # remove '.db'
        self._testData = _testDataForFeatureSetName[featureSetLocalName]
        super(FeatureSetTests, self).__init__(featureSetLocalName, dataPath)

    def getProtocolClass(self):
        return protocol.FeatureSet

    def getDataModelInstance(self, localId, dataPath):
        featureSet = sequenceAnnotations.Gff3DbFeatureSet(
            self._dataset, localId)
        featureSet.setOntology(self._ontology)
        featureSet.setReferenceSet(self._referenceSet)
        featureSet.populateFromFile(dataPath)
        return featureSet

    def testGetFeatureById(self):
        idString = _getFeatureCompoundId(
            _datasetName,
            self._testData["featureSetName"],
            self._testData["sampleFeatureId"])
        compoundId = datamodel.FeatureCompoundId.parse(idString)
        feature = self._gaObject.getFeature(compoundId)
        self.assertIsNotNone(feature)
        self.assertEqual(feature.id, idString)
        if self._testData["sampleParentId"] is not None:
            self.assertEqual(
                datamodel.FeatureCompoundId.parse(
                    feature.parent_id).featureId,
                str(self._testData["sampleParentId"]))
        else:
            self.assertEqual(feature.parent_id, '')
        self.assertEqual(
            feature.reference_name,
            self._testData["referenceName"])
        self.assertEqual(
            feature.start,
            self._testData["sampleStart"])
        self.assertEqual(
            feature.end,
            self._testData["sampleEnd"])
        self.assertEqual(
            feature.strand,
            self._testData["sampleStrand"])

    def testFetchAllFeaturesInRegion(self):
        features = []
        nextPageTokens = []
        for (feature, nextPageToken) in self._gaObject.getFeatures(
                self._testData["referenceName"],
                self._testData["region"][0],
                self._testData["region"][1],
                None, 1000):
            features.append(feature)
            nextPageTokens.append(nextPageToken)
        self.assertEqual(len(features), self._testData["totalFeatures"])
        self.assertIsNone(nextPageTokens[-1])

    def testFetchFeaturesRestrictedByOntology(self):
        features = []
        for (feature, _) in self._gaObject.getFeatures(
                self._testData["referenceName"],
                self._testData["region"][0],
                self._testData["region"][1],
                None, 1000,
                featureTypes=self._testData["ontologyRestriction"]):
            features.append(feature)
        self.assertEqual(len(features),
                         self._testData["featuresWithOntology"])

    def testFetchFeaturesRestrictedByParent(self):
        parentId = ""
        if self._testData["sampleParentId"] is not None:
            parentIdString = _getFeatureCompoundId(
                _datasetName,
                self._testData["featureSetName"],
                self._testData["sampleParentId"])
            parentId = datamodel.FeatureCompoundId.parse(
                parentIdString).featureId
        features = []
        for (feature, _) in self._gaObject.getFeatures(
                self._testData["referenceName"],
                self._testData["region"][0],
                self._testData["region"][1],
                None, 1000,
                parentId=parentId):
            features.append(feature)
        self.assertEqual(len(features),
                         self._testData["sampleSiblings"])
