"""Synthetic data generation for testing and demonstration."""

from typing import Optional, List, Tuple, Dict, Any
import numpy as np
import pandas as pd


def generate_dummy_data(
    n_days: int = 365,
    base_level: float = 10000,
    trend_slope: float = 0.001,
    trend_type: str = "linear",
    weekly_amplitude: float = 0.2,
    annual_amplitude: float = 0.1,
    noise_level: float = 0.05,
    event_dates: Optional[List[str]] = None,
    event_multiplier: float = 3.0,
    event_duration: int = 1,
    start_date: str = "2023-01-01",
    frequency: str = "D",
    random_seed: Optional[int] = None,
) -> pd.Series:
    """
    Generate synthetic ad traffic data with trend, seasonality, and events.

    Creates realistic ad impression data suitable for testing forecasting
    models. Includes:
    - Trend component (linear or exponential growth)
    - Weekly seasonality (higher traffic on weekdays)
    - Annual seasonality (optional)
    - Random noise
    - Event spikes (optional)

    Parameters
    ----------
    n_days : int, default 365
        Number of days to generate.
    base_level : float, default 10000
        Base level of impressions.
    trend_slope : float, default 0.001
        Daily trend growth rate. For linear: added per day.
        For exponential: percentage growth per day.
    trend_type : {'linear', 'exponential'}, default 'linear'
        Type of trend growth.
    weekly_amplitude : float, default 0.2
        Amplitude of weekly seasonality (as fraction of base).
    annual_amplitude : float, default 0.1
        Amplitude of annual seasonality (as fraction of base).
    noise_level : float, default 0.05
        Standard deviation of noise (as fraction of value).
    event_dates : List[str], optional
        Dates of special events (e.g., ['2023-02-12', '2023-11-24']).
    event_multiplier : float, default 3.0
        Multiplier for event traffic.
    event_duration : int, default 1
        Duration of each event in days.
    start_date : str, default '2023-01-01'
        Start date for the time series.
    frequency : str, default 'D'
        Frequency: 'D' for daily, 'W' for weekly, 'M' for monthly.
    random_seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    y : pd.Series
        Synthetic time series with DatetimeIndex.

    Examples
    --------
    >>> # Basic synthetic data
    >>> y = generate_dummy_data(n_days=365, trend_slope=0.001)
    >>>
    >>> # With events
    >>> y = generate_dummy_data(
    ...     n_days=730,
    ...     event_dates=['2023-02-12', '2024-02-11'],  # Super Bowl dates
    ...     event_multiplier=5.0,
    ... )

    Notes
    -----
    The generated data follows the pattern:
    Y(t) = Trend(t) × (1 + Weekly(t) + Annual(t)) × (1 + Noise(t)) × Event(t)
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    # Generate date index
    dates = pd.date_range(start=start_date, periods=n_days, freq=frequency)
    t = np.arange(n_days)

    # Generate trend
    if trend_type == "linear":
        trend = base_level + trend_slope * base_level * t
    elif trend_type == "exponential":
        trend = base_level * np.exp(trend_slope * t)
    else:
        raise ValueError(f"Unknown trend_type: {trend_type}")

    # Generate weekly seasonality
    # Monday=0, peak on Wed-Thu, low on weekend
    if frequency == "D":
        day_of_week = dates.dayofweek
        weekly_pattern = np.array([0.9, 1.0, 1.1, 1.1, 1.0, 0.7, 0.5])
        weekly = (weekly_pattern[day_of_week] - 1) * weekly_amplitude
    else:
        weekly = np.zeros(n_days)

    # Generate annual seasonality
    if frequency in ["D", "W"] and annual_amplitude > 0:
        # Peak in fall, low in summer
        day_of_year = dates.dayofyear if frequency == "D" else dates.isocalendar().week * 7
        annual = annual_amplitude * np.sin(2 * np.pi * (day_of_year - 90) / 365)
    else:
        annual = np.zeros(n_days)

    # Generate noise
    noise = np.random.normal(0, noise_level, n_days)

    # Combine components (multiplicative model)
    values = trend * (1 + weekly + annual) * (1 + noise)

    # Add event spikes
    if event_dates:
        event_spike = np.ones(n_days)
        for event_date in event_dates:
            event_ts = pd.Timestamp(event_date)
            if event_ts >= dates[0] and event_ts <= dates[-1]:
                event_idx = dates.get_indexer([event_ts], method="nearest")[0]
                for i in range(event_duration):
                    if event_idx + i < n_days:
                        event_spike[event_idx + i] = event_multiplier

        values = values * event_spike

    # Ensure non-negative
    values = np.maximum(values, 0)

    return pd.Series(values, index=dates, name="impressions")


def generate_multi_series_data(
    n_days: int = 365,
    n_series: int = 5,
    series_names: Optional[List[str]] = None,
    base_levels: Optional[List[float]] = None,
    trend_slopes: Optional[List[float]] = None,
    correlation: float = 0.5,
    start_date: str = "2023-01-01",
    random_seed: Optional[int] = None,
) -> pd.DataFrame:
    """
    Generate multiple correlated synthetic time series.

    Useful for testing hierarchical forecasting and demographic
    reconciliation.

    Parameters
    ----------
    n_days : int, default 365
        Number of days.
    n_series : int, default 5
        Number of series to generate.
    series_names : List[str], optional
        Names for each series. Default: ['series_0', 'series_1', ...].
    base_levels : List[float], optional
        Base levels for each series. Default: random between 5000-15000.
    trend_slopes : List[float], optional
        Trend slopes for each series. Default: random between 0-0.002.
    correlation : float, default 0.5
        Correlation between series (0 = independent, 1 = identical).
    start_date : str, default '2023-01-01'
        Start date.
    random_seed : int, optional
        Random seed.

    Returns
    -------
    df : pd.DataFrame
        DataFrame with multiple time series and 'total' column.
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    if series_names is None:
        series_names = [f"series_{i}" for i in range(n_series)]

    if base_levels is None:
        base_levels = np.random.uniform(5000, 15000, n_series)

    if trend_slopes is None:
        trend_slopes = np.random.uniform(0, 0.002, n_series)

    dates = pd.date_range(start=start_date, periods=n_days, freq="D")

    # Generate correlated noise
    common_noise = np.random.normal(0, 0.05, n_days)

    series_data = {}
    for i, (name, base, slope) in enumerate(zip(series_names, base_levels, trend_slopes)):
        # Individual noise component
        individual_noise = np.random.normal(0, 0.05, n_days)
        noise = correlation * common_noise + (1 - correlation) * individual_noise

        # Generate series with different characteristics
        y = generate_dummy_data(
            n_days=n_days,
            base_level=base,
            trend_slope=slope,
            noise_level=0,  # We add correlated noise separately
            start_date=start_date,
            random_seed=None,  # Don't reset seed
        )

        series_data[name] = y.values * (1 + noise)

    df = pd.DataFrame(series_data, index=dates)

    # Add total column
    df["total"] = df.sum(axis=1)

    return df


def generate_intermittent_demand(
    n_days: int = 365,
    demand_probability: float = 0.3,
    mean_demand: float = 100,
    demand_cv: float = 0.5,
    start_date: str = "2023-01-01",
    random_seed: Optional[int] = None,
) -> pd.Series:
    """
    Generate intermittent demand data (many zeros).

    Useful for testing Croston/TSB methods.

    Parameters
    ----------
    n_days : int, default 365
        Number of days.
    demand_probability : float, default 0.3
        Probability of demand occurring on any day.
    mean_demand : float, default 100
        Mean demand when it occurs.
    demand_cv : float, default 0.5
        Coefficient of variation of demand size.
    start_date : str, default '2023-01-01'
        Start date.
    random_seed : int, optional
        Random seed.

    Returns
    -------
    y : pd.Series
        Sparse time series with many zeros.
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    dates = pd.date_range(start=start_date, periods=n_days, freq="D")

    # Generate demand occurrences
    occurs = np.random.random(n_days) < demand_probability

    # Generate demand sizes (log-normal)
    sigma = np.sqrt(np.log(1 + demand_cv**2))
    mu = np.log(mean_demand) - sigma**2 / 2
    demand_sizes = np.random.lognormal(mu, sigma, n_days)

    # Combine
    values = np.where(occurs, demand_sizes, 0)

    return pd.Series(values, index=dates, name="demand")


def generate_event_pattern(
    event_duration: int = 5,
    peak_position: float = 0.3,
    peak_multiplier: float = 5.0,
    base_level: float = 1.0,
    noise_level: float = 0.1,
    random_seed: Optional[int] = None,
) -> np.ndarray:
    """
    Generate a realistic event traffic pattern.

    Creates a shape suitable for use as a source pattern in the
    analog/copy model.

    Parameters
    ----------
    event_duration : int, default 5
        Duration of event in days.
    peak_position : float, default 0.3
        Position of peak (0-1, as fraction of duration).
    peak_multiplier : float, default 5.0
        Peak value relative to base.
    base_level : float, default 1.0
        Base level of traffic.
    noise_level : float, default 0.1
        Noise level (fraction).
    random_seed : int, optional
        Random seed.

    Returns
    -------
    pattern : np.ndarray
        Event pattern values.

    Examples
    --------
    >>> pattern = generate_event_pattern(
    ...     event_duration=7,
    ...     peak_position=0.3,
    ...     peak_multiplier=8.0,
    ... )
    >>> # Scale to actual traffic levels
    >>> actual_pattern = pattern * 100000
    """
    if random_seed is not None:
        np.random.seed(random_seed)

    t = np.linspace(0, 1, event_duration)
    peak_idx = int(peak_position * event_duration)

    # Create asymmetric peak (faster rise, slower decay)
    pattern = np.zeros(event_duration)

    for i, ti in enumerate(t):
        if i <= peak_idx:
            # Rising phase
            pattern[i] = base_level + (peak_multiplier - base_level) * (i / peak_idx)
        else:
            # Decay phase
            decay_progress = (i - peak_idx) / (event_duration - peak_idx - 1)
            pattern[i] = peak_multiplier - (peak_multiplier - base_level) * decay_progress

    # Add noise
    noise = np.random.normal(0, noise_level, event_duration)
    pattern = pattern * (1 + noise)

    return np.maximum(pattern, 0)


def create_test_dataset(
    scenario: str = "basic",
    random_seed: int = 42,
) -> Dict[str, Any]:
    """
    Create a complete test dataset for a specific scenario.

    Parameters
    ----------
    scenario : str, default 'basic'
        Scenario to generate:
        - 'basic': Simple trend + seasonality
        - 'event': Includes event spikes
        - 'sparse': Intermittent demand
        - 'multi': Multiple series with total
        - 'complex': All features combined
    random_seed : int, default 42
        Random seed.

    Returns
    -------
    dataset : dict
        Dictionary with 'data' and 'metadata' keys.

    Examples
    --------
    >>> dataset = create_test_dataset('event')
    >>> y = dataset['data']
    >>> event_dates = dataset['metadata']['event_dates']
    """
    np.random.seed(random_seed)

    if scenario == "basic":
        y = generate_dummy_data(
            n_days=365,
            base_level=10000,
            trend_slope=0.001,
            weekly_amplitude=0.2,
            annual_amplitude=0.1,
            noise_level=0.05,
            random_seed=random_seed,
        )
        metadata = {
            "scenario": "basic",
            "has_trend": True,
            "has_weekly_seasonality": True,
            "has_annual_seasonality": True,
            "has_events": False,
        }

    elif scenario == "event":
        event_dates = ["2023-02-12", "2023-07-04", "2023-11-24"]
        y = generate_dummy_data(
            n_days=365,
            base_level=10000,
            trend_slope=0.001,
            event_dates=event_dates,
            event_multiplier=5.0,
            event_duration=2,
            random_seed=random_seed,
        )
        metadata = {
            "scenario": "event",
            "has_trend": True,
            "has_events": True,
            "event_dates": event_dates,
        }

    elif scenario == "sparse":
        y = generate_intermittent_demand(
            n_days=365,
            demand_probability=0.3,
            mean_demand=100,
            random_seed=random_seed,
        )
        metadata = {
            "scenario": "sparse",
            "is_sparse": True,
            "demand_probability": 0.3,
        }

    elif scenario == "multi":
        y = generate_multi_series_data(
            n_days=365,
            n_series=4,
            series_names=["age_18_24", "age_25_34", "age_35_44", "age_45_plus"],
            correlation=0.6,
            random_seed=random_seed,
        )
        metadata = {
            "scenario": "multi",
            "n_series": 4,
            "series_names": ["age_18_24", "age_25_34", "age_35_44", "age_45_plus"],
        }

    elif scenario == "complex":
        # Generate 2 years with multiple events
        event_dates = [
            "2023-02-12", "2023-07-04", "2023-11-24",
            "2024-02-11", "2024-07-04", "2024-11-28",
        ]
        y = generate_dummy_data(
            n_days=730,
            base_level=10000,
            trend_slope=0.0005,
            trend_type="exponential",
            weekly_amplitude=0.25,
            annual_amplitude=0.15,
            noise_level=0.08,
            event_dates=event_dates,
            event_multiplier=4.0,
            event_duration=3,
            random_seed=random_seed,
        )
        metadata = {
            "scenario": "complex",
            "has_trend": True,
            "trend_type": "exponential",
            "has_weekly_seasonality": True,
            "has_annual_seasonality": True,
            "has_events": True,
            "event_dates": event_dates,
        }

    else:
        raise ValueError(f"Unknown scenario: {scenario}")

    return {"data": y, "metadata": metadata}
