# Define data structure
class ExpenseItem(BaseModel):
    name: str
    amount: float
    currency: str
    payer: str
    emoji: str

class Expenses(BaseModel):
    expenses: List[ExpenseItem]

# Encode the images in base64
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")
images_content = [
    {
        "type": "image_url",
        "image_url": {"url": f"data:image/png;base64,{encode_image(ima)}"},
    }
    for ima in pathlib.Path("../tricount_pictures").glob("*.PNG")
]

# Classify
response = client.beta.chat.completions.parse(
    model=MODEL,
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "List all the expenses in a json format",
                },
                *images_content,
            ],
        },
    ],
    temperature=0.0,
    response_format=Expenses,
)

expenses = response.choices[0].message.parsed
# prints
#[ExpenseItem(name='Beers', amount=XXX, currency='DKK', payer='Dudu (me)', emoji='🍺'),
# ExpenseItem(name='Beers', amount=XXX, currency='DKK', payer='Dudu (me)', emoji='🍺'),
# ExpenseItem(name='Decatlon comun', amount=XXX, currency='DKK', payer='Vitto', emoji='💶'),

# define pydantic model 
class EmojiClassification(BaseModel):
    """
    A single emoji that describes a financial transaction (e.g., Airplane ticket -> ✈️)
    """

    emoji: str = Field(
        description="The emoji that describes the transaction (Must be a single character!)"
    )

    @field_validator("emoji")
    @classmethod
    def validate_emoji(cls, value):
        if not emoji.is_emoji(value):
            raise ValueError(f"The emoji must be a single character! Received: {value}")
        return value

# classification function 
def classify_openai(expense: str):
    client = OpenAI()
    messages = [
        {
            "role": "user",
            "content": f"Classify this transaction: {expense}",
        }
    ]
    completion = client.beta.chat.completions.parse(
        model="gpt-4o-mini-2024-07-18",
        messages=messages,
        response_format=EmojiClassification,
    )
    return completion.choices[0].message.parsed.emoji


classify_openai("Treno Milano Ancona")
# prints '🚆'

# Eats pydantic model, shoots 'json schema'
def schema(f) -> dict:
    kw = {
        n: (o.annotation, ... if o.default == Parameter.empty else o.default)
        for n, o in inspect.signature(f).parameters.items()
    }
    s = create_model(f"Input for `{f.__name__}`", **kw).schema()

    return {
        "type": "function",
        "function": {
            "name": f.__name__,
            "description": f.__doc__,
            "parameters": s,
        },
    }

# Don't even try without this (email me if you have a better option!)
SYSTEM_PROMPT = """ 
You are a genius expert. Your task is to classify the description of a financial transaction with an emoji. 
You must only use emojis that are a single character.
You must only reply with a single emoji.

Example 1:
User: Treno Milano Ancona
Response: EmojiClassification(emoji='🚆')

Example 2:
User: Dinner with friends
Response: EmojiClassification(emoji='🍽')

Example 3:
User: Plane ticket to New York
Response: EmojiClassification(emoji='✈️')

Use the EmojiClassification tool to help you classify the transactions.
""".strip()

# This is a bad example of retry logic - but you get the point
def classify_expense(
    expense_description: str,
    model: str,
    retries: int = 3,
    base_url: str,
    api_key: str,
) -> EmojiClassification:

    client = OpenAI(base_url=base_url, api_key=api_key)
    classification_prompt = f"User: {expense_description}"

        # poor man's retry
    for retry_num in range(retries + 1):
          # we use OpenAI's client
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": classification_prompt},
            ],
            tools=[schema(EmojiClassification)],
            temperature=0.0 + retry_num / 10,
        )

        tool_calls = getattr(response.choices[0].message, "tool_calls", None)
        if not tool_calls or tool_calls[0].function.name != "EmojiClassification":
            print(
                f"Failed to classify expense. Expected EmojiClassification but got {tool_calls[0].function.name if tool_calls else 'None'}"
            )
            continue
        try:
            return EmojiClassification(
                **json.loads(tool_calls[0].function.arguments)
            ).emoji
        except ValidationError as e:
            print(f"Validation error: {e}")
            continue
    return None

# we can now call this with ollama
classify_expense(
    expense_description="Donut",
    model="llama3.1",
    base_url="http://localhost:11434/v1",
    api_key="ollama", # dummy key - required
)
# prints '🍩'

# load a small db of emojis and descriptions 
df = pandas.read_csv("./emojis.csv", index_col=False, header=None)
df.columns = ["emoji", "description"]

# prepare column to be embedded (I hate this)
df["description"] = df["description"].str.replace('"', "")
df["to_embed"] = "passage: " + df["description"]

# embed it (multilingual embedding, since the description are in IT and PT)
model = sentence_transformers.SentenceTransformer(
    "intfloat/multilingual-e5-base"
)
embeddings = model.encode(df["to_embed"].tolist(), show_progress_bar=True, device="mps")
df = df.assign(vector=embeddings.tolist())

# create a lancedb table with the data 
db = lancedb.connect(".my_db")
table_name = "emojis"
table = db.create_table(table_name, data=df.to_dict(orient="records"))
print(f"Created table {table_name}")

# Why not try a re-ranker as well? 
ranker = Reranker("nreimers/mmarco-mMiniLMv2-L12-H384-v1")

# Create a classification function 
def embedding_classify(
    expense_description: str,
    top_k: int = 10, # number of closest emoji to re-rank
    use_reranker: bool = True,
    debug: bool = False,
) -> str:
    query_embedding = model.encode(
        f"query: {expense_description}", show_progress_bar=True, device="mps"
    )
    r = table.search(query_embedding).limit(top_k).to_pandas()

    if not use_reranker:
        return r.iloc[0]["emoji"]

    r = ranker.rank(
        query=expense_description,
        docs=r["description"].tolist(),
        doc_ids=r["emoji"].tolist(),
    )
    return r.top_k(1)[-1].doc_id

print(embedding_classify("Beers and pizza with friends", use_reranker=False))
# prints 🍕
print(embedding_classify("Almoco", use_reranker=False))
# prints ⚗

CACHE = {}

def benchmark(expense: ExpenseItem) -> dict:

    if expense.name in CACHE:
        return CACHE[expense.name]

    result = {
        "description": expense.name,
        "tricount_emoji": expense.emoji,
        "openai": classify_openai(expense.name),
        "llama3.1_7b: ": classify_expense(
            expense.name,
            model="llama3.1",
            **api_params,
        ),
        "qwen2.5_3b": classify_expense(
            expense.name,
            model="qwen2.5:3b",
            **api_params,
        ),
        "qwen2.5_1.5b": classify_expense(
            expense.name,
            model="qwen2.5:3b",
            **api_params,
        ),
        "llama3.2_3b": classify_expense(
            expense.name,
            model="llama3.2",
            **api_params,
        ),
        "embedding": embedding_classify(expense.name, use_reranker=False),
        "embedding_reranker": embedding_classify(expense.name, use_reranker=True),
    }

    CACHE[expense.name] = result
    return result


results = [benchmark(expense) for expense in expenses.expenses[:30]]
bench_df = pandas.DataFrame(results)
bench_df

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3.2-1B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# log probability of sentence 
def calculate_sentence_log_probability(sentence):
    inputs = tokenizer(sentence, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits[:, :-1, :]  # Align token predictions
    shifted_input_ids = inputs.input_ids[:, 1:]  # Drop first token

    log_probs = torch.log_softmax(logits, dim=-1)
    token_log_probs = log_probs.gather(2, shifted_input_ids.unsqueeze(-1)).squeeze(-1)

    return token_log_probs.sum().item()

# wrapper to run through the emojis
def find_most_likely_emoji(base_sentence, emoji_list):
    log_probs = {}
    for emoji in emoji_list:
        sentence = base_sentence.format(emoji=emoji)
        log_probs[emoji] = calculate_sentence_log_probability(sentence)

    most_likely_emoji = max(log_probs, key=log_probs.get)
    return most_likely_emoji, log_probs

# prompt
USER_MESSAGE = """ 
You are a genius expert. Your task is to classify the description of a financial transaction with an emoji. 
You must only use emojis that are a single character.
You must only reply with a single emoji.

Example 1:
User: Treno Milano Ancona
Response: 🚆
...

Now it's your turn! Classify the following transactions with an emoji: 
User: {description}
""".strip() 
USER_MESSAGE = USER_MESSAGE.format(description="Voli Tokyo") # 'Flights Tokyo'

# massaging the formats and tokenizing
ASSISTANT_MESSAGE = """Response: {emoji}"""
chat = [
    {"role": "user", "content": USER_MESSAGE},
    {"role": "assistant", "content": ASSISTANT_MESSAGE},
]
input_tokens = tokenizer.apply_chat_template(chat, tokenize=True)[:-1]
base_sentence = tokenizer.decode(input_tokens)

# classification
emoji_list = ["🎄", "🔔", "🚨", "💎", "🇬🇹", "🍹", "🚣", "⚽", "🥘", "🛩️", "🇯🇵", "✈️"] # example possibilities 
most_likely_emoji, log_probs = find_most_likely_emoji(base_sentence, emoji_list)
print(f"Most likely emoji: {most_likely_emoji}")
for emoji, log_prob in log_probs.items():
    print(f"Log probability of {emoji}: {log_prob}")

# Most likely emoji: ✈️
# Log probability of 🎄: -582.9215087890625
# Log probability of 🔔: -580.827392578125
# Log probability of 🚨: -577.4615478515625
# Log probability of 💎: -582.4413452148438
# Log probability of 🇬🇹: -580.4923095703125
# Log probability of 🍹: -581.0596923828125