"""Basic functionality."""

import contextlib
import functools
import inspect
import multiprocessing.pool
import re
import sys
import time
import traceback
from typing import (
    Any, Callable, Container, Dict, Hashable, Iterable, Iterator, List, Mapping, Optional, Set, Sequence, Type, Tuple,
    Union
)

import numpy as np

from .. import options


# define common types
Array = Any
RecArray = Any
Data = Dict[str, Array]
Options = Dict[str, Any]
Bounds = Tuple[Array, Array]

# define a pool managed by parallel and used by generate_items
pool = None


@contextlib.contextmanager
def parallel(processes: int) -> Iterator[None]:
    r"""Context manager used for parallel processing in a ``with`` statement context.

    This manager creates a context in which a pool of Python processes will be used by any method that requires
    market-by-market computation. These methods will distribute their work among the processes. After the context
    created by the ``with`` statement ends, all worker processes in the pool will be terminated. Outside of this
    context, such methods will not use multiprocessing.

    Importantly, multiprocessing will only improve speed if gains from parallelization outweigh overhead from
    serializing and passing data between processes. For example, if computation for a single market is very fast and
    there is a lot of data in each market that must be serialized and passed between processes, using multiprocessing
    may reduce overall speed.

    Arguments
    ---------
    processes : `int`
        Number of Python processes that will be created and used by any method that supports parallel processing.

    Examples
    --------
    .. raw:: latex

       \begin{examplenotebook}

    .. toctree::

       /_notebooks/api/parallel.ipynb

    .. raw:: latex

       \end{examplenotebook}

    """

    # validate the number of processes
    if not isinstance(processes, int):
        raise TypeError("processes must be an int.")
    if processes < 2:
        raise ValueError("processes must be at least 2.")

    # start the process pool, wait for work to be done, and then terminate it
    output(f"Starting a pool of {processes} processes ...")
    start_time = time.time()
    global pool
    try:
        with multiprocessing.pool.Pool(processes) as pool:
            output(f"Started the process pool after {format_seconds(time.time() - start_time)}.")
            yield
            output(f"Terminating the pool of {processes} processes ...")
            terminate_time = time.time()
    finally:
        pool = None
    output(f"Terminated the process pool after {format_seconds(time.time() - terminate_time)}.")


def generate_items(keys: Iterable, factory: Callable[[Any], tuple], method: Callable) -> Iterator:
    """Generate (key, method(*factory(key))) tuples for each key. The first element returned by factory is an instance
    of the class to which method is attached. If a process pool has been initialized, use multiprocessing; otherwise,
    use serial processing.
    """
    if pool is None:
        return (generate_items_worker((k, factory(k), method)) for k in keys)
    return pool.imap_unordered(generate_items_worker, ((k, factory(k), method) for k in keys))


def generate_items_worker(args: Tuple[Any, tuple, Callable]) -> Tuple[Any, Any]:
    """Call the the specified method of a class instance with any additional arguments. Return the associated key along
    with the returned object.
    """
    key, (instance, *method_args), method = args
    return key, method(instance, *method_args)


def structure_matrices(mapping: Mapping) -> RecArray:
    """Structure a mapping of keys to (array or None, type) tuples as a record array in which each sub-array is
    guaranteed to be at least two-dimensional.
    """

    # determine the number of rows in all matrices
    size = next(a.shape[0] for a, _ in mapping.values() if a is not None)

    # collect matrices and data types
    matrices: List[Array] = []
    dtypes: List[Tuple[Union[str, Tuple[Hashable, str]], Any, Tuple[int]]] = []
    for key, (array, dtype) in mapping.items():
        matrix = np.zeros((size, 0)) if array is None else np.c_[array]
        dtypes.append((key, dtype, (matrix.shape[1],)))
        matrices.append(matrix)

    # build the record array
    structured = np.recarray(size, dtypes)
    for dtype, matrix in zip(dtypes, matrices):
        structured[dtype[0] if isinstance(dtype[0], str) else dtype[0][1]] = matrix
    return structured


def update_matrices(matrices: RecArray, update_mapping: Dict) -> RecArray:
    """Update fields in a record array created by structure_matrices by re-structuring the matrices."""
    mapping = update_mapping.copy()
    for key in matrices.dtype.names:
        if key not in mapping:
            if len(matrices.dtype.fields[key]) > 2:
                mapping[(matrices.dtype.fields[key][2], key)] = (matrices[key], matrices[key].dtype)
            else:
                mapping[key] = (matrices[key], matrices[key].dtype)

    return structure_matrices(mapping)


def extract_matrix(structured_array_like: Mapping, key: Any) -> Optional[Array]:
    """Attempt to extract a field from a structured array-like object or horizontally stack field0, field1, and so on,
    into a full matrix. The extracted array will have at least two dimensions.
    """
    try:
        matrix = np.c_[structured_array_like[key]]
        return matrix if matrix.size > 0 else None
    except Exception:
        index = 0
        parts: List[Array] = []
        while True:
            try:
                part = np.c_[structured_array_like[f'{key}{index}']]
            except Exception:
                # output a warning if there's a 1 but no 0 (this is a common mistake)
                if index == 0:
                    try:
                        structured_array_like[f'{key}{index + 1}']
                    except Exception:
                        pass
                    else:
                        output("")
                        output(f"Warning: '{key}{index + 1}' was specified but not '{key}{index}'.")
                        output("")
                break
            index += 1
            if part.size > 0:
                parts.append(part)

        return np.hstack(parts) if parts else None


def extract_size(structured_array_like: Mapping) -> int:
    """Attempt to extract the number of rows from a structured array-like object."""
    size = 0
    getters = [
        lambda m: m.shape[0],
        lambda m: next(iter(structured_array_like.values())).shape[0],
        lambda m: len(next(iter(structured_array_like.values()))),
        lambda m: len(m)
    ]
    for get in getters:
        try:
            size = get(structured_array_like)
            break
        except Exception:
            pass
    if size > 0:
        return size
    raise TypeError(
        f"Failed to get the number of rows in the structured array-like object of type {type(structured_array_like)}. "
        f"Try using a dictionary, a NumPy structured array, a Pandas DataFrame, or any other standard type."
    )


def interact_ids(*columns: Array) -> Array:
    """Create interactions of ID columns."""
    interacted = columns[0].flatten().astype(np.object)
    if len(columns) > 1:
        interacted[:] = list(zip(*columns))
    return interacted


def output(message: Any) -> None:
    """Print a message if verbosity is turned on."""
    if options.verbose:
        if not callable(options.verbose_output):
            raise TypeError("options.verbose_output should be callable.")
        options.verbose_output(str(message))
        if options.flush_output:
            sys.stdout.flush()


def output_progress(iterable: Iterable, length: int, start_time: float) -> Iterator:
    """Yield results from an iterable while outputting progress updates at most every minute."""
    elapsed = time.time() - start_time
    next_minute = int(elapsed / 60) + 1
    for index, iterated in enumerate(iterable):
        yield iterated
        elapsed = time.time() - start_time
        if elapsed > 60 * next_minute:
            output(f"Finished {index + 1} out of {length} after {format_seconds(elapsed)}.")
            next_minute = int(elapsed / 60) + 1


def format_seconds(seconds: float) -> str:
    """Prepare a number of seconds to be displayed as a string."""
    hours, remainder = divmod(int(round(seconds)), 60**2)
    minutes, seconds = divmod(remainder, 60)
    return f'{hours:02}:{minutes:02}:{seconds:02}'


def format_number(number: Any) -> str:
    """Prepare a number to be displayed as a string."""
    if not isinstance(options.digits, int):
        raise TypeError("options.digits must be an int.")
    template = f"{{:^+{options.digits + 6}.{options.digits - 1}E}}"
    formatted = template.format(float(number))
    if "NAN" in formatted:
        formatted = formatted.replace("+", " ")
    return formatted


def format_se(se: Any) -> str:
    """Prepare a standard error to be displayed as a string."""
    formatted = format_number(se)
    for string in ["NAN", "-INF", "+INF"]:
        if string in formatted:
            return formatted.replace(string, f"({string})")

    return f"({formatted})"


def format_options(mapping: Options) -> str:
    """Prepare a mapping of options to be displayed as a string."""
    strings: List[str] = []
    for key, value in mapping.items():
        if callable(value):
            value = f'{value.__module__}.{value.__qualname__}'
        elif isinstance(value, float):
            value = format_number(value)
        strings.append(f'{key}: {value}')

    joined = ', '.join(strings)
    return f'{{{joined}}}'


def format_table(
        header: Sequence, *data: Sequence, title: Optional[str] = None, include_border: bool = True,
        include_header: bool = True, line_indices: Container[int] = ()) -> str:
    """Format table information as a string, which has fixed widths, vertical lines after any specified indices, and
    optionally a title, border, and header.
    """

    # construct the header rows
    row_index = -1
    header_rows: List[List[str]] = []
    header = [[c] if isinstance(c, str) else c for c in header]
    while True:
        header_row = ["" if len(c) < -row_index else c[row_index] for c in header]
        if not any(header_row):
            break
        header_rows.insert(0, header_row)
        row_index -= 1

    # construct the data rows
    data_rows = [[str(c) for c in r] + [""] * (len(header) - len(r)) for r in data]

    # compute column widths
    widths = []
    for column_index in range(len(header)):
        widths.append(max(len(r[column_index]) for r in header_rows + data_rows))

    # build the template
    template = "  " .join("{{:^{}}}{}".format(w, "  |" if i in line_indices else "") for i, w in enumerate(widths))

    # build the table
    lines = []
    if title is not None:
        lines.append(f"{title}:")
    if include_border:
        lines.append("=" * len(template.format(*[""] * len(widths))))
    if include_header:
        lines.extend([template.format(*r) for r in header_rows])
        lines.append(template.format(*("-" * w for w in widths)))
    lines.extend([template.format(*r) for r in data_rows])
    if include_border:
        lines.append("=" * len(template.format(*[""] * len(widths))))
    return "\n".join(lines)


def get_indices(ids: Array) -> Dict[Hashable, Array]:
    """From a one-dimensional array input, construct a dictionary with keys that are the unique values of the array
    and values that are the indices where the key appears in the array.
    """
    flat = ids.flatten()
    sort_indices = flat.argsort(kind='mergesort')
    sorted_ids = flat[sort_indices]
    changes = np.ones(flat.shape, np.bool)
    changes[1:] = sorted_ids[1:] != sorted_ids[:-1]
    reduce_indices = np.nonzero(changes)[0]
    return dict(zip(sorted_ids[reduce_indices], np.split(sort_indices, reduce_indices)[1:]))


def compute_finite_differences(f: Callable[[Array], Array], x: Array, epsilon_scale: float = 1.0) -> Array:
    """Approximate a matrix of derivatives with finite differences."""
    epsilon = epsilon_scale * options.finite_differences_epsilon

    columns = []
    for index in range(x.size):
        x1 = x.copy()
        x2 = x.copy()
        x1[index] += epsilon / 2
        x2[index] -= epsilon / 2
        columns.append((f(x1) - f(x2)) / epsilon)

    return np.column_stack(columns)


class SolverStats(object):
    """Structured statistics returned by a generic numerical solver."""

    converged: bool
    iterations: int
    evaluations: int

    def __init__(self, converged: bool = True, iterations: int = 0, evaluations: int = 0) -> None:
        """Structure the statistics."""
        self.converged = converged
        self.iterations = iterations
        self.evaluations = evaluations


class StringRepresentation(object):
    """Object that defers to its string representation."""

    def __repr__(self) -> str:
        """Defer to the string representation."""
        return str(self)


class Groups(object):
    """Computation of grouped statistics."""

    sort_indices: Array
    reduce_indices: Array
    unique: Array
    codes: Array
    counts: Array
    group_count: int

    def __init__(self, ids: Array) -> None:
        """Sort and index IDs that define groups."""

        # sort the IDs
        flat = ids.flatten()
        self.sort_indices = flat.argsort()
        sorted_ids = flat[self.sort_indices]

        # identify groups
        changes = np.ones(flat.shape, np.bool)
        changes[1:] = sorted_ids[1:] != sorted_ids[:-1]
        self.reduce_indices = np.nonzero(changes)[0]
        self.unique = sorted_ids[self.reduce_indices]

        # encode the groups
        sorted_codes = np.cumsum(changes) - 1
        self.codes = sorted_codes[self.sort_indices.argsort()]

        # compute counts
        self.group_count = self.reduce_indices.size
        self.counts = np.diff(np.append(self.reduce_indices, self.codes.size))

    def sum(self, matrix: Array) -> Array:
        """Compute the sum of each group."""
        return np.add.reduceat(matrix[self.sort_indices], self.reduce_indices)

    def mean(self, matrix: Array) -> Array:
        """Compute the mean of each group."""
        return self.sum(matrix) / self.counts[:, None]

    def expand(self, statistics: Array) -> Array:
        """Expand statistics for each group to the size of the original matrix."""
        return statistics[self.codes]


class Error(Exception):
    """Errors that are indistinguishable from others with the same message, which is parsed from the docstring."""

    stack: Optional[str]

    def __init__(self) -> None:
        """Optionally store the full current traceback for debugging purposes."""
        if options.verbose_tracebacks:
            self.stack = ''.join(traceback.format_stack())
        else:
            self.stack = None

    def __eq__(self, other: Any) -> bool:
        """Defer to hashes."""
        return hash(self) == hash(other)

    def __hash__(self) -> int:
        """Hash this instance such that in collections it is indistinguishable from others with the same message."""
        return hash((type(self).__name__, str(self)))

    def __repr__(self) -> str:
        """Defer to the string representation."""
        return str(self)

    def __str__(self) -> str:
        """Replace docstring markdown with simple text."""
        doc = inspect.getdoc(self)
        assert doc is not None

        # normalize LaTeX
        while True:
            match = re.search(r':math:`([^`]+)`', doc)
            if match is None:
                break
            start, end = match.span()
            doc = doc[:start] + re.sub(r'\s+', ' ', re.sub(r'[\\{}]', ' ', match.group(1))).lower() + doc[end:]

        # normalize references
        while True:
            match = re.search(r':ref:`[a-zA-Z0-9]+:([^`]+)`', doc)
            if match is None:
                break
            start, end = match.span()
            doc = doc[:start] + re.sub(r'<[^>]+>', '', match.group(1)) + doc[end:]

        # remove all remaining domains and compress whitespace
        doc = re.sub(r'[\s\n]+', ' ', re.sub(r':[a-z\-]+:|`', '', doc))

        # optionally add the full traceback
        if self.stack is not None:
            doc = f"{doc} Traceback:\n\n{self.stack}\n"
        return doc


class DerivedError(Error):
    """Error derived from another exception."""

    _exception: Exception

    def __init__(self, exception: Exception) -> None:
        """Store the exception from which this error is derived."""
        super().__init__()
        self._exception = exception

    def __str__(self) -> str:
        """Supplement the error with the exception's message."""
        return f"{super().__str__()} Exception encountered: '{self._exception}'."


class NumericalError(Error):
    """Floating point issues."""

    _messages: Set[str]

    def __init__(self) -> None:
        super().__init__()
        self._messages: Set[str] = set()

    def __str__(self) -> str:
        """Supplement the error with the messages."""
        combined = ", ".join(sorted(self._messages))
        return f"{super().__str__()} Errors encountered: {combined}."


class MultipleReversionError(Error):
    """Reversion of problematic elements."""

    _bad: int
    _total: int

    def __init__(self, bad_indices: Array) -> None:
        """Store element counts."""
        super().__init__()
        self._bad = bad_indices.sum()
        self._total = bad_indices.size

    def __str__(self) -> str:
        """Supplement the error with the counts."""
        return f"{super().__str__()} Number of reverted elements: {self._bad} out of {self._total}."


class InversionError(Error):
    """Problems with inverting a matrix."""

    _condition: float

    def __init__(self, matrix: Array) -> None:
        """Compute condition number of the matrix."""
        super().__init__()
        from .algebra import compute_condition_number
        self._condition = compute_condition_number(matrix)

    def __str__(self) -> str:
        """Supplement the error with the condition number."""
        return f"{super().__str__()} Condition number: {format_number(self._condition)}."


class InversionReplacementError(InversionError):
    """Problems with inverting a matrix led to the use of a replacement such as an approximation."""

    _replacement: str

    def __init__(self, matrix: Array, replacement: str) -> None:
        """Store the replacement description."""
        super().__init__(matrix)
        self._replacement = replacement

    def __str__(self) -> str:
        """Supplement the error with the description."""
        return f"{super().__str__()} The inverse was replaced with {self._replacement}."


class NumericalErrorHandler(object):
    """Decorator that appends errors to a function's returned list when numerical errors are encountered."""

    error: Type[NumericalError]

    def __init__(self, error: Type[NumericalError]) -> None:
        """Store the error class."""
        self.error = error

    def __call__(self, decorated: Callable) -> Callable:
        """Decorate the function."""
        @functools.wraps(decorated)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            """Configure NumPy to detect numerical errors."""
            detector = NumericalErrorDetector(self.error)
            with np.errstate(divide='call', over='call', under='ignore', invalid='call'):
                np.seterrcall(detector)
                returned = decorated(*args, **kwargs)
            if detector.detected is not None:
                returned[-1].append(detector.detected)
            return returned

        return wrapper


class NumericalErrorDetector(object):
    """Error detector to be passed to NumPy's error call function."""

    error: Type[NumericalError]
    detected: Optional[NumericalError]

    def __init__(self, error: Type[NumericalError]) -> None:
        """By default no error is detected."""
        self.error = error
        self.detected = None

    def __call__(self, message: str, _: int) -> None:
        """Initialize the error and store the error message."""
        if self.detected is None:
            self.detected = self.error()
        self.detected._messages.add(message)
