"""Pass 2: Dominant element service for telemetry generation.

Redistributes unit values so one element dominates each hour (60% share),
or creates varied distribution if no dominant element specified.
"""

import hashlib
from typing import Optional

import numpy as np
import polars as pl


class DominantElementService:
    """Redistributes telemetry so one element dominates per timestamp.

    Single responsibility: Redistribute unit values to create:
    - 60% share for dominant element (if specified)
    - Varied shares (10-30%) for all elements (if no dominant)
    """

    # Hardcoded constants (not exposed to config)
    _DOMINANT_SHARE_PCT = 60.0
    _MIN_NON_DOMINANT_SHARE = 4.0  # Minimum % for non-dominant elements
    _MAX_NON_DOMINANT_SHARE = 15.0  # Maximum % for non-dominant elements

    def __init__(
        self,
        dominant_element_id: Optional[str] = None,
        seed: int = 42
    ):
        """Initialize dominant element service.

        Args:
            dominant_element_id: Element that gets 60% share (None = varied)
            seed: Random seed for deterministic share generation
        """
        self.dominant_element_id = dominant_element_id
        self.seed = seed

    def apply_dominance(self, df: pl.DataFrame) -> pl.DataFrame:
        """Redistribute values to create dominant element or varied shares.

        Args:
            df: DataFrame with timestamp, element_id, unit_value columns

        Returns:
            DataFrame with redistributed unit_value per timestamp
        """
        if len(df) == 0:
            return df

        # Group by timestamp and redistribute within each group
        return df.group_by("timestamp").map_groups(
            self._redistribute_timestamp_group
        )

    def _redistribute_timestamp_group(self, group: pl.DataFrame) -> pl.DataFrame:
        """Redistribute unit values within a single timestamp group.

        Uses deterministic seeding based on timestamp for reproducibility.
        """
        if len(group) == 0:
            return group

        # Get timestamp for deterministic seeding
        timestamp = group["timestamp"][0]
        timestamp_seed = self._generate_timestamp_seed(timestamp)

        # Get element IDs in group
        element_ids = group["element_id"].to_list()

        # Calculate target shares
        if self.dominant_element_id:
            shares = self._calculate_dominant_shares(
                element_ids, timestamp_seed
            )
        else:
            shares = self._calculate_varied_shares(
                element_ids, timestamp_seed
            )

        # Calculate total current value
        total_value = group["unit_value"].sum()

        # Apply shares to create new unit values
        share_series = pl.Series("shares", shares)
        return group.with_columns([
            (pl.lit(total_value) * share_series).alias("unit_value")
        ])

    def _calculate_dominant_shares(
        self,
        element_ids: list,
        seed: int
    ) -> list:
        """Calculate shares with one dominant element (60% share)."""
        rng = np.random.Generator(np.random.PCG64(seed))

        # Find dominant element index
        try:
            dominant_idx = element_ids.index(self.dominant_element_id)
        except ValueError:
            # Dominant element not in this group, use varied shares
            return self._calculate_varied_shares(element_ids, seed)

        # Get non-dominant element IDs
        non_dominant_ids = [
            eid for i, eid in enumerate(element_ids) if i != dominant_idx
        ]

        # Dominant gets 60%, others share remaining 40%
        shares = self._distribute_remaining_share(
            non_dominant_ids,
            100.0 - self._DOMINANT_SHARE_PCT,
            rng
        )

        # Insert dominant share at correct position
        shares.insert(dominant_idx, self._DOMINANT_SHARE_PCT)

        # Convert from percentage to fraction
        return [s / 100.0 for s in shares]

    def _calculate_varied_shares(self, element_ids: list, seed: int) -> list:
        """Calculate varied shares when no dominant element (10-30% each)."""
        rng = np.random.Generator(np.random.PCG64(seed))

        shares = self._distribute_remaining_share(element_ids, 100.0, rng)

        # Convert from percentage to fraction
        return [s / 100.0 for s in shares]

    def _distribute_remaining_share(
        self,
        element_ids: list,
        total_pct: float,
        rng: np.random.Generator
    ) -> list:
        """Distribute total_pct across elements with persistent differentiation.

        Each element gets a deterministic base share (based on element_id hash)
        plus small per-timestamp variation for realism.
        """
        if len(element_ids) == 0:
            return []

        if len(element_ids) == 1:
            return [total_pct]

        # Get persistent base shares for each element (deterministic)
        base_shares = [
            self._get_element_base_share_pct(eid) for eid in element_ids
        ]

        # Add small per-timestamp variation (±2% around base share)
        variation_pct = 2.0
        variations = rng.uniform(-variation_pct, variation_pct, len(element_ids))
        varied_shares = [
            base + var for base, var in zip(base_shares, variations)
        ]

        # Normalize to sum to total_pct
        share_sum = sum(varied_shares)
        normalized_shares = [(s / share_sum) * total_pct for s in varied_shares]

        return normalized_shares

    def _get_element_base_share_pct(self, element_id: str) -> float:
        """Get persistent base share percentage for an element.

        Returns deterministic value in 4-15% range based on element_id hash.
        This ensures each element has a consistent base share across all timestamps.
        """
        # Hash element_id with seed for deterministic value
        hash_value = hashlib.md5(
            f"{self.seed}_{element_id}".encode()
        ).hexdigest()
        hash_int = int(hash_value[:8], 16)

        # Map to 4-15% range
        range_size = self._MAX_NON_DOMINANT_SHARE - self._MIN_NON_DOMINANT_SHARE
        normalized = (hash_int % 10000) / 10000.0  # 0.0 to 1.0
        base_share = self._MIN_NON_DOMINANT_SHARE + (normalized * range_size)

        return base_share

    def _generate_timestamp_seed(self, timestamp) -> int:
        """Generate deterministic seed from timestamp."""
        timestamp_str = timestamp.isoformat()
        combined = f"{self.seed}_{timestamp_str}"
        hash_value = hashlib.md5(combined.encode()).hexdigest()
        return int(hash_value[:8], 16) % (2**32)
