#!/usr/bin/env python3
import pytest
from numpy import isnan
from hashlib import sha256

from rdflib import Namespace
from probs_runner import PROBS, QUANTITYKIND
from numpy.testing import assert_allclose

PRODCOM2016 = Namespace("http://w3id.org/probs-lab/data/prodcom/2016/")
PRODCOM2017 = Namespace("http://w3id.org/probs-lab/data/prodcom/2017/")
GEONAMES = Namespace("https://sws.geonames.org/")

def assert_exact_results(result, expected_value):
    assert len(result) == 1
    assert result[0].bound == PROBS.ExactBound
    assert_allclose(result[0].measurement, expected_value, rtol=1e-3)

def assert_no_measurement(result):
    assert len(result) == 1
    assert isnan(result[0].measurement)

def assert_no_result(result):
    assert len(result) == 0


# Do most tests for 2016; then spot check a few values for the other years

class TestProdcom2016Data:
    """Test PRODCOM2016 DATA conversion"""

    endpoint_data = ["outputs/PRODCOM2016DATA.nt.gz",
                     "outputs/PRD_2016_20200617_185122.nt.gz"]

    @pytest.mark.parametrize("code,metric,expected_value", [
        ("10113250", QUANTITYKIND.Mass, 2717950),
        ("10841270", QUANTITYKIND.Mass, 622435736),
        ("15115100", QUANTITYKIND.Mass, 0),
        ("17121439", QUANTITYKIND.Mass, 43196000),
        ("20132455", QUANTITYKIND.MassAmountOfSubstance, 1396308),
    ])
    def test_expected_measurements(self, rdfox, code, metric, expected_value):
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2016,
            GEONAMES["2635167"],
            metric,
            PROBS.SoldProduction,
            object_code=code
        )
        assert_exact_results(result, expected_value)

    @pytest.mark.parametrize("code,metric", [
        ("14391090", QUANTITYKIND.AmountOfSubstance),
        ("20141325", QUANTITYKIND.Mass),
        # This one has no measurement or data at all, so is loaded with metric "Unknown"
        ("20421975", QUANTITYKIND.Unknown),
    ])
    def test_missing_measurements(self, rdfox, code, metric):
        # these codes have no data -- expect an observation with no measurement
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2016,
            GEONAMES["2635167"] ,
            metric,
            PROBS.SoldProduction,
            object_code=code
        )
        assert_no_measurement(result)

    @pytest.mark.parametrize("code,metric", [
        ("14391090", QUANTITYKIND.Mass),
        ("20141325", QUANTITYKIND.AmountOfSubstance),
    ])
    def test_no_result_with_wrong_metric(self, rdfox, code, metric):
        # same as above, but with the wrong metric
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2016,
            GEONAMES["2635167"] ,
            metric,
            PROBS.SoldProduction,
            object_code=code
        )
        assert_no_result(result)


class TestProdcom2017Data:
    """Test PRODCOM2017 DATA conversion.

    Just include a few expected values, since other behaviour is tested for the 2016 data.
    """

    endpoint_data = ["outputs/PRODCOM2017DATA.nt.gz",
                     "outputs/PRD_2017_20200617_185035.nt.gz"]

    @pytest.mark.parametrize("code,metric,expected_value", [
        ("10721253", QUANTITYKIND.Mass, 244871494),
        ("16212400", QUANTITYKIND.Volume, 0),
        ("23111290", QUANTITYKIND.Area, 28514781),
        ("25711430", QUANTITYKIND.AmountOfSubstance, 329685),
    ])
    def test_expected_measurements(self, rdfox, code, metric, expected_value):
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2017,
            GEONAMES["2635167"],
            metric,
            PROBS.SoldProduction,
            object_code=code
        )
        assert_exact_results(result, expected_value)


class TestProdcom2018Data:
    """Test PRODCOM2018 DATA conversion.

    Just include a few expected values, since other behaviour is tested for the 2016 data.
    """

    endpoint_data = ["outputs/PRODCOM2018DATA.nt.gz",
                     "outputs/PRD_2017_20200617_185035.nt.gz"]

    @pytest.mark.parametrize("code,metric,expected_value", [
        ("10131120", QUANTITYKIND.Mass, 104282828),
        ("11052000", QUANTITYKIND.Mass, 737526732),
        ("13921530", QUANTITYKIND.Area, 1197738),
        ("15112200", QUANTITYKIND.Area, 0),
    ])
    def test_expected_measurements(self, rdfox, code, metric, expected_value):
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2018,
            GEONAMES["2635167"],
            metric,
            PROBS.SoldProduction,
            object_code=code
        )
        assert_exact_results(result, expected_value)