from enum import Enum
from pydantic import BaseModel, constr

import pandas as pd
import outlines
import torch


examples = [
    ("대학로 복현오거리에서 경대교방향 원활.", "|대학로 복현오거리|, |경대교방향|, |원할|◉"),
    ("<긴급>구마로 본리네거리에서 동본리네거리방향 동본리네거리 130m 못간 지점 가장자리차로 택시, 승용차 사고", "|구미로 본리네거리|, |동본리네거리|, |사고|◉"),
    ("국채보상로 청구삼거리에서 MBC네거리방향 교통량은 많으나 원활.", "|국체보상로 청구 삼거리|, |MBC네거리방향|, |원할|◉"),
    ("서대구로 북부정류장에서 평리네거리방향 평리지하차도 오르막 지점 가장자리차로 지프차 고장으로 정차중.", "|서대구로 북부정류장|, |평리네거리|, |기타|◉"),
    ("달구벌대로 수성구청역네거리에서 만촌네거리방향 신호 주기에 따라 밀림, 풀림 반복.", "|달구벌대로 수성구청역네거리|, |만촌네거리|, |정체|◉"),
    ("동부로 동대구역네거리에서 경북수협네거리방향 경북수협네거리 80m 못간 지점 가장자리차로 크레인 동원한 작업중.", "|동부역 동대구역네거리|, |경북수협네거리|, |공사|◉"),
    ("국도 5호선 영주에서 안동구간 영주시 이산면 송리원주유소 부근 양방향 원활.", "|국도 5호선 영주|, |영주시 이산면 송리원 주유소|, |원할|◉"),
    ("<긴급>구마로 본리네거리에서 성당네거리방향 동본리네거리 지나자마자 가장자리차로 사고로 보이는 승용차, 승합차 정차중.", "|구마로 본리네거리|, |성당네거리방향 동본리네거리|, |정체|◉"),
    ("대학로 경북대학교 북문에서 복현오거리방향 복현오거리 신호 3,4번.", "|경북대학교 북문|, |복현오거리방향|, |정체|◉"),
    ("안들림", "|알수 없음|◉"),
]

# schema = """{
#     "$defs": {
#         "상태": {
#             "enum": ["원할", "사고", "정체", "공사", "기타"],
#             "title": "상태",
#             "type": "string"
#         }
#     },
#     "properties": {
#         "location_from": {"maxLength": 10, "title": "Name", "type": "string"},
#         "location_to": {"title": "Age", "type": "integer"},
#         "상태": {"$ref": "#/$defs/상태"},
#         "strength": {"title": "Strength", "type": "integer"}\
#     },
#     "required": ["name", "age", "armor", "strength"],
#     "title": "Character",
#     "type": "object"
# }"""


@outlines.prompt
def labelling(to_label, examples):
    """You are a sentiment-labelling assistant.
    신호 1번 == 원할
    신호 1,2번 == 원할
    신호 2번 == 정체
    신호 3번 == 정체
    신호 4번 == 정체
    {% for example in examples %}
    {{ example[0] }} // {{ example[1] }}
    {% endfor %}
    {{ to_label }} //
    """

model = outlines.models.transformers("mistralai/Mistral-7B-instruct-v0.2", device="cuda")

df = pd.read_csv("data/최종데이터.csv")

for row in df["내용"]:
    print(f"input : {row}")
    prompt = labelling(row, examples)
    answer = outlines.generate.text(model)(prompt, max_tokens=80, stop_at="◉")
    print(f"output : {answer}")