"""
Model Advisor module that analyzes dataset statistics and provides
recommendations for optimal preprocessing configurations.
"""

from typing import Dict, Any, Tuple
import numpy as np
from loguru import logger


class ModelAdvisor:
    """
    Analyzes dataset statistics and recommends optimal preprocessing configurations
    based on feature characteristics and distributions.
    """

    def __init__(self, features_stats: Dict[str, Any]):
        """
        Initialize the ModelAdvisor with dataset statistics.

        Args:
            features_stats: Dictionary containing feature statistics generated by DatasetStatistics
        """
        self.features_stats = features_stats
        self.recommendations = {}
        self.global_config = {}
        self.feature_interactions = {}
        self.distribution_thresholds = {
            "normal": {"skewness": 0.5, "kurtosis": 1.0},
            "uniform": {"variance": 0.1},
            "heavy_tailed": {"kurtosis": 4.0},
            "log_normal": {"skewness": 1.0},
            "periodic": {"autocorr": 0.7},
            "multimodal": {"bimodality": 0.55},
            "sparse": {"zero_ratio": 0.5},
        }

    def analyze_feature_stats(self) -> Dict[str, Any]:
        """
        Analyze feature statistics and generate recommendations.

        Returns:
            Dictionary containing recommendations for all features and global configuration
        """
        if not self.features_stats:
            logger.warning(
                "No feature statistics provided. Cannot generate recommendations."
            )
            return {}

        # Process numeric features
        self._analyze_numeric_features()

        # Process categorical features
        self._analyze_categorical_features()

        # Process text features
        self._analyze_text_features()

        # Process date features
        self._analyze_date_features()

        # Analyze feature interactions
        self._analyze_feature_interactions()

        # Generate global recommendations
        self._generate_global_recommendations()

        return {"features": self.recommendations, "global_config": self.global_config}

    def _analyze_numeric_features(self):
        """Analyze numeric features and generate specific recommendations."""
        numeric_stats = self.features_stats.get("numeric_stats", {})

        for feature, stats in numeric_stats.items():
            recommendation = {
                "feature_type": "NumericalFeature",
                "preprocessing": [],
                "config": {},
                "advanced_options": {},
            }

            # Extract statistics
            mean = stats.get("mean", 0)
            variance = stats.get("var", 1)
            skewness = stats.get("skewness", 0)
            kurtosis = stats.get("kurtosis", 3)
            zero_ratio = stats.get("zero_ratio", 0)
            autocorr = stats.get("autocorr", 0)
            bimodality = stats.get("bimodality", 0)

            # Check for zero variance
            if variance is not None and variance < 1e-10:
                recommendation["preprocessing"].append("CONSTANT")
                recommendation["config"]["constant_value"] = mean
                recommendation["notes"] = [
                    "Feature has near-zero variance, consider removing"
                ]
                self.recommendations[feature] = recommendation
                continue

            # Enhanced distribution detection
            dist_type, dist_confidence = self._detect_distribution_type(
                stats, skewness, kurtosis, zero_ratio, autocorr, bimodality
            )

            recommendation["detected_distribution"] = dist_type
            recommendation["distribution_confidence"] = dist_confidence

            # Smart preprocessing recommendations based on distribution
            self._recommend_numeric_preprocessing(
                recommendation, dist_type, stats, skewness, kurtosis
            )

            # Advanced feature engineering recommendations
            self._recommend_advanced_engineering(
                recommendation, feature, stats, dist_type
            )

            self.recommendations[feature] = recommendation

    def _detect_distribution_type(
        self,
        stats: Dict[str, Any],
        skewness: float,
        kurtosis: float,
        zero_ratio: float,
        autocorr: float,
        bimodality: float,
    ) -> Tuple[str, float]:
        """
        Enhanced distribution detection using multiple statistical tests.

        Returns:
            Tuple of (distribution_type, confidence_score)
        """
        # Check for sparse distribution first
        if zero_ratio > self.distribution_thresholds["sparse"]["zero_ratio"]:
            return "sparse", 0.9

        # Check for periodic distribution
        if autocorr > self.distribution_thresholds["periodic"]["autocorr"]:
            return "periodic", 0.8

        # Check for multimodal distribution
        if bimodality > self.distribution_thresholds["multimodal"]["bimodality"]:
            return "multimodal", 0.85

        # Check for normal distribution
        if (
            abs(skewness) < self.distribution_thresholds["normal"]["skewness"]
            and abs(kurtosis - 3) < self.distribution_thresholds["normal"]["kurtosis"]
        ):
            return "normal", 0.9

        # Check for uniform distribution
        if stats.get("var", 1) < self.distribution_thresholds["uniform"]["variance"]:
            return "uniform", 0.85

        # Check for heavy-tailed distribution
        if kurtosis > self.distribution_thresholds["heavy_tailed"]["kurtosis"]:
            return "heavy_tailed", 0.8

        # Check for log-normal distribution
        if skewness > self.distribution_thresholds["log_normal"]["skewness"]:
            return "log_normal", 0.8

        # Default to unknown with low confidence
        return "unknown", 0.5

    def _recommend_numeric_preprocessing(
        self,
        recommendation: Dict[str, Any],
        dist_type: str,
        stats: Dict[str, Any],
        skewness: float,
        kurtosis: float,
    ):
        """Recommend preprocessing steps based on distribution type."""
        if dist_type == "normal":
            recommendation["preprocessing"].append("FLOAT_NORMALIZED")
            recommendation["config"]["normalization"] = "z_score"
            recommendation["notes"] = [
                "Normal distribution detected, standard normalization recommended"
            ]

        elif dist_type == "uniform":
            recommendation["preprocessing"].append("FLOAT_RESCALED")
            recommendation["config"]["min"] = stats.get("min", 0)
            recommendation["config"]["max"] = stats.get("max", 1)
            recommendation["notes"] = [
                "Uniform distribution detected, rescaling recommended"
            ]

        elif dist_type == "heavy_tailed":
            recommendation["preprocessing"].append("DISTRIBUTION_AWARE")
            recommendation["config"]["prefered_distribution"] = "heavy_tailed"
            recommendation["config"]["robust_scaling"] = True
            recommendation["notes"] = [
                "Heavy-tailed distribution detected, robust scaling recommended"
            ]

        elif dist_type == "log_normal":
            recommendation["preprocessing"].append("DISTRIBUTION_AWARE")
            recommendation["config"]["prefered_distribution"] = "log_normal"
            recommendation["config"]["log_transform"] = True
            recommendation["notes"] = [
                "Log-normal distribution detected, logarithmic transformation recommended"
            ]

        elif dist_type == "periodic":
            recommendation["preprocessing"].append("DISTRIBUTION_AWARE")
            recommendation["config"]["prefered_distribution"] = "periodic"
            recommendation["config"]["trigonometric_features"] = True
            recommendation["notes"] = [
                "Periodic distribution detected, trigonometric features recommended"
            ]

        elif dist_type == "multimodal":
            recommendation["preprocessing"].append("DISTRIBUTION_AWARE")
            recommendation["config"]["prefered_distribution"] = "multimodal"
            recommendation["config"]["mixture_model"] = True
            recommendation["notes"] = [
                "Multimodal distribution detected, mixture model encoding recommended"
            ]

        elif dist_type == "sparse":
            recommendation["preprocessing"].append("DISTRIBUTION_AWARE")
            recommendation["config"]["prefered_distribution"] = "sparse"
            recommendation["config"]["zero_handling"] = "special"
            recommendation["notes"] = [
                "Sparse distribution detected, specialized zero handling recommended"
            ]

        else:
            # Default to robust normalization for unknown distributions
            recommendation["preprocessing"].append("FLOAT_NORMALIZED")
            recommendation["config"]["normalization"] = "robust"
            recommendation["notes"] = [
                f"Unknown distribution type: {dist_type}, using robust normalization"
            ]

    def _recommend_advanced_engineering(
        self,
        recommendation: Dict[str, Any],
        feature: str,
        stats: Dict[str, Any],
        dist_type: str,
    ):
        """Recommend advanced feature engineering based on feature characteristics."""
        # Check for complex distribution patterns
        if abs(stats.get("skewness", 0)) > 1.5 or abs(stats.get("kurtosis", 3) - 3) > 2:
            recommendation["advanced_options"][
                "use_advanced_numerical_embedding"
            ] = True
            recommendation["advanced_options"][
                "embedding_dim"
            ] = self._calculate_embedding_dim(stats.get("count", 1000))
            recommendation["advanced_options"]["num_bins"] = self._calculate_num_bins(
                stats.get("count", 1000)
            )
            recommendation["notes"].append(
                "Complex distribution detected, advanced numerical embedding recommended"
            )

        # Check for potential polynomial features
        if dist_type in ["normal", "uniform"]:
            recommendation["advanced_options"]["polynomial_features"] = True
            recommendation["advanced_options"]["degree"] = 2
            recommendation["notes"].append(
                "Linear distribution detected, polynomial features recommended"
            )

        # Check for potential interaction features
        if feature in self.feature_interactions:
            recommendation["advanced_options"]["interaction_features"] = True
            recommendation["advanced_options"][
                "interaction_pairs"
            ] = self.feature_interactions[feature]
            recommendation["notes"].append(
                "Strong feature interactions detected, interaction features recommended"
            )

    def _calculate_embedding_dim(self, n_samples: int) -> int:
        """Calculate optimal embedding dimension based on dataset size."""
        return min(16, max(4, int(np.log2(n_samples))))

    def _calculate_num_bins(self, n_samples: int) -> int:
        """Calculate optimal number of bins for numerical embedding."""
        return min(100, max(10, int(np.sqrt(n_samples))))

    def _analyze_feature_interactions(self):
        """Analyze interactions between features using correlation and mutual information."""
        numeric_stats = self.features_stats.get("numeric_stats", {})
        categorical_stats = self.features_stats.get("categorical_stats", {})

        # Analyze numeric-numeric interactions
        for feat1 in numeric_stats:
            for feat2 in numeric_stats:
                if feat1 < feat2:  # Avoid duplicate pairs
                    corr = self._calculate_correlation(feat1, feat2)
                    if abs(corr) > 0.7:  # Strong correlation threshold
                        self.feature_interactions[feat1] = (
                            self.feature_interactions.get(feat1, []) + [(feat2, corr)]
                        )
                        self.feature_interactions[feat2] = (
                            self.feature_interactions.get(feat2, []) + [(feat1, corr)]
                        )

        # Analyze numeric-categorical interactions
        for num_feat in numeric_stats:
            for cat_feat in categorical_stats:
                mi_score = self._calculate_mutual_information(num_feat, cat_feat)
                if mi_score > 0.5:  # Strong mutual information threshold
                    self.feature_interactions[num_feat] = self.feature_interactions.get(
                        num_feat, []
                    ) + [(cat_feat, mi_score)]
                    self.feature_interactions[cat_feat] = self.feature_interactions.get(
                        cat_feat, []
                    ) + [(num_feat, mi_score)]

    def _calculate_correlation(self, feat1: str, feat2: str) -> float:
        """Calculate correlation between two numeric features."""
        # Implementation would use the actual feature values from the dataset
        # For now, return a placeholder
        return 0.0

    def _calculate_mutual_information(self, num_feat: str, cat_feat: str) -> float:
        """Calculate mutual information between numeric and categorical features."""
        # Implementation would use the actual feature values from the dataset
        # For now, return a placeholder
        return 0.0

    def _analyze_categorical_features(self):
        """Analyze categorical features and generate recommendations."""
        # Try both "categorical" and "categorical_stats" keys
        categorical_features = self.features_stats.get("categorical", {})
        if not categorical_features:
            categorical_features = self.features_stats.get("categorical_stats", {})

        for feature, stats in categorical_features.items():
            vocabulary_size = stats.get("vocabulary_size", 0)
            rare_value_ratio = stats.get("rare_value_ratio", 0)

            if feature not in self.recommendations:
                self.recommendations[feature] = {}

            # Calculate unique count based on value counts if available
            unique_count = 0
            if "value_counts" in stats:
                unique_count = len(stats["value_counts"])
            # If size is available (used in categorical_stats)
            elif "size" in stats:
                unique_count = stats.get("size", 0)
            # If vocab is available (used in categorical_stats)
            elif "vocab" in stats:
                unique_count = len(stats.get("vocab", []))

            # If vocabulary size is 0, use unique count
            if vocabulary_size == 0 and unique_count > 0:
                vocabulary_size = unique_count

            # Set feature type and default configuration
            self.recommendations[feature]["feature_type"] = "CategoricalFeature"
            self.recommendations[feature]["preprocessing"] = []
            self.recommendations[feature]["config"] = {
                "feature_type": "STRING_CATEGORICAL"
            }
            self.recommendations[feature]["notes"] = []

            # Check if ordinal (ordered) data
            if stats.get("is_ordinal", False):
                self.recommendations[feature]["notes"].append(
                    "Ordinal feature, preserving order"
                )
                self.recommendations[feature]["config"]["is_ordinal"] = True

            # Generate encoding recommendations based on vocabulary size
            if vocabulary_size < 50:
                encoding = "ONE_HOT_ENCODING"
                self.recommendations[feature]["preprocessing"].append(
                    "ONE_HOT_ENCODING"
                )
                self.recommendations[feature]["config"][
                    "category_encoding"
                ] = "ONE_HOT_ENCODING"
                self.recommendations[feature]["notes"].append(
                    f"Small vocabulary ({vocabulary_size} categories), one-hot encoding recommended"
                )
            elif vocabulary_size < 100:
                encoding = "EMBEDDING"
                self.recommendations[feature]["preprocessing"].append("EMBEDDING")
                self.recommendations[feature]["config"][
                    "category_encoding"
                ] = "EMBEDDING"
                embedding_size = min(
                    max(4, vocabulary_size // 8), 64
                )  # Scale embedding size with vocabulary
                self.recommendations[feature]["config"][
                    "embedding_size"
                ] = embedding_size
                self.recommendations[feature]["notes"].append(
                    f"Medium vocabulary ({vocabulary_size} categories), embedding recommended"
                )
            else:
                encoding = "HASHING"
                self.recommendations[feature]["preprocessing"].append("HASHING")
                self.recommendations[feature]["config"]["category_encoding"] = "HASHING"

                # Calculate optimal hash bucket size based on vocabulary
                if vocabulary_size < 5000:
                    hash_bucket_size = (
                        vocabulary_size * 3 // 10 * 10
                    )  # Round to nearest 10
                    hash_bucket_size = max(
                        hash_bucket_size, 100
                    )  # At least 100 buckets
                elif vocabulary_size < 50000:
                    hash_bucket_size = (
                        vocabulary_size // 5 // 50 * 50
                    )  # Round to nearest 50
                    hash_bucket_size = max(
                        hash_bucket_size, 1000
                    )  # At least 1000 buckets
                else:
                    hash_bucket_size = (
                        vocabulary_size // 10 // 100 * 100
                    )  # Round to nearest 100
                    hash_bucket_size = max(
                        hash_bucket_size, 5000
                    )  # At least 5000 buckets

                self.recommendations[feature]["config"][
                    "hash_bucket_size"
                ] = hash_bucket_size

                # For very large vocabularies, recommend hash with embedding
                if vocabulary_size > 10000:
                    self.recommendations[feature]["config"][
                        "hash_with_embedding"
                    ] = True
                    embedding_size = min(max(8, hash_bucket_size // 128), 32)
                    self.recommendations[feature]["config"][
                        "embedding_size"
                    ] = embedding_size
                    self.recommendations[feature]["notes"].append(
                        f"Large vocabulary ({vocabulary_size} categories), hashing with embedding recommended"
                    )
                else:
                    self.recommendations[feature]["notes"].append(
                        f"Large vocabulary ({vocabulary_size} categories), hashing recommended"
                    )

                # Assign a salt value if multiple hashed features are present
                # This avoids hash collisions between different features
                salt_index = len(
                    [
                        r
                        for r in self.recommendations.values()
                        if r.get("preprocessing", [])
                        and "HASHING" in r["preprocessing"]
                    ]
                )
                if (
                    salt_index > 0
                ):  # Only add salt if there's more than one hashing feature
                    self.recommendations[feature]["config"]["salt"] = salt_index
                    self.recommendations[feature]["notes"].append(
                        f"Hash salt: {salt_index} to avoid collisions with other features"
                    )

                self.recommendations[feature]["notes"].append(
                    f"Hash bucket size: {hash_bucket_size}"
                )

            # Handle rare values if significant
            if rare_value_ratio > 0.1:
                self.recommendations[feature]["notes"].append(
                    f"High rare value ratio ({rare_value_ratio:.2f}), special handling recommended"
                )

                if encoding == "ONE_HOT_ENCODING" or encoding == "EMBEDDING":
                    self.recommendations[feature]["config"]["oov_buckets"] = min(
                        max(2, int(vocabulary_size * rare_value_ratio / 10)), 10
                    )

                    self.recommendations[feature]["advanced_options"] = {
                        "handle_rare_values": True,
                        "rare_value_threshold": max(0.001, 1.0 / vocabulary_size),
                    }

    def _analyze_text_features(self):
        """Analyze text features and generate specific recommendations."""
        text_stats = self.features_stats.get("text", {})

        for feature, stats in text_stats.items():
            recommendation = {
                "feature_type": "TextFeature",
                "preprocessing": [],
                "config": {},
                "advanced_options": {},
            }

            # Extract statistics
            vocab_size = stats.get("vocab_size", 10000)
            sequence_length = stats.get("sequence_length", 100)
            special_char_ratio = stats.get("special_char_ratio", 0)
            language = stats.get("language", "unknown")

            # Recommend tokenization and embedding strategy
            recommendation["preprocessing"].append("TEXT_VECTORIZATION")
            recommendation["config"]["max_tokens"] = self._calculate_max_tokens(
                vocab_size
            )
            recommendation["config"][
                "output_sequence_length"
            ] = self._calculate_sequence_length(sequence_length)
            recommendation["config"][
                "standardize"
            ] = self._determine_text_standardization(special_char_ratio, language)

            # Smart embedding dimension calculation
            embedding_dim = self._calculate_text_embedding_dim(
                vocab_size, sequence_length
            )
            recommendation["config"]["embedding_dim"] = embedding_dim

            # Add advanced text processing options
            if special_char_ratio > 0.1:
                recommendation["advanced_options"]["special_char_handling"] = True
                recommendation["notes"] = [
                    "High ratio of special characters detected, specialized handling recommended"
                ]

            if language != "unknown":
                recommendation["advanced_options"]["language_specific"] = True
                recommendation["config"]["language"] = language
                recommendation["notes"].append(
                    f"Language-specific processing recommended for {language}"
                )

            self.recommendations[feature] = recommendation

    def _calculate_max_tokens(self, vocab_size: int) -> int:
        """Calculate optimal number of tokens for text vectorization."""
        return min(20000, max(1000, vocab_size))

    def _calculate_sequence_length(self, avg_length: int) -> int:
        """Calculate optimal sequence length for text processing."""
        return min(200, max(50, avg_length * 2))

    def _calculate_text_embedding_dim(
        self, vocab_size: int, sequence_length: int
    ) -> int:
        """Calculate optimal embedding dimension for text features."""
        return min(300, max(16, vocab_size // 100))

    def _determine_text_standardization(
        self, special_char_ratio: float, language: str
    ) -> str:
        """Determine appropriate text standardization strategy."""
        if special_char_ratio > 0.1:
            return "lower_and_strip_punctuation"
        elif language != "unknown":
            return "language_specific"
        else:
            return "lower_and_strip_punctuation"

    def _analyze_date_features(self):
        """Analyze date features and generate specific recommendations."""
        date_stats = self.features_stats.get("date_stats", {})

        for feature, stats in date_stats.items():
            recommendation = {
                "feature_type": "DateFeature",
                "preprocessing": [],
                "config": {},
                "advanced_options": {},
                "notes": [],
            }

            # Extract statistics
            has_time = stats.get("has_time", False)
            timezone_info = stats.get("timezone_info", None)
            cyclical_patterns = stats.get("cyclical_patterns", [])

            # Basic date feature extraction
            recommendation["preprocessing"].append("DATE_FEATURES")
            recommendation["config"]["extract"] = [
                "year",
                "month",
                "day",
                "dayofweek",
                "quarter",
            ]

            if has_time:
                recommendation["config"]["extend"].extend(["hour", "minute", "second"])

            # Add timezone handling if needed
            if timezone_info:
                recommendation["config"]["timezone_aware"] = True
                recommendation["config"]["timezone"] = timezone_info

            # Add cyclical encoding for detected patterns
            if cyclical_patterns:
                recommendation["advanced_options"]["cyclical_encoding"] = True
                recommendation["config"]["cyclical_features"] = cyclical_patterns
                recommendation["notes"].append(
                    "Cyclical patterns detected, using cyclical encoding"
                )

            self.recommendations[feature] = recommendation

    def _generate_global_recommendations(self):
        """Generate global configuration recommendations based on feature analysis."""
        # Count features by type and encoding strategy
        categorical_features = self.features_stats.get("categorical_stats", {})
        num_high_cardinality = 0
        num_hashing_recommended = 0

        # Analyze recommended encoding strategies
        for feature, stats in categorical_features.items():
            vocab_size = stats.get("size", 0)
            if vocab_size > 100:
                num_high_cardinality += 1

            # Check if we recommended hashing for this feature
            if feature in self.recommendations:
                if "HASHING" in self.recommendations[feature].get("preprocessing", []):
                    num_hashing_recommended += 1

        self.global_config = {
            "output_mode": "CONCAT",
            "use_distribution_aware": True,
            "tabular_attention": True,
            "tabular_attention_heads": self._calculate_attention_heads(),
            "tabular_attention_placement": "multi_resolution",
            "notes": [],
        }

        # Add recommendations for high-cardinality features
        if num_high_cardinality > 0:
            self.global_config["notes"].append(
                f"Dataset contains {num_high_cardinality} high-cardinality categorical features"
            )

            if num_hashing_recommended > 0:
                self.global_config["notes"].append(
                    f"Hashing recommended for {num_hashing_recommended} features to efficiently handle high cardinality"
                )

                # Consider memory optimization when many hash features
                if num_hashing_recommended >= 3:
                    self.global_config["notes"].append(
                        "Consider using hash_with_embedding=True for memory optimization with multiple hashed features"
                    )

        # Add feature interaction recommendations
        num_features = len(self.features_stats.get("categorical_stats", {})) + len(
            self.features_stats.get("numeric_stats", {})
        )
        if num_features > 5:
            self.global_config["tabular_attention"] = True
            self.global_config["notes"].append(
                "Tabular attention recommended for capturing feature interactions"
            )

    def _calculate_attention_heads(self) -> int:
        """Calculate optimal number of attention heads based on feature count."""
        feature_count = len(self.recommendations)
        return min(8, max(2, feature_count // 4))

    def generate_code_snippet(self) -> str:
        """Generate a code snippet implementing the recommendations."""
        code = []
        code.append("from kdp.features import (")
        code.append(
            "    NumericalFeature, CategoricalFeature, TextFeature, DateFeature,"
        )
        code.append("    FeatureType, CategoryEncodingOptions")
        code.append(")")
        code.append("from kdp.processor import PreprocessingModel")
        code.append("")
        code.append("# Define features based on recommendations")
        code.append("features = {")

        # Generate feature definitions
        for feature_name, rec in self.recommendations.items():
            feature_type = rec.get("feature_type", "Feature")
            config = rec.get("config", {})
            preprocessing = rec.get("preprocessing", [])

            if feature_type == "NumericalFeature":
                code.append(f'    # {", ".join(rec.get("notes", [""]))}')
                code.append(f'    "{feature_name}": NumericalFeature(')
                code.append(f'        name="{feature_name}",')
                if "FLOAT_NORMALIZED" in preprocessing:
                    code.append("        feature_type=FeatureType.FLOAT_NORMALIZED,")
                elif "FLOAT_RESCALED" in preprocessing:
                    code.append("        feature_type=FeatureType.FLOAT_RESCALED,")
                elif "FLOAT_DISCRETIZED" in preprocessing:
                    code.append("        feature_type=FeatureType.FLOAT_DISCRETIZED,")
                if config.get("use_embedding", False):
                    code.append("        use_embedding=True,")
                    code.append(
                        f"        embedding_dim={config.get('embedding_dim', 8)},"
                    )
                if config.get("bin_boundaries"):
                    boundaries_str = (
                        "[" + ", ".join(map(str, config["bin_boundaries"])) + "]"
                    )
                    code.append(f"        bin_boundaries={boundaries_str},")
                code.append("    ),")

            elif feature_type == "CategoricalFeature":
                code.append(f'    # {", ".join(rec.get("notes", [""]))}')
                code.append(f'    "{feature_name}": CategoricalFeature(')
                code.append(f'        name="{feature_name}",')

                # Set feature type
                if (
                    config.get("feature_type", "STRING_CATEGORICAL")
                    == "STRING_CATEGORICAL"
                ):
                    code.append("        feature_type=FeatureType.STRING_CATEGORICAL,")
                else:
                    code.append("        feature_type=FeatureType.INTEGER_CATEGORICAL,")

                # Handle different encoding types
                encoding = config.get("category_encoding")
                if (
                    encoding == "ONE_HOT_ENCODING"
                    or "ONE_HOT_ENCODING" in preprocessing
                ):
                    code.append(
                        "        category_encoding=CategoryEncodingOptions.ONE_HOT_ENCODING,"
                    )
                elif encoding == "HASHING" or "HASHING" in preprocessing:
                    code.append(
                        "        category_encoding=CategoryEncodingOptions.HASHING,"
                    )

                    # Add hashing specific parameters
                    if "hash_bucket_size" in config:
                        code.append(
                            f"        hash_bucket_size={config['hash_bucket_size']},"
                        )
                    if config.get("hash_with_embedding"):
                        code.append("        hash_with_embedding=True,")
                        if "embedding_size" in config:
                            code.append(
                                f"        embedding_size={config['embedding_size']},"
                            )
                    if "salt" in config:
                        code.append(f"        salt={config['salt']},")
                else:
                    # Default to embedding
                    code.append(
                        "        category_encoding=CategoryEncodingOptions.EMBEDDING,"
                    )
                    if "embedding_size" in config:
                        code.append(
                            f"        embedding_size={config['embedding_size']},"
                        )

                code.append("    ),")

            elif feature_type == "TextFeature":
                code.append(f'    # {", ".join(rec.get("notes", [""]))}')
                code.append(f'    "{feature_name}": TextFeature(')
                code.append(f'        name="{feature_name}",')
                code.append("        feature_type=FeatureType.TEXT,")
                if "max_tokens" in config:
                    code.append(f"        max_tokens={config['max_tokens']},")
                if "output_sequence_length" in config:
                    code.append(
                        f"        output_sequence_length={config['output_sequence_length']},"
                    )
                if "embedding_dim" in config:
                    code.append(f"        embedding_dim={config['embedding_dim']},")
                code.append("    ),")

            elif feature_type == "DateFeature":
                code.append(f'    # {", ".join(rec.get("notes", [""]))}')
                code.append(f'    "{feature_name}": DateFeature(')
                code.append(f'        name="{feature_name}",')
                code.append("        feature_type=FeatureType.DATE,")
                if "date_format" in config:
                    code.append(f'        date_format="{config["date_format"]}",')
                if "output_format" in config:
                    code.append(f'        output_format="{config["output_format"]}",')
                if config.get("extract", []):
                    extract_str = (
                        "[" + ", ".join([f'"{e}"' for e in config["extract"]]) + "]"
                    )
                    code.append(f"        extract={extract_str},")
                code.append("    ),")

        code.append("}")
        code.append("")
        code.append("# Create preprocessing model with recommended configuration")
        code.append("model = PreprocessingModel(")
        code.append("    features_specs=features,")

        # Add global configuration
        if "output_mode" in self.global_config:
            code.append(f"    output_mode=\"{self.global_config['output_mode']}\",")
        if self.global_config.get("tabular_attention"):
            code.append("    tabular_attention=True,")
            code.append(
                f"    tabular_attention_heads={self.global_config.get('tabular_attention_heads', 4)},"
            )
            code.append(
                f"    tabular_attention_placement=\"{self.global_config.get('tabular_attention_placement', 'multi_resolution')}\","
            )
        if self.global_config.get("use_distribution_aware"):
            code.append("    use_distribution_aware=True,")

        code.append(")")

        return "\n".join(code)


def recommend_model_configuration(features_stats: Dict[str, Any]) -> Dict[str, Any]:
    """
    Analyze dataset statistics and recommend optimal preprocessing configuration.

    Args:
        features_stats: Dictionary containing feature statistics generated by DatasetStatistics

    Returns:
        Dictionary containing recommendations for features and global configuration
    """
    advisor = ModelAdvisor(features_stats)
    recommendations = advisor.analyze_feature_stats()

    # Add code snippet to recommendations
    recommendations["code_snippet"] = advisor.generate_code_snippet()

    return recommendations
