#!/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/total_production/DS-056121_006-2016.nt.gz",
                     "outputs/PRD_2017_20200617_185035.nt.gz"]

    @pytest.mark.parametrize("code,metric,expected_value", [
        ("13103100", QUANTITYKIND.Mass, 224299),
        ("13911910", QUANTITYKIND.Mass, 12566190),
        ("15115100", QUANTITYKIND.Mass, 0),
        ("20132455", QUANTITYKIND.MassAmountOfSubstance, 1369740),
        ("20141290", QUANTITYKIND.Mass, 208123),
    ])
    def test_expected_measurements(self, rdfox, code, metric, expected_value):
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2016,
            GEONAMES["2635167"],
            metric,
            PROBS.TotalProduction,
            object_code=code
        )
        assert_exact_results(result, expected_value)

    @pytest.mark.parametrize("code,metric", [
        ("13202042", QUANTITYKIND.Area),              # No PRDQNT, :C for PQNTFLAG 
        ("20141325", QUANTITYKIND.Mass),              # No PRDQNT, :C for PQNTFLAG
    ])
    def test_missing_measurements(self, rdfox, code, metric):
        # these codes have confidential PRDQNT (observation created with no measurement)
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2016,
            GEONAMES["2635167"] ,
            metric,
            PROBS.TotalProduction,
            object_code=code,
        )
        assert_no_measurement(result)

    @pytest.mark.parametrize("code,metric", [
        ("20132455", QUANTITYKIND.Mass),
        ("20141325", QUANTITYKIND.AmountOfSubstance),
    ])
    def test_no_result_with_wrong_metric(self, rdfox, code, metric):
        # query using wrong metric
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2016,
            GEONAMES["2635167"] ,
            metric,
            PROBS.TotalProduction,
            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/total_production/DS-056121_006-2017.nt.gz",
                     "outputs/PRD_2017_20200617_185035.nt.gz"]

    @pytest.mark.parametrize("code,metric,expected_value", [
        ("08111250", QUANTITYKIND.Mass, 604132000),
        ("20111120", QUANTITYKIND.Volume, 387410000),
        ("13201230", QUANTITYKIND.Area, 8038713),
        ("20122415", QUANTITYKIND.MassAmountOfSubstance, 31291132),
    ])
    def test_expected_measurements(self, rdfox, code, metric, expected_value):
        result = rdfox.get_observations(
            PROBS.TimePeriod_YearOf2017,
            GEONAMES["2635167"],
            metric,
            PROBS.TotalProduction,
            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/total_production/DS-056121_006-2018.nt.gz",
                     "outputs/PRD_2017_20200617_185035.nt.gz"]

    @pytest.mark.parametrize("code,metric,expected_value", [
        ("08114000", QUANTITYKIND.Mass, 41455000),
        ("13106135", QUANTITYKIND.Mass, 9193633),
        ("13203210", QUANTITYKIND.Area, 4300488),
        ("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.TotalProduction,
            object_code=code
        )
        assert_exact_results(result, expected_value)
