from datetime import datetime

import ee
import pytest

import openet.ndvi
import openet.ndvi.interpolate as interpolate
import openet.ndvi.utils as utils


def scene_coll(variables, ndvi=[0.6, 0.6, 0.6]):
    """Return a generic scene collection to test scene interpolation functions

    Parameters
    ----------
    variables : list
        The variables to return in the collection
    ndvi : list

    Returns
    -------
    ee.ImageCollection

    """
    img = (
        ee.Image('LANDSAT/LC08/C02/T1_L2/LC08_044033_20170716')
        .select(['SR_B3']).double().multiply(0)
    )
    mask = img.add(1).updateMask(1).uint8()

    # # The "date" is used for the time band since it needs to be the 0 UTC time
    # date1 = ee.Number(ee.Date.fromYMD(2017, 7, 8).millis())
    # date2 = ee.Number(ee.Date.fromYMD(2017, 7, 16).millis())
    # date3 = ee.Number(ee.Date.fromYMD(2017, 7, 24).millis())

    # The "time" is advanced to match the typical Landsat overpass time
    time1 = ee.Number(ee.Date.fromYMD(2017, 7, 8).advance(18, 'hours').millis())
    time2 = ee.Number(ee.Date.fromYMD(2017, 7, 16).advance(18, 'hours').millis())
    time3 = ee.Number(ee.Date.fromYMD(2017, 7, 24).advance(18, 'hours').millis())

    # Don't add mask or time band to scene collection
    # since they are now added in the interpolation calls
    scene_coll = ee.ImageCollection.fromImages([
        ee.Image([img.add(ndvi[0])]).rename(['ndvi'])
        .set({'system:index': 'LE07_044033_20170708', 'system:time_start': time1}),
        ee.Image([img.add(ndvi[1])]).rename(['ndvi'])
        .set({'system:index': 'LC08_044033_20170716', 'system:time_start': time2}),
        ee.Image([img.add(ndvi[2])]).rename(['ndvi'])
        .set({'system:index': 'LE07_044033_20170724', 'system:time_start': time3}),
    ])

    return scene_coll.select(variables)


def test_from_scene_ndvi_t_interval_daily_values(tol=0.0001):
    output_coll = interpolate.from_scene_ndvi(
        scene_coll(['ndvi'], ndvi=[0.2, 0.4, 0.6]),
        start_date='2017-07-01', end_date='2017-08-01',
        variables=['ndvi'],
        interp_args={'interp_method': 'linear', 'interp_days': 32},
        model_args={'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                    'et_reference_band': 'eto'},
        t_interval='daily',
    )

    TEST_POINT = (-121.5265, 38.7399)
    output = utils.point_coll_value(output_coll, TEST_POINT, scale=30)
    assert abs(output['ndvi']['2017-07-01'] - 0.2) <= tol
    assert abs(output['ndvi']['2017-07-08'] - 0.2) <= tol
    assert abs(output['ndvi']['2017-07-10'] - 0.25) <= tol
    assert abs(output['ndvi']['2017-07-12'] - 0.3) <= tol
    assert abs(output['ndvi']['2017-07-16'] - 0.4) <= tol
    assert abs(output['ndvi']['2017-07-24'] - 0.6) <= tol
    assert abs(output['ndvi']['2017-07-31'] - 0.6) <= tol
    assert '2017-08-01' not in output['ndvi'].keys()
    # assert output['count']['2017-07-01'] == 3


def test_from_scene_ndvi_t_interval_monthly_values(tol=0.0001):
    output_coll = interpolate.from_scene_ndvi(
        scene_coll(['ndvi']),
        start_date='2017-07-01', end_date='2017-08-01',
        variables=['ndvi', 'count'],
        interp_args={'interp_method': 'linear', 'interp_days': 32},
        model_args={'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                    'et_reference_band': 'eto'},
        t_interval='monthly',
    )

    TEST_POINT = (-121.5265, 38.7399)
    output = utils.point_coll_value(output_coll, TEST_POINT, scale=30)
    assert abs(output['ndvi']['2017-07-01'] - 0.6) <= tol
    assert output['count']['2017-07-01'] == 3


def test_from_scene_ndvi_t_interval_custom_values(tol=0.0001):
    output_coll = interpolate.from_scene_ndvi(
        scene_coll(['ndvi']),
        start_date='2017-07-01', end_date='2017-08-01',
        variables=['ndvi', 'count'],
        interp_args={'interp_method': 'linear', 'interp_days': 32},
        model_args={'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                    'et_reference_band': 'eto'},
        t_interval='custom',
    )

    TEST_POINT = (-121.5265, 38.7399)
    output = utils.point_coll_value(output_coll, TEST_POINT, scale=30)
    assert abs(output['ndvi']['2017-07-01'] - 0.6) <= tol
    assert output['count']['2017-07-01'] == 3


def test_from_scene_ndvi_t_interval_custom_daily_count(tol=0.0001):
    output_coll = interpolate.from_scene_ndvi(
        scene_coll(['ndvi']),
        start_date='2017-07-01', end_date='2017-08-01',
        variables=['ndvi', 'daily_count'],
        interp_args={'interp_method': 'linear', 'interp_days': 32},
        model_args={'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                    'et_reference_band': 'eto'},
        t_interval='custom',
    )

    TEST_POINT = (-121.5265, 38.7399)
    output = utils.point_coll_value(output_coll, TEST_POINT, scale=30)
    assert abs(output['ndvi']['2017-07-01'] - 0.6) <= tol
    assert output['daily_count']['2017-07-01'] == 31


def test_from_scene_ndvi_t_interval_monthly_interp_args_et_reference(tol=0.0001):
    # Check that the et_reference parameters can be set through the interp_args
    output_coll = interpolate.from_scene_ndvi(
        scene_coll(['ndvi']),
        start_date='2017-07-01', end_date='2017-08-01',
        variables=['ndvi', 'count'],
        interp_args={'interp_method': 'linear', 'interp_days': 32,
                     'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                     'et_reference_band': 'eto'},
        model_args={},
        t_interval='monthly',
    )

    TEST_POINT = (-121.5265, 38.7399)
    output = utils.point_coll_value(output_coll, TEST_POINT, scale=30)
    assert abs(output['ndvi']['2017-07-01'] - 0.6) <= tol
    assert output['count']['2017-07-01'] == 3


def test_from_scene_ndvi_t_interval_bad_value():
    # Function should raise a ValueError if t_interval is not supported
    with pytest.raises(ValueError):
        interpolate.from_scene_ndvi(
            scene_coll(['ndvi']),
            start_date='2017-07-01', end_date='2017-08-01', variables=['et'],
            interp_args={'interp_method': 'linear', 'interp_days': 32},
            model_args={'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                        'et_reference_band': 'etr'},
            t_interval='deadbeef',
        )


def test_from_scene_ndvi_t_interval_no_value():
    # Function should raise an Exception if t_interval is not set
    with pytest.raises(TypeError):
        interpolate.from_scene_ndvi(
            scene_coll(['ndvi']),
            start_date='2017-07-01', end_date='2017-08-01',
            variables=['ndvi', 'count'],
            interp_args={'interp_method': 'linear', 'interp_days': 32},
            model_args={'et_reference_source': 'IDAHO_EPSCOR/GRIDMET',
                        'et_reference_band': 'etr'},
        )
