"""
This module implements GB1 DMS,
using the open-available data from [1].

The black-box is a lookup of the combinatorically complete mutation landscape.
If the proposed sequence exists the associated value is returned, else a nan/inf value is returned.

The dataset reference is [2].

[1] "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
bioRxiv 2024.07.27.605457; doi: https://doi.org/10.1101/2024.07.27.605457.
[2] "Learning protein fitness landscapes with deep mutational scanning data from multiple sources"
Lin Chen, Zehong Zhang, Zhenghao Li, Rui Li, Ruifeng Huo, Lifan Chen, Dingyan Wang, Xiaomin Luo, Kaixian Chen, Cangsong Liao, Mingyue Zheng,
Cell Systems,
Volume 14, Issue 8, 2023,
ISSN 2405-4712; doi: https://doi.org/10.1016/j.cels.2023.07.003.

"""

from __future__ import annotations

from pathlib import Path
from time import time
from uuid import uuid4

import numpy as np
import pandas as pd

from poli.core.abstract_isolated_function import AbstractIsolatedFunction

THIS_DIR = Path(__file__).parent.resolve()


class DMSGB1IsolatedLogic(AbstractIsolatedFunction):
    """
    GB1 internal implementation.

    Parameters
    ----------
    alphabet : List[str], optional
        The alphabet for the problem, by default we use
        the amino acid list provided in poli.core.util.proteins.defaults.
    experiment_id : str, optional
        The experiment ID, by default None.

    Methods
    -------
    _black_box(x, context=None)
        The main black box method that performs the computation, i.e.
        it computes the stability of the mutant(s) in x.
    _load_dms_data()
        This function loads the DMS data under assets.

    Notes
    -----
    - If experiment_id is not provided, it is generated using the current timestamp and a random UUID.
    """

    def __init__(
        self,
        experiment_id: str = None,
    ):
        """
        Initialize the GB1 Register object.

        Parameters:
        -----------

        experiment_id : str, optional
            The experiment ID, by default None.

        Notes:
        ------
        - If experiment_id is not provided, it is generated using the current timestamp and a random UUID.
        """

        if experiment_id is None:
            experiment_id = f"{int(time())}_{str(uuid4())[:8]}"
        self.experiment_id = experiment_id

        # see [2] for reference:
        # "we investigated the fitness landscape of all variants
        # (204 = 160,000) at four amino acid sites (V39, D40, G41 and V54)
        # in an epistatic region of protein G domain B1"
        self.wt = "VDGV"
        self.positions = [39, 40, 41, 54]  # included for completeness

        self.dms_df = self._load_dms_data()

        self.x0 = np.array(list(self.wt))[None, :]

    def _load_dms_data(self) -> pd.DataFrame:
        return pd.read_csv(THIS_DIR / "assets" / "fitness.csv")

    def _return_dms_val(self, x) -> float:
        return self.dms_df[self.dms_df.Combo == x].fitness.values[0]

    def __call__(self, x, context=None):
        """
        Computes the fitness of the variant(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 longest sequence in the batch, and b is the batch size.

        Throws
        -----
        - ValueError exception if input is invalid.
        """
        results = []
        for i, x_i in enumerate(x):
            # Assuming x_i is an array of strings
            if len(x_i) != 4:
                raise ValueError(
                    f"Inputs of L=4 expected!\n Curren input length={len(x_i)} at index {i}"
                )
            mutant_residue_string = "".join(x_i)
            result = self._return_dms_val(mutant_residue_string)
            results.append(result)

        return np.array(results).reshape(-1, 1)


if __name__ == "__main__":
    from poli.core.registry import register_isolated_function

    register_isolated_function(
        DMSGB1IsolatedLogic,
        name="dmsgb1__isolated",
        conda_environment_name="poli__dms",
        force=True,
    )
