from pathlib import Path
import pandas as pd
from PIL import Image
import base64
import io
from sklearn.model_selection import train_test_split

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Extract pairs of (image, label)
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~

root_dir = Path() / "satellite_dataset" / "AID"
data = []
for d in root_dir.iterdir():
    if not d.is_dir():
        continue
    data.extend([{"label": d.name, "img_path": p} for p in d.iterdir()])

dataset_df = pd.DataFrame(data)
classes = [*dataset_df["label"].unique()]

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Load image and encode in base64 (this might take a few minutes)
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# NOTE: This is not needed here, but a nice additional step would be to resize
# the images into 1024 longest edge (if the image was too big)
# For more details see: https://docs.mistral.ai/capabilities/vision/

def encode_image_to_base64(image_path: str | Path) -> str:
    image = Image.open(image_path)

    # Resize the image by a factor of 0.5
    new_size = (image.width // 2, image.height // 2)
    image = image.resize(new_size, Image.Resampling.LANCZOS)

    buffer = io.BytesIO()
    image.save(buffer, format='JPEG')
    buffer.seek(0)

    encoded_string = (
        "data:image/jpeg;base64,"
        + base64.b64encode(buffer.read()).decode('utf-8')
    )

    return encoded_string


dataset_df["img_b64"] = [
    encode_image_to_base64(img_path)
    for img_path in dataset_df["img_path"]
]


# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Split dataset in train / test
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~

train_df, test_df = train_test_split(
    dataset_df,
    test_size=0.2,
    random_state=42,
    stratify=dataset_df["label"]
)
train_df = train_df.reset_index(drop=True)
test_df = test_df.reset_index(drop=True)

# Release a bit of memory
del dataset_df


# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Light check at the data
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
print("Classes:", classes)
print("Size train:", len(train_df))
print("Size test:", len(test_df))

train_df.head()

from IPython.display import display, HTML

def display_image(dataset_df: pd.DataFrame, idx: int) -> None:
    img_b64 = dataset_df["img_b64"].iloc[idx]
    label = dataset_df["label"].iloc[idx]
    display(HTML(f'<h2>{label}</h2><img src="{img_b64}" />'))

display_image(train_df, idx=0)
display_image(train_df, idx=2)

# NOTE: We had to engineer a slighlty complex system prompt as the current
# finetuning of Pixtral-12B (July 2nd 2025) misses a few features that more
# recent models offer:
# - The Classifier Factory in LaPlateforme only supports ministral-3b for now
#   See: https://docs.mistral.ai/capabilities/finetuning/classifier_factory/
# - Structured outputs are not supported by Pixtral-12B

from pydantic import BaseModel

class ClassifierOutput(BaseModel):
    image_description: str
    label: str

instruction_prompt = f"""
Classify the following image into the category it belongs to.
- These category labels are {'; '.join(classes)}.
- Output your result using exclusively the following schema: {ClassifierOutput.model_fields}
- Put your results between a json tag
  ```json
  ```
""".strip()

print("Instruction prompt:")
print(instruction_prompt)

_template_inference_message(train_df["img_b64"].iloc[0])

_template_finetuning_message(train_df["img_b64"].iloc[0], train_df["label"].iloc[0])

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# We will use the Batch API for inference on the test dataset,
# and thus apply the proper formatting
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~

inference_data = []
for row_id, row in test_df.iterrows():
    messages = _template_inference_message(img_b64=row["img_b64"])
    inference_data.append(
        {
            "custom_id": str(row_id),
            "body": {
                "messages": [m.model_dump() for m in messages]
            }
        }
    )

inference_data[0]

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# We will finetune Pixtral on the train dataset
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
finetuning_data = []
for row_id, row in train_df.iterrows():
    messages = _template_finetuning_message(img_b64=row["img_b64"], label=row["label"])
    finetuning_data.append([m.model_dump() for m in messages])

finetuning_data[0]

from getpass import getpass
from mistralai import Mistral

api_key= getpass("Type your API Key")
client = Mistral(api_key=api_key)
pixtral = "pixtral-12b-2409"

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# We call the "base" Pixtral-12B on an example that works
# well (right label is retrieved)
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
IDX = 0
user_message = _template_inference_message(train_df["img_b64"].iloc[IDX])

results = client.chat.complete(
    model=pixtral,
    messages=user_message,
    temperature=0
)
results = parse_content(results.choices[0].message.content)

display_image(train_df, idx=IDX)
print(results)

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# We call the "base" Pixtral-12B on an example that does not work
# well (wrong label is retrieved)
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
IDX = 4
user_message = _template_inference_message(train_df["img_b64"].iloc[IDX])

results = client.chat.complete(
    model=pixtral,
    messages=user_message,
    temperature=0
)
results = parse_content(results.choices[0].message.content)

display_image(train_df, idx=IDX)
print(results)

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Launch a batch inference job
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# NOTE: This can also be done via LaPlateforme UI by drag-and-drop of
# the input files
# NOTE: There might be some size limitation, both in terms of number
# of records and memory footprint of the model

# Upload the dataset
with open(Path() / "test_satellite_for_batch_inference.jsonl", "rb") as f:
    batch_inference_data = client.files.upload(
        file={
            "file_name": "test_satellite_for_batch_inference.jsonl",
            "content": f
        },
        purpose="batch"
    )

# Launch the job
created_job = client.batch.jobs.create(
    input_files=[batch_inference_data.id],
    model=pixtral,
    endpoint="/v1/chat/completions",
    metadata={"job_type": "testing"}
)

print("file ID:", batch_inference_data.id)
print("job ID:", created_job.id)

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Retrieve job results (with a timeout of 1h)
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
import time

with timeout(3600):
    while True:
        retrieved_job = client.batch.jobs.get(job_id=created_job.id)
        if retrieved_job.status not in ["RUNNING", "QUEUED"]:
            print(f"Job finished with status {retrieved_job.status}")
            break

        time.sleep(30)  # Pause for 30s


if not retrieved_job.status == "SUCCESS":
    raise RuntimeError("Batch job failed")

output_file_stream = client.files.download(file_id=retrieved_job.output_file)
results_df = pd.read_json(path_or_buf=output_file_stream, lines=True)
results_df.head()

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Apply processing to the retrieved results
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
parsed_results_df = results_df.copy()
parsed_results_df[f"predicted_{pixtral}"] = (
    parsed_results_df["response"]
    .apply(lambda response: parse_content(response["body"]["choices"][0]["message"]["content"]))
    .apply(lambda x: x.label if x else None)
)

parsed_results_df = (
    parsed_results_df.sort_values("custom_id")
    .reset_index(drop=True)
    [["custom_id", f"predicted_{pixtral}"]]
)

parsed_results_df.head()

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Let's compare against our "ground_truth"
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
comparison_df = test_df.copy()
comparison_df["custom_id"] = comparison_df.reset_index(drop=False)["index"]
comparison_df = comparison_df.merge(parsed_results_df)
comparison_df.head()

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Display traditional classification metrics
# ✅ Good baseline model: some classes are very well handled
#   (e.g. Desert, Railway station, etc.)
# ❌ But there are limitations
#   - some labels are "hallucinated" (i.e. not present in the system prompt)
#   - some classes are not well handled (e.g. River vs. Bridge)
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~

from sklearn.metrics import confusion_matrix, classification_report, accuracy_score, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

y_true = [*comparison_df["label"]]
y_pred = [
    # NOTE: Pixtral-12B might have hallucinated some categories
    # We'll put them under the "WRONG" label
    label if label in classes else "WRONG"
    for label in comparison_df[f"predicted_{pixtral}"]
]
labels = [*classes, "WRONG"]

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Confusion matrix
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
cm = confusion_matrix(y_true, y_pred, labels=labels)
print("Confusion Matrix:")

disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
fig, ax = plt.subplots(figsize=(8, 6))
disp.plot(ax=ax, cmap='Oranges', colorbar=False)
plt.title(f"Confusion Matrix ({pixtral})", fontsize=18, fontweight='bold')
plt.xticks(rotation=90)
plt.show()

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Classification report
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
print("Classification Report:")
print(classification_report(y_true, y_pred, target_names=labels))

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Launch a finetuning job
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# NOTE: This can also be done via LaPlateforme UI by drag-and-drop of
# the input files
# NOTE: There might be some size limitation, both in terms of number
# of records and memory footprint of the model

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Upload your training dataset
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~

with open(Path() / "train_satellite_for_finetuning.jsonl", "rb") as f:
    finetuning_data = client.files.upload(
        file={
            "file_name": "train_satellite_for_finetuning.jsonl",
            "content": f
        },
        purpose="fine-tune"
    )

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Create the job
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
#
# Fine-tuning is both an art and a science
#
# Our recommended approach for fine-tuning our models is to do so iteratively.
# Start with smaller dataset e.g 200-1000 examples (and potentially smaller
# models), larger learning rates, lower epochs and once you start to see the
# desired results in the correct direction start scaling these to be bigger
# and bigger.
#
# Once you are somewhat comfortable with the fine-tuning data used for the
# model, you want to start playing around with the different fine-tuning
# hyperparameters.
#
# Learning Rate - The learning rate defines how much the model adjusts for each
# training step. Bigger learning rates will make the model move faster so you
# potentially may need less training time, however it may also potentially
# overshoot and reach a local minima
#
# Epochs - This is the amount of full runs across the whole dataset the training
# will do. It’s somewhat proportional to the training time.
#
# Batch Size - This is the number of examples passed in each gradient step to
# updates the FT-weights. A higher batch size typically leads to smoother
# training, but with small datasets you typically use a smaller batch size.
#
# Here is an equation that summarize the relationship between the parameters:
# Epochs = Steps × Batch Size​ / Total Number of Training Samples

created_job = client.fine_tuning.jobs.create(
    model="pixtral-12b-latest",
    training_files=[{"file_id": finetuning_data.id, "weight": 1}],
    validation_files=[],
    hyperparameters={
        "epochs": 2,
        "learning_rate": 0.0001,
    },
    auto_start=True  # This could be set to False to let you manually validate the run
)


print("file ID:", finetuning_data.id)
print("job ID:", created_job.id)

with timeout(3600):
    while True:
        retrieved_job = client.fine_tuning.jobs.get(job_id=created_job.id)
        if retrieved_job.status not in ["VALIDATING", "RUNNING", "QUEUED"]:
            print(f"Job finished with status {retrieved_job.status}")
            break

        time.sleep(30)  # Pause for 30s

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# See all the details associated with a finetuning job
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
retrieved_job

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Visualize the evolution of the loss during FT
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
def plot_loss(job_id: str, model: str) -> None:
    # Extract step_number and train_loss from the job data
    retrieved_job = client.fine_tuning.jobs.get(job_id=job_id)
    step_numbers = [checkpoint.step_number for checkpoint in retrieved_job.checkpoints]
    train_losses = [checkpoint.metrics.train_loss for checkpoint in retrieved_job.checkpoints]

    # Plot the data
    plt.figure(figsize=(12, 8))
    plt.plot(step_numbers, train_losses, marker='o', linestyle='-', color='b', label='Train Loss')

    # Adding titles and labels
    plt.title(f'Finetune {model} on your data - Train Loss', fontsize=16, fontweight='bold')
    plt.xlabel('Step Number', fontsize=14)
    plt.ylabel('Train Loss', fontsize=14)

    # Adding grid for better readability
    plt.grid(True, linestyle='--', alpha=0.7)

    # Adding legend
    plt.legend(loc='best')

    # Customizing the appearance
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)

    # Adding a background color
    plt.gca().set_facecolor('#f0f0f0')

    # Display the plot
    plt.show()

plot_loss(job_id=retrieved_job.id, model=pixtral)

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Retrieve the ID of the fine-tuned model
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
fine_tuned_pixtral = client.fine_tuning.jobs.get(job_id=created_job.id).fine_tuned_model
fine_tuned_pixtral

# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
# Make a batch inference with fine-tuned model
# ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~

created_job = client.batch.jobs.create(
    input_files=[batch_inference_data.id],
    model=fine_tuned_pixtral,
    endpoint="/v1/chat/completions",
    metadata={"job_type": "testing"}
)

print("file ID:", batch_inference_data.id)
print("job ID:", created_job.id)

with timeout(3600):
    while True:
        retrieved_job = client.batch.jobs.get(job_id=created_job.id)
        if retrieved_job.status not in ["RUNNING", "QUEUED"]:
            print(f"Job finished with status {retrieved_job.status}")
            break

        time.sleep(30)  # Pause for 30s


if not retrieved_job.status == "SUCCESS":
    raise RuntimeError("Batch job failed")

output_file_stream = client.files.download(file_id=retrieved_job.output_file)
results_df = pd.read_json(path_or_buf=output_file_stream, lines=True)
results_df.head()