import numpy as np
from datasets import load_dataset
from sklearn.metrics.pairwise import cosine_similarity

palette_dataset = load_dataset("danielritchie/cinematic-mood-palette")["train"]

palette_vectors = []
palette_names = []

for row in palette_dataset:
    palette_vectors.append([row["V"], row["A"], row["D"], row["Cx"], row["Co"]])
    palette_names.append(row.get("name", "unknown"))

palette_vectors = np.array(palette_vectors)


def nearest_palette_vector(raw_vad):
    raw_vec = np.array([[raw_vad[k] for k in ["V","A","D","Cx","Co"]]])
    sims = cosine_similarity(raw_vec, palette_vectors)[0]
    idx = np.argmax(sims)
    anchor = palette_vectors[idx]

    return {
        "vector": {
            "V": anchor[0],
            "A": anchor[1],
            "D": anchor[2],
            "Cx": anchor[3],
            "Co": anchor[4],
        },
        "name": palette_names[idx]
    }


def amplify_with_palette(raw, drama):
    anchor_data = nearest_palette_vector(raw)
    anchor = anchor_data["vector"]

    amplified = {
        k: float(raw[k] + drama * (anchor[k] - raw[k]))
        for k in raw
    }

    return amplified, anchor_data["name"]