import datetime
from unittest.mock import ANY

import numpy as np
from numpy import nan
import pandas as pd
from .conftest import assert_frame_equal, assert_index_equal

import pytest

import pytz
from pytz.exceptions import UnknownTimeZoneError

import pvlib
from pvlib.location import Location
from pvlib.solarposition import declination_spencer71
from pvlib.solarposition import equation_of_time_spencer71
from .conftest import requires_ephem, requires_tables


def test_location_required():
    Location(32.2, -111)


def test_location_all():
    Location(32.2, -111, 'US/Arizona', 700, 'Tucson')


@pytest.mark.parametrize('tz', [
    pytz.timezone('US/Arizona'), 'America/Phoenix',  -7, -7.0,
    datetime.timezone.utc
])
def test_location_tz(tz):
    Location(32.2, -111, tz)


def test_location_invalid_tz():
    with pytest.raises(UnknownTimeZoneError):
        Location(32.2, -111, 'invalid')


def test_location_invalid_tz_type():
    with pytest.raises(TypeError):
        Location(32.2, -111, [5])


def test_location_print_all():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    expected_str = '\n'.join([
        'Location: ',
        '  name: Tucson',
        '  latitude: 32.2',
        '  longitude: -111',
        '  altitude: 700',
        '  tz: US/Arizona'
    ])
    assert tus.__str__() == expected_str


def test_location_print_pytz():
    tus = Location(32.2, -111, pytz.timezone('US/Arizona'), 700, 'Tucson')
    expected_str = '\n'.join([
        'Location: ',
        '  name: Tucson',
        '  latitude: 32.2',
        '  longitude: -111',
        '  altitude: 700',
        '  tz: US/Arizona'
    ])
    assert tus.__str__() == expected_str


@pytest.fixture
def times():
    return pd.date_range(start='20160101T0600-0700',
                         end='20160101T1800-0700',
                         freq='3H')


@requires_tables
def test_get_clearsky(mocker, times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    m = mocker.spy(pvlib.clearsky, 'ineichen')
    out = tus.get_clearsky(times)
    assert m.call_count == 1
    assert_index_equal(out.index, times)
    # check that values are 0 before sunrise and after sunset
    assert out.iloc[0, :].sum().sum() == 0
    assert out.iloc[-1:, :].sum().sum() == 0
    # check that values are > 0 during the day
    assert (out.iloc[1:-1, :] > 0).all().all()
    assert (out.columns.values == ['ghi', 'dni', 'dhi']).all()


def test_get_clearsky_ineichen_supply_linke(mocker):
    tus = Location(32.2, -111, 'US/Arizona', 700)
    times = pd.date_range(start='2014-06-24-0700', end='2014-06-25-0700',
                          freq='3h')
    mocker.spy(pvlib.clearsky, 'ineichen')
    out = tus.get_clearsky(times, linke_turbidity=3)
    # we only care that the LT is passed in this test
    pvlib.clearsky.ineichen.assert_called_once_with(ANY, ANY, 3, ANY, ANY)
    assert_index_equal(out.index, times)
    # check that values are 0 before sunrise and after sunset
    assert out.iloc[0:2, :].sum().sum() == 0
    assert out.iloc[-2:, :].sum().sum() == 0
    # check that values are > 0 during the day
    assert (out.iloc[2:-2, :] > 0).all().all()
    assert (out.columns.values == ['ghi', 'dni', 'dhi']).all()


def test_get_clearsky_haurwitz(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='haurwitz')
    expected = pd.DataFrame(data=np.array(
                            [[   0.        ],
                             [ 242.30085588],
                             [ 559.38247117],
                             [ 384.6873791 ],
                             [   0.        ]]),
                            columns=['ghi'],
                            index=times)
    assert_frame_equal(expected, clearsky)


def test_get_clearsky_simplified_solis(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis')
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  70.00146271,  638.01145669,  236.71136245],
               [ 101.69729217,  852.51950946,  577.1117803 ],
               [  86.1679965 ,  755.98048017,  385.59586091],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_simplified_solis_apparent_elevation(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    solar_position = {'apparent_elevation': pd.Series(80, index=times),
                      'apparent_zenith': pd.Series(10, index=times)}
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                solar_position=solar_position)
    expected = pd.DataFrame(data=np.
        array([[  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_simplified_solis_dni_extra(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                dni_extra=1370)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  67.82281485,  618.15469596,  229.34422063],
               [  98.53217848,  825.98663808,  559.15039353],
               [  83.48619937,  732.45218243,  373.59500313],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky)


def test_get_clearsky_simplified_solis_pressure(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                pressure=95000)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  70.20556637,  635.53091983,  236.17716435],
               [ 102.08954904,  850.49502085,  576.28465815],
               [  86.46561686,  753.70744638,  384.90537859],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_simplified_solis_aod_pw(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                aod700=0.25, precipitable_water=2.)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  85.77821205,  374.58084365,  179.48483117],
               [ 143.52743364,  625.91745295,  490.06254157],
               [ 114.63275842,  506.52275195,  312.24711495],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_valueerror(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    with pytest.raises(ValueError):
        tus.get_clearsky(times, model='invalid_model')


def test_from_tmy_3():
    from pvlib.tests.iotools.test_tmy import TMY3_TESTFILE
    from pvlib.iotools import read_tmy3
    data, meta = read_tmy3(TMY3_TESTFILE)
    loc = Location.from_tmy(meta, data)
    assert loc.name is not None
    assert loc.altitude != 0
    assert loc.tz != 'UTC'
    assert_frame_equal(loc.weather, data)


def test_from_tmy_2():
    from pvlib.tests.iotools.test_tmy import TMY2_TESTFILE
    from pvlib.iotools import read_tmy2
    data, meta = read_tmy2(TMY2_TESTFILE)
    loc = Location.from_tmy(meta, data)
    assert loc.name is not None
    assert loc.altitude != 0
    assert loc.tz != 'UTC'
    assert_frame_equal(loc.weather, data)


def test_from_epw():
    from pvlib.tests.iotools.test_epw import epw_testfile
    from pvlib.iotools import read_epw
    data, meta = read_epw(epw_testfile)
    loc = Location.from_epw(meta, data)
    assert loc.name is not None
    assert loc.altitude != 0
    assert loc.tz != 'UTC'
    assert_frame_equal(loc.weather, data)


def test_get_solarposition(expected_solpos, golden_mst):
    times = pd.date_range(datetime.datetime(2003, 10, 17, 12, 30, 30),
                          periods=1, freq='D', tz=golden_mst.tz)
    ephem_data = golden_mst.get_solarposition(times, temperature=11)
    ephem_data = np.round(ephem_data, 3)
    expected_solpos.index = times
    expected_solpos = np.round(expected_solpos, 3)
    assert_frame_equal(expected_solpos, ephem_data[expected_solpos.columns])


def test_get_airmass(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    airmass = tus.get_airmass(times)
    expected = pd.DataFrame(data=np.array(
                            [[        nan,         nan],
                             [ 3.61046506,  3.32072602],
                             [ 1.76470864,  1.62309115],
                             [ 2.45582153,  2.25874238],
                             [        nan,         nan]]),
                            columns=['airmass_relative', 'airmass_absolute'],
                            index=times)
    assert_frame_equal(expected, airmass)

    airmass = tus.get_airmass(times, model='young1994')
    expected = pd.DataFrame(data=np.array(
                            [[        nan,         nan],
                             [ 3.6075018 ,  3.31800056],
                             [ 1.7641033 ,  1.62253439],
                             [ 2.45413091,  2.25718744],
                             [        nan,         nan]]),
                            columns=['airmass_relative', 'airmass_absolute'],
                            index=times)
    assert_frame_equal(expected, airmass)


def test_get_airmass_valueerror(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    with pytest.raises(ValueError):
        tus.get_airmass(times, model='invalid_model')


def test_Location___repr__():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')

    expected = '\n'.join([
        'Location: ',
        '  name: Tucson',
        '  latitude: 32.2',
        '  longitude: -111',
        '  altitude: 700',
        '  tz: US/Arizona'
    ])
    assert tus.__repr__() == expected


@requires_ephem
def test_get_sun_rise_set_transit(golden):
    times = pd.DatetimeIndex(['2015-01-01 07:00:00', '2015-01-01 23:00:00'],
                             tz='MST')
    result = golden.get_sun_rise_set_transit(times, method='pyephem')
    assert all(result.columns == ['sunrise', 'sunset', 'transit'])

    result = golden.get_sun_rise_set_transit(times, method='spa')
    assert all(result.columns == ['sunrise', 'sunset', 'transit'])

    dayofyear = 1
    declination = declination_spencer71(dayofyear)
    eot = equation_of_time_spencer71(dayofyear)
    result = golden.get_sun_rise_set_transit(times, method='geometric',
                                             declination=declination,
                                             equation_of_time=eot)
    assert all(result.columns == ['sunrise', 'sunset', 'transit'])


def test_get_sun_rise_set_transit_valueerror(golden):
    times = pd.DatetimeIndex(['2015-01-01 07:00:00', '2015-01-01 23:00:00'],
                             tz='MST')
    with pytest.raises(ValueError):
        golden.get_sun_rise_set_transit(times, method='eyeball')


def test_extra_kwargs():
    with pytest.raises(TypeError, match='arbitrary_kwarg'):
        Location(32.2, -111, arbitrary_kwarg='value')
