#!/usr/bin/env python
# -*- coding: utf-8 -*-
from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from chromadb import HttpClient
import torch
import spacy
from spacy.matcher import PhraseMatcher
import re

# PARAMETERS
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
DATA_ID = "tapxc3/owast_new"
BOOK_TITLES = [
    "Nineteen Eighty-Four", "Animal Farm", "Shooting an Elephant", "Homage to Catalonia",
    "Coming Up for Air", "Down and Out in Paris and London", "Keep the Aspidistra Flying",
    "A Clergyman's Daughter", "Fifty Orwell Essays", "Politics and the English Language",
    "The Road to Wigan Pier", "Burmese Days"
]

spelled_to_digit = {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six": 6,
                   "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11, "twelve": 12,
                   "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16, "seventeen": 17,
                   "eighteen": 18, "nineteen": 19, "twenty": 20}

digit_to_spelled = {v: k for k, v in spelled_to_digit.items()}

digit_to_roman = {v: r for v, r in enumerate(
    ["i", "ii", "iii", "iv", "v", "vi", "vii", "viii", "ix", "x", "xi", "xii", "xiii", "xiv", "xv", "xvi", "xvii", "xviii", "xix", "xx"],
    start=1
)}

ordinal_to_digit = {
    "first": 1, "second": 2, "third": 3, "fourth": 4, "fifth": 5, "sixth": 6,
    "seventh": 7, "eighth": 8, "ninth": 9, "tenth": 10, "eleventh": 11, "twelfth": 12,
    "thirteenth": 13, "fourteenth": 14, "fifteenth": 15, "sixteenth": 16, "seventeenth": 17,
    "eighteenth": 18, "nineteenth": 19, "twentieth": 20
}

roman_to_int = {'I': 1, 'V': 5, 'X': 10, 'L': 50, 'C': 100, 'D': 500, 'M': 1000}

print("Connecting to ChromaDB...")
chroma_client = HttpClient(host="34.126.200.250", port=8000)
collection = chroma_client.get_collection(name="orwell_books")
print("Connected to ChromaDB.")