# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

"""
This module contains some utility functions and classes that are used in the chemenv package.
"""

__author__ = "David Waroquiers"
__copyright__ = "Copyright 2012, The Materials Project"
__credits__ = "Geoffroy Hautier"
__version__ = "2.0"
__maintainer__ = "David Waroquiers"
__email__ = "david.waroquiers@gmail.com"
__date__ = "Feb 20, 2016"


import numpy as np
from pymatgen.analysis.chemenv.utils.math_utils import power2_inverse_decreasing, power2_decreasing_exp, smootherstep
from pymatgen.analysis.chemenv.utils.math_utils import power2_inverse_power2_decreasing


class RatioFunction(object):

    ALLOWED_FUNCTIONS = {}

    def __init__(self, function, options_dict=None):
        if function not in self.ALLOWED_FUNCTIONS:
            raise ValueError('Function "{}" is not allowed in RatioFunction of '
                             'type "{}"'.format(function, self.__class__.__name__))
        self.eval = object.__getattribute__(self, function)
        self.function = function
        self.setup_parameters(options_dict=options_dict)

    def setup_parameters(self, options_dict):
        function_options = self.ALLOWED_FUNCTIONS[self.function]
        if len(function_options) > 0:
            # Check if there are missing options
            if options_dict is None:
                missing_options = True
            else:
                missing_options = False
                for op in function_options:
                    if op not in options_dict:
                        missing_options = True
                        break
            # If there are missing options, raise an error
            if missing_options:
                if len(function_options) == 1:
                    opts = 'Option "{}"'.format(function_options[0])
                else:
                    opts1 = ', '.join(['"{}"'.format(op) for op in function_options[:-1]])
                    opts = 'Options {}'.format(' and '.join([opts1,
                                                             '"{}"'.format(function_options[-1])]))
                if options_dict is None or len(options_dict) == 0:
                    missing = 'no option was provided.'
                else:
                    optgiven = list(options_dict.keys())
                    if len(options_dict) == 1:
                        missing = 'only "{}" was provided.'.format(optgiven[0])
                    else:
                        missing1 = ', '.join(['"{}"'.format(miss) for miss in optgiven[:-1]])
                        missing = 'only {} were provided.'.format(' and '.join([missing1,
                                                                                '"{}"'.format(optgiven[-1])]))
                raise ValueError('{} should be provided for function "{}" in RatioFunction of '
                                 'type "{}" while {}'.format(opts,
                                                             self.function,
                                                             self.__class__.__name__,
                                                             missing))
            # Setup the options and raise an error if a wrong option is provided
            for key, val in options_dict.items():
                if key not in function_options:
                    raise ValueError('Option "{}" not allowed for function "{}" in RatioFunction of '
                                     'type "{}"'.format(key, self.function, self.__class__.__name__))
                self.__setattr__(key, val)

    def evaluate(self, value):
        return self.eval(value)


class CSMFiniteRatioFunction(RatioFunction):

    ALLOWED_FUNCTIONS = {'power2_decreasing_exp': ['max_csm', 'alpha']}

    def power2_decreasing_exp(self, vals):
        return power2_decreasing_exp(vals, edges=[0.0, self.__dict__['max_csm']], alpha=self.__dict__['alpha'])

    def fractions(self, data):
        if len(data) == 0:
            return None
        total = np.sum([self.eval(dd) for dd in data])
        if total > 0.0:
            return [self.eval(dd) / total for dd in data]
        else:
            return None

    def mean_estimator(self, data):
        if len(data) == 0:
            return None
        elif len(data) == 1:
            return data[0]
        else:
            fractions = self.fractions(data)
            if fractions is None:
                return None
            return np.sum(np.array(fractions) * np.array(data))

    ratios = fractions


class CSMInfiniteRatioFunction(RatioFunction):

    ALLOWED_FUNCTIONS = {'power2_inverse_decreasing': ['max_csm'],
                         'power2_inverse_power2_decreasing': ['max_csm']}

    def power2_inverse_decreasing(self, vals):
        return power2_inverse_decreasing(vals, edges=[0.0, self.__dict__['max_csm']])

    def power2_inverse_power2_decreasing(self, vals):
        return power2_inverse_power2_decreasing(vals, edges=[0.0, self.__dict__['max_csm']])

    def fractions(self, data):
        if len(data) == 0:
            return None
        close_to_zero = np.isclose(data, 0.0, atol=1e-10).tolist()
        nzeros = close_to_zero.count(True)
        if nzeros == 1:
            fractions = [0.0] * len(data)
            fractions[close_to_zero.index(True)] = 1.0
            return fractions
        elif nzeros > 1:
            raise RuntimeError('Should not have more than one continuous symmetry measure with value equal to 0.0')
        else:
            fractions = self.eval(np.array(data))
            total = np.sum(fractions)
            if total > 0.0:
                return fractions / total
            else:
                return None

    def mean_estimator(self, data):
        if len(data) == 0:
            return None
        elif len(data) == 1:
            return data[0]
        else:
            fractions = self.fractions(data)
            if fractions is None:
                return None
            return np.sum(np.array(fractions) * np.array(data))

    ratios = fractions


class DeltaCSMRatioFunction(RatioFunction):

    ALLOWED_FUNCTIONS = {'smootherstep': ['delta_csm_min', 'delta_csm_max']}

    def smootherstep(self, vals):
        return smootherstep(vals, edges=[self.__dict__['delta_csm_min'], self.__dict__['delta_csm_max']])