"""
TrainOfflineUsecase
저장된 Experience 데이터로 Q-Table을 오프라인 학습
"""

from typing import List, Optional
from ..domain.model.q_table import QTable
from ..domain.model.experience import Experience
from ..domain.model.visit_table import VisitTable
from ..domain.repository.experience_repository import ExperienceRepository


class TrainOfflineUsecase:
    """
    수집된 Experience 데이터를 사용하여 Q-Table을 오프라인으로 학습

    배치 학습 방식:
    - 저장된 Experience를 로드
    - Q-Learning 알고리즘으로 Q-Table 업데이트
    - 여러 에포크 반복 가능
    """

    def __init__(
        self,
        q_table: QTable,
        experience_repository: ExperienceRepository,
        visit_table: Optional[VisitTable] = None,
    ):
        """
        Args:
            q_table: 학습할 Q-Table
            experience_repository: Experience 저장소
            visit_table: 방문 횟수를 함께 추적할 N-Table (선택)
        """
        self.q_table = q_table
        self.experience_repository = experience_repository
        self.visit_table = visit_table

    def train(
        self,
        filename: str = "experiences.jsonl",
        epochs: int = 1,
        batch_size: Optional[int] = None,
    ) -> dict:
        """
        저장된 Experience로 Q-Table 학습

        Args:
            filename: Experience 파일명
            epochs: 학습 반복 횟수
            batch_size: 배치 크기 (None이면 전체 데이터 사용)

        Returns:
            {
                'total_experiences': 1000,
                'epochs': 10,
                'updates': 10000,
                'avg_loss': 0.05
            }
        """
        # Experience 로드
        experiences = self.experience_repository.load_all(filename)

        if not experiences:
            return {
                "total_experiences": 0,
                "epochs": 0,
                "updates": 0,
                "avg_loss": 0.0,
                "message": "No experiences found",
            }

        total_updates = 0
        total_loss = 0.0

        # 에포크 반복
        for epoch in range(epochs):
            epoch_loss = 0.0

            # 배치 처리
            if batch_size:
                for i in range(0, len(experiences), batch_size):
                    batch = experiences[i : i + batch_size]
                    loss = self._train_batch(batch)
                    epoch_loss += loss
                    total_updates += len(batch)
            else:
                # 전체 데이터 한번에
                loss = self._train_batch(experiences)
                epoch_loss += loss
                total_updates += len(experiences)

            total_loss += epoch_loss

        avg_loss = total_loss / total_updates if total_updates > 0 else 0.0

        return {
            "total_experiences": len(experiences),
            "epochs": epochs,
            "updates": total_updates,
            "avg_loss": avg_loss,
        }

    def _train_batch(self, experiences: List[Experience]) -> float:
        """
        배치 Experience로 Q-Table 업데이트

        Args:
            experiences: Experience 리스트

        Returns:
            평균 손실값
        """
        total_loss = 0.0

        for exp in experiences:
            # State → state_index
            state_index = exp.state.to_index()

            # Q-value 업데이트 전 값
            old_q = self.q_table.get_q_value(state_index, exp.action_id)

            # Q-Table 업데이트
            if exp.next_state:
                next_state_index = exp.next_state.to_index()
                self.q_table.update(
                    state_index=state_index,
                    action_id=exp.action_id,
                    reward=exp.reward,
                    next_state_index=next_state_index,
                    done=exp.done,
                )
            else:
                # 종료 상태
                self.q_table.update(
                    state_index=state_index,
                    action_id=exp.action_id,
                    reward=exp.reward,
                    next_state_index=None,
                    done=True,
                )

            if self.visit_table is not None:
                self.visit_table.increment(state_index, exp.action_id)

            # Q-value 업데이트 후 값
            new_q = self.q_table.get_q_value(state_index, exp.action_id)

            # 손실 계산 (업데이트 크기)
            loss = abs(new_q - old_q)
            total_loss += loss

        return total_loss / len(experiences) if experiences else 0.0

    def train_by_episodes(
        self,
        episode_ids: List[str],
        filename: str = "experiences.jsonl",
        epochs: int = 1,
    ) -> dict:
        """
        특정 에피소드들만 선택하여 학습

        Args:
            episode_ids: 학습할 에피소드 ID 리스트
            filename: Experience 파일명
            epochs: 학습 반복 횟수

        Returns:
            학습 결과
        """
        # 해당 에피소드의 Experience만 로드
        all_experiences = self.experience_repository.load_all(filename)
        filtered_experiences = [
            exp for exp in all_experiences if exp.episode_id in episode_ids
        ]

        if not filtered_experiences:
            return {
                "total_experiences": 0,
                "episodes": 0,
                "epochs": 0,
                "updates": 0,
                "avg_loss": 0.0,
                "message": "No matching episodes found",
            }

        total_updates = 0
        total_loss = 0.0

        for epoch in range(epochs):
            loss = self._train_batch(filtered_experiences)
            total_loss += loss * len(filtered_experiences)
            total_updates += len(filtered_experiences)

        return {
            "total_experiences": len(filtered_experiences),
            "episodes": len(episode_ids),
            "epochs": epochs,
            "updates": total_updates,
            "avg_loss": total_loss / total_updates if total_updates > 0 else 0.0,
        }

