"""Reward function aligned with the Q-Table design specification (Version 3.0)."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Optional

from ..model.state import (
    AcceptanceRate,
    DistributionStructure,
    InputPriceZone,
    PartnerType,
    RevenueRange,
)


class NegotiationOutcome(Enum):
    ONGOING = "ongoing"
    SUCCESS = "success"
    FAILURE = "failure"


@dataclass
class RewardConfig:
    """보상 함수 설정"""

    beta: float = 0.2
    success_reward: float = 1.0
    ongoing_reward: float = 0.0
    failure_penalty: float = -0.5
    penalty_lambda: float = 0.02

    # 동적 가중치 계수 (Version 3.0)
    w1: float = 0.2  # 매출액 가격구간
    w2: float = 0.25  # 유통 구조
    w3: float = 0.2  # 파트너사 종류
    w4: float = 0.25  # 가격 수용률
    w5: float = 0.1  # 입력 금액 구간

    min_weight: float = 0.2
    max_weight: float = 0.8


@dataclass
class RewardBreakdown:
    price_reward: float
    end_reward: float
    penalty: float
    weight: float
    total: float


def calculate_reward(
    *,
    revenue_range: RevenueRange,
    distribution: DistributionStructure,
    partner_type: PartnerType,
    acceptance_rate: AcceptanceRate,
    input_price_zone: InputPriceZone,
    current_price: float,
    anchor_price: float,
    target_price: float,
    round_number: int,
    outcome: NegotiationOutcome,
    config: Optional[RewardConfig] = None,
) -> RewardBreakdown:
    """
    보상 계산 (Version 3.0)

    R = W × R_price + (1-W) × R_end - λ × t

    W = clip(W_raw, min_weight, max_weight)
    W_raw = w1×S_amount + w2×S_dist + w3×S_partner + w4×S_accept + w5×S_pricezone
    """
    cfg = config or RewardConfig()

    price_reward = _calculate_price_reward(
        current_price=current_price,
        anchor_price=anchor_price,
        target_price=target_price,
        beta=cfg.beta,
    )
    end_reward = _calculate_end_reward(outcome, cfg)
    weight = _calculate_weight(
        revenue_range=revenue_range,
        distribution=distribution,
        partner_type=partner_type,
        acceptance_rate=acceptance_rate,
        input_price_zone=input_price_zone,
        config=cfg,
    )
    penalty = _calculate_penalty(round_number, cfg)

    total = weight * price_reward + (1.0 - weight) * end_reward - penalty
    return RewardBreakdown(
        price_reward=price_reward,
        end_reward=end_reward,
        penalty=penalty,
        weight=weight,
        total=total,
    )


def _calculate_price_reward(
    *,
    current_price: float,
    anchor_price: float,
    target_price: float,
    beta: float,
) -> float:
    if anchor_price <= 0 or target_price <= 0:
        raise ValueError("anchor_price and target_price must be positive")
    if anchor_price > target_price:
        raise ValueError("anchor_price must not exceed target_price")

    if current_price < anchor_price:
        improvement = (anchor_price - current_price) / anchor_price
        return 1.0 + beta * improvement

    if anchor_price <= current_price <= target_price:
        span = target_price - anchor_price
        if span <= 0:
            return 0.0
        return (target_price - current_price) / span

    return 0.0


def _calculate_end_reward(outcome: NegotiationOutcome, cfg: RewardConfig) -> float:
    if outcome is NegotiationOutcome.SUCCESS:
        return cfg.success_reward
    if outcome is NegotiationOutcome.FAILURE:
        return cfg.failure_penalty
    return cfg.ongoing_reward


def _calculate_weight(
    *,
    revenue_range: RevenueRange,
    distribution: DistributionStructure,
    partner_type: PartnerType,
    acceptance_rate: AcceptanceRate,
    input_price_zone: InputPriceZone,
    config: RewardConfig,
) -> float:
    """
    동적 가중치 계산 (Version 3.0)

    W_raw = w1×S_amount + w2×S_dist + w3×S_partner + w4×S_accept + w5×S_pricezone
    W = clip(W_raw, min_weight, max_weight)
    """
    w_raw = (
        config.w1 * revenue_range.weight
        + config.w2 * distribution.weight
        + config.w3 * partner_type.weight
        + config.w4 * acceptance_rate.weight
        + config.w5 * input_price_zone.weight
    )
    return max(config.min_weight, min(config.max_weight, w_raw))


def _calculate_penalty(round_number: int, cfg: RewardConfig) -> float:
    if round_number < 0:
        raise ValueError("round_number must be non-negative")
    return cfg.penalty_lambda * float(round_number)
