#Utility functions for a memoized implementation of the likelihood calculation
import numpy as np
from functools import cache, wraps

class numpy_array_hashed:
    def __init__(self, x):
        self.val = x
        self.h = hash( x.data.tobytes() )
        return None
    
    def __hash__(self):
        return self.h
    
    def __eq__(self, other):
        return self.h == other.h
    
def memoize(function):
    @cache
    def cached_caller(*args, **kwargs):

        args = [a.val if isinstance(a, numpy_array_hashed) else a for a in args]
        kwargs = {
            k: v.val if isinstance(v, numpy_array_hashed) else v for k, v in kwargs.items()
        }

        return function(*args, **kwargs)

    @wraps(function)
    def wrapper(*args, **kwargs):
        args = [numpy_array_hashed(a) if isinstance(a, np.ndarray) else a for a in args]
        kwargs = {
            k: numpy_array_hashed(v) if isinstance(v, np.ndarray) else v for k, v in kwargs.items()
        }
        wrapper.calls += 1
        return cached_caller(*args, **kwargs)

    wrapper.cache_info = cached_caller.cache_info
    wrapper.cache_clear = cached_caller.cache_clear
    wrapper.calls = 0
    return wrapper