__all__ = ['HierarchyElement', 'HierarchicalClassifierModel']

import os
import io
import math
import pandas as pd
import numpy as np
import copy as cp
import pickle as pk
import datetime as dt
import time as ti
import itertools as it
import sklearn as sk
import joblib
from sklearn import preprocessing as skp
from sklearn import model_selection as skms
from sklearn import pipeline as skpl
from sklearn import decomposition as skd
from sklearn import linear_model as sklm
from sklearn import ensemble as skle
from sklearn import neighbors as skln
from sklearn import dummy as sky
from sklearn import metrics as skm
from sklearn import calibration as skc
from sklearn.utils import validation as skuv

from ..utilities import filehelpers as fh
from ..utilities import classificationhelpers as ch
from . import embedders as emb

import logging
import warnings
from typing import Any, Dict, List, Optional, Tuple

_log = logging.getLogger(__name__)


class HierarchyElement(object):

    def __init__(self):

        self.path = None
        self.depth = None
        self.classes = None
        self.filter_value = None
        self.parent_output_feature = None
        self.output_feature = None
        self.tune_params = {}
        self.estimator = None
        self.children = []
        self.data = {}

    @property
    def X(self):
        return self.data['training']['X']

    @X.setter
    def X(self, value):
        self.data['training']['X'] = value

    @property
    def y(self):
        return self.data['training']['y']

    @y.setter
    def y(self, value):
        self.data['training']['y'] = value

    def prepare_data_subset(self, subset):
        if subset not in self.data:
            self.data[subset] = {
                'X': None,
                'y': None
            }

    def subset_has_data(self, subset):
        if subset not in self.data:
            return False

        if 'X' not in self.data[subset] or self.data[subset]['X'] is None:
            return False
        
        if 'y' not in self.data[subset] or self.data[subset]['y'] is None:
            return False

        return True

    def save_configuration(self, include_data=False):
        configuration = {
            'path': self.path,
            'depth': self.depth,
            'classes': self.classes,
            'filter_value': self.filter_value,
            'parent_output_feature': self.parent_output_feature,
            'output_feature': self.output_feature,
            'tune_params': self.tune_params,
            'estimator': self.estimator,
            'children': []
        }

        if include_data:
            configuration['data'] = self.data

        for child in self.children:
            configuration['children'].append(
                child.save_configuration(
                    include_data=include_data))

        return configuration

    @staticmethod
    def load_configuration(configuration):
        element = HierarchyElement()
        element.path = configuration['path']
        element.depth = configuration['depth']
        element.classes = configuration['classes']
        element.filter_value = configuration['filter_value']
        element.parent_output_feature = configuration['parent_output_feature']
        element.output_feature = configuration['output_feature']
        element.tune_params = configuration['tune_params']
        element.estimator = configuration['estimator']

        if 'data' in configuration:
            element.data = configuration['data']

        for child in configuration['children']:
            element.children.append(
                HierarchyElement.load_configuration(child))

        return element

    def get_child_by_filter_value(self, filter_value):
        for child in self.children:
            if child.filter_value == filter_value:
                return child
        return None

    def get_class_name(self, y):
        return self.classes[y]

    def tune_classifiers(self, test_size, classifiers, min_pca, search_n_jobs, random_state: int = 0):

        self._tune_classifiers(
            test_size=test_size,
            classifiers=classifiers,
            min_pca=min_pca,
            search_n_jobs=search_n_jobs,
            random_state=random_state)

        for child in self.children:
            child.tune_classifiers(
                test_size=test_size,
                classifiers=classifiers,
                min_pca=min_pca,
                search_n_jobs=search_n_jobs,
                random_state=random_state)

    def score(self, X, y):
        if self.estimator is None:
            return None

        return self.estimator.score(X, y)

    def predict(self, X):
        if self.estimator is None:
            return None

        return self.estimator.predict(X)

    def get_accuracy(self):
        """Returns the weighted-average test score across all nodes in this subtree,
        weighted by each node's sample count. Each node contributes independently -
        scores are not multiplied across levels."""
        pairs = []
        self._collect_accuracy_pairs(pairs)
        total_weight = sum(w for _, w in pairs)
        if total_weight == 0:
            return 0.0
        return sum(s * w for s, w in pairs) / total_weight

    def _collect_accuracy_pairs(self, pairs):
        pairs.append((self._get_accuracy(), self.get_weight()))
        for child in self.children:
            child._collect_accuracy_pairs(pairs)

    def _get_accuracy(self):
        return self.tune_params['best_score']

    def get_weight(self):
        return self.X.shape[0]

    def get_summary(self, show_parameters, output=None):
        if output is None:
            output = []

        output.append(
            self._get_summary(
                show_parameters=show_parameters))

        for element in self.children:
            element.get_summary(
                show_parameters=show_parameters,
                output=output)

        return output

    def _get_summary(self, show_parameters):
        d = {
            'path': self.path,
            'depth': self.depth,
            'nb_classes': len(self.classes),
            'nb_samples': self.X.shape[0],
            'best_classifier': self.tune_params['best_classifier'],
            'best_score': self.tune_params['best_score'],
            'fit_time': self.tune_params['classifiers'][self.tune_params['best_classifier']]['fit_time']
        }

        if show_parameters:
            d.update(
                self.tune_params['classifiers'][self.tune_params['best_classifier']]['parameters'])

        return d

    @staticmethod
    def get_pca_nb_components(min_features, max_features, max_length):

        output = []
        value = min_features

        while value < max_features and len(output) < max_length:
            output.append(value)
            value = 2*value

        if len(output) < max_length:
            output.append(max_features)

        return output

    def _tune_classifiers(
            self,
            test_size,
            classifiers,
            min_pca,
            search_n_jobs,
            random_state: int = 0):
        """Tune all classifiers for this hierarchy node and pick the best.

        Trains each classifier in *classifiers* via GridSearchCV, evaluates on the
        held-out test split, selects the best estimator, and (if the model supports
        predict_proba) runs MulticlassClassifierOptimizer calibration and threshold
        optimisation.  Results are stored in ``self.tune_params`` and
        ``self.estimator``.
        """

        _log.info('Tuning classifiers for output-feature "{0}" for path "{1}".'.format(
            self.output_feature.feature_name,
            self.path))

        if not self.subset_has_data('training'):
            _log.info('    -> Missing training data...')
            raise KeyError('Missing training data')

        self.tune_params = {}
        self.tune_params['classifiers'] = {}

        if self.subset_has_data('testing'):
            X_tr = self.X
            y_tr = self.y
            X_te = self.data['testing']['X']
            y_te = self.data['testing']['y']
        else:
            # Create training and testing partition (validation
            # split is handled by the 'StratifiedKFold' object)
            X_tr, X_te, y_tr, y_te = skms.train_test_split(
                self.X,
                self.y,
                stratify=self.y,
                shuffle=True,
                random_state=random_state,
                test_size=test_size)

        _log.info('    -> Train size: ({0}); Test size: ({1}) Number of classes ({2}).'.format(
            X_tr.shape[0],
            X_te.shape[0],
            len(self.classes)))

        if len(self.classes) > 1:

            for classifier_name, classifier_value in classifiers.items():

                _log.info('    -> Tuning "{0}"'.format(classifier_name))

                # Search grid
                dict_grid = {}

                self.tune_params['classifiers'][classifier_name] = {
                    'results': {},
                    'parameters': {},
                    'best_estimator': None
                }

                # Add standard scaller
                steps = [
                    ('std', skp.StandardScaler())
                ]

                if min_pca is not None and X_tr.shape[0] > X_tr.shape[1] and X_tr.shape[1] > min_pca:
                    steps.append(
                        ('pca', skd.PCA(
                            random_state=random_state)))

                    dict_grid['pca__n_components'] = HierarchyElement.get_pca_nb_components(
                        min_pca, X_tr.shape[1], 3)

                if classifier_name == 'SGDClassifier':
                    steps.append(
                        (classifier_name, sklm.SGDClassifier(
                            shuffle=True,
                            random_state=random_state,
                            max_iter=1000,
                            penalty='l2',
                            loss='log_loss',
                            class_weight='balanced',
                            n_jobs=1)))

                elif classifier_name == 'RandomForestClassifier':
                    steps.append(
                        (classifier_name, skle.RandomForestClassifier(
                            random_state=random_state,
                            max_depth=None,
                            class_weight='balanced',
                            n_jobs=1)))

                elif classifier_name == 'KNeighborsClassifier':
                    steps.append(
                        (classifier_name, skln.KNeighborsClassifier(
                            n_jobs=1)))

                # Create a pipeline for the work to be done
                pipe = skpl.Pipeline(steps)

                for param_name, param_value in classifier_value.items():
                    # Add the search space to the grid
                    dict_grid['{0}__{1}'.format(
                        classifier_name, param_name)] = param_value

                # create the k-fold object
                kfold = skms.StratifiedKFold(
                    n_splits=5,
                    random_state=random_state,
                    shuffle=True)

                search = skms.GridSearchCV(
                    estimator=pipe,
                    param_grid=dict_grid,
                    scoring='f1_weighted',
                    refit=True,
                    cv=kfold,
                    n_jobs=search_n_jobs)

                # capture start time
                start_time = ti.time()

                search.fit(
                    X=X_tr,
                    y=y_tr)

                elapsed_time = dt.timedelta(
                    seconds=int(round(ti.time() - start_time)))

                # capture elapsed time
                self.tune_params['classifiers'][classifier_name]['fit_time'] = elapsed_time.total_seconds()

                # capture all tuning parameters
                self.tune_params['classifiers'][classifier_name]['parameters'].update(search.best_params_)

                # keep the best estimator
                self.tune_params['classifiers'][classifier_name]['best_estimator'] = search.best_estimator_

                # capture the scores
                self.tune_params['classifiers'][classifier_name]['results'] = {
                    'validation': search.best_score_,
                    'test': search.score(
                        X=X_te, 
                        y=y_te)
                }

                _log.info('        -> Best validation score: {0:.4%}'.format(
                    self.tune_params['classifiers'][classifier_name]['results']['validation']))

                _log.info('        -> Test score: {0:.4%}'.format(
                    self.tune_params['classifiers'][classifier_name]['results']['test']))

                _log.info(
                    '        -> Tuning time: {0} ({1}s)'.format(elapsed_time, elapsed_time.total_seconds()))

        else:
            classifier_name = 'DummyClassifier'

            _log.info('    -> Tuning "{0}"'.format(classifier_name))

            self.tune_params['classifiers'][classifier_name] = {}

            estimator = skpl.Pipeline(
                steps=[
                    (classifier_name, sky.DummyClassifier(
                        strategy='most_frequent'))
                ])

            estimator.fit(
                X=X_tr,
                y=y_tr)

            self.tune_params['classifiers'][classifier_name]['results'] = {
                'validation': 1.0,
                'test': 1.0
            }

            self.tune_params['classifiers'][classifier_name]['fit_time'] = 0
            self.tune_params['classifiers'][classifier_name]['parameters'] = {}
            self.tune_params['classifiers'][classifier_name]['best_estimator'] = estimator

        # find the model with the best results.
        all_classifiers = list(self.tune_params['classifiers'].keys())
        all_results = [self.tune_params['classifiers'][classifier]
                       ['results']['test'] for classifier in all_classifiers]

        best_estimator_index = np.argmax(all_results)
        best_estimator_name = all_classifiers[best_estimator_index]

        best_estimator = self.tune_params['classifiers'][best_estimator_name]['best_estimator']

        # We need to check whether the best estimator implements the 'predict_proba' method.
        # If it does, we can 1) calibrate the estimator and 2) compute the optimized thresholds. 
        if ch.MulticlassClassifierOptimizer.optimizable_model(best_estimator):

            _log.info('    -> Optimizing "{0}"'.format(best_estimator_name))

            # Create a calibrated estimator
            optimized_estimator = ch.MulticlassClassifierOptimizer(
                model=best_estimator,
                classes=self.classes,
                scoring_function=ch.BinaryClassifierHelper.f1_score,
                random_state=random_state)

            self.estimator = optimized_estimator.fit(
                X=X_tr,
                y=y_tr)

            self.tune_params['classifiers'][best_estimator_name]['results']['train_optimized'] = optimized_estimator.score(
                X=X_tr,
                y=y_tr)

            self.tune_params['classifiers'][best_estimator_name]['results']['test_optimized'] = optimized_estimator.score(
                X=X_te,
                y=y_te)
        else:
            self.estimator = self.tune_params['classifiers'][best_estimator_name]['best_estimator']

        self.tune_params['best_score'] = self.tune_params['classifiers'][best_estimator_name]['results']['test']
        self.tune_params['best_classifier'] = best_estimator_name

        _log.info('    -> Best classifier is "{0}" with a test score of {1:.4%}'.format(
            best_estimator_name,
            self.tune_params['best_score']))

        if 'test_optimized' in self.tune_params['classifiers'][best_estimator_name]['results']:

            _log.info('    -> Optimized test score: {0:.4%}'.format(
                self.tune_params['classifiers'][best_estimator_name]['results']['test_optimized']))

    def plot_curve(self, class_name, X, y, n_bins=10):

        if isinstance(self.estimator, ch.MulticlassClassifierOptimizer):
            self.estimator.plot_curve(
                class_name=class_name,
                X=X,
                y=y,
                n_bins=n_bins)

    def plot_curves(self, X, y, n_bins=10, n_top=5):

        if isinstance(self.estimator, ch.MulticlassClassifierOptimizer):
            self.estimator.plot_curves(
                X=X,
                y=y,
                n_bins=n_bins,
                n_top=n_top)

    def get_metrics_by_class(self, X, y):

        if isinstance(self.estimator, ch.MulticlassClassifierOptimizer):
            return self.estimator.get_metrics_by_class(
                X=X,
                y=y)

    def save_data(self, filename):
        fh.save_to_npz(
            filename,
            {
                'data': self.data,
                'classes': self.classes
            })

        return filename

    def load_data(self, filename):
        d = fh.load_from_npz(
            filename)

        self.data = d['data']
        self.classes = d['classes']

    def get_min_samples_per_class(self, min_samples_per_class=10):

        for subset in self.data:
            if self.subset_has_data(subset):
                indices = np.arange(len(self.data[subset]['y']))

                for i, name in enumerate(self.classes):
                    c_indices = indices[(self.data[subset]['y'] == i)]

                    if len(c_indices) < min_samples_per_class:
                        yield {
                            'path': self.path,
                            'subset': subset,
                            'class': name,
                            'n_classes': len(c_indices)
                        }

        for child in self.children:
            for result in child.get_min_samples_per_class(
                min_samples_per_class=min_samples_per_class):
                yield result

    def set_min_samples_per_class(self, min_samples_per_class=10, random_state: int = 0):
        """Duplicate minority-class samples to meet a minimum sample count per class.

        .. warning::
            This method only needs to be called in extreme class-imbalance situations.
            The classifiers already use ``class_weight='balanced'``, which handles
            imbalance without sample duplication.  Calling both simultaneously
            **double-weights** minority classes: the loss is over-penalised *and*
            the minority samples are memorised via exact duplicates, inflating
            training metrics.

        Args:
            min_samples_per_class: Target minimum number of samples per class.
                Must be between 1 and 100,000.
            random_state: Seed for the random number generator used when
                sampling extra minority-class rows. Default 0.

        Raises:
            ValueError: if ``min_samples_per_class`` is outside the allowed range.
        """
        if min_samples_per_class < 1:
            raise ValueError(
                'min_samples_per_class must be >= 1, got {0}.'.format(
                    min_samples_per_class))
        if min_samples_per_class > 100_000:
            raise ValueError(
                'min_samples_per_class={0} is suspiciously large (> 100,000). '
                'Pass a smaller value or ensure class imbalance is handled via '
                'class_weight="balanced" instead.'.format(min_samples_per_class))

        for subset in self.data:

            if self.subset_has_data(subset):

                _X = []
                _y = []

                indices = np.arange(len(self.data[subset]['y']))

                for i, _ in enumerate(self.classes):

                    c_indices = indices[(self.data[subset]['y'] == i)]

                    if len(c_indices) > 0:
                        samples_count = c_indices.shape[0]

                        if samples_count < min_samples_per_class:
                            missing = min_samples_per_class - samples_count
                            rng = np.random.default_rng(seed=random_state)
                            extra_indices = rng.choice(c_indices, size=missing, replace=True)
                            for new_idx in extra_indices:
                                _X.append(self.data[subset]['X'][new_idx])
                                _y.append(self.data[subset]['y'][new_idx])

                if len(_X) > 0:
                    self.data[subset]['X'] = np.concatenate(
                        [self.data[subset]['X'], np.array(_X)],
                        axis=0)

                    self.data[subset]['y'] = np.concatenate(
                        [self.data[subset]['y'], np.array(_y)],
                        axis=0)

    def get_dataset_balance_status(self, subset='training'):

        df = None

        if self.subset_has_data(subset):

            a, c = np.unique(
                self.data[subset]['y'],
                return_counts=True)

            df = pd.DataFrame.from_dict({
                'label_idx': a,
                'label_names': self.classes,
                'label_count': c
            })

            df.set_index('label_idx', inplace=True)

        return df


_DEFAULT_CLASSIFIERS = {
    'SGDClassifier': {
        'alpha': np.logspace(start=-8, stop=-1, num=8)
    },
    'RandomForestClassifier': {
        'n_estimators': np.linspace(start=64, stop=256, num=4, dtype=np.int32)
    },
    'KNeighborsClassifier': {
        'n_neighbors': [4, 6, 8, 10]
    },
}


class HierarchicalClassifierModel(object):

    def __init__(
        self,
        model_name: str,
        experiment_name: str,
        input_features: list,
        output_feature_hierarchy,
        cache_path: Optional[str] = None,
        max_nb_categories: int = 128,
        random_state: int = 0,
    ):
        self.embedders = {}
        self.model_name = model_name
        self.experiment_name = experiment_name
        self.input_features = input_features
        self.output_feature_hierarchy = output_feature_hierarchy
        self.cache_path = cache_path
        self.max_nb_categories = max_nb_categories
        self.random_state = random_state
        self.hierarchy = None

        self.register_default_embedders(
            cache_path=cache_path)

    def save_configuration(self, include_data=False):
        import pwml as _pwml
        import hashlib as _hashlib
        feature_hash = _hashlib.sha256(
            str([(f.feature_name, f.feature_type)
                 for f in self.input_features]).encode()
        ).hexdigest()[:16]
        configuration = {
            'model_name': self.model_name,
            'input_features': self.input_features,
            'output_feature_hierarchy': self.output_feature_hierarchy,
            'cache_path': self.cache_path,
            'max_nb_categories': self.max_nb_categories,
            'random_state': self.random_state,
            'hierarchy': self.hierarchy.save_configuration(
                include_data=include_data),
            'experiment_name': self.experiment_name,
            '_meta': {
                'pwml_version': _pwml.__version__,
                'sklearn_version': sk.__version__,
                'input_features_hash': feature_hash,
                'saved_at_utc': dt.datetime.utcnow().isoformat() + 'Z',
            },
        }

        return configuration

    @staticmethod
    def load_configuration(configuration):
        import pwml as _pwml
        meta = configuration.get('_meta', {})
        saved_version = meta.get('pwml_version')
        if saved_version and saved_version != _pwml.__version__:
            _log.warning(
                'Model was saved with pwml %s but the running version is %s. '
                'Predictions may differ if breaking changes were introduced.',
                saved_version, _pwml.__version__)

        model = HierarchicalClassifierModel(
            model_name=configuration['model_name'],
            input_features=configuration['input_features'],
            output_feature_hierarchy=configuration['output_feature_hierarchy'],
            experiment_name=configuration['experiment_name'],
            cache_path=configuration.get('cache_path', None),
            max_nb_categories=configuration['max_nb_categories'],
            random_state=configuration.get('random_state', 0))

        model.hierarchy = HierarchyElement.load_configuration(
            configuration['hierarchy'])

        return model

    def save_data(self, filename):
        self.hierarchy.save_data(filename)

    def save_model(self, filepath: str, include_data: bool = False) -> None:
        # joblib is more efficient than raw pickle for objects containing large
        # numpy arrays (sklearn estimators), and produces smaller files.
        joblib.dump(
            self.save_configuration(include_data=include_data),
            filepath)

        _log.info('Model saved in "{0}"'.format(filepath))

    @staticmethod
    def load_model(filepath: str) -> 'HierarchicalClassifierModel':
        """Load a model previously saved with :meth:`save_model`.

        .. warning::
            This method calls :func:`joblib.load`, which deserialises arbitrary
            Python objects.  **Only load files from trusted sources.**  A
            maliciously crafted file can execute arbitrary code during
            deserialisation.

        Args:
            filepath: Path to the ``.joblib`` file written by :meth:`save_model`.

        Returns:
            The restored :class:`HierarchicalClassifierModel` instance.
        """
        _log.warning(
            'load_model() deserialises "%s" via joblib.load(). '
            'Only load files from trusted sources - maliciously crafted files '
            'can execute arbitrary code.',
            filepath)
        return HierarchicalClassifierModel.load_configuration(
            joblib.load(filepath))

    def get_accuracy(self):
        if self.hierarchy is None:
            return None

        return self.hierarchy.get_accuracy()

    def get_dataset_balance_status(self):
        if self.hierarchy is None:
            return None

        return self.hierarchy.get_dataset_balance_status()

    def get_summary(self, show_parameters=False):
        if self.hierarchy is None:
            return None

        return pd.DataFrame(
            self.hierarchy.get_summary(
                show_parameters=show_parameters)).sort_values(
                    by=['depth', 'path'],
                    ascending=[True, True])

    def get_min_samples_per_class(self, min_samples_per_class=10):
        return pd.DataFrame.from_records(
            data=self.hierarchy.get_min_samples_per_class(
                min_samples_per_class=min_samples_per_class))

    def tune_classifiers(
        self,
        test_size=.2,
        min_pca=256,
        search_n_jobs=2,
        classifiers=None,
        random_state: Optional[int] = None):

        if self.hierarchy is None:
            return None

        if classifiers is None:
            classifiers = _DEFAULT_CLASSIFIERS

        _random_state = self.random_state if random_state is None else random_state

        self.hierarchy.tune_classifiers(
            test_size=test_size,
            classifiers=classifiers,
            min_pca=min_pca,
            search_n_jobs=search_n_jobs,
            random_state=_random_state)

    def register_embedder(self, embedder):
        self.embedders[embedder.input_type] = embedder

    def register_default_embedders(self, cache_path):
        self.register_embedder(emb.CategoryEmbedder(cache_path=cache_path))
        self.register_embedder(emb.NumericEmbedder(cache_path=cache_path))
        self.register_embedder(emb.SentenceTransformerEmbedder(
            input_type='text', model_name='all-MiniLM-L6-v2', cache_path=cache_path))
        self.register_embedder(emb.SentenceTransformerEmbedder(
            input_type='longtext', model_name='all-MiniLM-L6-v2', cache_path=cache_path))

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()

    def close(self):
        """Flush and close all embedder caches. Call when done to avoid shelve corruption."""
        for embedder in self.embedders.values():
            embedder.close()

    def flush_embedders_cache(self):
        for embedder in self.embedders.values():
            embedder.flush_cache()

    def load_from_csv(self, input_file, sep=',', header=0):
        self.load_from_dataframe(
            data=pd.read_csv(
                filepath_or_buffer=input_file,
                sep=sep,
                header=header))

    @staticmethod
    def analyze_csv(input_file, sep=',', header=0):
        return HierarchicalClassifierModel.analyze_dataframe(
            data=pd.read_csv(
                filepath_or_buffer=input_file,
                sep=sep,
                header=header))

    @staticmethod
    def analyze_dataframe(data):
        return data.describe(include=['object']).transpose()

    def _precompute_embeddings(self, data):
        """Embed all rows in data and return an (n_samples, n_features) float64 matrix.

        data must have a contiguous integer index (call reset_index(drop=True) first)
        so that filtered subsets can index into the returned matrix by position.
        """
        _log.info('Pre-computing embeddings for {0} rows...'.format(len(data)))
        X = []
        for i, r in enumerate(data.itertuples(index=True)):
            r_dict = r._asdict()
            if not self.valid_input(r_dict):
                raise KeyError(
                    'Row {0}: missing one or more input features.'.format(i))
            X.append(self._get_vector(r_dict))
            if (i + 1) % 10000 == 0:
                _log.info(' -> Embedded {0} rows so far.'.format(i + 1))
        _log.info(' -> Done embedding {0} rows.'.format(len(X)))
        return np.array(X, dtype=np.float64)

    def process_data(self, data):
        """Pre-compute all embeddings and return the feature matrix alongside
        all hierarchy label columns.

        Useful for external analysis, visualisation, or batch evaluation without
        running the full hierarchical inference.

        Returns:
            X (ndarray, shape (n_samples, n_features)): Embedding matrix.
            y (dict): Mapping from output-feature name to label array (str).
        """
        data = data.reset_index(drop=True)

        y_cols = []
        output_feature = self.output_feature_hierarchy
        while output_feature is not None:
            y_cols.append(output_feature.feature_name)
            output_feature = output_feature.child_feature

        X = self._precompute_embeddings(data)
        y = {col: data[col].values for col in y_cols if col in data.columns}

        return X, y

    def predict_dataframe(
        self,
        data: pd.DataFrame,
        append_proba: bool = False,
        last_level_only: bool = False,
        min_routing_confidence: Optional[float] = None,
        profile: bool = False,
    ):
        """Run hierarchical inference on every row of a DataFrame.

        Embeddings are pre-computed once before iterating, so this is
        substantially faster than calling predict() in a Python loop.

        Returns a DataFrame indexed like data (after reset_index), with one
        column per predicted level: <feature>_predicted, <feature>_confidence,
        and optionally <feature>_proba_<class> columns when append_proba=True.
        Rows where routing was stopped due to low confidence carry a
        <feature>_routing_uncertain=True column at the last reached level.

        Args:
            profile: when True, return ``(DataFrame, latency_dict)`` where
                ``latency_dict`` maps each visited node path to the total
                wall-clock inference time (seconds) summed across all rows.
        """
        data = data.reset_index(drop=True)
        X_all = self._precompute_embeddings(data)

        _profile_dict: Optional[Dict[str, float]] = {} if profile else None

        records = []
        for i, X in enumerate(X_all):
            row_output = []
            self._predict(
                X=X,
                output=row_output,
                append_proba=append_proba,
                last_level_only=last_level_only,
                min_routing_confidence=min_routing_confidence,
                _profile=_profile_dict)

            row = {}
            for pred in row_output:
                feat = pred['feature_name']
                row['{0}_predicted'.format(feat)] = pred['value']
                if 'confidence' in pred:
                    row['{0}_confidence'.format(feat)] = pred['confidence']
                if 'routing_uncertain' in pred:
                    row['{0}_routing_uncertain'.format(feat)] = pred['routing_uncertain']
                if append_proba and 'predicted_proba' in pred:
                    for cls, prob in pred['predicted_proba'].items():
                        row['{0}_proba_{1}'.format(feat, cls)] = prob
            records.append(row)

        result_df = pd.DataFrame(records, index=data.index)
        if profile:
            return result_df, _profile_dict
        return result_df

    def evaluate(self, data: pd.DataFrame) -> Tuple[Dict[str, float], pd.DataFrame]:
        """Evaluate the model on labeled data.

        Expects data to contain the output feature columns declared in the
        output_feature_hierarchy. Returns a dict mapping each output-feature
        name to its accuracy metrics, and a DataFrame of per-row predictions.

        Returns:
            metrics (dict): {feature_name: {'accuracy': float, 'correct': int, 'total': int}}
            predictions (DataFrame): output of predict_dataframe with label columns appended.
        """
        label_cols = []
        of = self.output_feature_hierarchy
        while of is not None:
            label_cols.append(of.feature_name)
            of = of.child_feature

        predictions = self.predict_dataframe(data, append_proba=False, last_level_only=False)

        data_reset = data.reset_index(drop=True)
        metrics = {}
        for col in label_cols:
            pred_col = '{0}_predicted'.format(col)
            if pred_col in predictions.columns and col in data_reset.columns:
                true_labels = data_reset[col]
                pred_labels = predictions[pred_col]
                correct = int((true_labels == pred_labels).sum())
                total = len(true_labels)
                metrics[col] = float(correct) / float(total) if total > 0 else 0.0

        return metrics, predictions

    def cross_validate(
        self,
        data: pd.DataFrame,
        n_splits: int = 5,
        test_size: float = 0.2,
        random_state: int = 0,
        classifiers: Optional[Dict] = None,
        search_n_jobs: int = 2,
    ) -> Tuple[Dict[str, Dict[str, float]], List[Dict[str, float]]]:
        """Run stratified K-fold cross-validation on the full hierarchical model.

        A fresh model is trained and evaluated for each fold. Stratification is
        done on the top-level output feature.

        Returns:
            summary (dict): Per-feature mean and std accuracy across folds.
            fold_results (list): Per-fold metrics dicts.
        """
        top_col = self.output_feature_hierarchy.feature_name

        # Warn if any top-level class lacks enough samples for reliable stratification.
        class_counts = data[top_col].value_counts()
        sparse_classes = class_counts[class_counts < n_splits]
        if len(sparse_classes) > 0:
            warnings.warn(
                'cross_validate: the following top-level classes have fewer than '
                'n_splits={0} samples and may produce unreliable fold splits: {1}. '
                'Child-node class distributions are not validated.'.format(
                    n_splits, sparse_classes.to_dict()),
                UserWarning,
                stacklevel=2)

        kf = skms.StratifiedKFold(
            n_splits=n_splits,
            shuffle=True,
            random_state=random_state)

        fold_results = []

        for fold_idx, (train_idx, test_idx) in enumerate(
                kf.split(data, data[top_col])):
            _log.info('Cross-validation fold {0}/{1}'.format(fold_idx + 1, n_splits))

            train_data = data.iloc[train_idx].reset_index(drop=True)
            test_data = data.iloc[test_idx].reset_index(drop=True)

            fold_model = HierarchicalClassifierModel(
                model_name=self.model_name,
                experiment_name=self.experiment_name,
                input_features=cp.deepcopy(self.input_features),
                output_feature_hierarchy=cp.deepcopy(self.output_feature_hierarchy),
                cache_path=self.cache_path,
                max_nb_categories=self.max_nb_categories,
                random_state=random_state)

            fold_model.load_from_dataframe(train_data, tune=False)

            if classifiers is not None:
                fold_model.tune_classifiers(
                    test_size=test_size,
                    classifiers=classifiers,
                    search_n_jobs=search_n_jobs,
                    random_state=random_state)
            else:
                fold_model.tune_classifiers(
                    test_size=test_size,
                    search_n_jobs=search_n_jobs,
                    random_state=random_state)

            fold_metrics, _ = fold_model.evaluate(test_data)
            fold_results.append(fold_metrics)

            _log.info('  Fold {0} results: {1}'.format(fold_idx + 1, {
                k: '{0:.4%}'.format(v)
                for k, v in fold_metrics.items()
            }))

        feature_cols = list(fold_results[0].keys())
        summary = {}
        for col in feature_cols:
            accs = [r[col] for r in fold_results if col in r]
            summary[col] = {
                'mean': float(np.mean(accs)),
                'std': float(np.std(accs))
            }

        return summary, fold_results

    def load_from_dataframe(self, data: pd.DataFrame, subset: str = 'training', tune: bool = True):
        """Load data and optionally train all hierarchy nodes.

        Parameters
        ----------
        tune : bool
            When True (default) and subset=='training', automatically calls
            tune_classifiers() after building the hierarchy so the model is
            immediately ready for inference. Set to False when you want to call
            tune_classifiers() separately with custom parameters (e.g. inside
            cross_validate).
        """
        # Reset to a contiguous integer index so that filtered subsets can
        # address the pre-computed X_all matrix by position.
        data = data.reset_index(drop=True)

        if self.hierarchy is None or subset == 'training':

            for input_feature in self.input_features:
                if input_feature.feature_type == 'category':
                    input_feature.classes = data[input_feature.feature_name].unique().tolist()

                    if len(input_feature.classes) > self.max_nb_categories:
                        raise OverflowError(
                            'Too many values ({0}) for category "{1}". Maximum number of categories allowed is {2}. Switch type to "text" instead.'.format(
                                len(input_feature.classes),
                                input_feature.feature_name,
                                self.max_nb_categories))

                elif input_feature.feature_type == 'numeric':
                    input_feature.range = (data[input_feature.feature_name].min(), data[input_feature.feature_name].max())

            # Pre-compute all embeddings once so that child nodes can slice
            # directly from this matrix instead of re-embedding the same rows.
            X_all = self._precompute_embeddings(data)

            self.hierarchy = self._load_from_dataframe(
                data=data,
                X_all=X_all,
                parent_output_feature=None,
                output_feature=self.output_feature_hierarchy,
                filter_value=None,
                depth=0,
                path='.')

            if tune:
                self.tune_classifiers()

        else:
            X_all = self._precompute_embeddings(data)

            self._dataload_from_dataframe(
                element=self.hierarchy,
                data=data,
                X_all=X_all,
                subset=subset)

    def _dataload_from_dataframe(self, element, data, X_all, subset='training'):

        _log.info('Processing output-feature "{0}" for path "{1}" for subset "{2}".'.format(
            element.output_feature.feature_name,
            element.path,
            subset))

        element.prepare_data_subset(subset=subset)

        if element.parent_output_feature is None:
            _data = data
        else:
            _data = data[
                data[element.parent_output_feature.feature_name] == element.filter_value
            ]

        class_to_idx = {v: i for i, v in enumerate(element.classes)}

        row_indices = _data.index.values
        labels = _data[element.output_feature.feature_name].values

        known_mask = np.array([label in class_to_idx for label in labels])

        for unknown in np.unique(labels[~known_mask]):
            _log.info('  -> Unknown class: {0}'.format(unknown))

        valid_indices = row_indices[known_mask]
        valid_labels = labels[known_mask]

        if len(valid_indices) > 0:
            element.data[subset]['X'] = X_all[valid_indices]
            element.data[subset]['y'] = np.array(
                [class_to_idx[l] for l in valid_labels],
                dtype=np.int32)

            _log.info(' -> There are {0} classes and {1} samples.'.format(
                len(element.classes),
                element.data[subset]['y'].shape[0]))

        for child_element in element.children:
            self._dataload_from_dataframe(
                element=child_element,
                data=_data,
                X_all=X_all,
                subset=subset)

    def _load_from_dataframe(self, data, X_all, parent_output_feature, output_feature, filter_value, depth, path):

        element = HierarchyElement()
        element.parent_output_feature = parent_output_feature
        element.output_feature = output_feature
        element.filter_value = filter_value
        element.depth = depth
        element.path = path

        _log.info('Processing output-feature "{0}" for path "{1}".'.format(
            element.output_feature.feature_name,
            element.path))

        if element.output_feature is None:
            return None

        element.prepare_data_subset(subset='training')

        if element.parent_output_feature is None:
            _data = data
        else:
            _data = data[
                data[element.parent_output_feature.feature_name] == element.filter_value
            ]

        element.classes = _data[element.output_feature.feature_name].unique().tolist()

        # O(1) class lookups during label encoding
        class_to_idx = {v: i for i, v in enumerate(element.classes)}

        _log.info(' -> There are {0} classes and {1} samples.'.format(
            len(element.classes),
            _data.shape[0]))

        # Slice the pre-computed embedding matrix; _data.index addresses X_all
        # directly because the top-level dataframe was reset to a 0-based index.
        row_indices = _data.index.values
        labels = _data[element.output_feature.feature_name].values

        element.X = X_all[row_indices]
        element.y = np.array([class_to_idx[l] for l in labels], dtype=np.int32)

        # class_weight='balanced' on the classifiers handles imbalance without
        # the overfitting risk of exact sample duplication.
        _log.info(' -> Loaded {0} samples.'.format(element.y.shape[0]))

        if element.output_feature.child_feature is not None:
            for class_value in element.classes:
                child_element = self._load_from_dataframe(
                    data=_data,
                    X_all=X_all,
                    parent_output_feature=element.output_feature,
                    output_feature=element.output_feature.child_feature,
                    filter_value=class_value,
                    depth=depth+1,
                    path='{0} / {1}'.format(element.path, class_value))

                if child_element is not None:
                    element.children.append(child_element)

        return element

    def valid_input(self, input):
        valid = True

        for input_feature in self.input_features:
            if input_feature.feature_name not in input:
                _log.debug(
                    'Missing required feature "%s" in input.',
                    input_feature.feature_name)
                valid = False
                break

        return valid

    def predict(
        self,
        input: Dict[str, Any],
        append_proba: bool = False,
        last_level_only: bool = False,
        min_routing_confidence: Optional[float] = None,
        profile: bool = False,
    ):
        """Predict the hierarchy labels for a single input dict.

        Args:
            input: dict of feature_name -> value.
            append_proba: include per-class probabilities in results.
            last_level_only: only return the deepest level reached.
            min_routing_confidence: if the predicted class probability at a node
                is below this threshold, stop descending and mark the last result
                with routing_uncertain=True. Requires a model that supports
                predict_proba (i.e. not a plain DummyClassifier node).
            profile: when True, return ``(output, latency_dict)`` where
                ``latency_dict`` maps each visited node path to its wall-clock
                inference time in seconds.

        Returns:
            list of result dicts (``profile=False``, default), or
            ``(list, dict)`` tuple when ``profile=True``.
        """
        output = []
        _profile_dict: Optional[Dict[str, float]] = {} if profile else None

        if self.valid_input(input):
            self._predict(
                X=self._get_vector(input),
                hierarchy_element=None,
                output=output,
                append_proba=append_proba,
                last_level_only=last_level_only,
                min_routing_confidence=min_routing_confidence,
                _profile=_profile_dict)

        if profile:
            return output, _profile_dict
        return output

    def _predict(self, X, hierarchy_element=None, output=None, append_proba=False,
                 last_level_only=False, min_routing_confidence=None, _profile=None):
        if hierarchy_element is None:
            if self.hierarchy is None:
                return None
            else:
                hierarchy_element = self.hierarchy

        if hierarchy_element.estimator is not None:
            node_path = (
                '/'.join(str(p) for p in hierarchy_element.path)
                if hierarchy_element.path else
                hierarchy_element.output_feature.feature_name)

            # Always predict so we can route to the correct child, even when
            # last_level_only=True suppresses this node's output.
            _t0 = ti.perf_counter() if _profile is not None else None
            y = hierarchy_element.predict([X])[0]
            y_name = hierarchy_element.get_class_name(y)

            # Fetch probabilities once; used for confidence reporting and routing.
            confidence = None
            proba_values = None
            if callable(getattr(hierarchy_element.estimator, 'predict_proba', None)):
                proba_values = hierarchy_element.estimator.predict_proba([X])[0]
                confidence = float(proba_values[y])

            if _profile is not None:
                elapsed = ti.perf_counter() - _t0
                _profile[node_path] = _profile.get(node_path, 0.0) + elapsed

            if not last_level_only or len(hierarchy_element.children) == 0:
                result = {
                    'feature_name': hierarchy_element.output_feature.feature_name,
                    'value': y_name,
                }

                if confidence is not None:
                    result['confidence'] = confidence

                if append_proba and proba_values is not None:
                    proba_results = {}
                    for idx, proba_value in enumerate(proba_values):
                        proba_results[hierarchy_element.get_class_name(idx)] = float(proba_value)
                    result['predicted_proba'] = proba_results

                output.append(result)

            if len(hierarchy_element.children) > 0:
                # Soft routing: if confidence is below the threshold, stop
                # descending and mark the last appended result as uncertain.
                if (min_routing_confidence is not None
                        and confidence is not None
                        and confidence < min_routing_confidence):
                    if output:
                        output[-1]['routing_uncertain'] = True
                    return

                child = hierarchy_element.get_child_by_filter_value(y_name)

                if child is not None:
                    self._predict(
                        X=X,
                        hierarchy_element=child,
                        output=output,
                        append_proba=append_proba,
                        last_level_only=last_level_only,
                        min_routing_confidence=min_routing_confidence,
                        _profile=_profile)

    def _get_vector(self, input):
        parts = [
            self.embedders[f.feature_type].embed_data(input[f.feature_name], f)
            for f in self.input_features
        ]
        return np.concatenate(parts)
