def test_get_airmass_valueerror():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    times = pd.DatetimeIndex(start='20160101T0600-0700',
                             end='20160101T1800-0700',
                             freq='3H')
    with pytest.raises(ValueError):
        clearsky = tus.get_airmass(times, model='invalid_model')

def test_get_airmass():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    times = pd.DatetimeIndex(start='20160101T0600-0700',
                             end='20160101T1800-0700',
                             freq='3H')
    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_irradiance():
    system = tracking.SingleAxisTracker(max_angle=90, axis_tilt=30,
                                        axis_azimuth=180, gcr=2.0/7.0,
                                        backtrack=True)
    times = pd.DatetimeIndex(start='20160101 1200-0700',
                             end='20160101 1800-0700', freq='6H')
    location = Location(latitude=32, longitude=-111)
    solar_position = location.get_solarposition(times)
    irrads = pd.DataFrame({'dni':[900,0], 'ghi':[600,0], 'dhi':[100,0]},
                          index=times)
    solar_zenith = solar_position['apparent_zenith']
    solar_azimuth = solar_position['azimuth']
    tracker_data = system.singleaxis(solar_zenith, solar_azimuth)

    irradiance = system.get_irradiance(irrads['dni'],
                                       irrads['ghi'],
                                       irrads['dhi'],
                                       solar_zenith=solar_zenith,
                                       solar_azimuth=solar_azimuth,
                                       surface_tilt=tracker_data['surface_tilt'],
                                       surface_azimuth=tracker_data['surface_azimuth'])

    expected = pd.DataFrame(data=np.array(
        [[ 961.80070,   815.94490,   145.85580,   135.32820,
          10.52757492],
       [          nan,           nan,           nan,           nan,
                  nan]]),
                            columns=['poa_global', 'poa_direct',
                                     'poa_diffuse', 'poa_sky_diffuse',
                                     'poa_ground_diffuse'],
                            index=times)

    assert_frame_equal(irradiance, expected, check_less_precise=2)

def test_PVSystem_get_irradiance():
    system = pvsystem.PVSystem(surface_tilt=32, surface_azimuth=135)
    times = pd.DatetimeIndex(start='20160101 1200-0700',
                             end='20160101 1800-0700', freq='6H')
    location = Location(latitude=32, longitude=-111)
    solar_position = location.get_solarposition(times)
    irrads = pd.DataFrame({'dni':[900,0], 'ghi':[600,0], 'dhi':[100,0]},
                          index=times)

    irradiance = system.get_irradiance(solar_position['apparent_zenith'],
                                       solar_position['azimuth'],
                                       irrads['dni'],
                                       irrads['ghi'],
                                       irrads['dhi'])

    expected = pd.DataFrame(data=np.array(
        [[ 883.65494055,  745.86141676,  137.79352379,  126.397131  ,
              11.39639279],
           [   0.        ,   -0.        ,    0.        ,    0.        ,    0.        ]]),
                            columns=['poa_global', 'poa_direct',
                                     'poa_diffuse', 'poa_sky_diffuse',
                                     'poa_ground_diffuse'],
                            index=times)

    irradiance = np.round(irradiance, 4)
    expected = np.round(expected, 4)
    assert_frame_equal(irradiance, expected)

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

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

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_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_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_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_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():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    times = pd.DatetimeIndex(start='20160101T0600-0700',
                             end='20160101T1800-0700',
                             freq='3H')
    clearsky = tus.get_clearsky(times)
    expected = pd.DataFrame(data=np.array([
        (  0.0,                0.0,              0.0),
        (262.77734276159333, 791.1972825869296, 46.18714900637892),
        (616.764693938387,   974.9610353623959, 65.44157429054201),
        (419.6512657626518,  901.6234995035793, 54.26016437839348),
        (  0.0,                0.0,              0.0)],
        dtype=[('ghi', '<f8'), ('dni', '<f8'), ('dhi', '<f8')]), index=times)
    assert_frame_equal(expected, clearsky, check_less_precise=2)

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():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    times = pd.DatetimeIndex(start='20160101T0600-0700',
                             end='20160101T1800-0700',
                             freq='3H')
    clearsky = tus.get_clearsky(times)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [ 258.60422702,  761.57329257,   50.1235982 ],
               [ 611.96347869,  956.95353414,   70.8232806 ],
               [ 415.10904044,  878.52649603,   59.07820922],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['ghi', 'dni', 'dhi'],
                            index=times)
    assert_frame_equal(expected, clearsky, check_less_precise=2)

def test_get_clearsky_haurwitz():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    times = pd.DatetimeIndex(start='20160101T0600-0700',
                             end='20160101T1800-0700',
                             freq='3H')
    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_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_ineichen_supply_linke():
    tus = Location(32.2, -111, 'US/Arizona', 700)
    times = pd.date_range(start='2014-06-24', end='2014-06-25', freq='3h')
    times_localized = times.tz_localize(tus.tz)
    expected = pd.DataFrame(np.
        array([[    0.        ,     0.        ,     0.        ],
               [    0.        ,     0.        ,     0.        ],
               [   79.73090244,   316.16436502,    40.45759009],
               [  703.43653498,   876.41452667,    95.15798252],
               [ 1042.37962396,   939.86391062,   118.44687715],
               [  851.32411813,   909.11186737,   105.36662462],
               [  257.18266827,   646.16644264,    62.02777094],
               [    0.        ,     0.        ,     0.        ],
               [    0.        ,     0.        ,     0.        ]]),
                            columns=['ghi', 'dni', 'dhi'],
                            index=times_localized)
    out = tus.get_clearsky(times_localized, linke_turbidity=3)
    assert_frame_equal(expected, out, check_less_precise=2)

def test_from_tmy_2():
    from 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.tmy_data, data)

    def set_location(self, time, latitude, longitude):
        '''
        Sets the location for the query.

        Parameters
        ----------
        time: datetime or DatetimeIndex
            Time range of the query.
        '''
        if isinstance(time, datetime.datetime):
            tzinfo = time.tzinfo
        else:
            tzinfo = time.tz

        if tzinfo is None:
            self.location = Location(latitude, longitude)
        else:
            self.location = Location(latitude, longitude, tz=tzinfo)