import json
import os
from typing import Optional, List

import numpy as np
from deepeval.metrics import BaseMetric, ContextualRelevancyMetric
from deepeval.metrics.contextual_relevancy.template import ContextualRelevancyTemplate
from deepeval.models import DeepEvalBaseLLM
from deepeval.test_case import LLMTestCase
from langchain_community.llms.tongyi import Tongyi


class ContextualRelevancyWrapper(BaseMetric):
    def __init__(self):
        qwen_model = QwenModel(
            model_name="qwen2-72b-instruct",
            temperature=0.2,
            top_p=0.8
        )
        self.delegate = ContextualRelevancyMetric(
            threshold=0.5,
            model=qwen_model,
            include_reason=True,
            async_mode=False
        )

    def response_detail(self) -> dict:
        verdicts = [
            dict(verdict=verdict.verdict, reason=verdict.reason) for verdict in self.delegate.verdicts
        ]
        ap = self.calculate_average_precision([verdict['verdict'] for verdict in verdicts])
        detail = dict(
            reason=self.delegate.reason,
            average_precision=ap,
            success=self.delegate.success,
            verdicts=json.dumps(verdicts, indent=4, ensure_ascii=False)
        )
        return detail

    def measure(self, test_case: LLMTestCase, *args, **kwargs) -> float:
        return self.delegate.measure(test_case=test_case, *args, **kwargs)

    async def a_measure(self, test_case: LLMTestCase, *args, **kwargs) -> float:
        return await self.delegate.a_measure(test_case=test_case, *args, **kwargs)

    def is_successful(self) -> bool:
        return self.delegate.is_successful()

    def calculate_average_precision(self, verdicts: List[str]):
        """
        計算mAP
        :param verdicts: ["yes", "no", "no"]
        :return:
        """
        verdict_list = np.array([1 if verdict == "yes" else 0 for verdict in verdicts])
        denominator = sum(verdict_list) + 1e-10

        numerator = sum(
            [
                (sum(verdict_list[: i + 1]) / (i + 1)) * verdict_list[i]
                for i in range(len(verdict_list))
            ]
        )

        score = numerator / denominator
        return score


class QwenModel(DeepEvalBaseLLM):
    def __init__(
            self,
            model_name: str = "qwen2-72b-instruct",
            dashscope_api_key: Optional[str] = None,
            **kwargs
    ):
        self.model_name = model_name

        dashscope_api_key = dashscope_api_key or os.environ.get("dashscope_sk")
        self.model = Tongyi(model=model_name, api_key=dashscope_api_key, **kwargs)
        self.kwargs = kwargs

    def load_model(self, *args, **kwargs):
        return self.model

    def generate(self, prompt_str: str) -> str:
        return self.model.invoke(prompt_str)

    async def a_generate(self, prompt_str: str) -> str:
        try:
            return await self.model.ainvoke(prompt_str)
        except Exception as e:
            print(f"Failed in invoke: {e}\nTraceback:\n{e}")
            return ""

    def get_model_name(self, *args, **kwargs) -> str:
        return self.model_name


class CustomContextualRelevancyTemplate:
    @staticmethod
    def generate_reason(input, irrelevancies, score):
        return f"""Based on the given input, reasons for why the retrieval context is irrelevant to the input, and the contextual relevancy score (the closer to 1 the better), please generate a CONCISE reason for the score.
In your reason, you should quote data provided in the reasons for irrelevancy to support your point.

** 
IMPORTANT: Please make sure to only return in JSON format, with the 'reason' key providing the reason, reason field must return in **Chinese**.
Example JSON:
{{
    "reason": "得分為 <contextual_relevancy_score>，因為 <YOUR_REASON>。"
}}

If the score is 1, keep it short and say something positive with an upbeat encouraging tone (but don't overdo it otherwise it gets annoying).
**


Contextual Relevancy Score:
{score}

Input:
{input}

Reasons for why the retrieval context is irrelevant to the input:
{irrelevancies}

JSON:
"""

    @staticmethod
    def generate_verdict(text, context):
        return f"""Based on the input and context, please generate a JSON object to indicate whether the context is relevant to the provided input. The JSON will have 1 mandatory field: 'verdict', and 1 optional field: 'reason'.
The 'verdict' key should STRICTLY be either 'yes' or 'no', and states whether the context is relevant to the input. 
Provide a 'reason' ONLY IF verdict is no. You MUST quote the irrelevant parts of the context to back up your reason.

**
IMPORTANT: Please make sure to only return in JSON format.reason field must return in **Chinese**
Example Context: "愛因斯坦因發現光電效應而獲得諾貝爾獎。他在1968年獲得了諾貝爾獎。那里有一只貓。"
Example Input: "愛因斯坦的一些成就是什么？"

Example:
{{
    "verdict": "no",
    "reason": "雖然上下文包含有關愛因斯坦獲得諾貝爾獎的信息，但不相關地提到“那里有一只貓”，這與愛因斯坦的成就無關。"
}}
**

Input:
{text}

Context:
{context}

JSON:
"""


def patch_contextual_relevancy():
    ContextualRelevancyTemplate.generate_reason = CustomContextualRelevancyTemplate.generate_reason
    ContextualRelevancyTemplate.generate_verdict = CustomContextualRelevancyTemplate.generate_verdict


if __name__ == '__main__':
    patch_contextual_relevancy()
    cr_metric = ContextualRelevancyWrapper()
    testcase = LLMTestCase(
        input="小狗貧血的表現",
        actual_output="",
        retrieval_context=[
            "狗狗貧血的主要癥狀包括：<br>1. 身體消瘦、毛發粗糙無光澤、體力較差，活動時間短、運動無力。<br>2. 皮膚和可視黏膜蒼白，如舌頭顏色變粉紅或蒼白，牙齦蒼白。<br>3. 精神不振、嗜睡、經常出現走路搖晃、倒地后起立困難、臥地不起等癥狀。<br>4. 食欲不振、日常可能會暈倒、身體虛弱等。<br>如果發現狗狗出現以上癥狀，建議及時帶狗狗去獸醫院進行檢查和治療。在日常生活中，可以通過改善飲食結構、補充造血物質和營養品等方法來幫助狗狗改善貧血狀況。",
            "狗狗一直餓可能是由以下幾種原因導致的：<br>1. 喂食量不足：如果狗狗每次的喂食量太少，即使一天喂5次，也可能導致狗狗吃不飽。狗狗一天的喂食總量一般是根據體重來確定的，一般按照狗狗體重的3%- 5%喂食。<br>2. 幼犬對吃飽沒有太大概念，少量多次喂食，且對于好吃的食物沒有抵抗力，此時并不是因為餓，只是饞。<br>3. 糖尿病問題：如果是中老年犬出現長期饑餓的情況，可能是糖尿病或其他疾病引起的多飲多食、體重減輕的癥狀，建議及時帶狗狗去寵物醫院檢查。<br>4. 飲食不均衡或者不足：如果狗狗的飲食缺乏必要的營養成分，或者狗狗沒有得到足夠的食物，那么他們就會變得更饑餓。<br>5. 狗狗處于生長期：如果狗狗還很小，或者正處于生長期，那么他們需要更多的食物來支持他們的身體發育。<br>6. 狗狗在鍛煉或者活動后饑餓：如果狗狗進行了大量的鍛煉或者活動，那么他們會消耗更多的能量，因此可能會感到更饑餓。<br>如果狗狗一直處于饑餓狀態，建議帶它去看獸醫，以確定是否存在任何潛在的健康問題。",
            "狗狗如果出現了貧血，<em>它會看上去很沒有精神，也沒有食欲，身體變得消瘦而且毛發還非常的枯燥，</em>引起狗狗貧血有三方面的原因。 導致狗狗貧血的原因是多方面的，主人一旦發現狗狗沒有食欲精神不好，就要及時采取措施給予狗狗營養的補充。",
            "1、牙齦、眼瞼、結膜等皮膚黏膜都會從粉紅色、淡粉色，變淡變白。 2、在<em>貧血</em>初期出現嗜睡，出門散步<em>表現</em>得過度喘息、有氣無力，容易疲倦。 3、食欲減退、廢絕，體重下降快，毛發干枯發黃，無光澤。 4、<em>貧血</em>后期會出現運動無力、搖晃、起身困難、呼吸困難、全身衰竭。",
            "小狗肚子餓了，<em>最多也就是發出比較輕微的嗚咽聲。</em>像個嬰兒一樣，聲音雖不大，但是聽上去很可憐。 幼犬的自理能力比較差，所以主人要做到準時喂食，一般3~5小時就要喂一次。 把食盆叼過來給主人 狗狗一大早起來就肚子餓了，但是看到主人還在睡，自己也不敢發出很大的聲響。 把食盆叼過來給主人 狗狗一大早起來就肚子餓了，但是看到主人還在睡，自己也不敢發出很大的聲響。 就會選擇叼著一個狗碗去找主人，示意你該喂狗糧了，主人再想賴床也不行了。",
            "<em>臨床表現為呼吸加快、全身無力等特征。</em>狗狗貧血和發炎，通過血檢的血球容積比、血色素可以判斷。數值低下時會引起貧血，數值升高時就是正在發炎，不過當免疫力降低時，就算白血球沒有增加，也有可能正在發炎。",
        ],
    )
    score = cr_metric.measure(testcase)
    detail = cr_metric.response_detail()
    print(f"score:{score}, mAP:{detail['average_precision']}")
    print(f"detail_verdicts:{json.dumps(json.loads(detail['verdicts']), indent=4, ensure_ascii=False)}")

score:0.6666666666666666, mAP:0.7708333333140625
detail_verdicts:[
    {
        "verdict": "yes",
        "reason": null
    },
    {
        "verdict": "no",
        "reason": "上下文中沒有提到小狗貧血的表現，而是詳細描述了狗狗可能一直感到饑餓的原因，這與輸入的小狗貧血的表現無關。"
    },
    {
        "verdict": "yes",
        "reason": null
    },
    {
        "verdict": "yes",
        "reason": null
    },
    {
        "verdict": "no",
        "reason": "上下文中沒有提到小狗貧血的表現，而是提到了小狗肚子餓的行為，這與輸入的問題無關。"
    },
    {
        "verdict": "yes",
        "reason": null
    }
]