import pandas as pd
from pyenergyaspects import eaapi as energyaspects_api


def test_get_all_timeseries_def():
    res = energyaspects_api.get_timeseries_def()
    assert 1 in res.index


def test_read_timeseries_1():
    res = energyaspects_api.read_timeseries(dataset_id=(1530, 1531), date_from=pd.to_datetime('2020-01-01'))
    assert isinstance(res.attrs['metadata'], dict)
    assert 1530 in res.columns
    assert 1531 in res.columns


def test_read_timeseries_2():
    d = (1800, 1801, 5237, 2920)
    f = pd.to_datetime('2020-01-01')
    ft = {'format': ('extract_region', 'aspect')}
    res = energyaspects_api.read_timeseries(dataset_id=d, date_from=f, column_format=ft)
    assert isinstance(res.attrs['metadata'], dict)
    assert 'US.demand' in res.columns
    assert 'AP.demand' in res.columns


def test_read_timeseries_3():
    res = energyaspects_api.read_timeseries(category="crude_oil", frequency='monthly',
                                            date_from=pd.to_datetime('2021-08-01'))
    assert isinstance(res.attrs['metadata'], dict)
    assert 30 in res.columns
    assert 67 in res.columns


def test_read_timeseries_4():
    fc = {'format': ('extract_region', 'aspect'), 'multiindex': True}
    res = energyaspects_api.read_timeseries(category_detail="gasoline", frequency='monthly',
                                            date_from=pd.to_datetime('2021-08-01'), geography='WORLD',
                                            column_format=fc)
    assert isinstance(res.attrs['metadata'], dict)
    assert ('WORLD', 'production') in res.columns
    assert ('WORLD', 'demand') in res.columns


def test_detailed_metadata():
    res = energyaspects_api.detailed_metadata(5236)
    assert res.source.iloc[0] == 'Energy Aspects'


def test_commod_balance():
    datef = pd.to_datetime('2021-08-01')
    res = energyaspects_api.read_commod_balance(commod='gasoline', date_from=datef)
    assert 5236 in res.attrs['metadata']
    assert 5233 in res.attrs['metadata']


def test_crude_quality_streams():
    f = pd.to_datetime('2020-01-01')
    ft = {'format': ('description',)}
    res = energyaspects_api.crude_quality_streams(date_from=f, column_format=ft)
    assert isinstance(res.attrs['metadata'], dict)
    assert 'Yearly Russia condensate production in kb/d' in res.columns
    assert 'Yearly Norway condensate production in kb/d' in res.columns


def test_crude_quality_streams_by_country():
    f = pd.to_datetime('2020-01-01')
    res = energyaspects_api.crude_quality_streams_by_country(date_from=f)
    assert 'SA' in res.columns
    assert 'TH' in res.columns


def test_extract_region():
    m = {'dataset_id': '1800',
         'country': '',
         'country_iso': '',
         'region': 'AP',
         'sub_region': '',
         'description': 'Monthly gasoline demand for Asia in kb/d'}
    res = energyaspects_api.extract_region(m)
    assert res == 'AP'

    m = {'dataset_id': '4031',
         'country': 'US',
         'country_iso': 'US',
         'region': 'NA',
         'sub_region': 'PADD_3',
         'description': 'Monthly natural gasoline supply for US PADD3 in kb/d', }
    res = energyaspects_api.extract_region(m)
    assert res == 'PADD_3'


def test_get_all_webfile_mappings():
    res = energyaspects_api.get_all_webfile_mappings()
    assert 'crude_oil' in res


def test_get_all_webfile_mappings_names():
    res = energyaspects_api.get_all_webfile_mappings_names()
    assert 'crude_oil' in res
    assert 'Price forecasts' in res['crude_oil']


def test_get_dataset_ids_for_dataset():
    res = energyaspects_api.get_dataset_ids_for_dataset(energyaspects_api.get_all_webfile_mappings(), 'Asia crude balance')
    assert isinstance(res, list)