# Copyright (c) 2017 The Verde Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
#
# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
#
import numpy as np
import pandas as pd

from .base import BaseGridder
from .base.base_classes import get_instance_region, project_coordinates
from .base.utils import check_data
from .coordinates import check_region, scatter_points


class CheckerBoard(BaseGridder):
    r"""
    Generate synthetic data in a checkerboard pattern.

    The mathematical model is:

    .. math::

        f(e, n) = a
        \sin\left(\frac{2\pi}{w_e} e\right)
        \cos\left(\frac{2\pi}{w_n} n\right)

    in which :math:`e` is the easting coordinate, :math:`n` is the northing
    coordinate, :math:`a` is the amplitude, and :math:`w_e` and :math:`w_n` are
    the wavelengths in the east and north directions, respectively.

    Parameters
    ----------
    amplitude : float
        The amplitude of the checkerboard undulations.
    region : tuple
        The boundaries (``[W, E, S, N]``) of the region used to generate the
        synthetic data.
    w_east : float
        The wavelength in the east direction. Defaults to half of the West-East
        size of the evaluating region.
    w_north : float
        The wavelength in the north direction. Defaults to half of the
        South-North size of the evaluating region.

    Examples
    --------

    >>> synth = CheckerBoard()
    >>> # Default values for the wavelengths are selected automatically
    >>> print(synth.w_east_, synth.w_north_)
    2500.0 2500.0
    >>> # CheckerBoard.grid produces an xarray.Dataset with data on a grid
    >>> grid = synth.grid(shape=(11, 6))
    >>> # scatter and profile generate pandas.DataFrame objects
    >>> table = synth.scatter()
    >>> print(sorted(table.columns))
    ['easting', 'northing', 'scalars']
    >>> profile = synth.profile(point1=(0, 0), point2=(2500, -2500), size=100)
    >>> print(sorted(profile.columns))
    ['distance', 'easting', 'northing', 'scalars']

    """

    def __init__(
        self, amplitude=1000, region=(0, 5000, -5000, 0), w_east=None, w_north=None
    ):
        super().__init__()
        self.amplitude = amplitude
        self.w_east = w_east
        self.w_north = w_north
        self.region = region

    @property
    def w_east_(self):
        "Use half of the E-W extent"
        if self.w_east is None:
            return (self.region[1] - self.region[0]) / 2
        return self.w_east

    @property
    def w_north_(self):
        "Use half of the N-S extent"
        if self.w_north is None:
            return (self.region[3] - self.region[2]) / 2
        return self.w_north

    @property
    def region_(self):
        "Used to fool the BaseGridder methods"
        check_region(self.region)
        return self.region

    def predict(self, coordinates):
        """
        Evaluate the checkerboard function on a given set of points.

        Parameters
        ----------
        coordinates : tuple of arrays
            Arrays with the coordinates of each data point. Should be in the
            following order: (easting, northing, vertical, ...). Only easting
            and northing will be used, all subsequent coordinates will be
            ignored.

        Returns
        -------
        data : array
            The evaluated checkerboard function.

        """
        easting, northing = coordinates[:2]
        data = (
            self.amplitude
            * np.sin((2 * np.pi / self.w_east_) * easting)
            * np.cos((2 * np.pi / self.w_north_) * northing)
        )
        return data

    def scatter(
        self,
        region=None,
        size=300,
        random_state=0,
        dims=None,
        data_names=None,
        projection=None,
        **kwargs,
    ):
        """
        Generate values on a random scatter of points.

        Point coordinates are generated by :func:`verde.scatter_points`. Other
        arguments for this function can be passed as extra keyword arguments
        (``kwargs``) to this method.

        By default, the region specified when creating the class instance will
        be used if ``region=None``.

        Use the *dims* and *data_names* arguments to set custom names for the
        dimensions and the data field(s) in the output
        :class:`pandas.DataFrame`. Default names are provided.

        Parameters
        ----------
        region : list = [W, E, S, N]
            The west, east, south, and north boundaries of a given region.
        size : int
            The number of points to generate.
        random_state : numpy.random.RandomState or an int seed
            A random number generator used to define the state of the random
            permutations. Use a fixed seed to make sure computations are
            reproducible. Use ``None`` to choose a seed automatically
            (resulting in different numbers with each run).
        dims : list or None
            The names of the northing and easting data dimensions,
            respectively, in the output dataframe. Default is determined from
            the ``dims`` attribute of the class. Must be defined in the
            following order: northing dimension, easting dimension.
            **NOTE: This is an exception to the "easting" then
            "northing" pattern but is required for compatibility with xarray.**
        data_names : str, list or None
            The name(s) of the data variables in the output dataframe. Defaults
            to ``'scalars'`` for scalar data,
            ``['east_component', 'north_component']`` for 2D vector data, and
            ``['east_component', 'north_component', 'vertical_component']`` for
            3D vector data.
        projection : callable or None
            If not None, then should be a callable object
            ``projection(easting, northing) -> (proj_easting, proj_northing)``
            that takes in easting and northing coordinate arrays and returns
            projected northing and easting coordinate arrays. This function
            will be used to project the generated scatter coordinates before
            passing them into ``predict``. For example, you can use this to
            generate a geographic scatter from a Cartesian gridder.

        Returns
        -------
        table : pandas.DataFrame
            The interpolated values on a random set of points.

        """
        dims = self._get_dims(dims)
        region = get_instance_region(self, region)
        coordinates = scatter_points(region, size, random_state=random_state, **kwargs)
        if projection is None:
            data = check_data(self.predict(coordinates))
        else:
            data = check_data(
                self.predict(project_coordinates(coordinates, projection))
            )
        data_names = self._get_data_names(data, data_names)
        columns = [(dims[0], coordinates[1]), (dims[1], coordinates[0])]
        extra_coords_names = self._get_extra_coords_names(coordinates)
        columns.extend(zip(extra_coords_names, coordinates[2:]))
        columns.extend(zip(data_names, data))
        return pd.DataFrame(dict(columns), columns=[c[0] for c in columns])
