# -*- mode: python; coding: utf-8 -*-
# Copyright 2013-2015 Peter Williams <peter@newton.cx> and collaborators.
# Licensed under the MIT license.

"""
pwkit.msmt - Working with uncertain measurements.

"""
from __future__ import absolute_import, division, print_function, unicode_literals

__all__ = str ('''uval_unary_math''').split ()

import operator, six
from six.moves import range
import numpy as np

from . import PKError
from .simpleenum import enumeration
from .numutil import broadcastize


# Datatypes for storage arrays.

n_samples = 1024 - 1 # each Uval takes 1024*itemsize bytes

class UvalDtypeGenerator (object):
    def __init__ (self):
        self.cache = {}

    def __call__ (self, sample_dtype):
        prev = self.cache.get (sample_dtype)
        if prev is not None:
            return prev

        # The `str` calls are needed for Python 2.x since unicode_literals is
        # activated.
        npad = np.dtype (sample_dtype).itemsize - 1
        dtype = np.dtype ([
            (str ('kind'), np.uint8, 1),
            (str ('_pad'), np.uint8, npad),
            (str ('samples'), sample_dtype, n_samples),
        ])
        self.cache[sample_dtype] = dtype
        return dtype

get_uval_dtype = UvalDtypeGenerator ()

def make_data (kinds, samples):
    data = np.empty (kinds.shape, dtype=get_uval_dtype (samples.dtype))
    data['kind'] = kinds
    data['samples'] = samples
    return data


uval_default_repval_method = 'pct'


@enumeration
class Domains (object):
    nonpositive = -1
    anything = 0
    nonnegative = 1

    names = ['nonpositive', 'anything', 'nonnegative']

    @classmethod
    def check (cls, domain):
        if not (domain >= cls.nonpositive and domain <= cls.nonnegative):
            raise ValueError ('illegal Uval domain %r' % domain)
        return domain

    negate = [nonnegative, anything, nonpositive]
    add = {
        (nonpositive, nonpositive): nonpositive,
        (nonpositive, anything): anything,
        (nonpositive, nonnegative): anything,
        (anything, anything): anything,
        (anything, nonnegative): anything,
        (nonnegative, nonnegative): nonnegative,
    }


@broadcastize(1)
def _all_in_domain (data, domain):
    if domain == Domains.anything:
        return True
    if domain == Domains.nonnegative:
        return np.all (data >= 0)
    if domain == Domains.nonpositive:
        return np.all (data <= 0)


@enumeration
class Kinds (object):
    undef = 0
    sampled = 1
    upper = 2
    lower = 3

    names = ['undef', 'sampled', 'upper', 'lower']

    @classmethod
    def check (cls, kind):
        if not (kind >= undef and kind <= lower):
            raise ValueError ('illegal Uval kind %r' % kind)
        return kind

    negate = np.array ([undef, sampled, lower, upper])

    # For binary operations, we need to be able to vectorize, and it would
    # be nice to be somewhat efficient about it. So, we need to make a bunch
    # of static arrays indexed by:
    #
    #    0 - undef   * undef
    #    1 - undef   * sampled
    #    2 - undef   * upper
    #    3 - undef   * lower
    #    4 - sampled * undef   (redundant with 1)
    #    5 - sampled * sampled
    #    6 - sampled * upper
    #    7 - sampled * lower
    #    8 - upper   * undef   (redundant with 2)
    #    9 - upper   * sampled (redundant with 6)
    #   10 - upper   * upper
    #   11 - upper   * lower
    #   12 - lower   * undef   (redundant with 3)
    #   13 - lower   * sampled (redundant with 7)
    #   14 - lower   * upper   (redundant with 11)
    #   15 - lower   * lower

    @classmethod
    def binop (cls, k1, k2):
        return k2 * 4 + k2

    add = np.array ([
        undef, undef, undef, undef, # undef * ...
        undef, sampled, upper, lower, # sampled * ...
        undef, upper, upper, undef, # upper * ...
        undef, lower, undef, lower, # lower * ...
    ])


class LimitError (PKError):
    def __init__ (self):
        super (LimitError, self).__init__ ('forbidden operation on a limit value')


def _ordpair (v1, v2):
    if v1 > v2:
        return (v2, v1)
    return (v1, v2)


class Uval (object):
    """An empirical uncertain value, represented by samples.

    """
    __slots__ = ('domain', 'data')

    def __init__ (self, domain, shape_or_data=None, sampledtype=np.double):
        self.domain = Domains.check (domain)

        if isinstance (shape_or_data, (tuple,) + six.integer_types):
            self.data = np.empty (shape_or_data, dtype=get_uval_dtype (sampledtype))
            self.data['kind'].fill (Kinds.undef)
        elif isinstance (shape_or_data, np.ndarray):
            # It's hard to check the array dtype thoroughly but let's do this:
            try:
                assert shape_or_data['kind'].dtype == np.uint8
            except Exception:
                raise ValueError ('illegal Uval initializer array %r' % shape_or_data)
            self.data = shape_or_data
        else:
            raise ValueError ('unrecognized Uval initializer %r' % shape_or_data)

    @staticmethod
    def from_other (v, copy=True):
        if isinstance (v, Uval):
            if copy:
                return v.copy ()
            return v

        return Uval.from_fixed (Domains.anything, Kinds.sampled, v)

    @staticmethod
    def from_fixed (domain, kind, v):
        Domains.check (domain)
        Kinds.check (kind)
        if not _all_in_domain (v, domain):
            raise ValueError ('illegal Uval initializer: data %r do not lie in '
                              'stated domain' % v)

        v = np.asarray (v)
        r = Uval (domain, v.shape)
        r.data['kind'] = kind
        r.data['samples'] = v[...,None]
        return r

    @staticmethod
    def from_norm (mean, std, shape=(), domain=Domains.anything):
        Domains.check (domain)
        if std < 0:
            raise ValueError ('std must be positive')

        r = Uval (domain, shape)
        r.data['kind'].fill (Kinds.sampled)
        r.data['samples'] = np.random.normal (mean, std, shape+(n_samples,))
        return r

    def copy (self):
        return self.__class__ (self.domain, self.data.copy ())

    # Basic array properties

    @property
    def shape (self):
        return self.data.shape

    @property
    def ndim (self):
        return self.data.ndim

    @property
    def size (self):
        return self.data.size

    # Math. We start with addition. It gets complicated!

    def __neg__ (self):
        return _uval_unary_negative (self)

    def __abs__ (self):
        return _uval_unary_absolute (self)


    def __add__ (self, other):
        v1 = self
        v2 = Uval.from_other (other, copy=False)

        dom = Domains.add[_ordpair (v1.domain, v2.domain)]
        kind = Kinds.add[Kinds.binop (v1.data['kind'], v2.data['kind'])]
        tot = v1.data['samples'] + v2.data['samples']
        return Uval (dom, make_data (kind, tot))

    __radd__ = __add__

    def __sub__ (self, other):
        return self + (-other)

    __rsub__ = __sub__


def _uval_unary_negative (v):
    r = v.copy ()
    r.domain = Domains.negate[v.domain]
    r.data['kind'] = Kinds.negate[v.data['kind']]
    np.negative (r.data['samples'], r.data['samples'])
    return r


def _uval_unary_absolute (v):
    r = v.copy ()
    r.domain = Domains.nonnegative
    np.absolute (r.data['samples'], r.data['samples'])

    i = np.nonzero (r.data['kind'] == Kinds.past_zero)
    r.data['samples'][i] = 0.
    r.data['kind'][i] = Kinds.to_inf

    return r
