"""Utility functions for GeoRF: synthetic data generation and parallel tree fitting.

Design
------
Each tree's expansion + fitting runs as a single unit inside a loky worker
process.  This maximises parallelism while remaining fully deterministic:

* HVRT.expand() passes ``self.random_state`` to the KDE sampler, so every
  call on the same instance returns identical output.  We shallow-copy the
  HVRT model and set a unique ``random_state`` on the copy before expanding;
  the original is never mutated.

* sklearn's DecisionTreeClassifier can produce non-deterministic tie-breaking
  when features have nearly equal impurity gains (due to CPU/BLAS state).
  Passing ``random_state=tree_seed`` to the tree eliminates this.

* Both mechanisms together make each worker function a pure deterministic
  transformation of (seed, fitted_hvrt_models) → fitted_tree, regardless
  of n_jobs value.

Generation strategy
-------------------
expand() is called without an explicit generation_strategy, deferring to
HVRT's bandwidth='auto' resolution.  Auto-mode selects a narrow Gaussian
kernel (bandwidth=0.1) per partition — no covariance inversion, no CDF
normalisation, robust to any feature correlation structure and ~40x faster
than the legacy univariate_kde_copula path.
"""
import copy

import numpy as np
from sklearn.neighbors import NearestNeighbors
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _sanitise_synthetic(X_syn, X_syn_z):
    """Remove rows from synthetic arrays that contain NaN or Inf.

    HVRT's univariate_kde_copula strategy can produce NaN values when the
    training set is very small (< ~150 samples) and a feature's CDF
    collapses to zero — the divide-by-zero propagates through ``_to_z()``.

    Both arrays share the same row indices, so the same mask is applied to
    each.  The caller is responsible for handling the (rare) case where all
    rows are dropped.

    Parameters
    ----------
    X_syn   : ndarray (n, p) — synthetic features in original space
    X_syn_z : ndarray (n, p) — z-transformed synthetic features

    Returns
    -------
    X_syn_clean   : ndarray (m, p), m <= n
    X_syn_z_clean : ndarray (m, p)
    """
    finite_mask = np.isfinite(X_syn).all(axis=1) & np.isfinite(X_syn_z).all(axis=1)
    if finite_mask.all():
        return X_syn, X_syn_z
    return X_syn[finite_mask], X_syn_z[finite_mask]


def _expand_with_seed(hvrt_model, n, seed):
    """Generate n synthetic samples using a specific random seed.

    Creates a shallow copy of *hvrt_model*, sets ``random_state=seed`` on
    the copy so the original is never mutated, then calls expand().

    No generation_strategy is specified; HVRT's bandwidth='auto' auto-selects
    a robust, fast per-partition Gaussian kernel.
    """
    h = copy.copy(hvrt_model)
    h.random_state = int(seed)
    return h.expand(n=n)


# ---------------------------------------------------------------------------
# Shared helper: synthetic dataset generation (expand + kNN y-assignment)
# ---------------------------------------------------------------------------

def _make_synthetic_dataset(hvrt_model, X_real_z, y_real, n_synthetic, seed, k_neighbors=5):
    """Generate (X_syn, y_syn) via HVRT expansion + kNN target assignment.

    Parameters
    ----------
    hvrt_model : HVRT
        A fitted HVRT instance.  Never mutated; a shallow copy is made internally.
    X_real_z : ndarray of shape (n_real, n_features)
        z-transformed real X — i.e. ``hvrt_model._to_z(X_real)``, precomputed
        by the caller to avoid redundant work.
    y_real : ndarray of shape (n_real,)
        Real target values.
    n_synthetic : int
        Number of synthetic samples to generate.
    seed : int
        Random seed forwarded to ``_expand_with_seed``.
    k_neighbors : int, default=5
        Nearest neighbours used for weighted target interpolation.

    Returns
    -------
    X_syn : ndarray of shape (n_synthetic, n_features)
        Synthetic features in original (un-z-transformed) space.
    y_syn : ndarray of shape (n_synthetic,)
        kNN-interpolated target values.
    """
    X_syn = _expand_with_seed(hvrt_model, n_synthetic, seed)
    X_syn_z = hvrt_model._to_z(X_syn)
    X_syn, X_syn_z = _sanitise_synthetic(X_syn, X_syn_z)

    nn = NearestNeighbors(n_neighbors=k_neighbors)
    nn.fit(X_real_z)
    dists, idxs = nn.kneighbors(X_syn_z)

    weights = 1.0 / (dists + 1e-10)
    weights /= weights.sum(axis=1, keepdims=True)
    y_syn = np.sum(weights * y_real[idxs], axis=1)

    return X_syn, y_syn


# ---------------------------------------------------------------------------
# Worker functions  (called inside loky subprocesses by joblib.Parallel)
# ---------------------------------------------------------------------------

def fit_one_tree_clf(
    tree_seed, hvrt_models, class_encoded_labels, n_per_class_list, tree_kwargs
):
    """Generate synthetic data + fit a DecisionTreeClassifier (classification).

    Determinism is guaranteed by:
    - HVRT shallow-copy + ``random_state=tree_seed`` (KDE sampling)
    - ``random_state=tree_seed`` on the DecisionTreeClassifier (tie-breaking)

    Parameters
    ----------
    tree_seed : int
        Unique seed for this tree.
    hvrt_models : dict {encoded_label -> HVRT}
        One fitted HVRT model per class.
    class_encoded_labels : list[int]
    n_per_class_list : list[int]
        Synthetic samples to draw per class.
    tree_kwargs : dict
        Forwarded to DecisionTreeClassifier (must NOT include random_state).
    """
    X_parts, y_parts = [], []
    for enc_label, n_cls in zip(class_encoded_labels, n_per_class_list):
        X_syn = _expand_with_seed(hvrt_models[enc_label], n_cls, tree_seed)
        X_parts.append(X_syn)
        y_parts.append(np.full(n_cls, enc_label, dtype=np.intp))

    X_syn = np.vstack(X_parts)
    y_syn = np.concatenate(y_parts)

    rng = np.random.default_rng(tree_seed)
    perm = rng.permutation(len(y_syn))
    X_syn, y_syn = X_syn[perm], y_syn[perm]

    tree = DecisionTreeClassifier(**tree_kwargs, random_state=tree_seed)
    tree.fit(X_syn, y_syn)
    return tree


def fit_one_tree_reg(
    tree_seed, hvrt_model, X_real_z, y_real, n_samples, k_neighbors, tree_kwargs
):
    """Generate synthetic data + fit a DecisionTreeRegressor (regression).

    Synthetic X from HVRT.expand(); synthetic y via k-NN weighted
    interpolation in z-space from the real (X, y) pairs.

    Parameters
    ----------
    tree_seed : int
    hvrt_model : HVRT
    X_real_z : ndarray (n_real, n_features)
        Precomputed z-transformed real X.
    y_real : ndarray (n_real,)
    n_samples : int
    k_neighbors : int
    tree_kwargs : dict
    """
    X_syn = _expand_with_seed(hvrt_model, n_samples, tree_seed)
    X_syn_z = hvrt_model._to_z(X_syn)
    X_syn, X_syn_z = _sanitise_synthetic(X_syn, X_syn_z)

    nn = NearestNeighbors(n_neighbors=k_neighbors)
    nn.fit(X_real_z)
    dists, idxs = nn.kneighbors(X_syn_z)

    weights = 1.0 / (dists + 1e-10)
    weights /= weights.sum(axis=1, keepdims=True)
    y_syn = np.sum(weights * y_real[idxs], axis=1)

    tree = DecisionTreeRegressor(**tree_kwargs, random_state=tree_seed)
    tree.fit(X_syn, y_syn)
    return tree
