import jiwer
from fuzzywuzzy import fuzz


class OCREvaluator:
    """
    정답(GT) 텍스트와 하나 이상의 예측(Hypothesis) 텍스트를 비교하여
    다양한 문자 오류율(CER) 및 단어 오류율(WER) 지표를 계산하는 클래스.
    """

    def __init__(self, ground_truth_text: str):
        """
        평가기 인스턴스를 초기화합니다.

        :param ground_truth_text: 비교의 기준이 되는 정답 텍스트.
        """
        self.ground_truth = ground_truth_text

    def evaluate(self, hypothesis_text: str) -> dict:
        """
        주어진 예측 텍스트에 대한 모든 평가 지표를 계산합니다.

        :param hypothesis_text: 평가할 OCR 예측 텍스트.
        :return: 평가 결과를 담은 딕셔너리.
        """
        cer_results = self._calculate_strict_cer(self.ground_truth, hypothesis_text)
        wer_results = self._calculate_strict_wer(self.ground_truth, hypothesis_text)
        flexible_cer = self._calculate_flexible_cer(self.ground_truth, hypothesis_text)
        flexible_wer = self._calculate_flexible_wer(self.ground_truth, hypothesis_text)

        results = {
            # Strict CER
            "strict_cer": cer_results["cer"],
            "char_substitutions": cer_results["S"],
            "char_deletions": cer_results["D"],
            "char_insertions": cer_results["I"],
            "char_hits": cer_results["H"],
            # Strict WER
            "strict_wer": wer_results["wer"],
            "word_substitutions": wer_results["S"],
            "word_deletions": wer_results["D"],
            "word_insertions": wer_results["I"],
            "word_hits": wer_results["H"],
            # Flexible Metrics
            "flexible_cer": flexible_cer,
            "flexible_wer": flexible_wer,
        }
        return results

    def _calculate_strict_cer(self, ref: str, hyp: str) -> dict:
        """
        jiwer를 사용하여 엄격한 순서의 CER을 계산합니다.
        """
        if not ref:
            return {"cer": 1.0 if hyp else 0.0, "S": 0, "D": 0, "I": len(hyp), "H": 0}
        output = jiwer.process_characters(ref, hyp)
        return {
            "cer": output.cer,
            "S": output.substitutions,
            "D": output.deletions,
            "I": output.insertions,
            "H": output.hits,
        }

    def _calculate_strict_wer(self, ref: str, hyp: str) -> dict:
        """
        jiwer를 사용하여 엄격한 순서의 WER을 계산합니다.
        """
        if not ref:
            return {"wer": 1.0 if hyp else 0.0, "S": 0, "D": 0, "I": len(hyp.split()), "H": 0}
        output = jiwer.process_words(ref, hyp)
        return {
            "wer": output.wer,
            "S": output.substitutions,
            "D": output.deletions,
            "I": output.insertions,
            "H": output.hits,
        }

    def _calculate_flexible_cer(self, ref: str, hyp: str) -> float:
        """
        fuzzywuzzy의 token_sort_ratio를 사용하여 순서에 유연한 CER을 계산합니다.
        """
        similarity_ratio = fuzz.token_sort_ratio(ref, hyp)
        return (100 - similarity_ratio) / 100.0

    def _calculate_flexible_wer(self, ref: str, hyp: str) -> float:
        """
        fuzzywuzzy의 token_set_ratio를 사용하여 순서에 유연한 WER을 계산합니다.
        """
        similarity_ratio = fuzz.token_set_ratio(ref, hyp)
        return (100 - similarity_ratio) / 100.0