<commit_before>"""Structure prediction implementation."""

from typing import Generator, List, Tuple, Optional

from .database import StructureDB
from .mutation import CationMutator
from .structure import SmactStructure


class StructurePredictor:
    """Provides structure prediction functionality.

    Implements a statistically-based model for determining
    likely structures of a given composition, based on a
    database of known compositions and a lambda table of
    weights.

    Based on the algorithm presented in:
        <NAME>., <NAME>., <NAME>., <NAME>., and <NAME>. (2011)
        Data Mined Ionic Substitutions for the Discovery of New Compounds.
        Inorganic Chemistry, 50(2), 656-663.
        `doi:10.1021/ic102031h <https://pubs.acs.org/doi/10.1021/ic102031h>`_

    """

    def __init__(self, mutator: CationMutator, struct_db: StructureDB):
        """Initialize class."""
        self.cm = mutator
        self.db = struct_db

    def predict_structs(
      self,
      species: List[Tuple[str, int]],
      thresh: Optional[float] = 1e-3, ) -> Generator[Tuple[SmactStructure, float], None, None]:
        """Predict structures for a combination of species.

        Args:
            species: A list of (element, charge). The constituent species
                of the target compound.
            thresh: The probability threshold, below which to discard
                predictions.

        Yields:
            Potential structures, as tuples of (structure, probability).

        """
        raise NotImplementedError<commit_msg>[WIP] `StructurePredictor` initial implementation<commit_after>"""Structure prediction implementation."""

import itertools
from typing import Generator, List, Tuple, Optional

from .database import StructureDB
from .mutation import CationMutator
from .structure import SmactStructure
from .utilities import unparse_spec


class StructurePredictor:
    """Provides structure prediction functionality.

    Implements a statistically-based model for determining
    likely structures of a given composition, based on a
    database of known compositions and a lambda table of
    weights.

    Based on the algorithm presented in:
        <NAME>., <NAME>., <NAME>., <NAME>., and <NAME>. (2011)
        Data Mined Ionic Substitutions for the Discovery of New Compounds.
        Inorganic Chemistry, 50(2), 656-663.
        `doi:10.1021/ic102031h <https://pubs.acs.org/doi/10.1021/ic102031h>`_

    """

    def __init__(self, mutator: CationMutator, struct_db: StructureDB):
        """Initialize class."""
        self.cm = mutator
        self.db = struct_db

    def predict_structs(
      self,
      species: List[Tuple[str, int]],
      thresh: Optional[float] = 1e-3, ) -> Generator[Tuple[SmactStructure, float], None, None]:
        """Predict structures for a combination of species.

        Args:
            species: A list of (element, charge). The constituent species
                of the target compound.
            thresh: The probability threshold, below which to discard
                predictions.

        Yields:
            Potential structures, as tuples of (structure, probability).

        """
        # For now, consider just structures with the same species, and unary substitutions.
        # This means we need only consider structures with a difference of 0 or 1 species.

        for identical in self.db.get_with_species(species):
            yield (identical, 1.0)

        sub_spec = itertools.combinations(species, len(species) - 1)

        potential_unary_parents: List[List[SmactStructure]] = list(
          self.db.get_with_species(specs) for specs in sub_spec
        )

        for spec_idx, parents in enumerate(potential_unary_parents):
            # Get missing ion
            (diff_spec, ) = set(species[spec_idx]) - set(sub_spec[spec_idx])
            diff_spec_str = unparse_spec(diff_spec)

            # Determine conditional substitution likelihoods
            diff_sub_probs = self.cm.cond_sub_probs(diff_spec_str)

            for parent in parents:
                # Filter out any structures with identical species
                if parent.has_species(diff_spec):
                    continue

                # Ensure parent has as many species as target
                if len(parent.species) != len(species):
                    continue

                ## Determine probability
                # Get species to be substituted
                (alt_spec, ) = (
                  set(parent.get_spec_strs()) - set(map(unparse_spec, species)) - {diff_spec_str}
                )
                try:
                    p = diff_sub_probs.loc[alt_spec]
                except:
                    # Not in the Series
                    continue

                if p > thresh:
                    yield (cm._mutate_structure(parent, alt_spec, diff_spec), p)
{'fraction_non_alphanumeric': 0.08186655751125665, 'fraction_numerical': 0.016987310683585756, 'mean_word_length': 2.3109756097560976, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 18, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '963349', 'n_tokens_mistral': 1394, 'n_tokens_neox': 1247, 'n_words': 420}