import pytest
from numpy.testing import assert_allclose, assert_approx_equal, assert_equal
import pandas as pd
import numpy as np
from pvlib import spectrum

from .conftest import DATA_DIR, assert_series_equal

SPECTRL2_TEST_DATA = DATA_DIR / 'spectrl2_example_spectra.csv'

@pytest.fixture
def spectrl2_data():
    # reference spectra generated with solar_utils==0.3
    """
    expected = solar_utils.spectrl2(
        units=1,
        location=[40, -80, -5],
        datetime=[2020, 3, 15, 10, 45, 59],
        weather=[1013, 15],
        orientation=[0, 180],
        atmospheric_conditions=[1.14, 0.65, 0.344, 0.1, 1.42],
        albedo=[0.3, 0.7, 0.8, 1.3, 2.5, 4.0] + [0.2]*6,
    )
    """
    kwargs = {
        'surface_tilt': 0,
        'relative_airmass': 1.4899535986910446,
        'apparent_zenith': 47.912086486816406,
        'aoi': 47.91208648681641,
        'ground_albedo': 0.2,
        'surface_pressure': 101300,
        'ozone': 0.344,
        'precipitable_water': 1.42,
        'aerosol_turbidity_500nm': 0.1,
        'dayofyear': 75
    }
    df = pd.read_csv(SPECTRL2_TEST_DATA, index_col=0)
    # convert um to nm
    df['wavelength'] = np.round(df['wavelength'] * 1000, 1)
    df[['specdif', 'specdir', 'specetr', 'specglo']] /= 1000
    return kwargs, df


def test_spectrl2(spectrl2_data):
    # compare against output from solar_utils wrapper around NREL spectrl2_2.c
    kwargs, expected = spectrl2_data
    actual = spectrum.spectrl2(**kwargs)
    assert_allclose(expected['wavelength'].values, actual['wavelength'])
    assert_allclose(expected['specdif'].values, actual['dhi'].ravel(),
                    atol=7e-5)
    assert_allclose(expected['specdir'].values, actual['dni'].ravel(),
                    atol=1.5e-4)
    assert_allclose(expected['specetr'], actual['dni_extra'].ravel(),
                    atol=2e-4)
    assert_allclose(expected['specglo'], actual['poa_global'].ravel(),
                    atol=1e-4)


def test_spectrl2_array(spectrl2_data):
    # test that supplying arrays instead of scalars works
    kwargs, expected = spectrl2_data
    kwargs = {k: np.array([v, v, v]) for k, v in kwargs.items()}
    actual = spectrum.spectrl2(**kwargs)

    assert actual['wavelength'].shape == (122,)

    keys = ['dni_extra', 'dhi', 'dni', 'poa_sky_diffuse', 'poa_ground_diffuse',
            'poa_direct', 'poa_global']
    for key in keys:
        assert actual[key].shape == (122, 3)


def test_spectrl2_series(spectrl2_data):
    # test that supplying Series instead of scalars works
    kwargs, expected = spectrl2_data
    kwargs.pop('dayofyear')
    index = pd.to_datetime(['2020-03-15 10:45:59']*3)
    kwargs = {k: pd.Series([v, v, v], index=index) for k, v in kwargs.items()}
    actual = spectrum.spectrl2(**kwargs)

    assert actual['wavelength'].shape == (122,)

    keys = ['dni_extra', 'dhi', 'dni', 'poa_sky_diffuse', 'poa_ground_diffuse',
            'poa_direct', 'poa_global']
    for key in keys:
        assert actual[key].shape == (122, 3)


def test_dayofyear_missing(spectrl2_data):
    # test that not specifying dayofyear with non-pandas inputs raises error
    kwargs, expected = spectrl2_data
    kwargs.pop('dayofyear')
    with pytest.raises(ValueError, match='dayofyear must be specified'):
        _ = spectrum.spectrl2(**kwargs)


def test_aoi_gt_90(spectrl2_data):
    # test that returned irradiance values are non-negative when aoi > 90
    # see GH #1348
    kwargs, _ = spectrl2_data
    kwargs['apparent_zenith'] = 70
    kwargs['aoi'] = 130
    kwargs['surface_tilt'] = 60

    spectra = spectrum.spectrl2(**kwargs)
    for key in ['poa_direct', 'poa_global']:
        message = f'{key} contains negative values for aoi>90'
        assert np.all(spectra[key] >= 0), message


def test_get_example_spectral_response():
    # test that the sample sr is read and interpolated correctly
    sr = spectrum.get_example_spectral_response()
    assert_equal(len(sr), 185)
    assert_equal(np.sum(sr.index), 136900)
    assert_approx_equal(np.sum(sr), 107.6116)

    wavelength = [270, 850, 950, 1200, 4001]
    expected = [0.0, 0.92778, 1.0, 0.0, 0.0]

    sr = spectrum.get_example_spectral_response(wavelength)
    assert_equal(len(sr), len(wavelength))
    assert_allclose(sr, expected, rtol=1e-5)


def test_get_am15g():
    # test that the reference spectrum is read and interpolated correctly
    e = spectrum.get_am15g()
    assert_equal(len(e), 2002)
    assert_equal(np.sum(e.index), 2761442)
    assert_approx_equal(np.sum(e), 1002.88, significant=6)

    wavelength = [270, 850, 950, 1200, 4001]
    expected = [0.0, 0.893720, 0.147260, 0.448250, 0.0]

    e = spectrum.get_am15g(wavelength)
    assert_equal(len(e), len(wavelength))
    assert_allclose(e, expected, rtol=1e-6)


def test_calc_spectral_mismatch_field(spectrl2_data):
    # test that the mismatch is calculated correctly with
    # - default and custom reference sepctrum
    # - single or multiple sun spectra

    # sample data
    _, e_sun = spectrl2_data
    e_sun = e_sun.set_index('wavelength')
    e_sun = e_sun.transpose()

    e_ref = spectrum.get_am15g()
    sr = spectrum.get_example_spectral_response()

    # test with single sun spectrum, same as ref spectrum
    mm = spectrum.calc_spectral_mismatch_field(sr, e_sun=e_ref)
    assert_approx_equal(mm, 1.0, significant=6)

    # test with single sun spectrum
    mm = spectrum.calc_spectral_mismatch_field(sr, e_sun=e_sun.loc['specglo'])
    assert_approx_equal(mm, 0.992397, significant=6)

    # test with single sun spectrum, also used as reference spectrum
    mm = spectrum.calc_spectral_mismatch_field(sr,
                                               e_sun=e_sun.loc['specglo'],
                                               e_ref=e_sun.loc['specglo'])
    assert_approx_equal(mm, 1.0, significant=6)

    # test with multiple sun spectra
    expected = [0.972982, 0.995581, 0.899782, 0.992397]

    mm = spectrum.calc_spectral_mismatch_field(sr, e_sun=e_sun)
    assert mm.index is e_sun.index
    assert_allclose(mm, expected, rtol=1e-6)


@pytest.mark.parametrize("module_type,expect", [
    ('cdte', np.array(
        [[ 0.99051020, 0.97640320, 0.93975028],
         [ 1.02928735, 1.01881074, 0.98578821],
         [ 1.04750335, 1.03814456, 1.00623986]])),
    ('monosi', np.array(
        [[ 0.97769770, 1.02043409, 1.03574032],
         [ 0.98630905, 1.03055092, 1.04736262],
         [ 0.98828494, 1.03299036, 1.05026561]])),
    ('polysi', np.array(
        [[ 0.97704080, 1.01705849, 1.02613202],
         [ 0.98992828, 1.03173953, 1.04260662],
         [ 0.99352435, 1.03588785, 1.04730718]])),
    ('cigs', np.array(
        [[ 0.97459190, 1.02821696, 1.05067895],
         [ 0.97529378, 1.02967497, 1.05289307],
         [ 0.97269159, 1.02730558, 1.05075651]])),
    ('asi', np.array(
        [[ 1.05552750, 0.87707583, 0.72243772],
         [ 1.11225204, 0.93665901, 0.78487953],
         [ 1.14555295, 0.97084011, 0.81994083]]))
])
def test_spectral_factor_firstsolar(module_type, expect):
    ams = np.array([1, 3, 5])
    pws = np.array([1, 3, 5])
    ams, pws = np.meshgrid(ams, pws)
    out = spectrum.spectral_factor_firstsolar(pws, ams, module_type)
    assert_allclose(out, expect, atol=0.001)


def test_spectral_factor_firstsolar_supplied():
    # use the cdte coeffs
    coeffs = (0.87102, -0.040543, -0.00929202, 0.10052, 0.073062, -0.0034187)
    out = spectrum.spectral_factor_firstsolar(1, 1, coefficients=coeffs)
    expected = 0.99134828
    assert_allclose(out, expected, atol=1e-3)


def test_spectral_factor_firstsolar_ambiguous():
    with pytest.raises(TypeError):
        spectrum.spectral_factor_firstsolar(1, 1)


def test_spectral_factor_firstsolar_ambiguous_both():
    # use the cdte coeffs
    coeffs = (0.87102, -0.040543, -0.00929202, 0.10052, 0.073062, -0.0034187)
    with pytest.raises(TypeError):
        spectrum.spectral_factor_firstsolar(1, 1, 'cdte', coefficients=coeffs)


def test_spectral_factor_firstsolar_large_airmass():
    # test that airmass > 10 is treated same as airmass==10
    m_eq10 = spectrum.spectral_factor_firstsolar(1, 10, 'monosi')
    m_gt10 = spectrum.spectral_factor_firstsolar(1, 15, 'monosi')
    assert_allclose(m_eq10, m_gt10)


def test_spectral_factor_firstsolar_low_airmass():
    with pytest.warns(UserWarning, match='Exceptionally low air mass'):
        _ = spectrum.spectral_factor_firstsolar(1, 0.1, 'monosi')


def test_spectral_factor_firstsolar_range():
    with pytest.warns(UserWarning, match='Exceptionally high pw values'):
        out = spectrum.spectral_factor_firstsolar(np.array([.1, 3, 10]),
                                                  np.array([1, 3, 5]),
                                                  module_type='monosi')
    expected = np.array([0.96080878, 1.03055092, np.nan])
    assert_allclose(out, expected, atol=1e-3)
    with pytest.warns(UserWarning, match='Exceptionally high pw values'):
        out = spectrum.spectral_factor_firstsolar(6, 1.5, max_pw=5,
                                                  module_type='monosi')
    with pytest.warns(UserWarning, match='Exceptionally low pw values'):
        out = spectrum.spectral_factor_firstsolar(np.array([0, 3, 8]),
                                                  np.array([1, 3, 5]),
                                                  module_type='monosi')
    expected = np.array([0.96080878, 1.03055092, 1.04932727])
    assert_allclose(out, expected, atol=1e-3)
    with pytest.warns(UserWarning, match='Exceptionally low pw values'):
        out = spectrum.spectral_factor_firstsolar(0.2, 1.5, min_pw=1,
                                                  module_type='monosi')


@pytest.mark.parametrize('airmass,expected', [
    (1.5, 1.00028714375),
    (np.array([[10, np.nan]]), np.array([[0.999535, 0]])),
    (pd.Series([5]), pd.Series([1.0387675]))
])
def test_spectral_factor_sapm(sapm_module_params, airmass, expected):

    out = spectrum.spectral_factor_sapm(airmass, sapm_module_params)

    if isinstance(airmass, pd.Series):
        assert_series_equal(out, expected, check_less_precise=4)
    else:
        assert_allclose(out, expected, atol=1e-4)
