#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Nov 25 15:27:28 2019

@author: jonasg
"""
import iris
import numpy as np
import numpy.testing as npt
import pandas as pd
import pytest
from cf_units import Unit

from pyaerocom import GriddedData, const, helpers
from pyaerocom.colocateddata import ColocatedData
from pyaerocom.colocation import (
    _colocate_site_data_helper,
    _colocate_site_data_helper_timecol,
    _regrid_gridded,
    colocate_gridded_gridded,
    colocate_gridded_ungridded,
)
from pyaerocom.io import ReadMscwCtm

from ._conftest_helpers import create_fake_station_data
from .conftest import TEST_RTOL, data_unavail


def test__regrid_gridded(data_tm5):
     one_way = _regrid_gridded(data_tm5, 'areaweighted', 5)
     another_way = _regrid_gridded(data_tm5, 'areaweighted',
                                   dict(lon_res_deg=5, lat_res_deg=5))

     assert one_way.shape == another_way.shape

S1 = create_fake_station_data('concpm10', {'concpm10': {'units' : 'ug m-3'}},
                              10, '2010-01-01', '2010-12-31', 'd',
                              {'ts_type' : 'daily'})

S1['concpm10'][10:20] = np.nan

S2 = create_fake_station_data('concpm10', {'concpm10': {'units' : 'ug m-3'}},
                              10, '2010-01-01', '2010-12-31', 'd',
                              {'ts_type' : 'daily'})

S3 = create_fake_station_data('concpm10', {'concpm10': {'units' : 'ug m-3'}},
                              10, '2010-01-01', '2010-12-31', '13d',
                              {'ts_type' : '13daily'})
S3['concpm10'][1] = np.nan
S3['concpm10'][3] = np.nan

S4 = create_fake_station_data('concpm10', {'concpm10': {'units' : 'ug m-3'}},
                              10, '2010-01-03', '2011-12-31', 'd',
                              {'ts_type' : 'daily'})

S4['concpm10'][0:5] = range(5)

@pytest.mark.parametrize(
    'stat_data,stat_data_ref,var,var_ref,ts_type,resample_how,min_num_obs, use_climatology_ref,num_valid', [
    (S4, S3, 'concpm10', 'concpm10', 'monthly', 'mean', {'monthly':{'daily':25}}, False, 10),
    (S3, S4, 'concpm10', 'concpm10', 'monthly', 'mean', {'monthly':{'daily':25}}, False, 24),
    (S1, S2, 'concpm10', 'concpm10', 'monthly', 'mean', 25, False, 12),
    (S2, S1, 'concpm10', 'concpm10', 'monthly', 'mean', 25, False, 11),
        ])
def test__colocate_site_data_helper_timecol(stat_data, stat_data_ref, var, var_ref,
                                            ts_type, resample_how,
                                            min_num_obs, use_climatology_ref,
                                            num_valid):
    result = _colocate_site_data_helper_timecol(
        stat_data, stat_data_ref, var, var_ref, ts_type, resample_how,
        min_num_obs, use_climatology_ref)

    assert isinstance(result, pd.DataFrame)
    assert result.data.isnull().sum() == result.ref.isnull().sum()
    assert len(result) - result.data.isnull().sum() == num_valid


def test__colocate_site_data_helper(aeronetsunv3lev2_subset):
    var = 'od550aer'
    stat1 = aeronetsunv3lev2_subset.to_station_data(3, var)
    stat2 = aeronetsunv3lev2_subset.to_station_data(4, var)
    df = _colocate_site_data_helper(stat1, stat2, var, var,
                                    'daily',None,None,False)

    assert isinstance(df, pd.DataFrame)
    assert len(df) == 9483
    means = [np.nanmean(df['data']),
             np.nanmean(df['ref'])]
    should_be = [0.31171085422102346,
                 0.07752743643132792]
    npt.assert_allclose(means, should_be, rtol=1e-5)

def test_colocate_gridded_ungridded_new_var(data_tm5, aeronetsunv3lev2_subset):
    data = data_tm5.copy()
    data.var_name='od550bc'
    coldata = colocate_gridded_ungridded(data, aeronetsunv3lev2_subset,
                                         var_ref='od550aer')

    assert coldata.metadata['var_name'] == ['od550aer', 'od550bc']

@data_unavail
@pytest.mark.parametrize('addargs,ts_type,shape,obsmean,modmean',[
    (dict(filter_name='WORLD-noMOUNTAINS',
          min_num_obs=const.OBS_MIN_NUM_RESAMPLE),
     'monthly', (2,12,8), 0.315930,0.275671),
    (dict(filter_name='WORLD-noMOUNTAINS'),
     'monthly', (2,12,8), 0.316924,0.275671),
    (dict(filter_name='WORLD-wMOUNTAINS',
          min_num_obs=const.OBS_MIN_NUM_RESAMPLE),
     'monthly', (2,12,11), 0.269707, 0.243861),
    (dict(filter_name='WORLD-noMOUNTAINS',
          use_climatology_ref=True,
          min_num_obs=const.OBS_MIN_NUM_RESAMPLE),
     'monthly', (2,12,13), 0.302636, 0.234147),
    (dict(filter_name='WORLD-noMOUNTAINS',
          regrid_res_deg=30, min_num_obs=const.OBS_MIN_NUM_RESAMPLE),
     'monthly', (2,12,8), 0.31593 , 0.1797),
    (dict(filter_name='WORLD-noMOUNTAINS', ts_type='yearly'),
     'yearly', (2,1,8), 0.417676, 0.275671)

    ])
def test_colocate_gridded_ungridded(data_tm5, aeronetsunv3lev2_subset,
                                    addargs, ts_type, shape,
                                    obsmean, modmean):

    coldata = colocate_gridded_ungridded(data_tm5, aeronetsunv3lev2_subset,
                                         **addargs)

    assert isinstance(coldata, ColocatedData)
    assert coldata.ts_type == ts_type
    assert coldata.shape == shape

    means = [np.nanmean(coldata.data.data[0]),
             np.nanmean(coldata.data.data[1])]

    npt.assert_allclose(means, [obsmean, modmean], rtol=TEST_RTOL)



@data_unavail
def test_colocate_gridded_ungridded_nonglobal(aeronetsunv3lev2_subset):
    times = [1,2]
    time_unit = Unit("days since 2010-1-1 0:0:0")
    cubes = iris.cube.CubeList()

    for time in times:
        time_coord = iris.coords.DimCoord(time, units=time_unit,
                                          standard_name='time')
        cube = helpers.make_dummy_cube_latlon(lat_res_deg=1, lon_res_deg=1,
                                              lat_range=[30.05,81.95],
                                              lon_range=[-29.5,89.95])
        cube.add_aux_coord(time_coord)
        cubes.append(cube)
    time_cube = cubes.merge_cube()
    gridded = GriddedData(time_cube)
    gridded.var_name = 'od550aer'
    gridded.units = Unit('1')

    coldata = colocate_gridded_ungridded(gridded,
                                         aeronetsunv3lev2_subset,
                                         colocate_time=False)
    assert isinstance(coldata, ColocatedData)
    assert coldata.shape == (2,2,2)


@data_unavail
def test_colocate_gridded_gridded_same_new_var(data_tm5):
    data = data_tm5.copy()
    data.var_name = 'Blaaa'
    coldata = colocate_gridded_gridded(data, data_tm5)

    assert coldata.metadata['var_name'] == ['od550aer', 'Blaaa']

@data_unavail
def test_colocate_gridded_gridded_same(data_tm5):
    coldata = colocate_gridded_gridded(data_tm5, data_tm5)

    assert isinstance(coldata, ColocatedData)
    stats = coldata.calc_statistics()
    # check mean value
    npt.assert_allclose(stats['data_mean'], 0.09825691)
    # check that mean value is same as in input GriddedData object
    npt.assert_allclose(stats['data_mean'], data_tm5.mean(areaweighted=False))
    assert stats['refdata_mean'] == stats['data_mean']
    assert stats['nmb'] == 0
    assert stats['mnmb'] == 0
    assert stats['R'] == 1
    assert stats['R_spearman'] == 1

@data_unavail
def test_read_emep_colocate_emep_tm5(data_tm5, path_emep):
    # tempfix
    with pytest.raises(ValueError):
        filepath = path_emep['monthly']
        r = ReadMscwCtm(path_emep['monthly'])
        data_emep = r.read_var('concpm10', ts_type='monthly')

        # Change units and year to match TM5 data
        data_emep.change_base_year(2010)
        data_emep.units = '1'
        col = colocate_gridded_gridded(data_emep, data_tm5)
        assert isinstance(col, ColocatedData)

if __name__ == '__main__':
    import sys
    pytest.main(sys.argv)
