"""
This module implements the DMS TrpB black-box,
using the data available from [1].

The black-box is a lookup of the combinatorically complete mutation landscape.
The proposed sequence has to exist and the associated value is returned.

The dataset reference is [2].

[1] "A combinatorially complete epistatic fitness landscape in an enzyme active site"
Kadina E. Johnston, Patrick J. Almhjell, Ella J. Watkins-Dulaney, Grace Liu, Nicholas J. Porter, Jason Yang, and Frances H. Arnold.
doi: https://doi.org/10.1073/pnas.2400439121 .
July 29, 2024.
[2] "Active learning-assisted directed evolution"
Jason Yang, Ravi G. Lal, James C. Bowden, Raul Astudillo, Mikhail A. Hameedi, Sukhvinder Kaur, Matthew Hill, Yisong Yue & Frances H. Arnold.
https://doi.org/10.1038/s41467-025-55987-8 .
January 16, 2025.

All rights of the data remain with the authors of the respective publications ([1], [2])
under the licenses of the provided resources.

"""

from __future__ import annotations

from poli.core.abstract_black_box import AbstractBlackBox
from poli.core.abstract_problem_factory import AbstractProblemFactory
from poli.core.black_box_information import BlackBoxInformation
from poli.core.problem import Problem
from poli.core.util.isolation.instancing import get_inner_function
from poli.core.util.proteins.defaults import AMINO_ACIDS
from poli.core.util.seeding import seed_python_numpy_and_torch


class DMSTrpBBlackBox(AbstractBlackBox):
    """
    TrpB Black Box implementation.

    Parameters
    ----------
    negative: bool, optional
        Apply sign-flip on observations for a minimizing black-box.
    experiment_id : str, optional
        The experiment ID, by default None.
    batch_size : int, optional
        The batch size for parallel evaluation, by default None.
    parallelize : bool, optional
        Flag indicating whether to parallelize evaluation, by default False.
    num_workers : int, optional
        The number of workers for parallel evaluation, by default None.
    evaluation_budget : int, optional
        The evaluation budget, by default float("inf").

    Methods
    -------
    _black_box(x, context=None)
        The main black box method, i.e.
        check x exists and return the DMS value.
    _load_dms_data()
        This function loads the datasets, found under assets.
    """

    def __init__(
        self,
        negative: bool = False,
        experiment_id: str = None,
        batch_size: int = None,
        parallelize: bool = False,
        num_workers: int = None,
        evaluation_budget: int = float("inf"),
        force_isolation: bool = False,
    ):
        """
        Initialize the DMSTrpB Register object.

        Parameters:
        -----------
        negative: bool, optional
            Apply sign flip to observations, ie. for minimizing black-box.
        experiment_id : str, optional
            The experiment ID, by default None.
        batch_size : int, optional
            The batch size for parallel evaluation, by default None.
        parallelize : bool, optional
            Flag indicating whether to parallelize evaluation, by default False.
        num_workers : int, optional
            The number of workers for parallel evaluation, by default None.
        evaluation_budget : int, optional
            The evaluation budget, by default float("inf").
        """
        if parallelize:
            print(
                "poli 🧪: BlackBox parallelization is handled by the isolated logic. Disabling it."
            )
            parallelize = False
        super().__init__(
            batch_size=batch_size,
            parallelize=parallelize,
            num_workers=num_workers,
            evaluation_budget=evaluation_budget,
        )
        self.negative = negative
        self.force_isolation = force_isolation
        self.experiment_id = experiment_id
        self.inner_function = get_inner_function(
            isolated_function_name="dms_trpb__isolated",
            class_name="DMSTrpBIsolatedLogic",
            module_to_import="poli.objective_repository.dms_trpb.isolated_function",
            force_isolation=self.force_isolation,
            experiment_id=self.experiment_id,
        )
        self.x0 = self.inner_function.x0

    def _black_box(self, x, context=None):
        """
        Computes the fitness of the mutant(s) in x.

        Parameters
        ----------
        x : np.ndarray
            Input array of shape [b, L] containing strings.
        context : dict, optional
            Additional context information (default is None).

        Returns
        -------
        y : np.ndarray
            The fitness of the mutant(s) in x.

        Notes
        -----
        - x is a np.array[str] of shape [b, L], where L is the length
          of the chain of residues, and b is the batch size.
          We process it by concantenating the array into a single string,
          where we assume the padding to be an empty string (if there was any).
          Each of these x_i's will be matched to the wildtype in self.  wildtype_residue_strings with the lowest Hamming distance.
        - negative flag inverts sign of observations, ie. for minimizing the black-box
        """
        y = self.inner_function(x, context=context)
        if self.negative:
            y = -y
        return y

    def get_black_box_info(self) -> BlackBoxInformation:
        """
        Returns the black box information for DMSGB1.
        """
        is_aligned = True
        is_fixed_length = True
        max_sequence_length = max([len("".join(x)) for x in self.x0])
        return BlackBoxInformation(
            name="dms_trpb",
            max_sequence_length=max_sequence_length,
            aligned=is_aligned,
            fixed_length=is_fixed_length,
            deterministic=True,
            alphabet=AMINO_ACIDS,
            log_transform_recommended=False,
            discrete=True,
            fidelity="low",
            padding_token="",
        )


class DMSTrpBProblemFactory(AbstractProblemFactory):
    def create(
        self,
        negative: bool = False,
        experiment_id: str = None,
        seed: int = None,
        batch_size: int = None,
        parallelize: bool = False,
        num_workers: int = None,
        evaluation_budget: int = float("inf"),
        force_isolation: bool = False,
    ) -> Problem:
        """
        Creates a DMS GB1 black box instance, alongside initial
        observations.

        Parameters
        ----------
        negative: bool, optional
            Invert signs of observations, by default False.
        experiment_id : str, optional
            The experiment ID, by default None.
        seed : int, optional
            The seed value for random number generation, by default None.
        batch_size : int, optional
            The batch size for parallel evaluation, by default None.
        parallelize : bool, optional
            Flag indicating whether to parallelize evaluation, by default False.
        num_workers : int, optional
            The number of workers for parallel evaluation, by default None.
        evaluation_budget : int, optional
            The evaluation budget, by default float("inf").

        Returns
        -------
        f : DMSTrpBBlackBox
            The DMS black box instance.
        x0 : np.ndarray
            The initial observations (i.e. the wildtype as a sequence
            of amino acids).
        y0 : np.ndarray
            The initial observations (i.e. the stability of the wildtypes).
        """
        if seed is not None:
            seed_python_numpy_and_torch(seed)

        f = DMSTrpBBlackBox(
            negative=negative,
            experiment_id=experiment_id,
            batch_size=batch_size,
            parallelize=parallelize,
            num_workers=num_workers,
            evaluation_budget=evaluation_budget,
            force_isolation=force_isolation,
        )

        # Constructing x0
        # (Moved to the isolated logic)
        x0 = f.inner_function.x0

        problem = Problem(f, x0)

        return problem
