afml/cache/robust_cache_keys.py

"""
Robust cache key generation for financial ML data structures.
Handles numpy arrays, pandas DataFrames, and time-series data properly.
"""

import hashlib
import pickle
from pathlib import Path
from typing import Any, Optional, Tuple

import numpy as np
import pandas as pd
from loguru import logger


class CacheKeyGenerator:
    """Generate robust, collision-resistant cache keys for ML data structures."""

    @staticmethod
    def generate_key(func, args: tuple, kwargs: dict) -> str:
        """
        Generate a robust cache key for a function call.

        Args:
            func: The function being cached
            args: Positional arguments
            kwargs: Keyword arguments

        Returns:
            MD5 hash string representing the unique call signature
        """
        key_parts = [
            func.__module__,
            func.__qualname__,
        ]

        # Process positional arguments
        for i, arg in enumerate(args):
            try:
                key_part = CacheKeyGenerator._hash_argument(arg, f"arg_{i}")
                key_parts.append(key_part)
            except Exception as e:
                logger.warning(f"Failed to hash argument {i} of type {type(arg)}: {e}")
                # Fallback to string representation
                key_parts.append(f"arg_{i}_{str(hash(str(arg)))}")

        # Process keyword arguments (sorted for consistency)
        for key, value in sorted(kwargs.items()):
            try:
                key_part = CacheKeyGenerator._hash_argument(value, key)
                key_parts.append(f"{key}={key_part}")
            except Exception as e:
                logger.warning(
                    f"Failed to hash kwarg '{key}' of type {type(value)}: {e}"
                )
                # Fallback
                key_parts.append(f"{key}={str(hash(str(value)))}")

        # Combine all parts and hash
        combined = "_".join(key_parts)
        return hashlib.md5(combined.encode()).hexdigest()

    @staticmethod
    def _hash_argument(arg: Any, name: str) -> str:
        """Hash a single argument based on its type."""
        try:
            from sklearn.base import BaseEstimator

            if isinstance(arg, BaseEstimator):
                return CacheKeyGenerator._hash_sklearn_estimator(arg, name)
        except ImportError:
            pass  # sklearn not available, continue with other types

        if isinstance(arg, np.ndarray):
            return CacheKeyGenerator._hash_numpy_array(arg, name)
        elif isinstance(arg, pd.DataFrame):
            return CacheKeyGenerator._hash_dataframe(arg, name)
        elif isinstance(arg, pd.Series):
            return CacheKeyGenerator._hash_series(arg, name)
        elif isinstance(arg, (list, tuple)):
            return CacheKeyGenerator._hash_sequence(arg, name)
        elif isinstance(arg, dict):
            return CacheKeyGenerator._hash_dict(arg, name)
        elif isinstance(arg, (int, float, str, bool, type(None))):
            return CacheKeyGenerator._hash_primitive(arg, name)
        else:
            # Fallback for unknown types
            return CacheKeyGenerator._hash_generic(arg, name)

    @staticmethod
    def _hash_numpy_array(arr: np.ndarray, name: str) -> str:
        """Hash numpy array including shape, dtype, and content."""
        # For large arrays, sample for performance
        if arr.size > 10000:
            # Hash shape, dtype, and a sample
            sample = arr.flat[:: max(1, arr.size // 1000)]  # Sample ~1000 points
            content_hash = hashlib.md5(sample.tobytes()).hexdigest()[:8]
        else:
            # Hash full content for small arrays
            content_hash = hashlib.md5(arr.tobytes()).hexdigest()[:8]

        return f"{name}_arr_{arr.shape}_{arr.dtype}_{content_hash}"

    @staticmethod
    def _hash_dataframe(df: pd.DataFrame, name: str) -> str:
        """Hash pandas DataFrame including index, columns, dtypes, and content."""
        parts = [
            f"shape_{df.shape}",
            f"cols_{hashlib.md5(str(tuple(df.columns)).encode()).hexdigest()[:8]}",
            f"dtypes_{hashlib.md5(str(tuple(df.dtypes)).encode()).hexdigest()[:8]}",
        ]

        # Hash index
        if isinstance(df.index, pd.DatetimeIndex):
            # For datetime index, hash start, end, and frequency
            parts.append(f"idx_dt_{df.index[0]}_{df.index[-1]}_{len(df.index)}")
        else:
            idx_hash = hashlib.md5(str(tuple(df.index)).encode()).hexdigest()[:8]
            parts.append(f"idx_{idx_hash}")

        # Hash content (sample for large DataFrames)
        if df.size > 10000:
            # Sample rows for hashing
            sample_rows = df.iloc[:: max(1, len(df) // 100)]  # ~100 rows
            content_hash = hashlib.md5(sample_rows.values.tobytes()).hexdigest()[:8]
        else:
            content_hash = hashlib.md5(df.values.tobytes()).hexdigest()[:8]

        parts.append(f"data_{content_hash}")

        return f"{name}_df_{'_'.join(parts)}"

    @staticmethod
    def _hash_series(series: pd.Series, name: str) -> str:
        """Hash pandas Series."""
        parts = [
            f"len_{len(series)}",
            f"dtype_{series.dtype}",
        ]

        # Hash index
        if isinstance(series.index, pd.DatetimeIndex):
            parts.append(f"idx_dt_{series.index[0]}_{series.index[-1]}")
        else:
            idx_hash = hashlib.md5(str(tuple(series.index)).encode()).hexdigest()[:8]
            parts.append(f"idx_{idx_hash}")

        # Hash values
        if len(series) > 1000:
            sample = series.iloc[:: max(1, len(series) // 100)]
            content_hash = hashlib.md5(sample.values.tobytes()).hexdigest()[:8]
        else:
            content_hash = hashlib.md5(series.values.tobytes()).hexdigest()[:8]

        parts.append(f"data_{content_hash}")

        return f"{name}_series_{'_'.join(parts)}"

    @staticmethod
    def _hash_sequence(seq: Tuple[Any, ...] | list, name: str) -> str:
        """Hash list or tuple recursively."""
        if len(seq) == 0:
            return f"{name}_empty_seq"

        # Hash each element
        element_hashes = []
        for i, item in enumerate(seq):
            elem_hash = CacheKeyGenerator._hash_argument(item, f"{name}_{i}")
            element_hashes.append(elem_hash)

        combined = "_".join(element_hashes)
        return hashlib.md5(combined.encode()).hexdigest()[:8]

    @staticmethod
    def _hash_dict(d: dict, name: str) -> str:
        """Hash dictionary recursively."""
        if len(d) == 0:
            return f"{name}_empty_dict"

        # Sort keys for consistency
        items_hash = []
        for key, value in sorted(d.items()):
            val_hash = CacheKeyGenerator._hash_argument(value, f"{name}_{key}")
            items_hash.append(f"{key}={val_hash}")

        combined = "_".join(items_hash)
        return hashlib.md5(combined.encode()).hexdigest()[:8]

    @staticmethod
    def _hash_primitive(value: Any, name: str) -> str:
        """Hash primitive types."""
        return f"{name}_{type(value).__name__}_{hash(value)}"

    @staticmethod
    def _hash_generic(obj: Any, name: str) -> str:
        """Fallback hashing for unknown types."""
        try:
            # Try to use object's __repr__
            return f"{name}_{type(obj).__name__}_{hash(repr(obj))}"
        except Exception:
            # Last resort: use id
            return f"{name}_{type(obj).__name__}_{id(obj)}"

    @staticmethod
    def _hash_sklearn_estimator(estimator: Any, name: str) -> str:
        """Hash sklearn estimator including nested estimators."""
        try:
            from sklearn.base import BaseEstimator

            if not isinstance(estimator, BaseEstimator):
                return CacheKeyGenerator._hash_generic(estimator, name)

            # Use the enhanced estimator hashing from cv_cache
            from .cv_cache import _hash_classifier

            estimator_hash = _hash_classifier(estimator)
            return f"{name}_estimator_{estimator_hash}"

        except ImportError:
            # Fallback if sklearn not available
            return CacheKeyGenerator._hash_generic(estimator, name)


class TimeSeriesCacheKey(CacheKeyGenerator):
    """
    Extended cache key generator with time-series awareness.
    Useful for financial data where lookback periods matter.
    """

    @staticmethod
    def generate_key_with_time_range(
        func,
        args: tuple,
        kwargs: dict,
        time_range: Tuple[pd.Timestamp, pd.Timestamp] = None,
    ) -> str:
        """
        Generate cache key that includes time range information.

        Args:
            func: Function being cached
            args: Positional arguments
            kwargs: Keyword arguments
            time_range: Optional (start, end) timestamp tuple

        Returns:
            Cache key string
        """
        base_key = CacheKeyGenerator.generate_key(func, args, kwargs)

        if time_range is None:
            # Try to extract time range from data
            time_range = TimeSeriesCacheKey._extract_time_range(args, kwargs)

        if time_range:
            start, end = time_range
            time_hash = f"time_{start}_{end}"
            return f"{base_key}_{time_hash}"

        return base_key

    @staticmethod
    def _extract_time_range(
        args: tuple, kwargs: dict
    ) -> Tuple[pd.Timestamp, pd.Timestamp] | None:
        """
        Attempt to extract time range from function arguments.
        Looks for DataFrames with DatetimeIndex or explicit start/end parameters.
        """
        # Check kwargs for explicit time parameters
        if "start_date" in kwargs and "end_date" in kwargs:
            return (
                pd.Timestamp(kwargs["start_date"]),
                pd.Timestamp(kwargs["end_date"]),
            )

        # Check for DataFrames with DatetimeIndex in args
        for arg in args:
            if isinstance(arg, pd.DataFrame) and isinstance(
                arg.index, pd.DatetimeIndex
            ):
                if len(arg.index) > 0:
                    return (arg.index[0], arg.index[-1])

            elif isinstance(arg, pd.Series) and isinstance(arg.index, pd.DatetimeIndex):
                if len(arg.index) > 0:
                    return (arg.index[0], arg.index[-1])

        return None


# =============================================================================
# Integration with existing cacheable decorator
# =============================================================================


def create_robust_cacheable(
    track_data_access: bool = False,
    dataset_name: Optional[str] = None,
    purpose: Optional[str] = None,
    use_time_awareness: bool = False,
):
    """
    Factory function to create robust cacheable decorators with data tracking.
    Args:
        track_data_access: Whether to track DataFrame accesses
        dataset_name: Name of the dataset for tracking
        purpose: One of 'train', 'test', 'validate', 'optimize', 'analyze'
        use_time_awareness: Whether to use time-series aware cache keys
    Returns:
        Decorator function
    """
    import time
    from functools import wraps

    from . import cache_stats, memory
    from .cache_monitoring import get_cache_monitor

    def decorator(func):
        func_name = f"{func.__module__}.{func.__qualname__}"
        cached_func = memory.cache(func)
        seen_signatures = set()
        monitor = get_cache_monitor()

        @wraps(func)
        def wrapper(*args, **kwargs):
            nonlocal seen_signatures

            # Track access time (ALWAYS do this first)
            monitor.track_access(func_name)

            # Generate cache key
            cache_key = None
            is_hit = False
            computation_start = None

            try:
                if use_time_awareness:
                    cache_key = TimeSeriesCacheKey.generate_key_with_time_range(
                        func, args, kwargs
                    )
                else:
                    cache_key = CacheKeyGenerator.generate_key(func, args, kwargs)

                # Track hit/miss
                try:
                    cached_func.check_call_in_cache(*args, **kwargs)
                    is_hit = True
                    cache_stats.record_hit(func_name)
                    logger.debug(f"Cache HIT for {func_name}")
                except:
                    cache_stats.record_miss(func_name)
                    is_hit = False
                    computation_start = time.time()  # Start timing for misses
                    logger.debug(f"Cache MISS for {func_name}")

                # Add to seen_signatures for this session
                seen_signatures.add(cache_key)

            except Exception as e:
                logger.warning(f"Cache key generation failed for {func_name}: {e}")
                cache_stats.record_miss(func_name)
                cache_key = None
                is_hit = False
                computation_start = time.time()  # Start timing for error case

            # Track data access if requested
            if track_data_access:
                try:
                    from .data_access_tracker import get_data_tracker

                    _track_dataframe_access(
                        get_data_tracker(), args, kwargs, dataset_name, purpose
                    )
                except Exception as e:
                    logger.warning(f"Data tracking failed for {func_name}: {e}")

            # Execute function
            try:
                if is_hit:
                    # For cache hits, just return cached result (no timing needed)
                    result = cached_func(*args, **kwargs)
                else:
                    # For cache misses, time the computation
                    result = cached_func(*args, **kwargs)
                    if computation_start:
                        computation_time = time.time() - computation_start
                        monitor.track_computation_time(func_name, computation_time)
                        logger.debug(
                            f"Computation time for {func_name}: {computation_time:.3f}s"
                        )

                return result

            except (EOFError, pickle.PickleError, OSError) as e:
                # Handle cache corruption
                logger.warning(
                    f"Cache corruption for {func_name}: {type(e).__name__} - recomputing"
                )

                # Clear corrupted cache if possible
                if cache_key is not None:
                    _clear_corrupted_cache(cached_func, args, kwargs, func_name)

                # Execute function directly and track time
                direct_start = time.time()
                result = func(*args, **kwargs)

                if computation_start:  # Track time if it was originally a miss
                    computation_time = time.time() - direct_start
                    monitor.track_computation_time(func_name, computation_time)
                    logger.debug(
                        f"Direct computation time for {func_name}: {computation_time:.3f}s"
                    )

                return result

            except Exception as e:
                # Other unexpected errors
                logger.error(f"Unexpected cache error for {func_name}: {e}")
                raise

        # Add cache info method for debugging
        def cache_info():
            return {
                "function_name": func_name,
                "seen_signatures": len(seen_signatures),
                "hits": cache_stats._stats.get(func_name, {}).get("hits", 0),
                "misses": cache_stats._stats.get(func_name, {}).get("misses", 0),
            }

        wrapper.cache_info = cache_info
        wrapper._afml_cacheable = True
        return wrapper

    return decorator


def _clear_corrupted_cache(cached_func, args, kwargs, func_name):
    """Helper to clear corrupted cache entries."""
    try:
        if hasattr(cached_func, "_get_cache_id"):
            joblib_cache_key = cached_func._get_cache_id(*args, **kwargs)
            cache_dir = Path(cached_func.store_backend.location)

            # Remove files matching this cache key
            removed_count = 0
            for cache_file in cache_dir.rglob("*"):
                if cache_file.is_file() and str(joblib_cache_key) in str(cache_file):
                    cache_file.unlink()
                    removed_count += 1
                    logger.debug(f"Removed corrupted file: {cache_file.name}")

            if removed_count > 0:
                logger.info(
                    f"Cleared {removed_count} corrupted cache files for {func_name}"
                )

    except Exception as clear_exc:
        logger.warning(f"Failed to clear corrupted cache for {func_name}: {clear_exc}")


def _track_dataframe_access(tracker, args, kwargs, dataset_name, purpose):
    """Track DataFrame accesses for data hygiene monitoring."""
    # Check all arguments for DataFrames with DatetimeIndex
    for i, arg in enumerate(args):
        if _is_trackable_dataframe(arg):
            _log_dataframe_access(tracker, arg, dataset_name or f"arg_{i}", purpose)

    for key, value in kwargs.items():
        if _is_trackable_dataframe(value):
            _log_dataframe_access(tracker, value, dataset_name or key, purpose)


def _is_trackable_dataframe(obj):
    """Check if object is a DataFrame with temporal index."""
    return (
        isinstance(obj, pd.DataFrame)
        and isinstance(obj.index, pd.DatetimeIndex)
        and len(obj) > 0
    )


def _log_dataframe_access(tracker, df, name, purpose):
    """Log DataFrame access to tracker."""
    tracker.log_access(
        dataset_name=name,
        start_date=df.index[0],
        end_date=df.index[-1],
        purpose=purpose or "unknown",
        data_shape=df.shape,
    )


# =============================================================================
# Final convenience exports
# =============================================================================

# Standard decorators (backward compatible)
robust_cacheable = create_robust_cacheable(use_time_awareness=False)
time_aware_cacheable = create_robust_cacheable(use_time_awareness=True)

# Data tracking decorators (new functionality)
data_tracking_cacheable = lambda dataset_name, purpose: create_robust_cacheable(
    track_data_access=True,
    dataset_name=dataset_name,
    purpose=purpose,
    use_time_awareness=False,
)

time_aware_data_tracking_cacheable = (
    lambda dataset_name, purpose: create_robust_cacheable(
        track_data_access=True,
        dataset_name=dataset_name,
        purpose=purpose,
        use_time_awareness=True,
    )
)

__all__ = [
    "CacheKeyGenerator",
    "TimeSeriesCacheKey",
    "data_tracking_cacheable",  # NEW
    "robust_cacheable",  # Backward compatible
    "time_aware_cacheable",  # Backward compatible
    "time_aware_data_tracking_cacheable",  # NEW
]