# Copyright 2021 NREL

# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

# See https://floris.readthedocs.io for documentation


from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Iterable

import attrs
from attrs import define, field
import numpy as np

from floris.utilities import Vec3, rotate_coordinates_rel_west, cosd, sind
from floris.type_dec import  (
    floris_float_type,
    floris_array_converter,
    NDArrayFloat,
    NDArrayInt
)

@define
class Grid(ABC):
    """
    Grid should establish domain bounds based on given criteria,
    and develop three arrays to contain components of the grid
    locations in space.

    This could be generalized to any number of dimensions to be
    used by perhaps a turbulence field.

    The grid will have to be reestablished for each wind direction since the planform
    area of the farm will be different.

    x are the locations in space in the primary direction (typically the direction of the wind)
    y are the locations in space in the lateral direction
    z are the locations in space in the vertical direction
    u are the velocity components at each point in space
    v are the velocity components at each point in space
    w are the velocity components at each point in space
    all of these arrays are the same size

    Args:
        turbine_coordinates (`list[Vec3]`): The collection of turbine coordinate (`Vec3`) objects.
        reference_turbine_diameter (:py:obj:`float`): The reference turbine's rotor diameter.
        grid_resolution (:py:obj:`int` | :py:obj:`Iterable(int,)`): Grid resolution specific to each grid type
    """
    turbine_coordinates: list[Vec3] = field()
    reference_turbine_diameter: float
    grid_resolution: int | Iterable = field()
    wind_directions: NDArrayFloat = field(converter=floris_array_converter)
    wind_speeds: NDArrayFloat = field(converter=floris_array_converter)

    n_turbines: int = field(init=False)
    n_wind_speeds: int = field(init=False)
    n_wind_directions: int = field(init=False)
    turbine_coordinates_array: NDArrayFloat = field(init=False)
    x: NDArrayFloat = field(init=False)
    y: NDArrayFloat = field(init=False)
    z: NDArrayFloat = field(init=False)

    def __attrs_post_init__(self) -> None:
        self.turbine_coordinates_array = np.array([c.elements for c in self.turbine_coordinates])

    @turbine_coordinates.validator
    def check_coordinates(self, instance: attrs.Attribute, value: list[Vec3]) -> None:
        """Ensures all elements are `Vec3` objects and keeps the `n_turbines` attribute up to date."""
        types = np.unique([isinstance(c, Vec3) for c in value])
        if not all(types):
            raise TypeError("'turbine_coordinates' must be `Vec3` objects.")

        self.n_turbines = len(value)

    @wind_speeds.validator
    def wind_speeds_validator(self, instance: attrs.Attribute, value: NDArrayFloat) -> None:
        """Using the validator method to keep the `n_wind_speeds` attribute up to date."""
        self.n_wind_speeds = value.size

    @wind_directions.validator
    def wind_directions_validator(self, instance: attrs.Attribute, value: NDArrayFloat) -> None:
        """Using the validator method to keep the `n_wind_directions` attribute up to date."""
        self.n_wind_directions = value.size

    @grid_resolution.validator
    def grid_resolution_validator(self, instance: attrs.Attribute, value: int | Iterable) -> None:
        # TODO move this to the grid types and off of the base class
        """Check that grid resolution is given as int or Vec3 with int components."""
        if isinstance(value, int) and type(self) is TurbineGrid:
            return
        elif isinstance(value, Iterable) and type(self) is FlowFieldPlanarGrid:
            assert type(value[0]) is int
            assert type(value[1]) is int
        elif isinstance(value, Iterable) and type(self) is FlowFieldGrid:
            assert type(value[0]) is int
            assert type(value[1]) is int
            assert type(value[2]) is int
        else:
            raise TypeError("`grid_resolution` must be of type int or Iterable(int,)")

    @abstractmethod
    def set_grid(self) -> None:
        raise NotImplementedError("Grid.set_grid")

    @abstractmethod
    def finalize(self) -> None:
        raise NotImplementedError("Grid.finalize")


@define
class TurbineGrid(Grid):
    """See `Grid` for more details.

    Args:
        turbine_coordinates (`list[Vec3]`): The collection of turbine coordinate (`Vec3`) objects.
        reference_turbine_diameter (:py:obj:`float`): The reference turbine's rotor diameter.
        wind_directions (`list[float]`): The input wind directions
        wind_speeds (`list[float]`): The input wind speeds
        grid_resolution (:py:obj:`int`): The number of points on each turbine
    """
    sorted_indices: NDArrayInt = field(init=False)
    unsorted_indices: NDArrayInt = field(init=False)

    def __attrs_post_init__(self) -> None:
        super().__attrs_post_init__()
        self.set_grid()

    def set_grid(self) -> None:
        """
        Create grid points at each turbine for each wind direction and wind speed in the simulation.
        This creates the underlying data structure for the calculation.

        arrays have shape (n wind directions, n wind speeds, n turbines, m grid spanwise, m grid vertically)
        - dimension 1: each wind direction
        - dimension 2: each wind speed
        - dimension 3: each turbine
        - dimension 4: number of points in the spanwise direction (ngrid)
        - dimension 5: number of points in the vertical dimension (ngrid)

        For example
        - x is [n wind direction, n wind speeds, n turbines, x-component of the points in the spanwise direction, x-component of the points in the vertical direction]
        - y is [n wind direction, n wind speeds, n turbines, y-component of the points in the spanwise direction, y-component of the points in the vertical direction]

        The x,y,z arrays contain the actual locations in that direction.

        # -   **self.grid_resolution** (*int*, optional): The square root of the number
        #             of points to use on the turbine grid. This number will be
        #             squared so that the points can be evenly distributed.
        #             Defaults to 5.

        If the grid conforms to the sequential solver interface, it must be sorted from upstream to downstream

        In a y-z plane on the rotor swept area, the -2 dimension is a column of points and the -1 dimension is the row number.
        So the following line prints the 0'th column of the the 0'th turbine's grid:
        print(grid.y[0,0,0,0,:])
        print(grid.z[0,0,0,0,:])
        And this line prints a single point
        print(grid.y[0,0,0,0,0])
        print(grid.z[0,0,0,0,0])
        Note that the x coordinates are all the same for the rotor plane.

        """
        # TODO: Where should we locate the coordinate system? Currently, its at
        # the foot of the turbine where the tower meets the ground.

        # These are the rotated coordinates of the wind turbines based on the wind direction
        x, y, z = rotate_coordinates_rel_west(self.wind_directions, self.turbine_coordinates_array)

        # -   **rloc** (*float, optional): A value, from 0 to 1, that determines
        #         the width/height of the grid of points on the rotor as a ratio of
        #         the rotor radius.
        #         Defaults to 0.5.

        # Create the data for the turbine grids
        radius_ratio = 0.5
        disc_area_radius = radius_ratio * self.reference_turbine_diameter / 2
        template_grid = np.ones(
            (
                self.n_wind_directions,
                self.n_wind_speeds,
                self.n_turbines,
                self.grid_resolution,
                self.grid_resolution,
            ),
            dtype=floris_float_type
        )

        # Calculate the radial distance from the center of the turbine rotor
        disc_grid = np.linspace(-1 * disc_area_radius, disc_area_radius, self.grid_resolution, dtype=floris_float_type)

        # Construct the turbine grids
        # Here, they are already rotated to the correct orientation for each wind direction
        _x = x[:, :, :, None, None] * template_grid

        ones_grid = np.ones((self.grid_resolution, self.grid_resolution), dtype=floris_float_type)
        # [:, None] on disc_grid below is effectively a transpose of this vector; compare with disc_grid.reshape(1,-1).T
        _y = y[:, :, :, None, None] + template_grid * ( disc_grid[:, None] * ones_grid )
        _z = z[:, :, :, None, None] + template_grid * ( disc_grid * ones_grid )

        # Sort the turbines at each wind direction

        # Get the sorted indices for the x coordinates. These are the indices
        # to sort the turbines from upstream to downstream for all wind directions.
        # Also, store the indices to sort them back for when the calculation finishes.
        self.sorted_indices = _x.argsort(axis=2)
        self.unsorted_indices = self.sorted_indices.argsort(axis=2)

        # Put the turbines into the final arrays in their sorted order
        self.x = np.take_along_axis(_x, self.sorted_indices, axis=2)
        self.y = np.take_along_axis(_y, self.sorted_indices, axis=2)
        self.z = np.take_along_axis(_z, self.sorted_indices, axis=2)

    def finalize(self):
        # Replace the turbines in their unsorted order so that
        # we can report values in a sane way.
        self.x = np.take_along_axis(self.x, self.unsorted_indices, axis=2)
        self.y = np.take_along_axis(self.y, self.unsorted_indices, axis=2)
        self.z = np.take_along_axis(self.z, self.unsorted_indices, axis=2)


@define
class FlowFieldGrid(Grid):
    """
    Args:
        grid_resolution (`Vec3`): The number of grid points to be created in each direction.
        turbine_coordinates (`list[Vec3]`): The collection of turbine coordinate (`Vec3`) objects.
        reference_turbine_diameter (:py:obj:`float`): The reference turbine's rotor diameter.
        grid_resolution (:py:obj:`int`): The number of points on each turbine
    """

    def __attrs_post_init__(self) -> None:
        super().__attrs_post_init__()
        self.set_grid()

    def set_grid(self) -> None:
        """
        Create a structured grid for the entire flow field domain.
        resolution: Vec3

        Calculates the domain bounds for the current wake model. The bounds
        are calculated based on preset extents from the
        given layout. The bounds consist of the minimum and maximum values
        in the x-, y-, and z-directions.

        If the Curl model is used, the predefined bounds are always set.

        First, sort the turbines so that we know the bounds in the correct orientation.
        Then, create the grid based on this wind-from-left orientation
        """

        # These are the rotated coordinates of the wind turbines based on the wind direction
        x, y, z = rotate_coordinates_rel_west(self.wind_directions, self.turbine_coordinates_array)

        # Construct the arrays storing the grid points
        eps = 0.01
        xmin = min(x[0,0]) - 2 * self.reference_turbine_diameter
        xmax = max(x[0,0]) + 10 * self.reference_turbine_diameter
        ymin = min(y[0,0]) - 2 * self.reference_turbine_diameter
        ymax = max(y[0,0]) + 2 * self.reference_turbine_diameter
        zmin = 0 + eps
        zmax = 6 * max(z[0,0])

        x_points, y_points, z_points = np.meshgrid(
            np.linspace(xmin, xmax, int(self.grid_resolution[0])),
            np.linspace(ymin, ymax, int(self.grid_resolution[1])),
            np.linspace(zmin, zmax, int(self.grid_resolution[2])),
            indexing="ij"
        )

        self.x = x_points[None, None, :, :, :]
        self.y = y_points[None, None, :, :, :]
        self.z = z_points[None, None, :, :, :]

    def finalize(self):
        pass


@define
class FlowFieldPlanarGrid(Grid):
    """
    Args:
        grid_resolution (`Vec3`): The number of grid points to be created in each direction.
        turbine_coordinates (`list[Vec3]`): The collection of turbine coordinate (`Vec3`) objects.
        reference_turbine_diameter (:py:obj:`float`): The reference turbine's rotor diameter.
        grid_resolution (:py:obj:`int`): The number of points on each turbine
    """
    normal_vector: str = field()
    planar_coordinate: float = field()
    x1_bounds: tuple = field(default=None)
    x2_bounds: tuple = field(default=None)

    sorted_indices: NDArrayInt = field(init=False)
    unsorted_indices: NDArrayInt = field(init=False)

    def __attrs_post_init__(self) -> None:
        super().__attrs_post_init__()
        self.set_grid()

    def set_grid(self) -> None:
        """
        Create a structured grid for the entire flow field domain.
        resolution: Vec3

        Calculates the domain bounds for the current wake model. The bounds
        are calculated based on preset extents from the
        given layout. The bounds consist of the minimum and maximum values
        in the x-, y-, and z-directions.

        If the Curl model is used, the predefined bounds are always set.

        First, sort the turbines so that we know the bounds in the correct orientation.
        Then, create the grid based on this wind-from-left orientation
        """
        # These are the rotated coordinates of the wind turbines based on the wind direction
        x, y, z = rotate_coordinates_rel_west(self.wind_directions, self.turbine_coordinates_array)

        max_diameter = self.reference_turbine_diameter

        if self.normal_vector == "z":  # Rules of thumb for horizontal plane
            if self.x1_bounds is None:
                self.x1_bounds = (np.min(x) - 2 * max_diameter, np.max(x) + 10 * max_diameter)

            if self.x2_bounds is None:
                self.x2_bounds = (np.min(y) - 2 * max_diameter, np.max(y) + 2 * max_diameter)

            # TODO figure out proper z spacing for GCH, currently set to +/- 10.0
            x_points, y_points, z_points = np.meshgrid(
                np.linspace(self.x1_bounds[0], self.x1_bounds[1], int(self.grid_resolution[0])),
                np.linspace(self.x2_bounds[0], self.x2_bounds[1], int(self.grid_resolution[1])),
                np.array([float(self.planar_coordinate) - 10.0, float(self.planar_coordinate), float(self.planar_coordinate) + 10.0]),
                indexing="ij"
            )

            self.x = x_points[None, None, :, :, :]
            self.y = y_points[None, None, :, :, :]
            self.z = z_points[None, None, :, :, :]

        elif self.normal_vector == "x":  # Rules of thumb for cross plane
            if self.x1_bounds is None:
                self.x1_bounds = (np.min(y) - 2 * max_diameter, np.max(y) + 2 * max_diameter)

            if self.x2_bounds is None:
                self.x2_bounds = (0.001, 6 * np.max(z))

            x_points, y_points, z_points = np.meshgrid(
                np.array([float(self.planar_coordinate)]),
                np.linspace(self.x1_bounds[0], self.x1_bounds[1], int(self.grid_resolution[0])),
                np.linspace(self.x2_bounds[0], self.x2_bounds[1], int(self.grid_resolution[1])),
                indexing="ij"
            )

            self.x = x_points[None, None, :, :, :]
            self.y = y_points[None, None, :, :, :]
            self.z = z_points[None, None, :, :, :]

        elif self.normal_vector == "y":  # Rules of thumb for y plane
            if self.x1_bounds is None:
                self.x1_bounds = (np.min(x) - 2 * max_diameter, np.max(x) + 10 * max_diameter)

            if self.x2_bounds is None:
                self.x2_bounds = (0.001, 6 * np.max(z))

            x_points, y_points, z_points = np.meshgrid(
                np.linspace(self.x1_bounds[0], self.x1_bounds[1], int(self.grid_resolution[0])),
                np.array([float(self.planar_coordinate)]),
                np.linspace(self.x2_bounds[0], self.x2_bounds[1], int(self.grid_resolution[1])),
                indexing="ij"
            )

            self.x = x_points[None, None, :, :, :]
            self.y = y_points[None, None, :, :, :]
            self.z = z_points[None, None, :, :, :]

        # self.sorted_indices = self.x.argsort(axis=2)
        # self.unsorted_indices = self.sorted_indices.argsort(axis=2)

        # Put the turbines into the final arrays in their sorted order
        # self.x = np.take_along_axis(self.x, self.sorted_indices, axis=2)
        # self.y = np.take_along_axis(self.y, self.sorted_indices, axis=2)
        # self.z = np.take_along_axis(self.z, self.sorted_indices, axis=2)

    def finalize(self):
        # sorted_indices = self.x.argsort(axis=2)
        # unsorted_indices = sorted_indices.argsort(axis=2)

        # # print(self.x)

        # x_coordinates, y_coordinates, _ = self.turbine_coordinates_array.T
        
        # x_center_of_rotation = (np.min(x_coordinates) + np.max(x_coordinates)) / 2
        # y_center_of_rotation = (np.min(y_coordinates) + np.max(y_coordinates)) / 2
        # # print(x_center_of_rotation)
        # # print(y_center_of_rotation)
        # # lkj

        # self.x = np.take_along_axis(self.x, self.unsorted_indices, axis=2)
        # self.y = np.take_along_axis(self.y, self.unsorted_indices, axis=2)
        # self.z = np.take_along_axis(self.z, self.unsorted_indices, axis=2)
        # # print(self.x)

        # self.x, self.y, self.z = self._rotated_grid(-1 * self.wind_directions, (x_center_of_rotation, y_center_of_rotation))
        # TODO figure out how to un-rotate grid for plotting after it has been solved
        pass

    # def _rotated_grid(self, angle, center_of_rotation):
    #     """
    #     Rotate the discrete flow field grid.
    #     """
    #     angle = ((angle - 270) % 360 + 360) % 360
    #     # angle = np.reshape(angle, (len(angle), 1, 1))
    #     xoffset = self.x - center_of_rotation[0]
    #     yoffset = self.y - center_of_rotation[1]
    #     rotated_x = (
    #         xoffset * cosd(angle) - yoffset * sind(angle) + center_of_rotation[0]
    #     )
    #     rotated_y = (
    #         xoffset * sind(angle) + yoffset * cosd(angle) + center_of_rotation[1]
    #     )
    #     return rotated_x, rotated_y, self.z