"""Multi-objective reward function for RL agent."""

import numpy as np
from typing import Dict


class RewardFunction:
    """Multi-objective reward function for adaptive mixing.

    Components:
    - Satisfaction (0.5): User feedback/preference
    - Quality (0.3): Audio quality metrics
    - Consistency (0.15): Character consistency
    - Latency (0.05): Generation latency penalty
    """

    def __init__(
        self,
        satisfaction_weight: float = 0.5,
        quality_weight: float = 0.3,
        consistency_weight: float = 0.15,
        latency_weight: float = 0.05,
        target_latency_ms: float = 200.0,
    ):
        """Initialize reward function.

        Args:
            satisfaction_weight: Weight for satisfaction component
            quality_weight: Weight for quality component
            consistency_weight: Weight for consistency component
            latency_weight: Weight for latency component
            target_latency_ms: Target latency in milliseconds
        """
        self.satisfaction_weight = satisfaction_weight
        self.quality_weight = quality_weight
        self.consistency_weight = consistency_weight
        self.latency_weight = latency_weight
        self.target_latency_ms = target_latency_ms

    def compute_reward(
        self,
        satisfaction_score: float,
        quality_metrics: Dict[str, float],
        consistency_score: float,
        latency_ms: float,
    ) -> float:
        """Compute total reward.

        Args:
            satisfaction_score: User satisfaction (0-1)
            quality_metrics: Quality metrics dict
            consistency_score: Character consistency (0-1)
            latency_ms: Generation latency in milliseconds

        Returns:
            Total reward
        """
        # Satisfaction component
        r_satisfaction = satisfaction_score

        # Quality component (average of metrics)
        r_quality = np.mean(list(quality_metrics.values()))

        # Consistency component
        r_consistency = consistency_score

        # Latency component (penalty for exceeding target)
        latency_ratio = latency_ms / self.target_latency_ms
        r_latency = max(0.0, 1.0 - (latency_ratio - 1.0))

        # Weighted sum
        total_reward = (
            self.satisfaction_weight * r_satisfaction
            + self.quality_weight * r_quality
            + self.consistency_weight * r_consistency
            + self.latency_weight * r_latency
        )

        return total_reward

    def compute_reward_components(
        self,
        satisfaction_score: float,
        quality_metrics: Dict[str, float],
        consistency_score: float,
        latency_ms: float,
    ) -> Dict[str, float]:
        """Compute reward components separately.

        Args:
            satisfaction_score: User satisfaction (0-1)
            quality_metrics: Quality metrics dict
            consistency_score: Character consistency (0-1)
            latency_ms: Generation latency in milliseconds

        Returns:
            Dictionary of reward components
        """
        r_satisfaction = satisfaction_score
        r_quality = np.mean(list(quality_metrics.values()))
        r_consistency = consistency_score
        latency_ratio = latency_ms / self.target_latency_ms
        r_latency = max(0.0, 1.0 - (latency_ratio - 1.0))

        return {
            "satisfaction": r_satisfaction * self.satisfaction_weight,
            "quality": r_quality * self.quality_weight,
            "consistency": r_consistency * self.consistency_weight,
            "latency": r_latency * self.latency_weight,
            "total": self.compute_reward(
                satisfaction_score, quality_metrics, consistency_score, latency_ms
            ),
        }
