"""
RandomRBF data generator.

Generates data using random Radial Basis Function centers.
"""

import numpy as np
from typing import Optional, Dict, Any

from tuiml.base.generators import ClassificationGenerator, GeneratedData

class RandomRBF(ClassificationGenerator):
    """RandomRBF synthetic data generator for classification.

    Generates data by first creating a random set of centers (centroids) for 
    each class. Each center is assigned a weight, a central position in the 
    feature space, and a standard deviation. Instances are generated by 
    selecting a center based on its weight and sampling from a Gaussian 
    distribution around it.

    Parameters
    ----------
    n_samples : int, default=100
        Number of data points to generate.
    n_features : int, default=10
        Number of attributes (features) for each instance.
    n_classes : int, default=2
        Number of distinct classes in the generated data.
    n_centroids : int, default=50
        Total number of RBF centers (centroids) to distribute among classes.
    random_state : int or None, default=None
        Random seed for reproducibility.

    Examples
    --------
    >>> from tuiml.datasets.generators.classification import RandomRBF
    >>> gen = RandomRBF(n_samples=1000, n_features=10, n_classes=3)
    >>> data = gen.generate()
    >>> print(data.X.shape)
    (1000, 10)
    """

    def __init__(
        self,
        n_samples: int = 100,
        n_features: int = 10,
        n_classes: int = 2,
        n_centroids: int = 50,
        random_state: Optional[int] = None
    ):
        """Initialize RandomRBF generator.

        Parameters
        ----------
        n_samples : int, default=100
            Number of samples to generate.
        n_features : int, default=10
            Number of features per instance.
        n_classes : int, default=2
            Number of distinct classes.
        n_centroids : int, default=50
            Total number of RBF centroids distributed among classes.
        random_state : int or None, default=None
            Random seed for reproducibility.
        """
        super().__init__(n_samples, n_features, n_classes, random_state)
        self.n_centroids = n_centroids

    def generate(self) -> GeneratedData:
        """Generate RandomRBF data.

        Returns
        -------
        GeneratedData
            Generated dataset with feature array X of shape
            (n_samples, n_features) and class labels y.
        """
        rng = self._rng

        # Generate centroids
        # Each centroid has: center point, standard deviation, weight, class
        centroids_per_class = self.n_centroids // self.n_classes
        remainder = self.n_centroids % self.n_classes

        centroids = []
        for c in range(self.n_classes):
            n_cent = centroids_per_class + (1 if c < remainder else 0)
            for _ in range(n_cent):
                center = rng.uniform(0, 1, self.n_features)
                std = rng.uniform(0.02, 0.15)
                weight = rng.uniform(0.1, 1.0)
                centroids.append({
                    'center': center,
                    'std': std,
                    'weight': weight,
                    'class': c
                })

        # Normalize weights
        total_weight = sum(c['weight'] for c in centroids)
        for c in centroids:
            c['weight'] /= total_weight

        # Generate samples
        X = np.zeros((self.n_samples, self.n_features))
        y = np.zeros(self.n_samples, dtype=int)

        weights = np.array([c['weight'] for c in centroids])
        cum_weights = np.cumsum(weights)

        for i in range(self.n_samples):
            # Select centroid based on weight
            r = rng.uniform()
            centroid_idx = np.searchsorted(cum_weights, r)
            centroid = centroids[centroid_idx]

            # Generate point around centroid
            direction = rng.standard_normal(self.n_features)
            direction /= np.linalg.norm(direction)
            distance = rng.normal(0, centroid['std'])

            X[i] = centroid['center'] + distance * direction
            y[i] = centroid['class']

        return GeneratedData(
            X=X,
            y=y,
            feature_names=[f"attr{i}" for i in range(self.n_features)],
            target_names=[f"class{i}" for i in range(self.n_classes)]
        )

    @classmethod
    def get_parameter_schema(cls) -> Dict[str, Any]:
        """Return JSON Schema for constructor parameters."""
        schema = super().get_parameter_schema()
        schema["n_centroids"] = {
            "type": "integer",
            "default": 50,
            "minimum": 1,
            "description": "Number of RBF centroids"
        }
        return schema
