import matplotlib.pyplot as plt
import numpy as np
from matplotlib.offsetbox import AnnotationBbox, OffsetImage
from PIL import Image

#### --- Plot data and logos --- ####

data_danish = {
    "comma-v0.1-2t": {"parameters (billions)": 7.0, "Danish Performance": 24.2},
    "Stage 1": {"parameters (billions)": 7.0, "Danish Performance": 35.6},
    "Stage 2": {"parameters (billions)": 7.0, "Danish Performance": 37.0},
    # stage 3:
    "dfm-decoder-open-v0-7b-pt": {
        "parameters (billions)": 7.0,
        "Danish Performance": 37.4,
    },
    "Pleias-350M": {"parameters (billions)": 0.35, "Danish Performance": 10.4},
    "Pleias-1.2B": {"parameters (billions)": 1.2, "Danish Performance": 17.7},
}

data_english = {
    "comma-v0.1-2t": {"parameters (billions)": 7.0, "English Performance": 51.6},
    "Stage 1": {"parameters (billions)": 7.0, "English Performance": 48.2},
    "Stage 2": {"parameters (billions)": 7.0, "English Performance": 50.2},
    # stage 3:
    "dfm-decoder-open-v0-7b-pt": {
        "parameters (billions)": 7.0,
        "English Performance": 50.1,
    },
    "Pleias-350M": {"parameters (billions)": 0.35, "English Performance": 18.9},
    "Pleias-1.2B": {"parameters (billions)": 1.2, "English Performance": 29.4},
}

data = {}
for model in data_danish.keys():
    data[model] = {
        "parameters (billions)": data_danish[model]["parameters (billions)"],
        "Danish Performance": data_danish[model]["Danish Performance"],
        "English Performance": data_english[model]["English Performance"],
        "Danish x English Performance": (
            data_danish[model]["Danish Performance"]
            + data_english[model]["English Performance"]
        )
        / 2,
    }


# Map models to logo files and sizes
logos = {
    "comma-v0.1-2t": {"path": "eleutherai.png", "zoom": 0.035},
    "dfm-decoder-open-v0-7b-pt": {"path": "dfm.png", "zoom": 0.018},
    "Pleias-350M": {"path": "pleias.png", "zoom": 0.30},
    "Pleias-1.2B": {"path": "pleias.png", "zoom": 0.30},
}


#### --- Create plot of Danish performance --- ####

model_names = list(data.keys())
sizes = [data[m]["parameters (billions)"] for m in model_names]
performance = [data[m]["Danish Performance"] for m in model_names]

# Create plot
plt.figure(figsize=(7, 7))

# Draw Pareto frontier
x_curve = np.linspace(0.0, 8, 100)
y_curve = 19 + 3.9 * np.log(x_curve)
plt.plot(x_curve, y_curve, "--", color="grey", linewidth=0.7, alpha=0.7)
plt.text(
    2,
    20,
    "Previous Pareto frontier for openly licensed data",
    fontsize=8,
    color="grey",
    style="italic",
    rotation=7.5,
)


# Get axis reference
ax = plt.gca()

# Add logos for specific models
for i, name in enumerate(model_names):
    if name in logos:
        img = Image.open(logos[name]["path"])
        imagebox = OffsetImage(img, zoom=logos[name]["zoom"])
        ab = AnnotationBbox(imagebox, (sizes[i], performance[i]), frameon=False, pad=0)
        ax.add_artist(ab)

        # Add label below/beside logo
        plt.annotate(
            name,
            (sizes[i], performance[i]),
            xytext=(5, -15),
            textcoords="offset points",
            fontsize=9,
        )

plt.xlabel("Parameters (billions)", fontsize=12)
plt.ylabel("Danish Performance", fontsize=12)

# Remove top and right spines
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
plt.xlim(0, 10)
plt.ylim(0, 40)
plt.savefig("danish-perf.png", dpi=300)


#### --- Create plot of Danish x English performance --- ####

model_names = list(data.keys())
sizes = [data[m]["parameters (billions)"] for m in model_names]
performance = [data[m]["Danish x English Performance"] for m in model_names]

# Create plot
plt.figure(figsize=(7, 7))

# Draw Pareto frontier
x_curve = np.linspace(0.2, 8, 100)
y_curve = 26 + 7 * np.log(x_curve)
plt.plot(x_curve, y_curve, "--", color="grey", linewidth=0.7, alpha=0.7)

# Get axis reference
ax = plt.gca()

# Add logos for specific models
for i, name in enumerate(model_names):
    if name in logos:
        img = Image.open(logos[name]["path"])
        imagebox = OffsetImage(img, zoom=logos[name]["zoom"])
        ab = AnnotationBbox(imagebox, (sizes[i], performance[i]), frameon=False, pad=0)
        ax.add_artist(ab)

        # Add label below/beside logo
        plt.annotate(
            name,
            (sizes[i], performance[i]),
            xytext=(5, -15),
            textcoords="offset points",
            fontsize=9,
        )

plt.xlabel("Parameters (billions)", fontsize=12)
plt.ylabel("Danish x English Performance", fontsize=12)

# Remove top and right spines
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
plt.xlim(0, 10)
plt.ylim(0, 45)
plt.savefig("danish-english-perf.png", dpi=300)


#### --- Create plot of English performance --- ####

model_names = list(data.keys())
sizes = [data[m]["parameters (billions)"] for m in model_names]
performance = [data[m]["English Performance"] for m in model_names]

# Create plot
plt.figure(figsize=(7, 7))


# Get axis reference
ax = plt.gca()

# Add logos for specific models
for i, name in enumerate(model_names):
    if name in logos:
        img = Image.open(logos[name]["path"])
        imagebox = OffsetImage(img, zoom=logos[name]["zoom"])
        ab = AnnotationBbox(imagebox, (sizes[i], performance[i]), frameon=False, pad=0)
        ax.add_artist(ab)

        # Add label below/beside logo
        plt.annotate(
            name,
            (sizes[i], performance[i]),
            xytext=(5, -15),
            textcoords="offset points",
            fontsize=9,
        )

plt.xlabel("Parameters (billions)", fontsize=12)
plt.ylabel("English Performance", fontsize=12)

# Remove top and right spines
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
plt.xlim(0, 10)
plt.ylim(0, 55)
plt.savefig("english-perf.png", dpi=300)