# -*- coding: utf-8 -*-
"""
Components that require Floris.

@author: ricriv
"""

# %% Import.

import sys
import warnings
from copy import deepcopy

import numpy as np
import xarray as xr

from wind_farm_loads.tool_agnostic import (
    preallocate_ilktn,
    rotate_grid,
)

# %% Utils.


def monkey_patch_floris_solver():
    """
    Replace Floris solver with the ones that are modified to provide no deficit around the current turbine. Must be called before importing Floris.

    Returns
    -------
    None.

    """
    import wind_farm_loads._floris_solver

    sys.modules["floris.core.solver"] = wind_farm_loads._floris_solver
    import floris  # noqa: F401


def _wt_wd_ws_from_fmodel(fmodel):
    """
    Extract the turbine index, wind direction and wind speed from the Floris model.

    We assume that their cartesian product generates the findex - turbine table.

    Parameters
    ----------
    fmodel : FlorisModel
        Floris model.

    Returns
    -------
    wt : (I, ) numpy array
        Turbine indices
    wd : (L, ) numpy array
        Wind directions.
    ws : (K, ) numpy array
        Wind speeds.

    """
    # Get number of turbines.
    I = fmodel.n_turbines
    wt = np.arange(I)

    # Map wind direction and speed to findex.
    # unique() is needed because fmodel should contain all combinations of wind direction and speed.
    wd = np.unique(fmodel.core.flow_field.wind_directions)
    ws = np.unique(fmodel.core.flow_field.wind_speeds)
    L = wd.size
    K = ws.size
    assert (
        fmodel.core.flow_field.n_findex == L * K
    ), f"There are {fmodel.core.flow_field.n_findex} flow cases, but I have found {L} wind directions and {K} wind speeds."

    return wt, wd, ws


def findex_to_ilk(v, I, L, K):
    """
    Convert a variable from shape `(findex, I)` to `(I, L, K)`.

    Assumes that the flow conditions have been generated as

    .. code-block:: python

       wd, ws = np.meshgrid(wd_ambient, ws_ambient, indexing="ij")
       wind_directions = wd.ravel()
       wind_speeds = ws.ravel()

    Parameters
    ----------
    v : float, array_like
        Input variable that needs to be converted to `(I, L, K)` shape.
    I : int
        Number of turbines.
    L : int
        Number of wind directions.
    K : int
        Number of wind speeds.

    Returns
    -------
    r : (I, L, K) ndarray
        Input variable converted to `(I, L, K)` shape.

    """
    return v.T.reshape(I, L, K)


# %% Functions to extract the inflow.


def get_rotor_averaged_wind_speed_and_turbulence_intensity(fmodel):
    """
    Get rotor-averaged effective wind speed and turbulence intensity.

    The wind speed and TI computed by Floris are converted to the PyWake convention.
    Thus, this function attempts to convert from Floris (`findex`, `turbine`) to PyWake `ilk` form.

    Parameters
    ----------
    fmodel : FlorisModel
        Floris model. Must follow a call to `run()`.

    Returns
    -------
    ws_eff : xarray DataArray
        Effective wind speed.
    ti_eff : xarray DataArray
        Effective turbulence intensity.

    """
    # Get ws and TI and convert to DataArray.
    wt, wd, ws = _wt_wd_ws_from_fmodel(fmodel)
    I, L, K = wt.size, wd.size, ws.size
    ws_eff = preallocate_ilktn(
        wt=wt,
        wd=wd,
        ws=ws,
        data=findex_to_ilk(fmodel.turbine_average_velocities, I, L, K),
    )
    ti_eff = preallocate_ilktn(
        wt=wt, wd=wd, ws=ws, data=findex_to_ilk(fmodel.get_turbine_TIs(), I, L, K)
    )

    return ws_eff, ti_eff


def compute_flow_map(
    fmodel,
    grid,
    align_in_yaw=True,
    align_in_tilt=True,
    axial_wind=False,
    wt=None,
    wd=None,
    ws=None,
    time=None,
    dtype=np.float32,
    save_grid=False,
):
    r"""
    Compute the effective wind speed and Turbulence Intensity over a rotor.

    This function receives a grid, and then rotates it by the wind direction. Optionally,
    the grid is also rotated by the yaw and tilt of each turbine to align it with the rotor plane.
    Finally, the grid is translated to each rotor center and the flow map is computed.

    Parameters
    ----------
    fmodel : FlorisModel
        Floris model. Must follow a call to `run()`.
    grid : (N, M, 3) or (N, M, 3, Type) ndarray
        x, y and z coordinate of the grid points, before rotation by yaw and tilt.
        Typically generated by `make_rectangular_grid` or `make_polar_grid`.
        The first 2 dimensions cover the rotor, then there are x, y, z and finally (optionally) the turbine type.
        If the user passes a 3D array, the grid is assumed to be the same for all turbine types.
    align_in_yaw : bool, optional
        If `True` (default) the grid is aligned in yaw with the rotor plane.
    align_in_tilt : bool, optional
        If `True` (default) the grid is aligned in tilt with the rotor plane.
    axial_wind : bool, optional
        If `True` the axial wind speed and TI are returned. That is, the downstream wind speed computed by PyWake
        is multiplied by :math:`\cos(\mathrm{yaw}) \cos(\mathrm{tilt})`. The default is `False`.
    wt : int, (I) array_like, optional
        Wind turbines. Must be a subset of the one contained in `sim_res`.
        The default is `None`, which means to use all available wind turbines.
    wd : float, (L) array_like, optional
        Wind direction, in deg. Must be a subset of the one contained in `sim_res`.
        The default is `None`, which means to use all available wind directions.
    ws : float, (K) array_like, optional
        Wind speed. Must be a subset of the one contained in `sim_res`.
        The default is `None`, which means to use all available wind speeds.
    time : float, (Time) array_like, optional
        Time. Must be a subset of the one contained in `sim_res`.
        The default is `None`, which means to use all available time instants.
    dtype : data-type, optional
        The desired data-type for the result. The default is single precision,
        which should be enough for all outputs. The properties of each type can
        be checked with `np.finfo(np.float32(1.0))`.
    save_grid : bool, optional
        If `True` the grid will be saved for all inflow conditions. Since this comes at a significant
        memory cost, it is recommended to switch it on only for debug purposes.
        The default is `False`.

    Returns
    -------
    flow_map : xarray DataSet
        Effective wind speed, effective turbulence intensity and corresponding grid points
        for each turbine and flow case.

    """
    n_types = len(fmodel.configuration["farm"]["turbine_type"])

    # The grid must be a numpy array with 3 dimensions.
    # The first 2 cover the rotor, while the last is over the types.
    # This implies that all turbine types must have the same number of grid points.
    # If the user passes a 2D array, the grid is assumed to be the same for all types.
    if grid.ndim == 3:
        grid_t = np.broadcast_to(grid[:, :, :, np.newaxis], (*grid.shape, n_types))
    elif grid.ndim == 4:
        grid_t = grid
        # Check that there is 1 grid per turbine type.
        if grid_t.shape[3] != n_types:
            raise ValueError(
                f"{grid_t.shape[3]} grid types provided, but {n_types} were expected."
            )
    else:
        raise ValueError("The grid must be a 3D or 4D array.")

    # Extract wind turbine, wind direction and wind speed from fmodel.
    wt_, wd_, ws_ = _wt_wd_ws_from_fmodel(fmodel)
    I, L, K = wt_.size, wd_.size, ws_.size

    findex_table = xr.DataArray(
        data=np.arange(fmodel.core.flow_field.n_findex).reshape(L, K),
        coords={
            "wd": wd_,
            "ws": ws_,
        },
    )

    # The default value of wt, wd, ws and time is the one in fmodel.
    if wt is not None:
        # Test that the input for wt is contained in fmodel.
        if not np.all(np.isin(np.atleast_1d(wt), wt_)):
            raise ValueError("Input 'wt' must be contained in the Floris model.")
        wt_ = np.atleast_1d(wt)

    if wd is not None:
        wd_ = np.atleast_1d(wd)
        # Test that the input for wd is contained in fmodel.
        if not np.all(np.isin(wd_, fmodel.core.flow_field.wind_directions)):
            raise ValueError("Input 'wd' must be contained in the Floris model.")

    if ws is not None:
        ws_ = np.atleast_1d(ws)
        # Test that the input for ws is contained in fmodel.
        if not np.all(np.isin(ws_, fmodel.core.flow_field.wind_speeds)):
            raise ValueError("Input 'ws' must be contained in the Floris model.")

    # Convert yaw and tilt to arrays.
    # If time is not present the result has shape (I, L, K), i.e. (turbines, wind directions, wind speeds).
    # Instead, if time is present, the result has shape (I, Time), i.e. (turbines, time).
    # These arrays are contained in fmodel, therefore all turbines, directions and speeds and times must be used.
    # In Floris yaw and tilt have shape (fmodel.core.flow_field.n_findex, I).
    # Transpose it to bring it to (I, L*K) and then reshape to (I, L, K).
    # Here we are assuming that the flow conditions have been generated as
    #    wd, ws = np.meshgrid(wd_ambient, ws_ambient, indexing="ij")
    #    wind_directions = wd.ravel()
    #    wind_speeds = ws.ravel()
    if align_in_yaw:
        yaw_turbines_ = fmodel.core.farm.yaw_angles
    else:
        yaw_turbines_ = 0.0
    yaw_turbines_ = findex_to_ilk(yaw_turbines_, I, L, K)

    # Set tilt angle.
    if align_in_tilt:
        # Compute tilt angle for the effective wind speed at each findex and turbine.
        tilt_turbines_ = fmodel.core.farm.calculate_tilt_for_eff_velocities(
            fmodel.turbine_average_velocities
        )
        tilt_turbines_ = tilt_turbines_
    else:
        tilt_turbines_ = 0.0
    tilt_turbines_ = findex_to_ilk(tilt_turbines_, I, L, K)

    # Conveniently access turbine position.
    x_turbines_ = fmodel.core.farm.layout_x  # Shape: (I, )
    y_turbines_ = fmodel.core.farm.layout_y  # Shape: (I, )
    z_turbines_ = fmodel.core.farm.hub_heights  # Shape: (n_findex, I)

    # Preallocate DataSet for effective wind speed, turbulence intensity and grid points.
    # In the flow map returned by this function wt, wd and ws, or time, are placed first, followed by the quantity and grid dimensions.
    # This order enables vectorization in predict_loads_pod().
    # Each turbine type is allowed to have a different grid, but all grids must have the same number of points.
    # The grid dimensions are labeled q0 and q1 because they might either be y and z or radius and azimuth.
    xr_dict = {}
    xr_dict["flow"] = xr.DataArray(
        data=np.empty(
            (
                wt_.size,
                wd_.size,
                ws_.size,
                grid_t.shape[0],
                grid_t.shape[1],
                2,  # Effective wind speed and TI.
            ),
            dtype=dtype,
        ),
        coords={
            "wt": wt_,
            "wd": wd_,
            "ws": ws_,
            "quantity": ["WS_eff", "TI_eff"],
        },
        dims=["wt", "wd", "ws", "q0", "q1", "quantity"],
    )

    if save_grid:
        xr_dict["grid"] = xr.DataArray(
            data=np.full(
                (
                    wt_.size,
                    wd_.size,
                    ws_.size,
                    grid_t.shape[0],
                    grid_t.shape[1],
                    3,  # x, y, z
                ),
                np.nan,
                dtype=dtype,
            ),
            coords={
                "wt": wt_,
                "wd": wd_,
                "ws": ws_,
                "axis": ["x", "y", "z"],
            },
            dims=["wt", "wd", "ws", "q0", "q1", "axis"],
        )
    ds = xr.Dataset(xr_dict)

    # Convert all angles from deg to rad.
    wd_rad = np.deg2rad(wd_)
    yaw_turbines_ = np.deg2rad(yaw_turbines_)
    tilt_turbines_ = np.deg2rad(tilt_turbines_)

    angle_ref = np.deg2rad(90.0)

    # Make a local copy of fmodel.
    fmodel_ = deepcopy(fmodel)

    # Save the grid for all turbines. We are going to update it at each wind direction and speed.
    # It is used to vectorize Floris flow map over the turbines.
    grid_all_turbines = np.zeros(
        (
            wt_.size,
            grid_t.shape[0],
            grid_t.shape[1],
            3,  # x, y, z
        ),
        dtype=dtype,
    )

    # Loop over wind directions and wind speeds.
    for l in range(wd_.size):
        for k in range(ws_.size):
            # Get flow index.
            findex = int(findex_table[l, k])
            lk = {"wd": wd_[l], "ws": ws_[k]}

            # Set the current flow condition and run the model.
            # fmodel_.reset_operation()
            fmodel_.set(
                wind_directions=[wd_[l]],
                wind_speeds=[ws_[k]],
                turbulence_intensities=[
                    fmodel.core.flow_field.turbulence_intensities[findex]
                ],
                yaw_angles=fmodel.core.farm.yaw_angles[[findex], :],
            )
            fmodel_.run()

            # Loop over the turbines.
            for i in wt_:
                # Get type of current turbine.
                # The type is just a string in fmodel.configuration["farm"]["turbine_type"]. How do we convert it to int?
                i_type = 0

                # Convert grid from downwind-crosswind-z to east-north-z.
                # While doing that, also rotate by yaw and tilt.
                # This can be done because the order of rotations is first yaw and then tilt.
                # It will NOT work for a floating turbine.
                # We rely on this function to create new arrays, so that the following
                # translation will not affect the original ones.
                # The formula for the yaw is taken from py_wake.wind_turbines._wind_turbines.WindTurbines.plot_xy()
                grid_current = rotate_grid(
                    grid_t[:, :, :, i_type],
                    yaw=angle_ref - wd_rad[l] + yaw_turbines_[i, l, k],  # [rad]
                    tilt=-tilt_turbines_[i, l, k],  # [rad]
                    degrees=False,
                )

                # rotate_grid() provides a JAX array. Let us convert it to numpy.
                grid_all_turbines[i, :, :, :] = np.asarray(grid_current)

            # Translate grids to each rotor center. The turbine position is in east-north-z coordinates.
            grid_all_turbines[:, :, :, 0] += x_turbines_[wt_, np.newaxis, np.newaxis]
            grid_all_turbines[:, :, :, 1] += y_turbines_[wt_, np.newaxis, np.newaxis]
            grid_all_turbines[:, :, :, 2] += z_turbines_[
                findex, wt_, np.newaxis, np.newaxis
            ]

            if save_grid:
                ds["grid"].loc[lk] = grid_all_turbines

            # Now that the grid is available for all turbines, compute the wind speed over the rotor disk.
            # There is only 1 flow index on the rows. Points are on the columns.
            flow_map_ws = fmodel_.sample_flow_at_points(
                grid_all_turbines[:, :, :, 0].ravel(),
                grid_all_turbines[:, :, :, 1].ravel(),
                grid_all_turbines[:, :, :, 2].ravel(),
            )[0, :]
            ds["flow"].loc[{**lk, "quantity": "WS_eff"}] = flow_map_ws.reshape(
                wt_.size, grid_t.shape[0], grid_t.shape[1]
            )

            try:
                # Same for the TI. This feature is available since Floris 4.5.
                # Calling sample_ti_at_points() would run again the solver.
                # Instead, we do exactly what sample_ti_at_points() does.
                # That is, we leverage sample_flow_at_points() and then access a variable.
                # There is only 1 flow index on the rows. Points are on the columns.
                flow_map_ti = fmodel_.core.flow_field.turbulence_intensity_field_sorted[
                    0, :, 0, 0
                ]
                ds["flow"].loc[{**lk, "quantity": "TI_eff"}] = flow_map_ti.reshape(
                    wt_.size, grid_t.shape[0], grid_t.shape[1]
                )
            except AttributeError:  # pragma: no cover
                # If sample_ti_at_points() is not available, raise a warning and provide a uniform TI.
                # This cannot be tested because we lock the version of Floris.
                warnings.warn(
                    "Floris sample_ti_at_points() is not available. TI will be uniform over the rotor disk."
                )  # pragma: no cover
                turbine_ti = fmodel_.get_turbine_TIs()[0, :, 0, 0]  # pragma: no cover
                ds["flow"].loc[{**lk, "quantity": "TI_eff"}] = turbine_ti[
                    :, np.newaxis, np.newaxis
                ]  # pragma: no cover

    # Project wind speed.
    if axial_wind:
        cos_yaw_cos_tilt = np.cos(yaw_turbines_) * np.cos(tilt_turbines_)
        ds["flow"] *= cos_yaw_cos_tilt[:, :, :, np.newaxis, np.newaxis, np.newaxis]

    return ds
