from tqdm import tqdm
import os
from suno_utils.utils.text import read_jsonl, read_json, write_jsonl

"""
This script assumes you've already run modal encode and have the ditto embeddings saved to COVER_PATH
DITTO_PATH is the json created by save_ditto_json.py

Creates a metas.jsonl file for training that is a subset of COVER_PATH
    filters cover pairs by ditto self_sim scores
    option to include default generations (equal to number of cover pairs)
    option to filter based on whether a cover is in discogs (professional filter)
"""

TASK = "self_sim"  # update this
DITTO_PATH = "/app/suno/sara/cover_filter/raw_ditto_mappings_v0.json"
COVER_PATH = "/app/suno/sara/cover_filter/filtered_cover_name_view_detailed.jsonl"
DISCOGS_PATH = "/app/suno/sara/metas_v0_discogs.jsonl"
OUT_FOLDER = "/app/suno/sara/cover_filter_v0/"


min_sim = 0.4
max_sim = 0.8
discogs_filter_type = "all"
max_covers_per_source = 10

if discogs_filter_type == "all":
    discogs_ids = {}
else:
    print("Loading discogs...")
    discogs = read_jsonl(DISCOGS_PATH, progress=True)
    discogs_ids = {}
    for data in discogs:
        discogs_ids[data["id"]] = 0  # data
    del discogs

print("Loading covers..")
test = read_jsonl(COVER_PATH, progress=True)

print("loading ditto scores...")
ditto_scores = read_json(DITTO_PATH)

print("Filtering...")
parent_to_data = {}
parent_to_covers = {}
skipped = 0
for row in tqdm(test):
    if "parent_id" not in row:
        parent_to_data[row["id"]] = row
        parent_to_covers[row["id"]] = []
for row in tqdm(test):
    if "parent_id" in row:
        parent = row["parent_id"]
        child = row["id"]

        if parent not in ditto_scores or child not in ditto_scores[parent]:
            skipped += 1
            continue
        similarity_score = float(ditto_scores[parent][child])
        if similarity_score > min_sim and similarity_score < max_sim:
            if discogs_filter_type == "professional":
                if child in discogs_ids:
                    parent_to_covers[parent].append(row)
            elif discogs_filter_type == "amateur":
                if child not in discogs_ids:
                    parent_to_covers[parent].append(row)
            else:
                parent_to_covers[parent].append(row)
print(f"{skipped} missing ditto scores out of {len(test)}.")

final_data = []
sources = 0
num_covers = 0
for parent, covers in parent_to_covers.items():
    if len(covers) > 0:
        final_data.append(parent_to_data[parent])
        sources += 1
        trimmed_covers = covers[:10]
        for cover in trimmed_covers:
            final_data.append(cover)
            num_covers += 1

print(f"Found {num_covers} from {sources} sources. Writing to file...")
write_jsonl(
    final_data,
    os.path.join(
        OUT_FOLDER,
        f"covers_filtered_{min_sim}_{max_sim}_{discogs_filter_type}.jsonl",
    ),
)