from __future__ import division

from itertools import chain

import numpy as np

from ._graph_matrix import DistanceMatrix3D
from ._atomic_property import AtomicProperty
from ._base.descriptor import Descriptor

__all__ = 'MoRSE',


class MoRSE(Descriptor):
    __slots__ = ('_prop', '_distance')

    require_3D = True

    @classmethod
    def preset(cls):
        return chain(
            (cls(None, i) for i in range(1, 33)),
            (cls('m', i) for i in range(1, 33)),
            (cls('v', i) for i in range(1, 33)),
            (cls('se', i) for i in range(1, 33)),
            (cls('p', i) for i in range(1, 33)),
        )

    def __str__(self):
        p = '' if self._prop is None else self._prop.as_argument
        return 'Mor{:02d}{}'.format(self._distance, p)

    def parameters(self):
        p = None if self._prop is None else self._prop.as_argument
        return p, self._distance

    def __init__(self, prop=None, distance=2):
        if prop is None:
            self._prop = None
        else:
            self._prop = AtomicProperty(self.explicit_hydrogens, prop)

        self._distance = distance

    def dependencies(self):
        d = {'D': DistanceMatrix3D(self.explicit_hydrogens)}

        if self._prop is not None:
            d['A'] = self._prop

        return d

    def calculate(self, D, A=None):
        if D.shape[0] <= 1:
            self.fail(ValueError('require 2 or more atoms'))

        N = D.shape[0]

        if A is None:
            A = np.ones(N)
        else:
            A = A / self._prop.carbon

        A = A.reshape(1, -1)

        if self._distance == 1:
            n = np.ones((N, N), dtype='float')

        else:
            with self.rethrow_zerodiv():
                sr = (self._distance - 1) * D
                np.fill_diagonal(sr, 1)
                n = np.sin(sr) / sr

        np.fill_diagonal(n, 0)

        return np.float(0.5 * A.dot(n).dot(A.T))

    rtype = float
