from pydantic import Field, model_validator, PrivateAttr
from typing import List, Optional, Dict, Any, Union

from ...core.base.hyperparameters_base import ModelHyperparameters


class XGBoostModelHyperparameters(ModelHyperparameters):
    """
    Hyperparameters for the XGBoost model training, extending the base ModelHyperparameters.

    Fields are organized into three tiers:
    1. Tier 1: Essential User Inputs - fields that users must explicitly provide
    2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
    3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
    """

    # ===== Essential User Inputs (Tier 1) =====
    # These are fields that users must explicitly provide

    # Most essential XGBoost hyperparameters
    num_round: int = Field(description="The number of boosting rounds for XGBoost.")

    max_depth: int = Field(description="Maximum depth of a tree.")

    # ===== System Inputs with Defaults (Tier 2) =====
    # These are fields with reasonable defaults that users can override

    # Override model_class from base
    model_class: str = Field(
        default="xgboost", description="Model class identifier, set to XGBoost."
    )

    min_child_weight: float = Field(
        default=1.0,
        description="Minimum sum of instance weight (hessian) needed in a child.",
    )

    # General Parameters
    booster: str = Field(
        default="gbtree",
        description="Specify which booster to use: gbtree, gblinear or dart.",
    )

    # Booster Parameters
    eta: float = Field(
        default=0.3,
        ge=0.0,
        le=1.0,
        description="Step size shrinkage used in update to prevents overfitting. Alias: learning_rate.",
    )

    gamma: float = Field(
        default=0.0,
        ge=0.0,
        description="Minimum loss reduction required to make a further partition on a leaf node of the tree.",
    )

    max_delta_step: float = Field(
        default=0.0,
        description="Maximum delta step we allow each tree's weight estimation to be. If 0, no constraint.",
    )

    subsample: float = Field(
        default=1.0,
        gt=0.0,
        le=1.0,
        description="Subsample ratio of the training instances.",
    )

    colsample_bytree: float = Field(
        default=1.0,
        gt=0.0,
        le=1.0,
        description="Subsample ratio of columns when constructing each tree.",
    )

    colsample_bylevel: float = Field(
        default=1.0,
        gt=0.0,
        le=1.0,
        description="Subsample ratio of columns for each level.",
    )

    colsample_bynode: float = Field(
        default=1.0,
        gt=0.0,
        le=1.0,
        description="Subsample ratio of columns for each split.",
    )

    lambda_xgb: float = Field(
        default=1.0,
        ge=0.0,
        description="L2 regularization term on weights. Alias: reg_lambda.",
    )

    alpha_xgb: float = Field(
        default=0.0,
        ge=0.0,
        description="L1 regularization term on weights. Alias: reg_alpha.",
    )

    tree_method: str = Field(
        default="auto", description="The tree construction algorithm used in XGBoost."
    )

    sketch_eps: Optional[float] = Field(
        default=None,
        ge=0.0,
        le=1.0,
        description="For tree_method 'approx'. Approximately (1 / sketch_eps) buckets are made.",
    )

    scale_pos_weight: float = Field(
        default=1.0,
        description="Control the balance of positive and negative weights, useful for unbalanced classes.",
    )

    num_parallel_tree: Optional[int] = Field(
        default=None,
        ge=1,
        description="Number of parallel trees constructed during each iteration. Used for random forests.",
    )

    # Learning Task Parameters
    base_score: Optional[float] = Field(
        default=None,
        description="The initial prediction score of all instances, global bias.",
    )

    seed: Optional[int] = Field(default=None, description="Random number seed.")

    early_stopping_rounds: Optional[int] = Field(
        default=None,
        ge=1,
        description="Activates early stopping. Requires eval_metric.",
    )

    # ===== Derived Fields (Tier 3) =====
    # These are fields calculated from other fields

    _objective: Optional[str] = PrivateAttr(default=None)
    _eval_metric: Optional[Union[str, List[str]]] = PrivateAttr(default=None)

    model_config = ModelHyperparameters.model_config.copy()
    model_config.update(
        {"extra": "allow"}
    )  # Changed from "forbid" to "allow" to fix circular reference handling

    # Public read-only properties for derived fields

    @property
    def objective(self) -> str:
        """Get objective derived from is_binary."""
        if self._objective is None:
            self._objective = "binary:logistic" if self.is_binary else "multi:softmax"
        return self._objective

    @property
    def eval_metric(self) -> List[str]:
        """Get evaluation metrics derived from is_binary."""
        if self._eval_metric is None:
            self._eval_metric = (
                ["logloss", "auc"] if self.is_binary else ["mlogloss", "merror"]
            )
        return self._eval_metric

    @model_validator(mode="after")
    def validate_xgboost_hyperparameters(self) -> "XGBoostModelHyperparameters":
        """Validate XGBoost-specific hyperparameters"""
        # Call the base model validator first
        super().validate_dimensions()

        # Initialize derived fields
        self._objective = "binary:logistic" if self.is_binary else "multi:softmax"
        self._eval_metric = (
            ["logloss", "auc"] if self.is_binary else ["mlogloss", "merror"]
        )

        # Validate multiclass parameters
        if self._objective.startswith("multi:") and self.num_classes < 2:
            raise ValueError(
                f"For multiclass objective '{self._objective}', 'num_classes' must be >= 2. "
                f"Current num_classes: {self.num_classes}"
            )

        # Validate early stopping configuration
        if self.early_stopping_rounds is not None and not self._eval_metric:
            raise ValueError(
                "'early_stopping_rounds' requires 'eval_metric' to be set."
            )

        # Validate GPU usage
        if self.tree_method == "gpu_hist" and self.device == -1:
            print(
                f"Warning: tree_method is '{self.tree_method}' but device is CPU (-1). "
                f"Ensure SageMaker instance is GPU for gpu_hist."
            )

        return self

    def get_public_init_fields(self) -> Dict[str, Any]:
        """
        Override get_public_init_fields to include XGBoost-specific fields.
        Gets a dictionary of public fields suitable for initializing a child config.
        """
        # Get fields from parent class
        base_fields = super().get_public_init_fields()

        # Add derived fields that should be exposed
        derived_fields = {"objective": self.objective, "eval_metric": self.eval_metric}

        # Combine (derived fields take precedence if overlap)
        return {**base_fields, **derived_fields}
