# Source code for pymatgen.analysis.structure_prediction.substitution_probability

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

    """
    This module provides classes for representing species substitution
    probabilities.
    """

    from collections import defaultdict
    from operator import mul
    from pymatgen.core.periodic_table import Specie, get_el_sp
    from monty.design_patterns import cached_class

    import functools
    import itertools
    import json
    import logging
    import math
    import os

    __author__ = "Will Richards, Geoffroy Hautier"
    __copyright__ = "Copyright 2012, The Materials Project"
    __version__ = "1.2"
    __maintainer__ = "Will Richards"
    __email__ = "wrichard@mit.edu"
    __date__ = "Aug 31, 2012"

    @cached_class
    class SubstitutionProbability:
        """
        This class finds substitution probabilities given lists of atoms
        to substitute. The inputs make more sense if you look through the
        from_defaults static method.

        The substitution prediction algorithm is presented in:
        Hautier, G., Fischer, C., Ehrlacher, V., Jain, A., and Ceder, G. (2011)
        Data Mined Ionic Substitutions for the Discovery of New Compounds.
        Inorganic Chemistry, 50(2), 656-663. doi:10.1021/ic102031h
        """

        def __init__(self, lambda_table=None, alpha=-5):
            """
            Args:
                lambda_table:
                    json table of the weight functions lambda if None,
                    will use the default lambda.json table
                alpha:
                    weight function for never observed substitutions
            """
            if lambda_table is not None:
                self._lambda_table = lambda_table
            else:
                module_dir = os.path.dirname(__file__)
                json_file = os.path.join(module_dir, 'data', 'lambda.json')
                with open(json_file) as f:
                    self._lambda_table = json.load(f)

            # build map of specie pairs to lambdas
            self.alpha = alpha
            self._l = {}
            self.species = set()
            for row in self._lambda_table:
                if 'D1+' not in row:
                    s1 = Specie.from_string(row[0])
                    s2 = Specie.from_string(row[1])
                    self.species.add(s1)
                    self.species.add(s2)
                    self._l[frozenset([s1, s2])] = float(row[2])

            # create Z and px
            self.Z = 0
            self._px = defaultdict(float)
            for s1, s2 in itertools.product(self.species, repeat=2):
                value = math.exp(self.get_lambda(s1, s2))
                self._px[s1] += value / 2
                self._px[s2] += value / 2
                self.Z += value

        def get_lambda(self, s1, s2):
            """
            Args:
                s1 (Structure): 1st Structure
                s2 (Structure): 2nd Structure

            Returns:
                Lambda values
            """
            k = frozenset([get_el_sp(s1),
                           get_el_sp(s2)])
            return self._l.get(k, self.alpha)

        def get_px(self, sp):
            """
            Args:
                sp (Specie/Element): Species

            Returns:
                Probability
            """
            return self._px[get_el_sp(sp)]

        def prob(self, s1, s2):
            """
            Gets the probability of 2 species substitution. Not used by the
            structure predictor.

            Returns:
                Probability of s1 and s2 substitution.
            """
            return math.exp(self.get_lambda(s1, s2)) / self.Z

        def cond_prob(self, s1, s2):
            """
            Conditional probability of substituting s1 for s2.

            Args:
                s1:
                    The *variable* specie
                s2:
                    The *fixed* specie

            Returns:
                Conditional probability used by structure predictor.
            """
            return math.exp(self.get_lambda(s1, s2)) / self.get_px(s2)

        def pair_corr(self, s1, s2):
            """
            Pair correlation of two species.

            Returns:
                The pair correlation of 2 species
            """
            return math.exp(self.get_lambda(s1, s2)) * self.Z / (self.get_px(s1) * self.get_px(s2))

        def cond_prob_list(self, l1, l2):
            """
            Find the probabilities of 2 lists. These should include ALL species.
            This is the probability conditional on l2

            Args:
                l1, l2:
                    lists of species

            Returns:
                The conditional probability (assuming these species are in
                l2)
            """
            assert len(l1) == len(l2)
            p = 1
            for s1, s2 in zip(l1, l2):
                p *= self.cond_prob(s1, s2)
            return p

        def as_dict(self):
            """
            Returns: MSONAble dict
            """
            return {"name": self.__class__.__name__, "version": __version__,
                    "init_args": {"lambda_table": self._lambda_table,
                                  "alpha": self._alpha},
                    "@module": self.__class__.__module__,
                    "@class": self.__class__.__name__}

        @classmethod
        def from_dict(cls, d):
            """
            Args:
                d(dict): Dict representation

            Returns:
                Class
            """
            return cls(**d['init_args'])

    class SubstitutionPredictor:
        """
        Predicts likely substitutions either to or from a given composition
        or species list using the SubstitutionProbability
        """

        def __init__(self, lambda_table=None, alpha=-5, threshold=1e-3):
            """
            Args:
                lambda_table (): Input lambda table.
                alpha (float): weight function for never observed substitutions
                threshold (float): Threshold to use to identify high probability structures.
            """
            self.p = SubstitutionProbability(lambda_table, alpha)
            self.threshold = threshold

        def list_prediction(self, species, to_this_composition=True):
            """
            Args:
                species:
                    list of species
                to_this_composition:
                    If true, substitutions with this as a final composition
                    will be found. If false, substitutions with this as a
                    starting composition will be found (these are slightly
                    different)
            Returns:
                List of predictions in the form of dictionaries.
                If to_this_composition is true, the values of the dictionary
                will be from the list species. If false, the keys will be
                from that list.
            """
            for sp in species:
                if get_el_sp(sp) not in self.p.species:
                    raise ValueError("the species {} is not allowed for the"
                                     "probability model you are using".format(sp))
            max_probabilities = []
            for s1 in species:
                if to_this_composition:
                    max_p = max([self.p.cond_prob(s2, s1) for s2 in self.p.species])
                else:
                    max_p = max([self.p.cond_prob(s1, s2) for s2 in self.p.species])
                max_probabilities.append(max_p)

            output = []

            def _recurse(output_prob, output_species):
                best_case_prob = list(max_probabilities)
                best_case_prob[:len(output_prob)] = output_prob
                if functools.reduce(mul, best_case_prob) > self.threshold:
                    if len(output_species) == len(species):
                        odict = {
                            'probability': functools.reduce(mul, best_case_prob)}
                        if to_this_composition:
                            odict['substitutions'] = dict(
                                zip(output_species, species))
                        else:
                            odict['substitutions'] = dict(
                                zip(species, output_species))
                        if len(output_species) == len(set(output_species)):
                            output.append(odict)
                        return
                    for sp in self.p.species:
                        i = len(output_prob)
                        if to_this_composition:
                            prob = self.p.cond_prob(sp, species[i])
                        else:
                            prob = self.p.cond_prob(species[i], sp)
                        _recurse(output_prob + [prob], output_species + [sp])

            _recurse([], [])
            logging.info('{} substitutions found'.format(len(output)))
            return output

        def composition_prediction(self, composition, to_this_composition=True):
            """
            Returns charged balanced substitutions from a starting or ending
            composition.

            Args:
                composition:
                    starting or ending composition
                to_this_composition:
                    If true, substitutions with this as a final composition
                    will be found. If false, substitutions with this as a
                    starting composition will be found (these are slightly
                    different)

            Returns:
                List of predictions in the form of dictionaries.
                If to_this_composition is true, the values of the dictionary
                will be from the list species. If false, the keys will be
                from that list.
            """
            preds = self.list_prediction(list(composition.keys()),
                                         to_this_composition)
            output = []
            for p in preds:
                if to_this_composition:
                    subs = {v: k for k, v in p['substitutions'].items()}
                else:
                    subs = p['substitutions']
                charge = 0
                for k, v in composition.items():
                    charge += subs[k].oxi_state * v
                if abs(charge) < 1e-8:
                    output.append(p)
            logging.info('{} charge balanced substitutions found'
                         .format(len(output)))
            return output