"""高斯混合数据生成工具

提供灵活的链式API来生成具有多种特性的混合分布数据，支持：
- 多分量高斯混合
- 均匀分布组件
- 指数分布组件
- 局部噪声和偏斜扰动
- 全局偏斜度控制
- 噪声添加
- 数据范围约束
- 可复现的随机生成
"""

import typing

import warnings
from typing import List, Literal, Optional, Tuple, Union

import numpy as np
from numpy.typing import NDArray
from scipy import stats

# 类型别名
ComponentConfig = dict
TransformConfig = dict
GenerationResult = Union[
    NDArray,
    Tuple[NDArray, NDArray],
    Tuple[NDArray, List[NDArray]],
    Tuple[NDArray, NDArray, List[NDArray]],
]


def _generate_data_only(generator: 'GaussianMixtureGenerator') -> NDArray:
    result = generator.generate()
    if not isinstance(result, np.ndarray):
        raise AssertionError('Expected ndarray result when no extras are requested')
    return result


class GaussianMixtureGenerator:
    """混合分布数据生成器

    支持链式调用的灵活API，可以逐步添加分量、噪声、偏斜等特性。
    支持高斯分布、均匀分布、指数分布等多种分布类型的混合。

    使用示例
    ----------
    ```python
    # 简单双峰分布
    data = (GaussianMixtureGenerator(1000)
            .add_bimodal(separation=6.0)
            .add_noise(0.05)
            .generate())

    # 复杂多分量with偏斜
    data = (GaussianMixtureGenerator(2000, random_state=42)
            .add_component(mean=-2, std=0.8, weight=0.3)
            .add_component(mean=1, std=1.2, weight=0.5)
            .add_component(mean=4, std=0.6, weight=0.2)
            .add_skewness(1.5)
            .add_noise(level=0.1, noise_type='uniform')
            .set_range(-10, 10)
            .generate())

    # 混合不同分布类型
    data = (GaussianMixtureGenerator(1500, random_state=42)
            .add_component(mean=0, std=1, weight=0.4, local_noise=0.1)
            .add_uniform_component(low=2, high=5, weight=0.3)
            .add_exponential_component(scale=1.5, offset=6, weight=0.3)
            .generate())
    ```
    """

    def __init__(self, n_samples: int, random_state: Optional[int] = None):
        """初始化高斯混合生成器

        Parameters
        ----------
        - `n_samples` : `int`
            - 要生成的样本数量
        - `random_state` : `Optional[int]`, optional
            - 随机种子，用于可复现的结果, by default None
        """
        if n_samples <= 0:
            raise ValueError('样本数量必须大于0')

        self.n_samples = n_samples
        self.random_state = random_state
        self.components: List[ComponentConfig] = []
        self.transforms: List[TransformConfig] = []
        self._rng = np.random.RandomState(random_state)

    def add_component(
        self,
        mean: float,
        std: float,
        weight: float = 1.0,
        local_noise: float = 0.0,
        local_skewness: float = 0.0,
    ) -> 'GaussianMixtureGenerator':
        """添加一个高斯分量

        Parameters
        ----------
        - `mean` : `float`
            - 高斯分量的均值
        - `std` : `float`
            - 高斯分量的标准差
        - `weight` : `float`, optional
            - 分量的权重（相对重要性）, by default 1.0
        - `local_noise` : `float`, optional
            - 该分量的局部噪声强度（相对于该分量标准差的比例）, by default 0.0
        - `local_skewness` : `float`, optional
            - 该分量的局部偏斜度, by default 0.0

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if std <= 0:
            raise ValueError('标准差必须大于0')
        if weight <= 0:
            raise ValueError('权重必须大于0')
        if local_noise < 0:
            raise ValueError('局部噪声强度必须非负')

        self.components.append(
            {
                'type': 'gaussian',
                'mean': mean,
                'std': std,
                'weight': weight,
                'local_noise': local_noise,
                'local_skewness': local_skewness,
            }
        )
        return self

    def add_bimodal(
        self,
        separation: float = 5.0,
        std: float = 1.0,
        weights: Tuple[float, float] = (0.5, 0.5),
    ) -> 'GaussianMixtureGenerator':
        """快速添加双峰分布

        Parameters
        ----------
        - `separation` : `float`, optional
            - 两个峰之间的距离, by default 5.0
        - `std` : `float`, optional
            - 两个峰的标准差, by default 1.0
        - `weights` : `Tuple[float, float]`, optional
            - 两个峰的相对权重, by default (0.5, 0.5)

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if std <= 0:
            raise ValueError('标准差必须大于0')
        if len(weights) != 2 or any(w <= 0 for w in weights):
            raise ValueError('权重必须是两个正数')

        mean1 = -separation / 2
        mean2 = separation / 2

        self.add_component(mean1, std, weights[0])
        self.add_component(mean2, std, weights[1])
        return self

    def add_multimodal(
        self, n_modes: int, spread: float = 10.0, std: float = 1.0
    ) -> 'GaussianMixtureGenerator':
        """快速添加多峰分布

        Parameters
        ----------
        - `n_modes` : `int`
            - 峰的数量
        - `spread` : `float`, optional
            - 峰之间的总分布范围, by default 10.0
        - `std` : `float`, optional
            - 每个峰的标准差, by default 1.0

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if n_modes < 2:
            raise ValueError('峰的数量必须至少为2')
        if std <= 0:
            raise ValueError('标准差必须大于0')

        means = np.linspace(-spread / 2, spread / 2, n_modes)
        weight = 1.0 / n_modes

        for mean in means:
            self.add_component(mean, std, weight)
        return self

    def add_exponential_component(
        self, scale: float, weight: float = 1.0, offset: float = 0.0
    ) -> 'GaussianMixtureGenerator':
        """
        添加指数分布分量

        Parameters
        ----------
        - `scale` : `float`
            - 指数分布的尺度参数
        - `weight` : `float`, optional
            - 该分量的权重, by default 1.0
        - `offset` : `float`, optional
            - 分布的偏移量, by default 0.0

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if scale <= 0:
            raise ValueError('指数分布尺度参数必须大于0')
        if weight <= 0:
            raise ValueError('权重必须大于0')

        component = {
            'type': 'exponential',
            'scale': scale,
            'weight': weight,
            'offset': offset,
        }
        self.components.append(component)
        return self

    def add_uniform_component(
        self,
        low: float,
        high: float,
        weight: float = 1.0,
        local_noise: float = 0.0,
        local_skewness: float = 0.0,
    ) -> 'GaussianMixtureGenerator':
        """添加均匀分布分量

        Parameters
        ----------
        - `low` : `float`
            - 均匀分布的下界
        - `high` : `float`
            - 均匀分布的上界
        - `weight` : `float`, optional
            - 该分量的权重, by default 1.0
        - `local_noise` : `float`, optional
            - 该分量的局部噪声强度（相对于该分量标准差的比例）, by default 0.0
        - `local_skewness` : `float`, optional
            - 该分量的局部偏斜度, by default 0.0

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if low >= high:
            raise ValueError('均匀分布下界必须小于上界')
        if weight <= 0:
            raise ValueError('权重必须大于0')
        if local_noise < 0:
            raise ValueError('局部噪声强度必须非负')

        component = {
            'type': 'uniform',
            'low': low,
            'high': high,
            'weight': weight,
            'local_noise': local_noise,
            'local_skewness': local_skewness,
        }
        self.components.append(component)
        return self

    def add_skewness(
        self, skewness: float, method: Literal['transform', 'mixture'] = 'transform'
    ) -> 'GaussianMixtureGenerator':
        """添加偏斜度

        Parameters
        ----------
        - `skewness` : `float`
            - 目标偏斜度值。正值表示右偏，负值表示左偏
        - `method` : `Literal['transform', 'mixture']`, optional
            - 偏斜实现方法, by default 'transform'

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        self.transforms.append(
            {'type': 'skewness', 'skewness': skewness, 'method': method}
        )
        return self

    def add_noise(
        self,
        level: float,
        noise_type: Literal['gaussian', 'uniform', 'laplace'] = 'gaussian',
        relative: bool = True,
    ) -> 'GaussianMixtureGenerator':
        """添加噪声

        Parameters
        ----------
        - `level` : `float`
            - 噪声强度。如果relative=True，则为相对于数据标准差的比例
        - `noise_type` : `Literal['gaussian', 'uniform', 'laplace']`, optional
            - 噪声类型, by default 'gaussian'
        - `relative` : `bool`, optional
            - 是否为相对噪声强度, by default True

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if level < 0:
            raise ValueError('噪声强度必须非负')

        self.transforms.append(
            {
                'type': 'noise',
                'level': level,
                'noise_type': noise_type,
                'relative': relative,
            }
        )
        return self

    def set_range(
        self, min_val: float, max_val: float, method: Literal['clip', 'rescale'] = 'clip'
    ) -> 'GaussianMixtureGenerator':
        """设置数据范围约束

        Parameters
        ----------
        - `min_val` : `float`
            - 最小值
        - `max_val` : `float`
            - 最大值
        - `method` : `Literal['clip', 'rescale']`, optional
            - 范围约束方法, by default 'clip'

        Returns
        ----------
        - `GaussianMixtureGenerator`
            - 返回自身以支持链式调用
        """
        if min_val >= max_val:
            raise ValueError('最小值必须小于最大值')

        self.transforms.append(
            {'type': 'range', 'min_val': min_val, 'max_val': max_val, 'method': method}
        )
        return self

    def generate(
        self, return_components: bool = False, return_labels: bool = False
    ) -> GenerationResult:
        """生成最终数据

        Parameters
        ----------
        - `return_components` : `bool`, optional
            - 是否返回各分量的独立数据, by default False
        - `return_labels` : `bool`, optional
            - 是否返回每个样本所属的分量标签, by default False

        Returns
        ----------
        - `GenerationResult`
            - 根据参数返回生成的数据及可选的附加信息
        """
        if not self.components:
            raise ValueError('必须至少添加一个分量')

        # 1. 生成基础GMM
        data, labels, components = self._generate_base_mixture()

        # 2. 应用各种变换
        for transform in self.transforms:
            data = self._apply_transform(data, transform)

        # 3. 构造返回结果
        if return_labels and return_components:
            return data, labels, components
        if return_labels:
            return data, labels
        if return_components:
            return data, components
        return data

    def _generate_base_mixture(self) -> Tuple[NDArray, NDArray, List[NDArray]]:
        """生成基础高斯混合数据"""
        # 计算权重归一化
        total_weight = sum(comp['weight'] for comp in self.components)
        normalized_weights = [comp['weight'] / total_weight for comp in self.components]

        # 分配样本数给各分量
        samples_per_component = self._rng.multinomial(self.n_samples, normalized_weights)

        all_data = []
        all_labels = []
        component_data = []

        for i, (comp, n_samples_comp) in enumerate(
            zip(self.components, samples_per_component)
        ):
            if n_samples_comp > 0:
                if comp['type'] == 'gaussian':
                    comp_data = self._rng.normal(
                        comp['mean'], comp['std'], n_samples_comp
                    )
                elif comp['type'] == 'exponential':
                    comp_data = (
                        self._rng.exponential(comp['scale'], n_samples_comp)
                        + comp['offset']
                    )
                elif comp['type'] == 'uniform':
                    comp_data = self._rng.uniform(
                        comp['low'], comp['high'], n_samples_comp
                    )
                else:
                    raise ValueError(f"未知的分量类型: {comp['type']}")

                # 应用局部扰动
                comp_data = self._apply_local_disturbances(comp_data, comp)

                all_data.append(comp_data)
                all_labels.extend([i] * n_samples_comp)
                component_data.append(comp_data)
            else:
                component_data.append(np.array([]))

        # 合并并打乱
        data = np.concatenate(all_data) if all_data else np.array([])
        labels = np.array(all_labels)

        # 打乱数据和标签（保持对应关系）
        if len(data) > 0:
            shuffle_indices = self._rng.permutation(len(data))
            data = data[shuffle_indices]
            labels = labels[shuffle_indices]

        return data, labels, component_data

    def _apply_local_disturbances(self, data: NDArray, comp: ComponentConfig) -> NDArray:
        """对单个分量应用局部扰动（噪声和偏斜）"""
        # 应用局部偏斜
        local_skewness = comp.get('local_skewness', 0.0)
        if local_skewness != 0:
            data = self._apply_local_skewness(data, local_skewness)

        # 应用局部噪声
        local_noise = comp.get('local_noise', 0.0)
        if local_noise > 0:
            data = self._apply_local_noise(data, local_noise)

        return data

    def _apply_local_skewness(self, data: NDArray, skewness: float) -> NDArray:
        """对数据应用局部偏斜变换"""
        if skewness == 0:
            return data

        # 使用Box-Cox类似的变换
        data_shifted = data - np.min(data) + 1

        if skewness > 0:
            # 右偏：使用幂变换
            power = 1 / (1 + skewness)
            transformed = np.power(data_shifted, power)
        else:
            # 左偏：使用对数变换的变体
            transformed = np.log(data_shifted + abs(skewness))

        # 重新缩放到原始数据的统计特性
        transformed = (transformed - np.mean(transformed)) / np.std(transformed)
        transformed = transformed * np.std(data) + np.mean(data)

        return transformed

    def _apply_local_noise(self, data: NDArray, noise_level: float) -> NDArray:
        """对数据应用局部噪声"""
        # 噪声强度相对于该分量的标准差
        noise_std = noise_level * np.std(data)
        noise = self._rng.normal(0, noise_std, len(data))
        return data + noise

    def _apply_transform(self, data: NDArray, transform: TransformConfig) -> NDArray:
        """应用指定的变换"""
        if transform['type'] == 'skewness':
            return self._apply_skewness(data, transform)
        elif transform['type'] == 'noise':
            return self._apply_noise(data, transform)
        elif transform['type'] == 'range':
            return self._apply_range(data, transform)
        else:
            raise ValueError(f"未知的变换类型: {transform['type']}")

    def _apply_skewness(self, data: NDArray, transform: TransformConfig) -> NDArray:
        """应用偏斜度变换"""
        skewness = transform['skewness']
        method = transform['method']

        if method == 'transform':
            # 使用Box-Cox类似的变换
            if skewness == 0:
                return data

            # 先标准化数据到正数范围
            data_shifted = data - np.min(data) + 1

            if skewness > 0:
                # 右偏：使用幂变换
                power = 1 / (1 + skewness)
                transformed = np.power(data_shifted, power)
            else:
                # 左偏：使用对数变换的变体
                transformed = np.log(data_shifted + abs(skewness))

            # 重新缩放到原始数据的统计特性
            transformed = (transformed - np.mean(transformed)) / np.std(transformed)
            transformed = transformed * np.std(data) + np.mean(data)

            return transformed

        elif method == 'mixture':
            # 通过不对称混合实现偏斜
            warnings.warn('mixture方法的偏斜度实现尚未完全优化', UserWarning)
            return data
        else:
            raise ValueError(f'未知的偏斜方法: {method}')

    def _apply_noise(self, data: NDArray, transform: TransformConfig) -> NDArray:
        """应用噪声"""
        level = transform['level']
        noise_type = transform['noise_type']
        relative = transform['relative']

        if level == 0:
            return data

        # 计算噪声强度
        if relative:
            noise_std = level * np.std(data)
        else:
            noise_std = level

        # 生成噪声
        if noise_type == 'gaussian':
            noise = self._rng.normal(0, noise_std, len(data))
        elif noise_type == 'uniform':
            # 均匀噪声：范围设为±sqrt(3)*std以匹配方差
            noise_range = noise_std * np.sqrt(3)
            noise = self._rng.uniform(-noise_range, noise_range, len(data))
        elif noise_type == 'laplace':
            # 拉普拉斯噪声：scale参数设为std/sqrt(2)
            noise_scale = noise_std / np.sqrt(2)
            noise = self._rng.laplace(0, noise_scale, len(data))
        else:
            raise ValueError(f'未知的噪声类型: {noise_type}')

        return data + noise

    def _apply_range(self, data: NDArray, transform: TransformConfig) -> NDArray:
        """应用范围约束"""
        min_val = transform['min_val']
        max_val = transform['max_val']
        method = transform['method']

        if method == 'clip':
            return np.clip(data, min_val, max_val)
        elif method == 'rescale':
            # 线性重缩放到指定范围
            data_min, data_max = np.min(data), np.max(data)
            if data_max == data_min:
                return np.full_like(data, (min_val + max_val) / 2)

            scaled = (data - data_min) / (data_max - data_min)
            return scaled * (max_val - min_val) + min_val
        else:
            raise ValueError(f'未知的范围方法: {method}')


# 便捷函数
def create_bimodal_data(
    n_samples: int,
    separation: float = 5.0,
    std: float = 1.0,
    weights: Tuple[float, float] = (0.5, 0.5),
    noise_level: float = 0.0,
    random_state: Optional[int] = None,
) -> NDArray:
    """快速创建双峰分布数据

    Parameters
    ----------
    - `n_samples` : `int`
        - 样本数量
    - `separation` : `float`, optional
        - 峰间距离, by default 5.0
    - `std` : `float`, optional
        - 峰的标准差, by default 1.0
    - `weights` : `Tuple[float, float]`, optional
        - 峰的权重, by default (0.5, 0.5)
    - `noise_level` : `float`, optional
        - 噪声强度, by default 0.0
    - `random_state` : `Optional[int]`, optional
        - 随机种子, by default None

    Returns
    ----------
    - `NDArray`
        - 生成的双峰分布数据
    """
    generator = GaussianMixtureGenerator(n_samples, random_state)
    generator.add_bimodal(separation, std, weights)

    if noise_level > 0:
        generator.add_noise(noise_level)

    return _generate_data_only(generator)


def create_skewed_gaussian_data(
    n_samples: int,
    mean: float = 0.0,
    std: float = 1.0,
    skewness: float = 0.0,
    noise_level: float = 0.0,
    random_state: Optional[int] = None,
) -> NDArray:
    """快速创建偏态高斯数据

    Parameters
    ----------
    - `n_samples` : `int`
        - 样本数量
    - `mean` : `float`, optional
        - 均值, by default 0.0
    - `std` : `float`, optional
        - 标准差, by default 1.0
    - `skewness` : `float`, optional
        - 偏斜度, by default 0.0
    - `noise_level` : `float`, optional
        - 噪声强度, by default 0.0
    - `random_state` : `Optional[int]`, optional
        - 随机种子, by default None

    Returns
    ----------
    - `NDArray`
        - 生成的偏态高斯数据
    """
    generator = GaussianMixtureGenerator(n_samples, random_state)
    generator.add_component(mean, std)

    if skewness != 0:
        generator.add_skewness(skewness)

    if noise_level > 0:
        generator.add_noise(noise_level)

    return _generate_data_only(generator)


def create_mixed_distribution_data(
    n_samples: int,
    components: typing.Sequence[typing.Mapping],
    random_state: Optional[int] = None,
) -> NDArray:
    """快速创建混合分布数据

    Parameters
    ----------
    - `n_samples` : `int`
        - 样本数量
    - `components` : `List[dict]`
        - 分量配置列表，每个字典包含分量参数
    - `random_state` : `Optional[int]`, optional
        - 随机种子, by default None

    Returns
    ----------
    - `NDArray`
        - 生成的混合分布数据

    Examples
    ----------
    ```python
    # 创建包含高斯和均匀分布的混合数据
    components = [
        {'type': 'gaussian', 'mean': 0, 'std': 1, 'weight': 0.6},
        {'type': 'uniform', 'low': -3, 'high': 3, 'weight': 0.4, 'local_noise': 0.1}
    ]
    data = create_mixed_distribution_data(1000, components, random_state=42)
    ```
    """
    generator = GaussianMixtureGenerator(n_samples, random_state)

    for comp in components:
        comp_type = comp.get('type', 'gaussian')

        if comp_type == 'gaussian':
            generator.add_component(
                mean=comp['mean'],
                std=comp['std'],
                weight=comp.get('weight', 1.0),
                local_noise=comp.get('local_noise', 0.0),
                local_skewness=comp.get('local_skewness', 0.0),
            )
        elif comp_type == 'uniform':
            generator.add_uniform_component(
                low=comp['low'],
                high=comp['high'],
                weight=comp.get('weight', 1.0),
                local_noise=comp.get('local_noise', 0.0),
                local_skewness=comp.get('local_skewness', 0.0),
            )
        elif comp_type == 'exponential':
            generator.add_exponential_component(
                scale=comp['scale'],
                weight=comp.get('weight', 1.0),
                offset=comp.get('offset', 0.0),
            )
        else:
            raise ValueError(f'不支持的分量类型: {comp_type}')

    return _generate_data_only(generator)
