import numpy as np
from .artifacts import AbstractArtifact
from .utils import _softmax
class Cluster:
    """ An implementation of a Cluster of an island.  This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)

    **Citation**:
    
    @Article{FunSearch2023,
        author  = {Romera-Paredes, Bernardino and Barekatain, Mohammadamin and Novikov, Alexander and Balog, Matej and Kumar, M. Pawan and Dupont, Emilien and Ruiz, Francisco J. R. and Ellenberg, Jordan and Wang, Pengming and Fawzi, Omar and Kohli, Pushmeet and Fawzi, Alhussein},
        journal = {Nature},
        title   = {Mathematical discoveries from program search with large language models},
        year    = {2023},
        doi     = {10.1038/s41586-023-06924-6}
    }
    """
    def __init__(self,score: float,first_program: AbstractArtifact,epsilon=1e-6,sample_with_replacement=False, default_program_temperature=0.1):
        self.score: float = score
        self.programs: list[AbstractArtifact] = [first_program]
        self.lengths = np.array([len(str(first_program))],dtype=np.float32)
        self.epsilon = epsilon
        self.sample_with_replacement = sample_with_replacement
        self.default_program_temperature = default_program_temperature
        
        
    def compute_length_probs(self,program_temperature: float):
        """ Compute the probability of each program given the length of the program. The probability is computed as the softmax of the negative length of the program. The temperature of the softmax is controlled by the program_temperature parameter.
        
        :param program_temperature: The temperature of the softmax
        :type program_temperature: float
        :return: The probability of each program given the length of the program
        :rtype: np.array
        """
        min_length = np.min(self.lengths)
        max_length = np.max(self.lengths)
        

        length_logits = (self.lengths - min_length)/(max_length  + self.epsilon)
        
        probs =  _softmax(-length_logits,program_temperature)
        return probs
    
    def register_program(self,program: str):
        """ Register a program on the cluster.
        
        :param program: The program to register
        :type program: str
        """
        self.programs.append(program)
        self.lengths = np.append(self.lengths,len(str(program)))
        
    def sample_program(self,program_temperature=None):
        """ Sample a program from the cluster given the program temperature.
        
        :param program_temperature: The temperature of the program
        :type program_temperature: float, optional
        :return: The sampled program
        :rtype: str
        """
        if program_temperature is None:
            program_temperature = self.default_program_temperature
        
        
        probs = self.compute_length_probs(program_temperature)
        #sample an index of probs randomly givent the probs
        index = np.random.choice(len(probs),p=probs,replace=self.sample_with_replacement)
        
        return self.programs[index]