# The StringAccessor class defined below is an adaptation of the
# pandas string methods source code (see pd.core.strings)

# For reference, here is a copy of the pandas copyright notice:

# (c) 2011-2012, Lambda Foundry, Inc. and PyData Development Team
# All rights reserved.

# Copyright (c) 2008-2011 AQR Capital Management, LLC
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:

#     * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.

#     * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided
#        with the distribution.

#     * Neither the name of the copyright holder nor the names of any
#        contributors may be used to endorse or promote products derived
#        from this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import codecs
import re
import textwrap

import numpy as np

from .computation import apply_ufunc

_cpython_optimized_encoders = (
    "utf-8",
    "utf8",
    "latin-1",
    "latin1",
    "iso-8859-1",
    "mbcs",
    "ascii",
)
_cpython_optimized_decoders = _cpython_optimized_encoders + ("utf-16", "utf-32")


def _is_str_like(x):
    return isinstance(x, str) or isinstance(x, bytes)


class StringAccessor:
    """Vectorized string functions for string-like arrays.

    Similar to pandas, fields can be accessed through the `.str` attribute
    for applicable DataArrays.

        >>> da = xr.DataArray(['some', 'text', 'in', 'an', 'array'])
        >>> ds.str.len()
        <xarray.DataArray (dim_0: 5)>
        array([4, 4, 2, 2, 5])
        Dimensions without coordinates: dim_0

    """

    __slots__ = ("_obj",)

    def __init__(self, obj):
        self._obj = obj

    def _apply(self, f, dtype=None):
        # TODO handling of na values ?
        if dtype is None:
            dtype = self._obj.dtype

        g = np.vectorize(f, otypes=[dtype])
        return apply_ufunc(g, self._obj, dask="parallelized", output_dtypes=[dtype])

    def len(self):
        """
        Compute the length of each element in the array.

        Returns
        -------
        lengths array : array of int
        """
        return self._apply(len, dtype=int)

    def __getitem__(self, key):
        if isinstance(key, slice):
            return self.slice(start=key.start, stop=key.stop, step=key.step)
        else:
            return self.get(key)

    def get(self, i):
        """
        Extract element from indexable in each element in the array.

        Parameters
        ----------
        i : int
            Position of element to extract.
        default : optional
            Value for out-of-range index. If not specified (None) defaults to
            an empty string.

        Returns
        -------
        items : array of objects
        """
        obj = slice(-1, None) if i == -1 else slice(i, i + 1)
        return self._apply(lambda x: x[obj])

    def slice(self, start=None, stop=None, step=None):
        """
        Slice substrings from each element in the array.

        Parameters
        ----------
        start : int, optional
            Start position for slice operation.
        stop : int, optional
            Stop position for slice operation.
        step : int, optional
            Step size for slice operation.

        Returns
        -------
        sliced strings : same type as values
        """
        s = slice(start, stop, step)
        f = lambda x: x[s]
        return self._apply(f)

    def slice_replace(self, start=None, stop=None, repl=""):
        """
        Replace a positional slice of a string with another value.

        Parameters
        ----------
        start : int, optional
            Left index position to use for the slice. If not specified (None),
            the slice is unbounded on the left, i.e. slice from the start
            of the string.
        stop : int, optional
            Right index position to use for the slice. If not specified (None),
            the slice is unbounded on the right, i.e. slice until the
            end of the string.
        repl : str, optional
            String for replacement. If not specified, the sliced region
            is replaced with an empty string.

        Returns
        -------
        replaced : same type as values
        """
        repl = self._obj.dtype.type(repl)

        def f(x):
            if len(x[start:stop]) == 0:
                local_stop = start
            else:
                local_stop = stop
            y = self._obj.dtype.type("")
            if start is not None:
                y += x[:start]
            y += repl
            if stop is not None:
                y += x[local_stop:]
            return y

        return self._apply(f)

    def capitalize(self):
        """
        Convert strings in the array to be capitalized.

        Returns
        -------
        capitalized : same type as values
        """
        return self._apply(lambda x: x.capitalize())

    def lower(self):
        """
        Convert strings in the array to lowercase.

        Returns
        -------
        lowerd : same type as values
        """
        return self._apply(lambda x: x.lower())

    def swapcase(self):
        """
        Convert strings in the array to be swapcased.

        Returns
        -------
        swapcased : same type as values
        """
        return self._apply(lambda x: x.swapcase())

    def title(self):
        """
        Convert strings in the array to titlecase.

        Returns
        -------
        titled : same type as values
        """
        return self._apply(lambda x: x.title())

    def upper(self):
        """
        Convert strings in the array to uppercase.

        Returns
        -------
        uppered : same type as values
        """
        return self._apply(lambda x: x.upper())

    def isalnum(self):
        """
        Check whether all characters in each string are alphanumeric.

        Returns
        -------
        isalnum : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isalnum(), dtype=bool)

    def isalpha(self):
        """
        Check whether all characters in each string are alphabetic.

        Returns
        -------
        isalpha : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isalpha(), dtype=bool)

    def isdecimal(self):
        """
        Check whether all characters in each string are decimal.

        Returns
        -------
        isdecimal : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isdecimal(), dtype=bool)

    def isdigit(self):
        """
        Check whether all characters in each string are digits.

        Returns
        -------
        isdigit : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isdigit(), dtype=bool)

    def islower(self):
        """
        Check whether all characters in each string are lowercase.

        Returns
        -------
        islower : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.islower(), dtype=bool)

    def isnumeric(self):
        """
        Check whether all characters in each string are numeric.

        Returns
        -------
        isnumeric : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isnumeric(), dtype=bool)

    def isspace(self):
        """
        Check whether all characters in each string are spaces.

        Returns
        -------
        isspace : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isspace(), dtype=bool)

    def istitle(self):
        """
        Check whether all characters in each string are titlecase.

        Returns
        -------
        istitle : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.istitle(), dtype=bool)

    def isupper(self):
        """
        Check whether all characters in each string are uppercase.

        Returns
        -------
        isupper : array of bool
            Array of boolean values with the same shape as the original array.
        """
        return self._apply(lambda x: x.isupper(), dtype=bool)

    def count(self, pat, flags=0):
        """
        Count occurrences of pattern in each string of the array.

        This function is used to count the number of times a particular regex
        pattern is repeated in each of the string elements of the
        :class:`~xarray.DatArray`.

        Parameters
        ----------
        pat : str
            Valid regular expression.
        flags : int, default 0, meaning no flags
            Flags for the `re` module. For a complete list, `see here
            <https://docs.python.org/3/howto/regex.html#compilation-flags>`_.

        Returns
        -------
        counts : array of int
        """
        pat = self._obj.dtype.type(pat)
        regex = re.compile(pat, flags=flags)
        f = lambda x: len(regex.findall(x))
        return self._apply(f, dtype=int)

    def startswith(self, pat):
        """
        Test if the start of each string element matches a pattern.

        Parameters
        ----------
        pat : str
            Character sequence. Regular expressions are not accepted.

        Returns
        -------
        startswith : array of bool
            An array of booleans indicating whether the given pattern matches
            the start of each string element.
        """
        pat = self._obj.dtype.type(pat)
        f = lambda x: x.startswith(pat)
        return self._apply(f, dtype=bool)

    def endswith(self, pat):
        """
        Test if the end of each string element matches a pattern.

        Parameters
        ----------
        pat : str
            Character sequence. Regular expressions are not accepted.

        Returns
        -------
        endswith : array of bool
            A Series of booleans indicating whether the given pattern matches
            the end of each string element.
        """
        pat = self._obj.dtype.type(pat)
        f = lambda x: x.endswith(pat)
        return self._apply(f, dtype=bool)

    def pad(self, width, side="left", fillchar=" "):
        """
        Pad strings in the array up to width.

        Parameters
        ----------
        width : int
            Minimum width of resulting string; additional characters will be
            filled with character defined in `fillchar`.
        side : {'left', 'right', 'both'}, default 'left'
            Side from which to fill resulting string.
        fillchar : str, default ' '
            Additional character for filling, default is whitespace.

        Returns
        -------
        filled : same type as values
            Array with a minimum number of char in each element.
        """
        width = int(width)
        fillchar = self._obj.dtype.type(fillchar)
        if len(fillchar) != 1:
            raise TypeError("fillchar must be a character, not str")

        if side == "left":
            f = lambda s: s.rjust(width, fillchar)
        elif side == "right":
            f = lambda s: s.ljust(width, fillchar)
        elif side == "both":
            f = lambda s: s.center(width, fillchar)
        else:  # pragma: no cover
            raise ValueError("Invalid side")

        return self._apply(f)

    def center(self, width, fillchar=" "):
        """
        Filling left and right side of strings in the array with an
        additional character.

        Parameters
        ----------
        width : int
            Minimum width of resulting string; additional characters will be
            filled with ``fillchar``
        fillchar : str
            Additional character for filling, default is whitespace

        Returns
        -------
        filled : same type as values
        """
        return self.pad(width, side="both", fillchar=fillchar)

    def ljust(self, width, fillchar=" "):
        """
        Filling right side of strings in the array with an additional
        character.

        Parameters
        ----------
        width : int
            Minimum width of resulting string; additional characters will be
            filled with ``fillchar``
        fillchar : str
            Additional character for filling, default is whitespace

        Returns
        -------
        filled : same type as values
        """
        return self.pad(width, side="right", fillchar=fillchar)

    def rjust(self, width, fillchar=" "):
        """
        Filling left side of strings in the array with an additional character.

        Parameters
        ----------
        width : int
            Minimum width of resulting string; additional characters will be
            filled with ``fillchar``
        fillchar : str
            Additional character for filling, default is whitespace

        Returns
        -------
        filled : same type as values
        """
        return self.pad(width, side="left", fillchar=fillchar)

    def zfill(self, width):
        """
        Pad strings in the array by prepending '0' characters.

        Strings in the array are padded with '0' characters on the
        left of the string to reach a total string length  `width`. Strings
        in the array  with length greater or equal to `width` are unchanged.

        Parameters
        ----------
        width : int
            Minimum length of resulting string; strings with length less
            than `width` be prepended with '0' characters.

        Returns
        -------
        filled : same type as values
        """
        return self.pad(width, side="left", fillchar="0")

    def contains(self, pat, case=True, flags=0, regex=True):
        """
        Test if pattern or regex is contained within a string of the array.

        Return boolean array based on whether a given pattern or regex is
        contained within a string of the array.

        Parameters
        ----------
        pat : str
            Character sequence or regular expression.
        case : bool, default True
            If True, case sensitive.
        flags : int, default 0 (no flags)
            Flags to pass through to the re module, e.g. re.IGNORECASE.
        regex : bool, default True
            If True, assumes the pat is a regular expression.
            If False, treats the pat as a literal string.

        Returns
        -------
        contains : array of bool
            An array of boolean values indicating whether the
            given pattern is contained within the string of each element
            of the array.
        """
        pat = self._obj.dtype.type(pat)
        if regex:
            if not case:
                flags |= re.IGNORECASE

            regex = re.compile(pat, flags=flags)

            if regex.groups > 0:  # pragma: no cover
                raise ValueError("This pattern has match groups.")

            f = lambda x: bool(regex.search(x))
        else:
            if case:
                f = lambda x: pat in x
            else:
                uppered = self._obj.str.upper()
                return uppered.str.contains(pat.upper(), regex=False)

        return self._apply(f, dtype=bool)

    def match(self, pat, case=True, flags=0):
        """
        Determine if each string matches a regular expression.

        Parameters
        ----------
        pat : string
            Character sequence or regular expression
        case : boolean, default True
            If True, case sensitive
        flags : int, default 0 (no flags)
            re module flags, e.g. re.IGNORECASE

        Returns
        -------
        matched : array of bool
        """
        if not case:
            flags |= re.IGNORECASE

        pat = self._obj.dtype.type(pat)
        regex = re.compile(pat, flags=flags)
        f = lambda x: bool(regex.match(x))
        return self._apply(f, dtype=bool)

    def strip(self, to_strip=None, side="both"):
        """
        Remove leading and trailing characters.

        Strip whitespaces (including newlines) or a set of specified characters
        from each string in the array from left and/or right sides.

        Parameters
        ----------
        to_strip : str or None, default None
            Specifying the set of characters to be removed.
            All combinations of this set of characters will be stripped.
            If None then whitespaces are removed.
        side : {'left', 'right', 'both'}, default 'left'
            Side from which to strip.

        Returns
        -------
        stripped : same type as values
        """
        if to_strip is not None:
            to_strip = self._obj.dtype.type(to_strip)

        if side == "both":
            f = lambda x: x.strip(to_strip)
        elif side == "left":
            f = lambda x: x.lstrip(to_strip)
        elif side == "right":
            f = lambda x: x.rstrip(to_strip)
        else:  # pragma: no cover
            raise ValueError("Invalid side")

        return self._apply(f)

    def lstrip(self, to_strip=None):
        """
        Remove leading and trailing characters.

        Strip whitespaces (including newlines) or a set of specified characters
        from each string in the array from the left side.

        Parameters
        ----------
        to_strip : str or None, default None
            Specifying the set of characters to be removed.
            All combinations of this set of characters will be stripped.
            If None then whitespaces are removed.

        Returns
        -------
        stripped : same type as values
        """
        return self.strip(to_strip, side="left")

    def rstrip(self, to_strip=None):
        """
        Remove leading and trailing characters.

        Strip whitespaces (including newlines) or a set of specified characters
        from each string in the array from the right side.

        Parameters
        ----------
        to_strip : str or None, default None
            Specifying the set of characters to be removed.
            All combinations of this set of characters will be stripped.
            If None then whitespaces are removed.

        Returns
        -------
        stripped : same type as values
        """
        return self.strip(to_strip, side="right")

    def wrap(self, width, **kwargs):
        """
        Wrap long strings in the array to be formatted in paragraphs with
        length less than a given width.

        This method has the same keyword parameters and defaults as
        :class:`textwrap.TextWrapper`.

        Parameters
        ----------
        width : int
            Maximum line-width
        expand_tabs : bool, optional
            If true, tab characters will be expanded to spaces (default: True)
        replace_whitespace : bool, optional
            If true, each whitespace character (as defined by
            string.whitespace) remaining after tab expansion will be replaced
            by a single space (default: True)
        drop_whitespace : bool, optional
            If true, whitespace that, after wrapping, happens to end up at the
            beginning or end of a line is dropped (default: True)
        break_long_words : bool, optional
            If true, then words longer than width will be broken in order to
            ensure that no lines are longer than width. If it is false, long
            words will not be broken, and some lines may be longer than width.
            (default: True)
        break_on_hyphens : bool, optional
            If true, wrapping will occur preferably on whitespace and right
            after hyphens in compound words, as it is customary in English. If
            false, only whitespaces will be considered as potentially good
            places for line breaks, but you need to set break_long_words to
            false if you want truly insecable words. (default: True)

        Returns
        -------
        wrapped : same type as values
        """
        tw = textwrap.TextWrapper(width=width)
        f = lambda x: "\n".join(tw.wrap(x))
        return self._apply(f)

    def translate(self, table):
        """
        Map all characters in the string through the given mapping table.

        Parameters
        ----------
        table : dict
            A a mapping of Unicode ordinals to Unicode ordinals, strings,
            or None. Unmapped characters are left untouched. Characters mapped
            to None are deleted. :meth:`str.maketrans` is a helper function for
            making translation tables.

        Returns
        -------
        translated : same type as values
        """
        f = lambda x: x.translate(table)
        return self._apply(f)

    def repeat(self, repeats):
        """
        Duplicate each string in the array.

        Parameters
        ----------
        repeats : int
            Number of repetitions.

        Returns
        -------
        repeated : same type as values
            Array of repeated string objects.
        """
        f = lambda x: repeats * x
        return self._apply(f)

    def find(self, sub, start=0, end=None, side="left"):
        """
        Return lowest or highest indexes in each strings in the array
        where the substring is fully contained between [start:end].
        Return -1 on failure.

        Parameters
        ----------
        sub : str
            Substring being searched
        start : int
            Left edge index
        end : int
            Right edge index
        side : {'left', 'right'}, default 'left'
            Starting side for search.

        Returns
        -------
        found : array of integer values
        """
        sub = self._obj.dtype.type(sub)

        if side == "left":
            method = "find"
        elif side == "right":
            method = "rfind"
        else:  # pragma: no cover
            raise ValueError("Invalid side")

        if end is None:
            f = lambda x: getattr(x, method)(sub, start)
        else:
            f = lambda x: getattr(x, method)(sub, start, end)

        return self._apply(f, dtype=int)

    def rfind(self, sub, start=0, end=None):
        """
        Return highest indexes in each strings in the array
        where the substring is fully contained between [start:end].
        Return -1 on failure.

        Parameters
        ----------
        sub : str
            Substring being searched
        start : int
            Left edge index
        end : int
            Right edge index

        Returns
        -------
        found : array of integer values
        """
        return self.find(sub, start=start, end=end, side="right")

    def index(self, sub, start=0, end=None, side="left"):
        """
        Return lowest or highest indexes in each strings where the substring is
        fully contained between [start:end]. This is the same as
        ``str.find`` except instead of returning -1, it raises a ValueError
        when the substring is not found.

        Parameters
        ----------
        sub : str
            Substring being searched
        start : int
            Left edge index
        end : int
            Right edge index
        side : {'left', 'right'}, default 'left'
            Starting side for search.

        Returns
        -------
        found : array of integer values
        """
        sub = self._obj.dtype.type(sub)

        if side == "left":
            method = "index"
        elif side == "right":
            method = "rindex"
        else:  # pragma: no cover
            raise ValueError("Invalid side")

        if end is None:
            f = lambda x: getattr(x, method)(sub, start)
        else:
            f = lambda x: getattr(x, method)(sub, start, end)

        return self._apply(f, dtype=int)

    def rindex(self, sub, start=0, end=None):
        """
        Return highest indexes in each strings where the substring is
        fully contained between [start:end]. This is the same as
        ``str.rfind`` except instead of returning -1, it raises a ValueError
        when the substring is not found.

        Parameters
        ----------
        sub : str
            Substring being searched
        start : int
            Left edge index
        end : int
            Right edge index

        Returns
        -------
        found : array of integer values
        """
        return self.index(sub, start=start, end=end, side="right")

    def replace(self, pat, repl, n=-1, case=None, flags=0, regex=True):
        """
        Replace occurrences of pattern/regex in the array with some string.

        Parameters
        ----------
        pat : string or compiled regex
            String can be a character sequence or regular expression.

        repl : string or callable
            Replacement string or a callable. The callable is passed the regex
            match object and must return a replacement string to be used.
            See :func:`re.sub`.

        n : int, default -1 (all)
            Number of replacements to make from start
        case : boolean, default None
            - If True, case sensitive (the default if `pat` is a string)
            - Set to False for case insensitive
            - Cannot be set if `pat` is a compiled regex
        flags : int, default 0 (no flags)
            - re module flags, e.g. re.IGNORECASE
            - Cannot be set if `pat` is a compiled regex
        regex : boolean, default True
            - If True, assumes the passed-in pattern is a regular expression.
            - If False, treats the pattern as a literal string
            - Cannot be set to False if `pat` is a compiled regex or `repl` is
            a callable.

        Returns
        -------
        replaced : same type as values
            A copy of the object with all matching occurrences of `pat`
            replaced by `repl`.
        """
        if not (_is_str_like(repl) or callable(repl)):  # pragma: no cover
            raise TypeError("repl must be a string or callable")

        if _is_str_like(pat):
            pat = self._obj.dtype.type(pat)

        if _is_str_like(repl):
            repl = self._obj.dtype.type(repl)

        is_compiled_re = isinstance(pat, type(re.compile("")))
        if regex:
            if is_compiled_re:
                if (case is not None) or (flags != 0):
                    raise ValueError(
                        "case and flags cannot be set" " when pat is a compiled regex"
                    )
            else:
                # not a compiled regex
                # set default case
                if case is None:
                    case = True

                # add case flag, if provided
                if case is False:
                    flags |= re.IGNORECASE
            if is_compiled_re or len(pat) > 1 or flags or callable(repl):
                n = n if n >= 0 else 0
                compiled = re.compile(pat, flags=flags)
                f = lambda x: compiled.sub(repl=repl, string=x, count=n)
            else:
                f = lambda x: x.replace(pat, repl, n)
        else:
            if is_compiled_re:
                raise ValueError(
                    "Cannot use a compiled regex as replacement "
                    "pattern with regex=False"
                )
            if callable(repl):
                raise ValueError(
                    "Cannot use a callable replacement when " "regex=False"
                )
            f = lambda x: x.replace(pat, repl, n)
        return self._apply(f)

    def decode(self, encoding, errors="strict"):
        """
        Decode character string in the array using indicated encoding.

        Parameters
        ----------
        encoding : str
        errors : str, optional

        Returns
        -------
        decoded : same type as values
        """
        if encoding in _cpython_optimized_decoders:
            f = lambda x: x.decode(encoding, errors)
        else:
            decoder = codecs.getdecoder(encoding)
            f = lambda x: decoder(x, errors)[0]
        return self._apply(f, dtype=np.str_)

    def encode(self, encoding, errors="strict"):
        """
        Encode character string in the array using indicated encoding.

        Parameters
        ----------
        encoding : str
        errors : str, optional

        Returns
        -------
        encoded : same type as values
        """
        if encoding in _cpython_optimized_encoders:
            f = lambda x: x.encode(encoding, errors)
        else:
            encoder = codecs.getencoder(encoding)
            f = lambda x: encoder(x, errors)[0]
        return self._apply(f, dtype=np.bytes_)
