from unsloth import FastLanguageModel
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-14B",
    max_seq_length = 2048,   # Context length - can be longer, but uses more memory
    load_in_4bit = True,     # 4bit uses much less memory
)

# 1. Load the PEFT config (LoRA metadata)
peft_model_id = "junnei/qwen3-14b-tft"
config = PeftConfig.from_pretrained(peft_model_id)

# 3. Apply the LoRA weights
model = PeftModel.from_pretrained(model, peft_model_id)

# 4. Tokenizer도 로딩
tokenizer = AutoTokenizer.from_pretrained(peft_model_id, trust_remote_code=True)

# (Optional) 모델 evaluation 모드로 전환
model.eval()

import json
from datasets import load_dataset, Dataset
from tqdm import tqdm

prompt_template = """다음은 TFT 게임의 보드 상태입니다. 아래는 각 플레이어의 보드가 가진 특성들과 유닛 구성입니다. 유닛의 위치는 슬롯번호로 표현됩니다.

[Player]
- Traits: {player_traits}
- Units:
{player_units}

[Opponent]
- Traits: {opponent_traits}
- Units:
{opponent_units}

이 보드를 바탕으로 'player'가 승리할지 'opponent'가 승리할지 예측해주세요."""

def trait_to_str(traits):
    return ', '.join([
        f"{t['Grade']} {t['Count']} {t['Trait']}"
        for t in traits
    ])

def units_to_str(units):
    lines = []
    for pos_str in sorted(units.keys(), key=lambda x: int(x)):
        u = units[pos_str]
        if not u or not u.get("Unit"):
            continue
        unit_name = u["Unit"]
        tier = u.get("Tier", "").replace("Tier ", "").replace("tier ", "")
        tier_star = f"Tier {tier}" if tier else ""
        items = u.get("Items", [])
        item_str = f" with {', '.join(items)}" if items else ""
        lines.append(f"  - Slot {pos_str}: {unit_name} ({tier_star}){item_str}")
    return '\n'.join(lines) if lines else "  (no units)"

def convert_example(example):
    player = example["board"]["player"]
    opponent = example["board"]["opponent"]

    player_traits = trait_to_str(player["traits"])
    opponent_traits = trait_to_str(opponent["traits"])

    player_units = units_to_str(player["units"])
    opponent_units = units_to_str(opponent["units"])

    prompt = prompt_template.format(
        player_traits=player_traits,
        player_units=player_units,
        opponent_traits=opponent_traits,
        opponent_units=opponent_units
    )

    label = "player" if example["is_win"] else "opponent"

    return {
        "conversations": [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": f"winner : {label}"}
        ]
    }

# 5. 데이터 로딩 및 huggingface Dataset 객체로 변환
with open("data.jsonl", "r", encoding="utf-8") as f:
    raw_data = [json.loads(line) for line in f]
dataset = Dataset.from_list(raw_data)

# 6. 대화 데이터셋으로 변환
conversation_dataset = dataset.map(convert_example)

print(conversation_dataset[0]['conversations'])

import torch
import torch.nn.functional as F
from tqdm import tqdm

def score_sequence(prefix_ids, candidate_ids, model):
    """
    prefix_ids: List[int], candidate_ids: List[int]
    → prefix에 이어 candidate를 한 토큰씩 feed하면서 log-prob 합산
    """
    input_ids = torch.tensor([prefix_ids], device="cuda")
    total_logprob = 0.0

    for token_id in candidate_ids:
        with torch.no_grad():
            logits = model(input_ids=input_ids).logits[0, -1]
            logprob = F.log_softmax(logits, dim=-1)[token_id].item()
            total_logprob += logprob

        # 다음 단계 계산을 위해 토큰을 append
        input_ids = torch.cat([input_ids, torch.tensor([[token_id]], device="cuda")], dim=1)

    return total_logprob

# ────────────────────────────────────────────────────────────────────────────
correct = 0
total = 0

# 후보 시퀀스 토큰 ID
player_ids   = tokenizer.encode(" player",   add_special_tokens=False)
opponent_ids = tokenizer.encode(" opponent", add_special_tokens=False)
max_cand_len = max(len(player_ids), len(opponent_ids))

for text in tqdm(conversation_dataset['conversations'][:1000]):
    # 1) prompt 생성
    full_prompt = tokenizer.apply_chat_template(text, tokenize=False)
    assistant_content = 'winner :'
    assistant_start = full_prompt.find(assistant_content) + len(assistant_content)
    prompt_text = full_prompt[:assistant_start]

    # 2) prefix 토크나이즈
    inputs = tokenizer(prompt_text, return_tensors="pt").to("cuda")
    prefix_ids = inputs["input_ids"][0].tolist()

    # 3) 두 후보 시퀀스의 누적 log-prob 계산
    player_score   = score_sequence(prefix_ids, player_ids,   model)
    opponent_score = score_sequence(prefix_ids, opponent_ids, model)
    pred_ids       = player_ids if player_score > opponent_score else opponent_ids
    pred_token     = "player"    if player_score > opponent_score else "opponent"
    
    log_probs = torch.tensor([player_score, opponent_score])
    probs = torch.softmax(log_probs, dim=0)
    
    pred_score     = probs[0].item() if player_score > opponent_score else probs[1].item()
    
    # 4) generate() 결과 토큰 시퀀스 획득 (최대 후보 길이만큼)
    output_ids = model.generate(
        **inputs,
        max_new_tokens=max_cand_len,
        do_sample=False,
    )[0]
    gen_token_ids = output_ids[len(prefix_ids): len(prefix_ids) + len(pred_ids)]

    # 5) 매칭 검사
    is_match = gen_token_ids.tolist() == pred_ids
    # 디버깅용 출력 (일치하지 않을 때만)
    if not is_match:
        decoded = tokenizer.decode(gen_token_ids, skip_special_tokens=True)
        print(f"[Mismatch] Score : {pred_score*100:.2f}%, gen_tokens={gen_token_ids.tolist()} ({decoded!r}), "
              f"pred_tokens={pred_ids} ({pred_token})")
    else:
        correct += 1
        if pred_score < 0.9:
            print(f"[Match] Score : {pred_score*100:.2f}%")

    total += 1

print(f"\nAccuracy: {correct}/{total} = {correct/total:.2%}")