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

"""
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 an info.json file to specify which indexes in COVER_PATH to include in training
    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)
"""

SPLIT = "val"  # val or tr
DITTO_PATH = "/app/suno/sara/cover_filter/raw_ditto_mappings_v0.json"
COVER_PATH = f"/app/suno/sara/cover_filter_v0/metas_{SPLIT}.jsonl"
DISCOGS_PATH = "/app/suno/sara/metas_v0_discogs.jsonl"
OUT_FOLDER = "/app/suno/sara/cover_filter_v0/"


# update these as needed
min_sim = 0.4
max_sim = 0.9
discogs_filter_type = "all"
max_covers_per_source = 15
include_defaults = True

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_idx = {}
parent_to_covers = {}
skipped = 0
for idx, row in tqdm(enumerate(test)):
    if row["task"] == "covers" and "parent_id" not in row:
        parent_to_idx[row["id"]] = idx
        parent_to_covers[row["id"]] = []
for idx, row in tqdm(enumerate(test)):
    if row["task"] == "covers" and "parent_id" in row:
        parent = row["parent_id"]
        child = row["id"]

        if (
            parent not in ditto_scores
            or child not in ditto_scores[parent]
            or parent not in parent_to_covers
        ):
            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(idx)
            elif discogs_filter_type == "amateur":
                if child not in discogs_ids:
                    parent_to_covers[parent].append(idx)
            else:
                parent_to_covers[parent].append(idx)
print(f"{skipped} missing ditto scores out of {len(test)}.")

final_data_covers = {}
sources = 0
num_covers = 0
for parent, covers in parent_to_covers.items():
    if len(covers) > 0:
        final_data_covers[str(parent_to_idx[parent])] = []
        sources += 1
        trimmed_covers = covers[:max_covers_per_source]
        for cover in trimmed_covers:
            final_data_covers[str(parent_to_idx[parent])].append(cover)
            num_covers += 1

final_data = {"covers": {"task": "covers", "idx_map": final_data_covers}}


if include_defaults:
    datasets_to_idx = {}
    total_default = 0
    for idx, row in tqdm(enumerate(test)):
        if row["task"] == "default":
            ds = row["dataset"]
            if ds not in datasets_to_idx:
                datasets_to_idx[ds] = []
            datasets_to_idx[ds].append(idx)
            total_default += 1

    num_to_keep = {
        ds: int(len(idxes) / total_default * num_covers) for ds, idxes in datasets_to_idx.items()
    }

    for k, v in datasets_to_idx.items():
        final_data[k] = {"idx_list": v[: num_to_keep[k]]}

print(f"Found {num_covers} from {sources} sources. Writing to file...")
defaults = "w_def" if include_defaults else "no_def"
write_json(
    final_data,
    os.path.join(
        OUT_FOLDER,
        f"info_{SPLIT}_{min_sim}_{max_sim}_{defaults}_{discogs_filter_type}.json",
    ),
)