# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
#  OpenModes - An eigenmode solver for open electromagnetic resonantors
#  Copyright (C) 2013-2015 David Powell
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
# -----------------------------------------------------------------------------
"Classes for arrays where one or more axes can be indexed by Part objects"

from __future__ import division

import numpy as np
from openmodes.parts import Part
import numbers
import collections
import six
import warnings


def part_ranges_lowest(parent_part, basis_container):
    """Construct the slice objects for the parent part, iterating down only
    to the specified lowest level"""

    lowest = getattr(basis_container, 'lowest_parts', None)

    # get the size of each child part
    sizes = [len(basis_container[part]) for part in parent_part.iter_lowest(lowest)]

    offsets = np.cumsum([0]+sizes)

    ranges = {}

    # index of single parts to get the offsets
    lowest_part_num = 0

    # iterate with parents last, so they can get data from their child objects
    for part in parent_part.iter_lowest(lowest, parent_order='after'):
        if part in lowest:
            ranges[part] = slice(offsets[lowest_part_num], offsets[lowest_part_num+1])
            lowest_part_num += 1
        else:
            # take slice information from first and last child
            start = ranges[part.children[0]].start
            stop = ranges[part.children[-1]].stop
            ranges[part] = slice(start, stop)

    return ranges


def part_ranges(parent_part, basis_container):
    "Construct the slice objects for the parent part and all of its children"
    if hasattr(basis_container, 'lowest_parts'):
        return part_ranges_lowest(parent_part, basis_container)

    # get the size of each child part
    sizes = [len(basis_container[part]) for part in parent_part.iter_single()]
    offsets = np.cumsum([0]+sizes)

    ranges = {}

    # index of single parts to get the offsets
    single_part_num = 0

    # iterate with parents last, so they can get data from their child objects
    for part in parent_part.iter_all(parent_first=False):
        if hasattr(part, 'children'):
            # take slice information from first and last child
            start = ranges[part.children[0]].start
            stop = ranges[part.children[-1]].stop
            ranges[part] = slice(start, stop)
        else:
            ranges[part] = slice(offsets[single_part_num], offsets[single_part_num+1])
            single_part_num += 1

    return ranges


def build_lookup(index_data):
    "Create the lookup table for a LookupArray"
    lookup = []
    shape = []
    for index in index_data:
        if isinstance(index, collections.Iterable) and isinstance(index[0], Part):
            # a hierarchy of parts
            basis_container = index[1]
            part = index[0]
            ranges = part_ranges(part, basis_container)
            lookup.append((ranges, basis_container, part))
            shape.append(ranges[part].stop)
        elif isinstance(index, numbers.Integral):
            # an integer for a specific length
            lookup.append(None)
            shape.append(index)
        else:
            # A tuple of strings representing quantities
            this_lookup = {}
            for (x, y) in enumerate(index):
                if not isinstance(y, six.string_types):
                    raise ValueError("Unknown index type %s" % str(y))
                this_lookup[y] = x

            lookup.append(this_lookup)
            shape.append(len(index))

    return lookup, shape


def compatible_quantity_part(first, second):
    "Determine if two parts of lookup lists are compatible quantities and parts"

    a = first[0]
    b = second[0]
    if not (isinstance(a, dict) and isinstance(b, dict)):
        return False
    if set(a.values()) != set(b.values()):
        return False

    a = first[1]
    b = second[1]
    if not (isinstance(a[0], dict) and isinstance(b[0], dict)):
        return False

    return a[:2] == b[:2]


class LookupArray(np.ndarray):
    """
    A subclass of a numpy array, where for certain dimensions, Part objects
    or strings can be used to index array elements.

    For explanation of subclassing numpy arrays, see
    http://docs.scipy.org/doc/numpy/user/basics.subclassing.html

    The following functionality of numpy arrays may cause problems, so they
    should be used with "extreme caution":
        - transpose
        - adding new axes by indexing with np.newaxis/None
        - flattening
        - anything other than C ordering
        - Functions which reduce dimensions
        - Indexing with ...
    """

    def __new__(subtype, index_data=None, lookup=None, shape=None, dtype=float):
        """Construct an empty vector which can be indexed by parts

        Parameters
        ----------
        index_data : tuple, optional
            Tuple elements can be integer, for a fixed length,
            a tuple (Part, BasisContainer), for hierarchical indexing by Parts,
            or a tuple of strings, for quantities
        lookup, shape: tuple, optional
            Instead of providing index_data, these elements can be provided
            directly if the lookup table is known in advance
        dtype : dtype, optional
            The numpy data type of the vector
        """

        if lookup is None or shape is None:
            lookup, shape = build_lookup(index_data)
        obj = np.ndarray.__new__(subtype, shape, dtype)
        obj.lookup = lookup

        # Finally, we must return the newly created object:
        return obj

    def __array_finalize__(self, obj):
        "Function is called when creating array from view as well"
        if obj is None:
            return

        # set default values for the custom attributes
        self.lookup = getattr(obj, 'lookup', None)

    def __setstate__(self, state):
        """Allow additional attributes of this array type to be unpickled

        Note that some metadata may be lost when unpickling."""
        base_state, extended_state = state
        super(LookupArray, self).__setstate__(base_state)
        self.lookup, = extended_state

    def __reduce__(self):
        """Allow additional attributes of this array type to be pickled

        Note that some metadata may be lost when unpickling."""
        base_reduce = list(super(LookupArray, self).__reduce__(self))
        full_state = (base_reduce[2], (self.lookup,))
        base_reduce[2] = full_state
        return tuple(base_reduce)

    # Under python 3.x, these members will not be called. However, they should
    # not cause any trouble.
    def __getslice__(self, start, stop):
        "Needed due to CPython bug"
        return self.__getitem__(slice(start, stop))

    def __setslice__(self, start, stop, val):
        "Needed due to CPython bug"
        self.__setitem__(slice(start, stop), val)

    def __getitem__(self, idx):
        """Gets an item or items from the array. Any of the indices may be the
        name of a range, in addition to all the usual fancy indexing options"""

        if not isinstance(idx, tuple):
            # force a single index to be a tuple
            idx = idx,

        new_idx = []
        sub_lookup = []
        entry_num = 0

        # try to lookup every part of the index to convert to a range
        for entry in idx:
            if isinstance(entry, Part):
                # Need to pass this metadata to the sub-array for its
                # lookup table
                this_lookup, container, parent_part = self.lookup[entry_num]
                sub_lookup.append((part_ranges(entry, container), container, entry))
                new_idx.append(this_lookup[entry])
            elif isinstance(entry, six.string_types):
                # If a string has been passed, then this dimension will have
                # been flattened out, so no metadata is needed
                this_lookup = self.lookup[entry_num]
                new_idx.append(this_lookup[entry])
            else:
                new_idx.append(entry)

                if not isinstance(entry, numbers.Integral):
                    # Integers mean a dimension is dropped, in all other
                    # cases it is kept
                    if entry is None:
                        # Need to record that a new dimension is added, so
                        # keep the place in the lookup of the original
                        sub_lookup.append(None)
                        entry_num -= 1
                    elif isinstance(entry, slice) and entry == slice(None):
                        # If slicing the whole dimension, metadata can be kept
                        sub_lookup.append(self.lookup[entry_num])
                    elif isinstance(entry, collections.Iterable):
                        # TODO: find a better solution to avoid this probelm
                        # warnings.warn("Indexing LookupArray with iterable is unreliable")
                        pass
                    else:
                        # In all other cases metadata is lost
                        sub_lookup.append(None)

            entry_num += 1

        # now add lookup data for all the non-indexed dimensions
        sub_lookup = sub_lookup+self.lookup[entry_num:]

        try:
            result = super(LookupArray, self).__getitem__(tuple(new_idx))
        except IndexError as exc:
            message = "Invalid index %s" % str(idx)
            exc.args = (message,)+tuple(str(n) for n in exc.args[1:])
            raise

        # May get a LookupArray or an array scalar back
        if isinstance(result, LookupArray):
            result.lookup = sub_lookup

        return result

    def __setitem__(self, idx, value):
        """Gets an item or items in the array. Any of the indices may be the
        name of a range, in addition to all the usual fancy indexing options"""
        if not isinstance(idx, tuple):
            # force a single index to be a tuple
            idx = idx,

        new_idx = []

        # try to lookup every part of the index to convert to a range
        for entry_num, entry in enumerate(idx):
            if isinstance(entry, Part):
                this_lookup, container, parent_part = self.lookup[entry_num]
                new_idx.append(this_lookup[entry])
            elif isinstance(entry, six.string_types):
                this_lookup = self.lookup[entry_num]
                new_idx.append(this_lookup[entry])
            else:
                new_idx.append(entry)

        try:
            super(LookupArray, self).__setitem__(tuple(new_idx), value)
        except IndexError as exc:
            message = "Invalid index %s" % idx
            exc.args = (message,)+tuple(str(n) for n in exc.args[1:])
            raise

    def transpose(self, **args):
        raise NotImplementedError

    @property
    def T(self):
        result = super(LookupArray, self).T
        assert(type(result) == LookupArray)
        result.lookup = list(reversed(self.lookup))
        return result

    def simple_view(self):
        """Return a view where quantity dimensions (with string keys) are
        collapsed into the subsequent dimension. View is of type ndarray."""
        new_shape = []

        for dim_n, lu_n in zip(reversed(self.shape), reversed(self.lookup)):
            if type(lu_n) == dict and type(list(lu_n.keys())[0]) == str:
                new_shape[-1] *= dim_n
            else:
                new_shape.append(dim_n)
        new_shape.reverse()
        return self.reshape(new_shape).view(np.ndarray)

    def dot(self, other):
        """Matrix/vector multiplication with another LookupArray"""
        if not isinstance(other, LookupArray):
            assert(self.shape[-1] == other.shape[0])
            new_lookup = self.lookup[:-1]+[None,]*(other.ndim-1)
            new_shape = self.shape[:-1]+other.shape[1:]
        elif compatible_quantity_part(self.lookup[-2:], other.lookup[:2]):
            new_lookup = self.lookup[:-2]+other.lookup[2:]
            new_shape = self.shape[:-2]+other.shape[2:]
            other = other.simple_view()
        else:
            raise NotImplementedError

        if len(new_shape) == 0:
            # handle the case of a scalar result
            new_array = np.dot(self.simple_view(), other)
        else:
            new_array = LookupArray(lookup=new_lookup, shape=new_shape,
                                    dtype=np.promote_types(self.dtype, other.dtype))
            new_array.simple_view()[:] = np.dot(self.simple_view(), other)
        return new_array

    def vdot(self, other):
        "Conjugated dot product"
        return self.conj().dot(other)


def view_lookuparray(original, index_data=None, lookup=None, shape=None):
    """Convert an array to a LookupArray, where possible avoiding copying"""
    if lookup is None:
        lookup, shape = build_lookup(index_data)
    result = original.reshape(shape).view(LookupArray)
    result.lookup = lookup
    return result


def loop_star_indices(x):
    """Return the indices into the array corresponding to the loop and star
    parts. The array must have been constructed using loop/star basis
    functions.

    Parameters
    ----------
    x: LookupArray
        Must have either 2 dimensions, both of which must be indexable
        by Parts

    Returns
    -------
    indices_loop, indices_star: list of ndarray
        For each dimension n, indices_loop[n] is an array indexing the loop
        part, and indices_star[n] is an array indexing the star part.
    """

    indices_loop = []
    indices_star = []

    for lookup_num, lookup in enumerate(x.lookup):
        if isinstance(list(lookup[0].keys())[0], Part):
            # This index is a lookup for Parts, so find all the SingleParts
            # along this index and add the relevant ranges to the indexing
            # array

            loop_list = []
            star_list = []
            # First find the parent part, the one covering the largest range
            part_list = list(lookup[0].keys())
            parent_part = part_list[np.argmax(lookup[n].stop-lookup[n].start
                                              for n in part_list)]

            # now iterate over all SingleParts of this parent part
            for part in parent_part.iter_single():
                part_start = lookup[0][part].start

                bf = lookup[1][part]
                loop_range = bf.loop_range
                loop_list.append(np.arange(loop_range.start+part_start,
                                           loop_range.stop+part_start))

                star_range = bf.star_range
                star_list.append(np.arange(star_range.start+part_start,
                                           star_range.stop+part_start))

            # If this is not the last axis, then add the necessary number of
            # extra dimensions to each array so that they will be broadcast
            # correctly when the caller goes to use them
            new_shape = (-1,)+(1,)*(x.ndim-lookup_num-1)
            loop_array = np.hstack(loop_list).reshape(new_shape)
            star_array = np.hstack(star_list).reshape(new_shape)

            indices_loop.append(loop_array)
            indices_star.append(star_array)
        else:
            # This index is not for parts, so just take the whole axis
            indices_loop.append(slice(None))
            indices_star.append(slice(None))

    return indices_loop, indices_star
