from typing import Dict, List, Optional
"""Allocation telemetry data generator."""

import random
from datetime import datetime, timedelta

import numpy as np
import polars as pl

from ..patterns import AnomalyConfig, MetricScale, SeasonalityConfig, TrendConfig
from .config import (
    AllocationConfig,
    ElementPatternConfig,
)


class AllocationGenerator:
    """Generator for allocation telemetry data with element-based usage distribution."""

    def __init__(self, config: AllocationConfig):
        self.config = config
        self.elements = self._generate_element_list()
        self.element_patterns = self._build_element_patterns_map()

    def _generate_element_list(self) -> List[str]:
        """Generate list of element IDs based on configuration."""
        if self.config.elements.type == "provided":
            return self.config.elements.list.copy()
        elif self.config.elements.type == "generated":
            elements = []
            for i in range(1, self.config.elements.count + 1):
                element_id = f"{self.config.elements.prefix}_{self.config.elements.naming_pattern.format(i)}"
                elements.append(element_id)
            return elements
        else:
            raise ValueError(f"Unsupported element type: {self.config.elements.type}")

    def _build_element_patterns_map(self) -> Dict[str, ElementPatternConfig]:
        """Build a map of element IDs to their specific pattern configurations."""
        patterns_map = {}
        if self.config.elements.patterns:
            for pattern in self.config.elements.patterns:
                patterns_map[pattern.element_id] = pattern
        return patterns_map

    def _create_datetime_range(
        self, start_date: str, end_date: str, frequency: str
    ) -> List[datetime]:
        """Create datetime range with specified frequency."""
        start = datetime.fromisoformat(start_date)
        end = datetime.fromisoformat(end_date)

        dates = []
        current = start
        delta = timedelta(hours=1) if frequency == "H" else timedelta(days=1)

        while current <= end:
            dates.append(current)
            current += delta

        return dates

    def _get_base_value(self, scale: MetricScale) -> float:
        """Get a base value within the specified scale range."""
        min_val, max_val = scale.value
        return random.uniform(min_val, max_val)

    def _apply_seasonality(
        self, value: float, date: datetime, config: SeasonalityConfig
    ) -> float:
        """Apply seasonal patterns to a value."""
        import math

        if config.pattern == "weekly":
            # Weekly pattern: 0=Monday, 6=Sunday
            weekday = date.weekday()
            # Create sine wave pattern with peak on Wednesday (weekday 2)
            cycle_position = (weekday + config.phase) % 7
            seasonal_factor = math.sin(2 * math.pi * cycle_position / 7)
            day_factor = 1 + config.amplitude * seasonal_factor
        elif config.pattern == "monthly":
            # Monthly pattern based on day of month
            day_of_month = date.day
            # Create pattern with peak mid-month
            cycle_position = (day_of_month + config.phase) % 31
            seasonal_factor = math.sin(2 * math.pi * cycle_position / 31)
            day_factor = 1 + config.amplitude * seasonal_factor
        elif config.pattern == "yearly":
            # Yearly pattern based on day of year
            day_of_year = date.timetuple().tm_yday
            cycle_position = (day_of_year + config.phase) % 365
            seasonal_factor = math.sin(2 * math.pi * cycle_position / 365)
            day_factor = 1 + config.amplitude * seasonal_factor
        else:
            day_factor = 1

        return value * day_factor

    def _apply_trend(
        self, value: float, days_from_start: int, config: TrendConfig
    ) -> float:
        """Apply trend patterns to a value."""
        if config.type == "linear":
            return value * (1 + config.growth_rate * days_from_start / 30)
        elif config.type == "exponential":
            return value * ((1 + config.growth_rate) ** (days_from_start / 30))
        elif config.type == "logistic":
            if config.cap is None:
                raise ValueError("Cap must be specified for logistic trend")
            growth_factor = (1 + config.growth_rate) ** (days_from_start / 30)
            return min(value * growth_factor, config.cap)
        else:
            return value

    def _should_add_anomaly(self, date: datetime, config: AnomalyConfig) -> bool:
        """Determine if an anomaly should be added on this date."""
        random.seed(config.seed + date.toordinal())
        monthly_probability = config.frequency / 30.0
        return random.random() < monthly_probability

    def _apply_anomaly(self, value: float, config: AnomalyConfig) -> float:
        """Apply anomaly multiplier to a value."""
        return value * config.magnitude

    def _apply_variance(self, value: float, factor: float) -> float:
        """Apply random variance to a value."""
        variance = random.uniform(-factor, factor)
        return max(0, value * (1 + variance))

    def _generate_element_usage_distribution(
        self, total_usage: float
    ) -> Dict[str, float]:
        """Generate usage distribution across elements using configured strategy."""
        distribution_config = self.config.elements.distribution
        random.seed(distribution_config.seed)
        np.random.seed(distribution_config.seed)

        if distribution_config.strategy == "random":
            return self._random_distribution(total_usage)
        elif distribution_config.strategy == "weighted":
            return self._weighted_distribution(total_usage, distribution_config.weights)
        elif distribution_config.strategy == "clustered":
            return self._clustered_distribution(
                total_usage, distribution_config.cluster_factor
            )
        elif distribution_config.strategy == "pareto":
            return self._pareto_distribution(
                total_usage, distribution_config.pareto_factor
            )
        else:
            raise ValueError(
                f"Unsupported distribution strategy: {distribution_config.strategy}"
            )

    def _random_distribution(self, total_usage: float) -> Dict[str, float]:
        """Simple random distribution - each element gets a random portion."""
        raw_portions = [random.random() for _ in self.elements]
        total_portions = sum(raw_portions)

        distribution = {}
        for i, element in enumerate(self.elements):
            portion = (raw_portions[i] / total_portions) * total_usage
            distribution[element] = max(0.001, portion)  # Minimum usage to avoid zero

        return distribution

    def _weighted_distribution(
        self, total_usage: float, weights: Optional[List[float]]
    ) -> Dict[str, float]:
        """Weighted distribution based on provided weights."""
        if not weights:
            return self._random_distribution(total_usage)

        if len(weights) != len(self.elements):
            raise ValueError(
                f"Weights count ({len(weights)}) must match elements count ({len(self.elements)})"
            )

        total_weight = sum(weights)
        distribution = {}

        for i, element in enumerate(self.elements):
            portion = (weights[i] / total_weight) * total_usage
            distribution[element] = max(0.001, portion)

        return distribution

    def _clustered_distribution(
        self, total_usage: float, cluster_factor: float
    ) -> Dict[str, float]:
        """Clustered distribution - some elements get much more usage than others."""
        # Create clusters: high usage (20%), medium usage (30%), low usage (50%)
        element_count = len(self.elements)
        high_count = max(1, int(element_count * 0.2))
        medium_count = max(1, int(element_count * 0.3))
        low_count = element_count - high_count - medium_count

        # Assign usage levels based on cluster factor
        high_usage_factor = (
            1.0 + cluster_factor * 4
        )  # High users get 5x more at max cluster_factor
        medium_usage_factor = 1.0 + cluster_factor * 1.5  # Medium users get 2.5x more
        low_usage_factor = (
            1.0 - cluster_factor * 0.8
        )  # Low users get 20% at max cluster_factor

        # Create weights for each element
        weights = []
        elements_shuffled = self.elements.copy()
        random.shuffle(elements_shuffled)

        for i, element in enumerate(elements_shuffled):
            if i < high_count:
                weights.append(high_usage_factor)
            elif i < high_count + medium_count:
                weights.append(medium_usage_factor)
            else:
                weights.append(low_usage_factor)

        # Map back to original element order
        element_weights = dict(zip(elements_shuffled, weights, strict=False))
        ordered_weights = [element_weights[element] for element in self.elements]

        return self._weighted_distribution(total_usage, ordered_weights)

    def _pareto_distribution(
        self, total_usage: float, pareto_factor: float
    ) -> Dict[str, float]:
        """Pareto distribution - 80/20 rule or similar."""
        element_count = len(self.elements)

        # Determine how many elements are in the "high usage" group
        high_usage_count = max(1, int(element_count * (1 - pareto_factor)))
        low_usage_count = element_count - high_usage_count

        # High usage elements get pareto_factor of total usage
        # Low usage elements get (1 - pareto_factor) of total usage
        high_usage_total = total_usage * pareto_factor
        low_usage_total = total_usage * (1 - pareto_factor)

        # Randomly select which elements are high usage
        elements_shuffled = self.elements.copy()
        random.shuffle(elements_shuffled)

        distribution = {}

        # Distribute among high usage elements
        for i, element in enumerate(elements_shuffled[:high_usage_count]):
            portion = high_usage_total / high_usage_count
            # Add some variance within the high usage group
            variance = random.uniform(0.7, 1.3)
            distribution[element] = max(0.001, portion * variance)

        # Distribute among low usage elements
        for i, element in enumerate(elements_shuffled[high_usage_count:]):
            portion = low_usage_total / low_usage_count
            # Add some variance within the low usage group
            variance = random.uniform(0.5, 1.5)
            distribution[element] = max(0.001, portion * variance)

        # Normalize to ensure total equals total_usage
        actual_total = sum(distribution.values())
        normalization_factor = total_usage / actual_total

        for element in distribution:
            distribution[element] *= normalization_factor

        return distribution

    def _apply_element_patterns(
        self, element_id: str, base_usage: float, date: datetime, start_date: datetime
    ) -> float:
        """Apply element-specific patterns to usage value."""
        usage = base_usage

        # Check if this element has specific patterns configured
        if element_id in self.element_patterns:
            pattern = self.element_patterns[element_id]

            # Apply usage multiplier first
            usage *= pattern.usage_multiplier

            # Apply element-specific seasonality
            if pattern.seasonality:
                usage = self._apply_seasonality(usage, date, pattern.seasonality)

            # Apply element-specific trend
            if pattern.trend:
                days_from_start = (date - start_date).days
                usage = self._apply_trend(usage, days_from_start, pattern.trend)

            # Apply element-specific anomalies
            if pattern.anomaly and self._should_add_anomaly(date, pattern.anomaly):
                usage = self._apply_anomaly(usage, pattern.anomaly)

            # Apply element-specific variance if configured
            if pattern.variance_factor is not None:
                usage = self._apply_variance(usage, pattern.variance_factor)

        return max(0.001, usage)  # Ensure minimum usage

    def generate(self) -> pl.DataFrame:
        """Generate allocation telemetry data."""
        dates = self._create_datetime_range(
            self.config.start_date, self.config.end_date, self.config.frequency
        )

        records = []
        start_date = datetime.fromisoformat(self.config.start_date)

        for date in dates:
            # Generate total usage for this time period
            base_usage = self._get_base_value(self.config.scale)

            # Apply patterns to total usage
            total_usage = base_usage

            if self.config.seasonality:
                total_usage = self._apply_seasonality(
                    total_usage, date, self.config.seasonality
                )

            if self.config.trend:
                days_from_start = (date - start_date).days
                total_usage = self._apply_trend(
                    total_usage, days_from_start, self.config.trend
                )

            if self.config.anomaly and self._should_add_anomaly(
                date, self.config.anomaly
            ):
                total_usage = self._apply_anomaly(total_usage, self.config.anomaly)

            # Apply variance
            total_usage = self._apply_variance(total_usage, self.config.variance_factor)

            # Distribute usage across elements
            element_usage = self._generate_element_usage_distribution(total_usage)

            # Create records for each element with individual patterns applied
            for element_id, base_element_usage in element_usage.items():
                # Apply element-specific patterns if configured
                final_usage = self._apply_element_patterns(
                    element_id, base_element_usage, date, start_date
                )

                record = {
                    "timestamp": date,
                    "element_id": element_id,  # Standardized column name
                    "unit_value": final_usage,  # Standardized column name
                }

                # Add cost filters as additional columns
                if self.config.cost_filters:
                    record.update(self.config.cost_filters)

                records.append(record)

        return pl.DataFrame(records)
