'''Home of the CrunchCube class.

This module contains the definition of the CrunchCube class. It represents
the open-source library used for manipulating the crunch cubes (JSON responses
from the Crunch.io platform).
'''
from __future__ import division

import json

import numpy as np
from scipy.stats import norm

from .dimension import Dimension


class CrunchCube(object):
    '''Implementation of the CrunchCube API class.

    Class is used for the implementation of the main API functions that are
    needed for seamless integration with the crunch cube responses
    (from Crunch.io platform).

    Main API functions are:
      - as_array
      - margin
      - proportions
      - percentages

    These functions are used to retrieve statistical information of interest,
    from the JSON like crunch cubes. Complete usage of each API function is
    described within the appropriate docstring.

    Crunch Cubes contain richer metadata than standart Python objects, and
    they also conceal certain complexity in the data structures from the user.
    In particular, Multiple Response variables are generally represented as
    single dimensions in result tables, but in the actual data, they may
    comprise of two dimensions. These methods (API) understand the subtleties
    in the Crunch data types, and correctly compute margins and
    percentages off of them.
    '''

    def __init__(self, response):
        '''Initializes the CrunchCube class with the cube JSON response.

        Class can be initialized with both JSON string, and dict types.
        There's no real parsing of the data at the point of initialization,
        just storing. The functionality is implemented in the
        appropriate API functions.
        '''
        # @cube: Represents the cube response object, as generated by
        # the zz9d cube factory. Cube can come in as a JSON or as a dictionary,
        # So we need to check its type, and convert it to a dictionary if
        # it's JSON, if possible.

        # If the provided response is dict, create cube immediately
        if isinstance(response, dict):
            self._cube = response.get('value', response)
            return

        try:
            response = json.loads(response)
            self._cube = response.get('value', response)
        except TypeError:
            # If an unexpected type is provided raise descriptive exception.
            if not isinstance(response, dict):
                raise TypeError((
                    'Unsupported type provided: {}. '
                    'A `cube` must be JSON or `dict`.'
                ).format(type(response)))

    @classmethod
    def _get_dimensions(cls, cube):
        '''Gets the dimensions of the crunch cube.

        This function is internal, and is not mean to be used by ouside users
        of the CrunchCube class. The main reason for this is the internal
        representation of the different variable types (namely the MR and the
        CA). These types have two dimensions each, but in the case of MR, the
        second dimensions shouldn't be visible to the user. This function
        returns such dimensions as well, since they're necessary for the
        correct implementation of the functionality for the MR type.
        The version that is mentioned to be used by users is the
        property 'dimensions'.
        '''
        entries = cube['result']['dimensions']
        return [
            (
                # Multiple Response and Categorical Array variables have
                # two subsequent dimensions (elements and selections). For
                # this reason it's necessary to pass in both of them in the
                # Dimension class init method. This is needed in order to
                # determine the correct type (CA or MR). We only skip the
                # two-argument constructor for the last dimension in the list
                # (where it's not possible to fetch the subsequent one).
                Dimension(entry)
                if i + 1 >= len(entries)
                else Dimension(entry, entries[i + 1])
            )
            for (i, entry) in enumerate(entries)
        ]

    @classmethod
    def _get_mr_selections_indices(cls, dimensions):
        '''Gets indices of each 'selection' dim, for corresponding MR dim.

        Multiple Response (MR) and Categorical Array (CA) variables are
        represented by two dimensions each. These dimensions can be thought of
        as 'elements' and 'selections'. This function returns the indices of
        the 'selections' dimension for each MR variable.
        '''
        mr_dimensions_indices = [
            i for (i, dim) in enumerate(dimensions)
            if (i + 1 < len(dimensions) and
                dim.type == 'multiple_response')
        ]
        # For each MR and CA dimension, the 'selections' dimension
        # follows right after it (in the originating cube).
        # Here we increase the MR index by 1, which gives us
        # the index of the corresponding 'selections' dimension.
        return [i + 1 for i in mr_dimensions_indices]

    @classmethod
    def _get_valid_indices(cls, dimensions, include_missing, get_non_selected):
        '''Gets valid indices for each dimension.

        Main criterion for a valid index is most often the information about
        whether the corresponding value of the dimension is missing or not.
        For MR variables, since they use two dimensions, the valid index
        for the 'selections' dimensions is [0], except in the case of
        non-selected slices calculation, where it needs to be [1].
        '''
        valid_indices = [dim.valid_indices(include_missing)
                         for dim in dimensions]

        mr_selections_indices = cls._get_mr_selections_indices(dimensions)
        mr_slice = [1] if get_non_selected else [0]
        if mr_selections_indices:
            # In the case of MR variables, we only need to select the
            # 'selected' slice of the 'selections' dimension.
            valid_indices = [
                (
                    valid_indices[i]
                    if i not in mr_selections_indices
                    else mr_slice
                )
                for (i, _) in enumerate(valid_indices)
            ]

        return valid_indices

    @classmethod
    def _fix_shape(cls, array):
        '''Fixes shape of MR variables.
        For MR variables, where 'selections' dims are dropped, the ndarray
        needs to be reshaped, in order to seem as if those dims never existed.
        '''

        # If a first dimension only has one element, we don't want to
        # remove it. Hence the i == 0 part.
        new_shape = [dim for (i, dim) in enumerate(array.shape)
                     if dim != 1 or i == 0]
        return array.reshape(new_shape)

    @property
    def _has_means(self):
        return self._cube['result']['measures'].get('mean', None) is not None

    def _get_values(self, weighted):
        if self._has_means:
            values = self._cube['result']['measures']['mean']['data']
        elif not weighted or not self.is_weighted:
            values = self._cube['result']['counts']
        else:
            values = self._cube['result']['measures']['count']['data']

        values = [(val if type(val) != dict else np.nan) for val in values]
        return values

    def _as_array(self, include_missing=False, get_non_selected=False,
                  weighted=True, adjusted=False, include_transforms_for_dims=False):
        '''Get crunch cube as ndarray.

        Args
            include_missing (bool): Include rows/cols for missing values
            get_non_selected (bool): Get non-selected slices for MR vars
        Returns
            res (ndarray): Tabular representation of crunch cube
        '''
        values = self._get_values(weighted)
        all_dimensions = self._get_dimensions(self._cube)
        shape = [len(dim.elements(include_missing=True)) for dim in all_dimensions]
        valid_indices = self._get_valid_indices(
            all_dimensions,
            include_missing,
            get_non_selected
        )
        res = np.array(values).reshape(shape)
        if include_transforms_for_dims:
            for (i, dim) in enumerate(all_dimensions):
                if (not (dim._has_transforms and i in include_transforms_for_dims) or
                        dim.type == 'categorical_array'):
                    continue
                ind_offset = 0
                for indices in dim._subtotals_indices():
                    ind_subtotal_elements = np.array(indices['inds'])
                    ind_insertion = indices['anchor_ind'] + ind_offset
                    ind_subtotal_elements = np.array(ind_subtotal_elements) + ind_offset
                    if i == 0:
                        value = sum(res[ind_subtotal_elements])
                        res = np.insert(res, ind_insertion + 1, value, axis=i)
                    else:
                        value = np.sum(res[:, ind_subtotal_elements], axis=1)
                        res = np.insert(res, ind_insertion + 1, value, axis=i)
                    valid_indices = self._fix_valid_indices(valid_indices, ind_insertion, i)
                    ind_offset += 1

        res = res[np.ix_(*valid_indices)]

        adjustment = 1 if adjusted else 0
        return self._fix_shape(res) + adjustment

    @classmethod
    def _fix_valid_indices(cls, valid_indices, insertion_index, dim):
        indices = np.array(sorted(valid_indices[dim]))
        slice_index = np.sum(indices <= insertion_index)
        indices[slice_index:] += 1
        indices = np.insert(indices, slice_index, insertion_index + 1)
        valid_indices[dim] = indices.tolist()
        return valid_indices

    @classmethod
    def _calculate_constraints_sum(cls, prop_table, prop_margin, axis):
        '''Calculate sum of constraints (part of the standard error equation).

        This method calculates the sum of the cell proportions multiplied by
        row (or column) marginal proportions (margins divide by the total
        count). It does this by utilizing the matrix multiplication, which
        directly translates to the mathematical definition (the sum
        across i and j indices).
        '''
        if axis not in [0, 1]:
            raise ValueError('Unexpected value for `axis`: {}'.format(axis))

        V = prop_table * (1 - prop_table)
        if axis == 0:
            # If axis is 0, sumation is performed across the 'i' index, which
            # requires the matrix to be multiplied from the right
            # (because of the inner matrix dimensions).
            return np.dot(V, prop_margin)
        elif axis == 1:
            # If axis is 1, sumation is performed across the 'j' index, which
            # requires the matrix to be multiplied from the left
            # (because of the inner matrix dimensions).
            return np.dot(prop_margin, V)

    def _calculate_standard_error(self, axis):
        total = self.margin(weighted=False, adjusted=True)
        # Calculate margin across axis, as percentages of the total count
        margin = self.margin(axis=axis, weighted=False, adjusted=True) / total
        # Adjusted proportions table, necessary for the standard error,
        # because of the division by it.
        props = self.proportions(axis=axis, weighted=False, adjusted=True)

        constraints = self._calculate_constraints_sum(props, margin, axis)
        if axis == 0:
            # If the s.e. is calculated across rows, the addition of the
            # 'constraints' member must be done element-wise, for 'd' and for
            # each column of the props variance matrix, thus the transformation
            # to the column vector.
            constraints = constraints[:, np.newaxis]

        d = (1 - 2 * margin) / margin
        if axis == 1:
            # If the s.e. is calculated across rows, the multiplication of
            # the 'd' member must be done element-wise, for 'd' and for each
            # row of the props variance matrix, thus the transformation to
            # the column vector.
            d = d[:, np.newaxis]

        return np.sqrt((d * props * (1 - props) + constraints) / total)

    def _calculate_statistics(self, axis):
        if axis not in [0, 1]:
            raise ValueError('Unexpected value for `axis`: {}'.format(axis))

        props = self.proportions(axis=axis)

        # Statistics are calculated by operating on both axes' margins. In this
        # function, we need to determine the cross-axis (other than the one
        # we're doing the calculation for), in order to be able to calculate
        # *that* margin, which will serve as the basis for the
        # statistics calculation.
        cross_axis = 0 if axis == 1 else 1
        cross_margin = self.margin(axis=cross_axis) / self.margin()

        if cross_axis == 1:
            # If the row proportional margins are required, they also need to
            # be broadcast into the vector column shape, in order to be able to
            # perform the subtration from the matrix.
            cross_margin = cross_margin[:, np.newaxis]

        return (props - cross_margin) / self._calculate_standard_error(axis)

    def _is_double_multiple_response(self):
        types = [dim.type for dim in self.dimensions]
        if types == ['multiple_response', 'multiple_response']:
            return True
        return False

    def _double_mr_proportions(self, axis, weighted):
        if axis is None:
            ind_sel, ind_non = 1, 1
        else:
            ind_sel, ind_non = 1 - axis, axis

        all_dimensions = self._get_dimensions(self._cube)
        shape = [len(dim.elements(include_missing=True)) for dim in all_dimensions]
        values = np.array(self._get_values(weighted)).reshape(shape)

        return (
            values[:, 0, :, 0] /
            (values[:, 0, :, 0] + values[:, ind_sel, :, ind_non])
        )

    # API Functions

    @property
    def has_multiple_response(self):
        '''Determines if a cube has MR dimensions.'''
        all_dimensions = self._get_dimensions(self._cube)
        mr_indices = self._get_mr_selections_indices(all_dimensions)
        return bool(mr_indices)

    @property
    def name(self):
        '''Return the name of the cube.

        If the cube has 2 diensions, return the name of the second one. In case
        of a different number of dimensions, default to returning the name of
        the last one. In case of no dimensions, return the empty string.
        '''
        if not self.dimensions:
            return None
        return self.dimensions[0].name

    @property
    def description(self):
        '''Return the description of the cube.'''
        if not self.dimensions:
            return None
        return self.dimensions[0].description

    def labels(self, include_missing=False, include_transforms_for_dims=False):
        '''Gets labels for each cube's dimension.

        Args
            include_missing (bool): Include labels for missing values

        Returns
            labels (list of lists): Labels for each dimension
        '''
        return [dim.labels(include_missing, include_transforms_for_dims) for dim in self.dimensions]

    @property
    def dimensions(self):
        '''Dimensions of the crunch cube.'''
        all_dimensions = self._get_dimensions(self._cube)
        mr_selections = self._get_mr_selections_indices(all_dimensions)
        return [
            dim for (i, dim) in enumerate(all_dimensions)
            if i not in mr_selections
        ]

    def as_array(self, include_missing=False, weighted=True, adjusted=False,
                 include_transforms_for_dims=None):
        '''Get crunch cube as ndarray.

        Returns the tabular representation of the crunch cube. The returning
        value has as many dimensions, as there are dimensions in the crunch
        cube itself. E.g. for a cross-tab representation of a categorical and
        numerical variable, the resulting cube will have two dimensions.

        Args
            include_missing (bool): Include rows/cols for missing values
        Returns
            (ndarray): Tabular representation of the crunch cube

        Example 1 (Categorical x Categorical):
            >>> cube = CrunchCube(response)
            >>> cube.as_array()
            np.array([
                [5, 2],
                [5, 3],
            ])

        Example 2 (Categorical x Categorical, include missing values):
            >>> cube = CrunchCube(response)
            >>> cube.as_array(include_missing=True)
            np.array([
                [5, 3, 2, 0],
                [5, 2, 3, 0],
                [0, 0, 0, 0],
            ])
        '''
        return self._as_array(
            include_missing=include_missing,
            weighted=weighted,
            adjusted=adjusted,
            include_transforms_for_dims=include_transforms_for_dims,
        )

    def margin(self, axis=None, weighted=True, adjusted=False,
               include_transforms_for_dims=None):
        '''Get margin for the selected axis.

        the selected axis. For MR variables, this is the sum of the selected
        and non-selected slices.

        Args
            axis (int): Axis across the margin is calculated. If no axis is
                        provided the margin is calculated across all axis.
                        For Categoricals, Num, Datetime, and Text, this
                        translates to sumation of all elements.
        Returns
            Calculated margin for the selected axis

        Example 1:
            >>> cube = CrunchCube(fixt_cat_x_cat)
            np.array([
               [5, 2],
               [5, 3],
            ])

            >>> cube.margin(axis=0)
            np.array([10, 5])

        Example 2:
            >>> cube = CrunchCube(fixt_cat_x_num_x_datetime)
            np.array([
                [[1, 1],
                 [0, 0],
                 [0, 0],
                 [0, 0]],
                [[2, 1],
                 [1, 1],
                 [0, 0],
                 [0, 0]],
                [[0, 0],
                 [2, 3],
                 [0, 0],
                 [0, 0]],
                [[0, 0],
                 [0, 0],
                 [3, 2],
                 [0, 0]],
                [[0, 0],
                 [0, 0],
                 [1, 1],
                 [0, 1]]
            ])

            >>> cube.margin(axis=0)
            np.array([
                [3, 2],
                [3, 4],
                [4, 3],
                [0, 1],
            ])
        '''
        transform_dims = (
            [(1 - axis)]
            if axis is not None and isinstance(axis, int)
            else None
        )
        array = self.as_array(
            weighted=weighted,
            adjusted=adjusted,
            include_transforms_for_dims=transform_dims,
        )

        all_dimensions = self._get_dimensions(self._cube)
        mr_indices = self._get_mr_selections_indices(all_dimensions)
        if mr_indices:
            # Special case treatment for MR variables, that are set as the
            # first dimension of a cube.
            margin = array + self._as_array(
                get_non_selected=True,
                weighted=weighted,
                adjusted=adjusted
            )
            if axis is None and len(margin.shape) > 1 and 1 in mr_indices:
                # If MR margin is being calculated as total, we need the
                # combination of selected + non-selected, and that's why
                # we're using margin and not array.
                return np.sum(margin, 1)[:, np.newaxis]
            if axis is None and len(margin.shape) > 1:
                # This covers the case when there's a MR variable, but
                # not as a first dimension.
                return np.sum(margin, 0)
            if axis == 1:
                # If MR margin is calculated by rows, we only need the counts
                # and that's why we use array and not margin.
                return np.sum(array, axis)
            return margin

        return np.sum(array, axis)

    def proportions(self, axis=None, weighted=True, adjusted=False,
                    include_transforms_for_dims=None, include_missing=False):
        '''Get proportions of a crunch cube.

        This function calculates the proportions across the selected axis
        of a crunch cube. For most variable types, it means the value divided
        by the margin value. For Multiple Response types, the value is divied
        by the sum of selected and non-selected slices.

        Args
            axis (int): Base axis of proportions calculation. If no axis is
                        provided, calculations are done accros entire table.

        Returns
            (nparray): Calculated array of crunch cube proportions.

        Example 1:
            >>> cube = CrunchCube(fixt_cat_x_cat)
            np.array([
               [5, 2],
               [5, 3],
            ])

            >>> cube.proportions()
            np.array([
                [0.3333333, 0.1333333],
                [0.3333333, 0.2000000],
            ])

        Example 2:
            >>> cube = CrunchCube(fixt_cat_x_cat)
            np.array([
               [5, 2],
               [5, 3],
            ])

            >>> cube.proportions(axis=0)
            np.array([
                [0.5, 0.4],
                [0.5, 0.6],
            ])
        '''

        if self._is_double_multiple_response():
            # For the case of MR x MR cube, proportions are calculated in a
            # specific way, which needs to be handled speparately.
            return self._double_mr_proportions(axis, weighted)

        margin = self.margin(axis=axis, weighted=weighted, adjusted=adjusted)
        if axis == 1:
            margin = margin[:, np.newaxis]

        return self.as_array(
            include_missing=include_missing,
            weighted=weighted,
            adjusted=adjusted,
            include_transforms_for_dims=include_transforms_for_dims,
        ) / margin

    def percentages(self, axis=None):
        '''Get the percentages for crunch cube values.

        This function calculates the percentages for crunch cube values. The
        percentages are based on the values of the 'proportions'.

        Args
            axis (int): Base axis of percentages calculation. If no axis is
                        provided, calculations are done accros entire table.

        Returns
            (nparray): Calculated array of crunch cube percentages.

        Example 1:
            >>> cube = CrunchCube(fixt_cat_x_cat)
            np.array([
               [5, 2],
               [5, 3],
            ])

            >>> cube.percentages()
            np.array([
                [33.33333, 13.33333],
                [33.33333, 20.00000],
            ])

        Example 2:
            >>> cube = CrunchCube(fixt_cat_x_cat)
            np.array([
               [5, 2],
               [5, 3],
            ])

            >>> cube.percentages(axis=0)
            np.array([
                [50., 40.],
                [50., 60.],
            ])
        '''
        return self.proportions(axis) * 100

    def pvals(self, axis):
        '''Calculate p-vals.

        This function calculates statistically signifficant results for
        categorical contingency tables. The values can be calculated across
        columns (axis = 0), or across rows (axis = 1).

        Args
            axis (int): Dimension across which to calculate the p-values.
                        0 - calculate across columns
                        1 - calculate across rows
        Returns
            (ndarray): 2-Dimensional array, representing the p-values for each
                       cell of the table-like representation of the
                       crunch cube.
        '''
        stats = self._calculate_statistics(axis)
        sign = np.sign(stats)

        p_values = 2 * (1 - norm.cdf(np.abs(stats)))
        p_values *= sign

        return p_values

    @property
    def missing(self):
        '''Get missing count of a cube.'''
        if self._has_means:
            return self._cube['result']['measures']['mean']['n_missing']
        return self._cube['result'].get('missing')

    def y_offset(self, expand=False):
        '''Gets y offset for sheet manipulation.

        Args:
            - expand (bool): If a cube is a categorical array, and it's also a
                0-index cube of multiple cube sequence, it needs to be
                displayed in a specific manner (sliced across categories
                dimension). This argument enables such behavior. If a CA cube
                is the only one (not a part of the multitable), it is output
                normally, without expansion.
        '''
        if not self.dimensions:
            return 4

        first_dim_length = self.as_array().shape[0]

        # Special case of CA as a 0-ind cube
        if expand and self.dimensions[0].type == 'categorical_array':
            return first_dim_length * (len(self.dimensions[1].elements()) + 4)

        if len(self.dimensions) <= 2 and self.dimensions[0].type:
            return first_dim_length + 4
        return first_dim_length * (self.as_array().shape[1] + 4)

    @property
    def is_weighted(self):
        '''Check if the cube dataset is weighted.'''
        weighted = self._cube.get('query', {}).get('weight', None) is not None
        weighted = weighted or self._cube.get('weight_var', None) is not None
        weighted = weighted or self._cube.get('weight_url', None) is not None
        weighted = weighted or (
            self._cube['result']['counts'] !=
            self._cube['result']['measures'].get('count', {}).get('data')
        )
        return weighted

    @property
    def has_means(self):
        '''Check if cube has means.'''
        return self._has_means

    @property
    def filter_annotation(self):
        '''Get cube's filter annotation.'''
        return self._cube.get('filter_names', [])
