"""Node and SplitRecord dataclasses for tree construction."""

from dataclasses import dataclass, field
from typing import Optional

import numpy as np


@dataclass
class SurrogateRecord:
    """A surrogate split that can substitute for the primary split."""

    feature: int
    threshold: float
    mirror: bool  # True if surrogate sends samples opposite direction
    lambda_: float  # Predictive measure of association
    missing_direction: int  # -1 left, +1 right (fallback if surrogate feature also missing)


@dataclass
class SplitRecord:
    """Records all information about a node's primary split."""

    feature: int = -1
    threshold: float = 0.0
    gain: float = -np.inf
    missing_direction: int = -1  # -1 left, +1 right
    n_left: int = 0
    n_right: int = 0
    impurity_left: float = 0.0
    impurity_right: float = 0.0
    surrogates: list = field(default_factory=list)  # list[SurrogateRecord]


@dataclass
class Node:
    """A node in the decision tree."""

    node_id: int = 0
    depth: int = 0
    is_leaf: bool = True
    n_samples: int = 0
    impurity: float = 0.0
    weighted_n_samples: float = 0.0

    # Leaf info
    value: Optional[np.ndarray] = None  # class counts (clf) or [sum, count] (reg)

    # Internal node info
    split: Optional[SplitRecord] = None
    left_child: Optional["Node"] = None
    right_child: Optional["Node"] = None

    def predict_single(self, x: np.ndarray) -> np.ndarray:
        """Route a single sample through the tree, handling missing values."""
        if self.is_leaf:
            return self.value

        feature = self.split.feature
        val = x[feature]

        if np.isnan(val):
            # Try surrogates in order
            for surr in self.split.surrogates:
                surr_val = x[surr.feature]
                if not np.isnan(surr_val):
                    if surr.mirror:
                        go_left = surr_val > surr.threshold
                    else:
                        go_left = surr_val <= surr.threshold
                    if go_left:
                        return self.left_child.predict_single(x)
                    else:
                        return self.right_child.predict_single(x)

            # All surrogates also missing → use learned missing direction
            if self.split.missing_direction <= 0:
                return self.left_child.predict_single(x)
            else:
                return self.right_child.predict_single(x)

        # Primary feature available
        if val <= self.split.threshold:
            return self.left_child.predict_single(x)
        else:
            return self.right_child.predict_single(x)
