from h5py.h5t import np

from py_wake.aep import AEP
from py_wake.examples.data.iea37 import iea37_path
from py_wake.examples.data.iea37._iea37 import IEA37_WindTurbines
from py_wake.examples.data.iea37.iea37_reader import read_iea37_windrose,\
    read_iea37_windfarm
from py_wake.site._site import UniformSite
from py_wake.tests import npt
from py_wake.wake_models.noj import NOJ
from py_wake.examples.data import hornsrev1
from py_wake.wake_models.gaussian import IEA37SimpleBastankhahGaussian


def test_wake_map():
    _, _, freq = read_iea37_windrose(iea37_path + "iea37-windrose.yaml")
    n_wt = 16
    x, y, _ = read_iea37_windfarm(iea37_path + 'iea37-ex%d.yaml' % n_wt)

    site = UniformSite(freq, ti=0.75)
    windTurbines = IEA37_WindTurbines(iea37_path + 'iea37-335mw.yaml')
    wake_model = NOJ(windTurbines)
    aep = AEP(site, windTurbines, wake_model, np.arange(0, 360, 22.5), [9.8])
    x_j = np.linspace(-1500, 1500, 200)
    y_j = np.linspace(-1500, 1500, 100)
    X, Y, Z = aep.wake_map(x_j, y_j, 110, x, y, wd=[0], ws=[9])

    if 0:
        import matplotlib.pyplot as plt
        c = plt.contourf(X, Y, Z)  # , np.arange(2, 10, .01))
        plt.colorbar(c)

        plt.plot(x, y, '2k')
        for i, (x_, y_) in enumerate(zip(x, y)):
            plt.annotate(i, (x_, y_))
        plt.axis('equal')

        plt.show()

    ref = [3.27, 3.27, 9.0, 7.46, 7.46, 7.46, 7.46, 7.31, 7.31, 7.31, 7.31, 8.3, 8.3, 8.3, 8.3, 8.3, 8.3]
    npt.assert_array_almost_equal(Z[49, 100:133:2], ref, 2)


def test_aep_map():

    _, _, freq = read_iea37_windrose(iea37_path + "iea37-windrose.yaml")
    n_wt = 16
    x, y, _ = read_iea37_windfarm(iea37_path + 'iea37-ex%d.yaml' % n_wt)

    x = [0, 0]
    y = [0, 200]
    site = UniformSite(freq, ti=0.75)
    windTurbines = IEA37_WindTurbines(iea37_path + 'iea37-335mw.yaml')
    wake_model = IEA37SimpleBastankhahGaussian(windTurbines)
    aep = AEP(site, windTurbines, wake_model)
    print(aep.calculate_AEP([0], [0]).sum())
    x_j = np.arange(-150, 150, 20)
    y_j = np.arange(-250, 250, 20)
    X, Y, Z = aep.aep_map(x_j, y_j, 0, x, y, wd=[0])
    print(y_j)
    if 0:
        import matplotlib.pyplot as plt
        c = plt.contourf(X, Y, Z, 100)  # , np.arange(2, 10, .01))
        plt.colorbar(c)

        plt.plot(x, y, '2k')
        for i, (x_, y_) in enumerate(zip(x, y)):
            plt.annotate(i, (x_, y_))
        plt.axis('equal')

        plt.show()
    print(np.round(Z[17], 2).tolist())
    ref = [21.5, 21.4, 21.02, 20.34, 18.95, 16.54, 13.17, 10.17, 10.17, 13.17, 16.54, 18.95, 20.34, 21.02, 21.4]
    npt.assert_array_almost_equal(Z[17], ref, 2)


def test_aep_no_wake_loss():
    _, _, freq = read_iea37_windrose(iea37_path + "iea37-windrose.yaml")
    n_wt = 16
    x, y, _ = read_iea37_windfarm(iea37_path + 'iea37-ex%d.yaml' % n_wt)

    x = [0, 0]
    y = [0, 200]
    site = UniformSite(freq, ti=0.75)
    windTurbines = IEA37_WindTurbines(iea37_path + 'iea37-335mw.yaml')
    wake_model = IEA37SimpleBastankhahGaussian(windTurbines)
    aep = AEP(site, windTurbines, wake_model)
    print(aep.calculate_AEP_no_wake_loss(x, y))