"""
Modal Runner for model evaluation. Launches EvalConductorStub to send requests to an engine worker, then on generation
completion launches an EvalEncoderStub to run (optional) source separation and ditto embedding calculation.
"""

import json
import logging
import random
import time
import os
import re
from datetime import datetime
import tempfile
from uuid import uuid4
import modal
from enum import Enum

from suno_utils.utils.s3 import check_s3_file_exists
from suno_utils.worker.settings import s3_client
from suno_utils.worker.modal_base import MOUNT_PATH
from suno_utils.worker.feature_eval import (
    Prompt,
    load_prompt_arr,
    FeatureEval,
    InlineTagEval,
    GenreEval,
    TEST_SETS,
    DITTO_S3_PATH,
)
from suno_utils.gpt.chirp_v2_5 import _get_model_if_needed
from suno_utils.worker.loader import S3Loader
from suno_utils.worker.tracing import distributed_trace
from suno_utils.worker.schema import QueueItem, HistoryPrompt
from suno_utils.tasks import ss_vad
from suno_utils.tasks.ditto_v2 import preload_models as preload_ditto_models
from suno_utils.worker.modal_base import get_modal_base_image_with_flash_attention
import suno_utils.gpt.chirp_v2_5 as chirp
from suno_utils.tasks.hoot import load_model_list as hoot_load_model_list


class EvalTestType(Enum):
    COVER_AND_ARTIST = "cover_artist"
    BASE_GENERATION = "base_gen"
    INLINE_TAG_EVAL = "inline_tag_eval"
    VOCALIST_GENDER_EVAL = "VOCALIST_GENDER_EVAL"
    GENRE_EVAL = "genre_eval"


############## CHANGE THESE ##############
TEST_TYPE = EvalTestType.BASE_GENERATION

MODEL_NAME = "chirp-auk-eval"  # model to test, must have an engine modal app running
EXTRA_TAGS = ""  # additional experiment labels, like config changes

if TEST_TYPE == EvalTestType.COVER_AND_ARTIST:
    EVAL_NAME = "all_genres"  # hard or all_genres
    GENS_PER_SOURCE = 5
    FEATURE_EVAL_DATA_PATH = TEST_SETS[EVAL_NAME]
elif TEST_TYPE == EvalTestType.BASE_GENERATION:
    EVAL_NAME = ""  # leave blank for this case
    GENS_PER_SOURCE = 20
    FEATURE_EVAL_DATA_PATH = "../../task_eval/labelbox/labelbox_sources.json"
elif TEST_TYPE == EvalTestType.INLINE_TAG_EVAL:
    EVAL_NAME = ""  # leave blank for this case
    GENS_PER_SOURCE = 20
    FEATURE_EVAL_DATA_PATH = "../../task_eval/genres_instruments.json"
    LYRICS_PATH = "../../task_eval/lyrics_with_inline_tags.json"
elif TEST_TYPE == EvalTestType.VOCALIST_GENDER_EVAL:
    EVAL_NAME = ""  # leave blank for this case
    GENS_PER_SOURCE = 5
    FEATURE_EVAL_DATA_PATH = "../../task_eval/genre_descriptions.json"
    LYRICS_PATH = "../../task_eval/lyrics_with_inline_tags.json"
elif TEST_TYPE == EvalTestType.GENRE_EVAL:
    EVAL_NAME = ""  # leave blank for this case
    GENS_PER_SOURCE = 4
    FEATURE_EVAL_DATA_PATH = "../../task_eval/rate_your_music_filtered.json"

INFER_CONFIG = None  # if you want to force inference parameter changes
##########################################

DEPLOYMENT_TYPE = "dev"
EVAL_CONDUCTOR_DDOG_SERVICE = "eval-conductor-worker"
EVAL_ENCODER_DDOG_SERVICE = "eval-encoder-worker"

logger = logging.getLogger(__name__)
logging.basicConfig()
logger.setLevel(logging.INFO)

OUTPUT_FOLDER = "../../task_eval/modal_runs"
MOUNT_PATH = "/suno/models"
APP_NAME = f"eval-orchestrator-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)

if DEPLOYMENT_TYPE == "dev":
    from suno_utils.worker.modal_model_configuration import MODEL_DEV_FNS as MODAL_FNS
elif DEPLOYMENT_TYPE == "prod":
    from suno_utils.worker.modal_model_configuration import MODEL_PROD_FNS as MODAL_FNS
else:
    err_msg = f"DEPLOYMENT_TYPE {DEPLOYMENT_TYPE} unrecognized"
    raise AssertionError(err_msg)

# Use random seed based on time
random_seed = int((time.time() * 1000) % 100000000)
print("Random seed set to:", random_seed)
random.seed(random_seed)

DD_SAMPLE_RATE = "0.0001"

# modal traffic lookup
MODAL_TRAFFIC_WORKER = "engine-chirpv2_engine_13b_special_8"

# looks up the H100 traffic
_ = modal.Cls.from_name(
    f"{MODAL_TRAFFIC_WORKER}_" + f"{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}_240s_H100",
    "ChirpV2Stub",
)().generate


class LoadEvalModelsWorker(S3Loader):
    """Load models for source separation and ditto"""

    @staticmethod
    def download_models_ditto(dir_path: str = MOUNT_PATH) -> None:
        _ = preload_ditto_models(model_filepath=DITTO_S3_PATH)

        ss_vad_config_path = _get_model_if_needed(ss_vad.YAML_PATH, cache_dir=dir_path)
        ss_vad_model_path = _get_model_if_needed(ss_vad.MODEL_PATH, cache_dir=dir_path)
        ss_vad.preload_models(
            checkpoint_filepath=ss_vad_model_path,
            config_path=ss_vad_config_path,
        )

    @staticmethod
    def download_models_tagging(dir_path: str = MOUNT_PATH) -> None:
        hoot_ckpt_path = chirp._get_model_if_needed(chirp.HOOT_CKPT_PATH, cache_dir=dir_path)
        hoot_tokenizer_path = chirp._get_model_if_needed(chirp.HOOT_TOKENIZER_PATH, cache_dir=dir_path)
        # load hoot on cuda as well
        chirp.preload_hoot_models(
            hoot_ckpt_path,
            hoot_tokenizer_path,
        )


aws_secret = modal.Secret.from_name("studio-aws")
sample_rule = json.dumps([{"sample_rate": DD_SAMPLE_RATE}])
SECRETS = [
    aws_secret,
    modal.Secret.from_dict(
        {
            "DD_SITE": "datadoghq.com",
            "DD_ENV": DEPLOYMENT_TYPE,
            "DD_SERVICE": "chatgpt-worker",
            "DD_LOGS_ENABLED": "true",
            "DD_TRACE_ENABLED": "true",
            "DD_TRACE_SAMPLING_RULES": sample_rule,
        },
    ),
    modal.Secret.from_name("datadog-metrics"),
    modal.Secret.from_name("api-callback-token"),
]


def download_model_wrapper_a() -> None:
    if TEST_TYPE == EvalTestType.INLINE_TAG_EVAL:
        LoadEvalModelsWorker.download_models_tagging()
    else:
        LoadEvalModelsWorker.download_models_ditto()


base_image = get_modal_base_image_with_flash_attention().add_local_python_source(
    "suno_utils", copy=False
)
image = base_image.run_function(download_model_wrapper_a, secrets=SECRETS)
app = modal.App(APP_NAME, image=image)


@app.cls(
    cpu=1.0,
    secrets=SECRETS,
    min_containers=5,
    timeout=600,
    retries=modal.Retries(
        max_retries=2,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    scaledown_window=400,
)
@modal.concurrent(max_inputs=50)
class EvalConductorStub(S3Loader):
    """Expose EvalConductorStub worker as a modal app

    Responsible for distributing queue items to the appropriate chirp worker based on the model name.
    When a generation completes, it triggers an EvalEncoderStub to save the evaluation results.

    Unlike the ConductorStub, this worker does not apply any model or parameter A/B test experiments.
    """

    def __init__(self):
        # looks up the H100 traffic
        self.modal_f_chirp_main_H100 = modal.Cls.lookup(
            f"{MODAL_TRAFFIC_WORKER}_" + f"{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}_240s_H100",
            "ChirpV2Stub",
        )().generate
        self.modal_f_history_loader = modal.Cls.lookup(
            f"history_encoder-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "HistoryEncoderStub",
        )().encode_history

        # key is a tuple of worker name and function name
        self.modal_chirp_worker_lookup_cache = {}

    def _get_engine_worker(self, item: QueueItem) -> modal.Function | None:
        """Get appropriate engine worker based on model name"""
        if item.model_name is None:
            raise ValueError("Model name is required to get the GPT worker.")
        stub_name, stub_function = MODAL_FNS[item.model_name]
        stub_name_worker = f"{stub_name}_H100"
        # we don't need to lookup it again if we have already done it
        if (stub_name_worker, stub_function) in self.modal_chirp_worker_lookup_cache:
            return self.modal_chirp_worker_lookup_cache[(stub_name_worker, stub_function)]
        else:
            try:
                cls_name = stub_function.split(".")[0]
                func_name = stub_function.split(".")[1]
                chirp_f = getattr(modal.Cls.lookup(stub_name_worker, cls_name)(), func_name)
                self.modal_chirp_worker_lookup_cache[(stub_name_worker, stub_function)] = chirp_f
                return chirp_f
            except Exception as e:
                print(
                    f"Error getting chirp worker for {item.model_name}: {e}. stub_name, stub_function: {stub_name_worker}, {stub_function}"
                )

    def _validate_input_queue_item(self, item: QueueItem) -> bool:
        """Validate the input queue item to check if the request if valid."""
        if (
            item.metadata.get("task") in ["infill", "infill_intro", "infill_outro"]
        ) and not item.is_infill:
            print(f"WTF is happening with {item.id} -- task is infill but not infill arguments.")
            return False
        # add task validations checks -- we should not be passing in clips if task isn't specified
        if not item.is_cover_condition and item.metadata.get("cover_clip_id", None) is not None:
            print(f"WTF is happening with {item.id} -- not cover but without cover id passed in.")
            return False
        if not item.is_artist_condition and item.metadata.get("artist_clip_id", None) is not None:
            print(f"WTF is happening with {item.id} -- not artist but without artist id passed in.")
            return False
        if (
            item.metadata.get("task") == "extend"
            or item.is_cover_extend
            or item.is_cover_infill
            or item.is_artist_extend
            or item.is_artist_infill
        ) and item.prompt_audio is None:
            print(f"WTF is happening with {item.id} -- task is extend but prompt audio id passed in.")
            return False
        if item.metadata.get("task") == "upsample" and not item.is_upsample:
            print(f"WTF is happening with {item.id} -- task is upsample but without upsample arguments.")
            return False
        return True

    @modal.method()
    @distributed_trace("redirect_and_generate", EVAL_CONDUCTOR_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def redirect_and_generate(
        self,
        queue_item_json: str,
        history_prompt: HistoryPrompt | None = None,
        time_prefix: str = "",
        task: str | None = None,
    ) -> None:
        """This worker will distribute the queue item to the appropriate model worker, then kick
        off ditto evaluation once the generation is complete.

        If it is a continue and needs to fetch the history, it will go to the history loader.
            History loader will pass the item back into redirect_and_generate.
        Otherwise, it will go through augmentation and to the GPT worker directly.

        Args:
            queue_item_json: the queue item in json format. Clip-level only.
            history_prompt: the history prompt in HistoryPrompt.
            time_prefix: timestamp label for the evaluation
            task: ditto task to trigger for evaluation
        """
        item = QueueItem(**json.loads(queue_item_json))
        logger.info(f"  {item}")

        # validate the queue item -- should meet expected conditions
        if not self._validate_input_queue_item(item):
            raise ValueError("Invalid generation request.")

        if (
            item.prompt_audio or item.is_cover_condition or item.is_artist_condition or item.is_infill
        ) and history_prompt is None:
            # This is a continue job, without history, we need to fetch the history first
            _ = self.modal_f_history_loader.spawn(item.json())
            logger.info(f"Request {item.id}: Spawned modal job into history loader")
            # nothing to be done for here, job isn't in queue yet, except upsample
            if (item.is_upsample or item.is_stem) and item.ids:
                for clip_id in item.ids:
                    curr_item = item.model_copy(deep=True, update={"id": clip_id, "ids": None})
                    # we need to copy the image -- note that this is independent of the job queue
                    self.copy_and_upload_image(curr_item)
            return

        if item.model_name not in MODAL_FNS:
            # this is a bad model name, we should notify at the request level
            raise ValueError(f"Invalid model name {item.model_name}")

        # start the audio runner, do a look up first
        engine_worker = self._get_engine_worker(item)

        # directly goes to the GPT worker, with empty HistoryPrompt
        _ = engine_worker.spawn(item.json(), history_prompt)
        logger.info(f"Request {item.id}: Spawned modal job {engine_worker} item id {item.id}")

        file_found = False
        logger.info(f"Waiting for S3 file for {item.id} to be created...")
        while not file_found:
            if check_s3_file_exists(f"s3://suno-data-uploads/studio/uploads/{item.id}.mp3"):
                file_found = True
                if task is not None:
                    logger.info(f"Found S3 file for {item.id}, launching evaluation.")
                    if task == "tags_inline":
                        EvalInlineTagStub.process.spawn(
                            json.dumps(dict(id=item.id, prompt_text=item.prompt_text, metadata={})),
                            time_prefix=time_prefix,
                        )
                    elif task == "gender_inline":
                        EvalInlineTagStub.process.spawn(
                            json.dumps(dict(id=item.id, prompt_text=item.prompt_text, metadata={})),
                            time_prefix=time_prefix,
                        )
                    elif task == "genre_encode":
                        EvalEncodeGenreStub.process.spawn(
                            json.dumps(dict(id=item.id, metadata={})), time_prefix=time_prefix
                        )
                    else:
                        EvalEncodeStub.process.spawn(
                            json.dumps(dict(id=item.id, metadata={})), time_prefix=time_prefix, task=task
                        )
            else:
                # keep waiting for generation to finish
                time.sleep(5)


@app.cls(
    cpu=12,
    gpu="a100",
    secrets=SECRETS,
    timeout=600,
    scaledown_window=400,
    memory=32000,
    retries=modal.Retries(
        max_retries=1,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    min_containers=1,
    max_containers=5,
)
@modal.concurrent(max_inputs=10)
class EvalEncodeStub:
    def __init__(self):
        import torch

        torch.set_num_threads(8)

        self.feature_eval = FeatureEval(
            model_name=MODEL_NAME + "_" + EVAL_NAME + EXTRA_TAGS,
        )

    @modal.method()
    @distributed_trace("eval_process", EVAL_ENCODER_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def process(
        self,
        queue_item_json: str,
        time_prefix="",
        task="self_sim",
    ):
        item = QueueItem(**json.loads(queue_item_json))
        s3_id = item.id

        with tempfile.NamedTemporaryFile(suffix=".mp3") as temp_file:
            s3_client.download_file("suno-data-uploads", f"studio/uploads/{s3_id}.mp3", temp_file.name)
            self.feature_eval.run_feature_eval(s3_id, temp_file.name, task, time_prefix)


@app.cls(
    cpu=12,
    gpu="a100",
    secrets=SECRETS,
    timeout=600,
    scaledown_window=400,
    memory=32000,
    retries=modal.Retries(
        max_retries=1,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    min_containers=1,
    max_containers=5,
)
@modal.concurrent(max_inputs=10)
class EvalEncodeGenreStub:
    def __init__(self):
        import torch

        torch.set_num_threads(8)
        _ = preload_ditto_models(model_filepath=DITTO_S3_PATH)

        self.feature_eval = GenreEval(
            model_name=MODEL_NAME + "_" + EVAL_NAME + EXTRA_TAGS,
        )

    @modal.method()
    @distributed_trace("eval_process", EVAL_ENCODER_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def process(
        self,
        queue_item_json: str,
        time_prefix="",
    ):
        item = QueueItem(**json.loads(queue_item_json))
        s3_id = item.id

        with tempfile.NamedTemporaryFile(suffix=".mp3") as temp_file:
            self.feature_eval.run_feature_eval(s3_id, time_prefix)


@app.cls(
    cpu=12,
    gpu="a100",
    secrets=SECRETS,
    timeout=600,
    scaledown_window=400,
    memory=32000,
    retries=modal.Retries(
        max_retries=1,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    min_containers=1,
    max_containers=5,
)
@modal.concurrent(max_inputs=10)
class EvalInlineTagStub:
    def __init__(self):
        import torch

        torch.set_num_threads(8)

        hoot_ckpt_path = chirp._get_model_if_needed(chirp.HOOT_CKPT_PATH, cache_dir=MOUNT_PATH)
        hoot_tokenizer_path = chirp._get_model_if_needed(chirp.HOOT_TOKENIZER_PATH, cache_dir=MOUNT_PATH)
        model_dict = hoot_load_model_list(
            "s3://suno-data/checkpoints/hoot_v3/hoot_ckpt.pt",
            "s3://suno-data/checkpoints/hoot_v3/tokenizer.model",
        )[0]
        tokenizer = model_dict["tokenizer"]

        self.tag_eval = InlineTagEval(
            model_name=MODEL_NAME + "_" + EVAL_NAME + EXTRA_TAGS,
            tokenizer=tokenizer,
            threshold=0.6,
            tag_type="instrument" if TEST_TYPE == EvalTestType.INLINE_TAG_EVAL else "gender",
        )

    @modal.method()
    @distributed_trace("eval_process", EVAL_ENCODER_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def process(
        self,
        queue_item_json: str,
        time_prefix="",
        task="",
    ):
        item = QueueItem(**json.loads(queue_item_json))
        s3_id = item.id
        lyrics = item.prompt_text

        self.tag_eval.run_feature_eval(s3_id, lyrics, None, time_prefix)


def test_cover_test(source_data, num_gens_per_source=5, time_prefix="", task="self_sim"):
    cover_mappings = {}
    conductor = EvalConductorStub()

    for cover_source in source_data:
        inputs = []
        prompt = Prompt(**cover_source)
        prompt.cover_arr = load_prompt_arr(prompt.s3_id)
        cover_mappings[prompt.s3_id] = []
        for _ in range(num_gens_per_source):
            uid = str(uuid4())
            cover_mappings[prompt.s3_id].append(uid)
            inputs.append(
                json.dumps(
                    dict(
                        id=uid,
                        prompt_text=prompt.lyrics,
                        model_name=MODEL_NAME,
                        metadata={
                            "lang": "English",
                            "tags": prompt.test_tags,
                            "task": "cover",
                            "type": "gen",
                            "is_bot": False,
                            "artist_end_s": 120,
                            "feature_flags": 3,
                            "cover_clip_id": prompt.s3_id,
                            "edited_clip_id": prompt.s3_id,
                            "forced_infer_config": INFER_CONFIG,
                        },
                    ),
                ),
            )

        history = HistoryPrompt(cover_audio=prompt.cover_arr)
        for input in inputs:
            conductor.redirect_and_generate.spawn(
                input, history, time_prefix=time_prefix, task="self_sim"
            )
            time.sleep(2)
        source_input = json.dumps(dict(id=prompt.s3_id, metadata={}))
        EvalEncodeStub.process.spawn(source_input, time_prefix=time_prefix, task=task)

    with open(os.path.join(OUTPUT_FOLDER, f"cover_mappings_{time_prefix}.json"), "w") as f:
        json.dump(cover_mappings, f, indent=4)


def test_artist_test(source_data, num_gens_per_source=5, time_prefix="", task="artist_vox_sim"):
    artist_mappings = {}
    conductor = EvalConductorStub()

    for artist_source in source_data:
        inputs = []
        prompt = Prompt(**artist_source)
        prompt.artist_arr = load_prompt_arr(prompt.s3_id)
        artist_mappings[prompt.s3_id] = []
        for _ in range(num_gens_per_source):
            uid = str(uuid4())
            artist_mappings[prompt.s3_id].append(uid)
            inputs.append(
                json.dumps(
                    dict(
                        id=uid,
                        prompt_text=prompt.test_lyrics,
                        model_name=MODEL_NAME,
                        metadata={
                            "lang": "English",
                            "tags": prompt.tags,
                            "task": "artist_consistency",
                            "type": "gen",
                            "is_bot": False,
                            "artist_end_s": 120,
                            "feature_flags": 3,
                            "artist_clip_id": prompt.s3_id,
                            "edited_clip_id": prompt.s3_id,
                            "forced_infer_config": INFER_CONFIG,
                        },
                    ),
                ),
            )

        history = HistoryPrompt(artist_audio=prompt.artist_arr)
        for input in inputs:
            conductor.redirect_and_generate.spawn(input, history, time_prefix=time_prefix, task=task)
            time.sleep(2)
        source_input = json.dumps(dict(id=prompt.s3_id, metadata={}))
        EvalEncodeStub.process.spawn(source_input, time_prefix=time_prefix, task=task)

    with open(os.path.join(OUTPUT_FOLDER, f"artist_mappings_{time_prefix}.json"), "w") as f:
        json.dump(artist_mappings, f, indent=4)


def test_generation_test(source_data, num_gens_per_source=1, time_prefix=""):
    genre_mappings = {}

    conductor = EvalConductorStub()

    for genre in source_data:
        genre_mappings[genre] = []
        inputs = []
        for data in source_data[genre]:
            for _ in range(num_gens_per_source):
                style = data["tags"]
                lyrics = data["lyrics"]
                uid = str(uuid4())
                genre_mappings[genre].append(dict(s3_id=uid, tags=style, lyrics=lyrics))
                inputs.append(
                    json.dumps(
                        dict(
                            id=uid,
                            prompt_text=lyrics,
                            model_name=MODEL_NAME,
                            metadata={
                                "lang": "English",
                                "tags": style,
                                "type": "gen",
                                "is_bot": False,
                                "forced_infer_config": INFER_CONFIG,
                            },
                        ),
                    ),
                )

        history = None
        for input in inputs:
            conductor.redirect_and_generate.spawn(
                input, history, time_prefix=time_prefix, task="self_sim"
            )
            time.sleep(2)

    with open(
        os.path.join(OUTPUT_FOLDER, f"genre_mappings_{MODEL_NAME}{EXTRA_TAGS}_{time_prefix}.json"), "w"
    ) as f:
        json.dump(genre_mappings, f, indent=4)


def replace_placeholders(template_string, replacements):
    # Find all placeholders in the format {something}
    placeholders = re.findall(r"\{[^}]*\}", template_string)

    # Replace each placeholder with the corresponding element from the replacements list
    result = template_string
    for i, placeholder in enumerate(placeholders):
        if i < len(replacements):
            result = result.replace(placeholder, replacements[i])

    return result


def test_inline_tag_test(source_data, eval_lyrics, num_gens_per_source=1, time_prefix=""):
    genre_mappings = {}

    conductor = EvalConductorStub()

    section_names = ["Instrumental Intro", "Verse 1", "Chorus", "Verse 2"]

    for genre, genre_data in source_data.items():
        genre_mappings[genre] = []
        inputs = []

        instrument_groups = genre_data["instruments"]
        descriptions = genre_data["descriptions"]
        genre_lyrics = eval_lyrics[genre]
        for _ in range(num_gens_per_source):
            # Create a list to hold all tuples for this iteration
            iteration_tuples = []

            # First, select one random description for each instrument
            # to use consistently in this iteration
            instrument_to_description = {}
            for instrument, desc_options in descriptions.items():
                instrument_to_description[instrument] = random.choice(desc_options)

            # Then generate a prompt for each instrument group using the consistent descriptions
            for idx, instrument_group in enumerate(instrument_groups):
                description_list = []

                # Add the consistent description for each instrument in the group
                for instrument in instrument_group:
                    if instrument in instrument_to_description:
                        description_list.append(instrument_to_description[instrument])

                # Create the prompt string
                prompt = f"{', '.join(description_list)}"

                # Add the tuple to the iteration list
                iteration_tuples.append([section_names[idx], instrument_group, prompt])

            # now we built the prompt
            style = genre
            lyrics = replace_placeholders(random.choice(genre_lyrics), [x[2] for x in iteration_tuples])
            uid = str(uuid4())
            genre_mappings[genre].append(
                dict(s3_id=uid, tags=style, lyrics=lyrics, mappings=iteration_tuples)
            )

            inputs.append(
                json.dumps(
                    dict(
                        id=uid,
                        prompt_text=lyrics,
                        model_name=f"{MODEL_NAME}",
                        metadata={
                            "lang": "English",
                            "tags": style,
                            "type": "gen",
                            "is_bot": False,
                            "forced_infer_config": INFER_CONFIG,
                        },
                    ),
                ),
            )

        history = None
        for input in inputs:
            conductor.redirect_and_generate.spawn(
                input, history, time_prefix=time_prefix, task="tags_inline"
            )
            time.sleep(2)

    with open(
        os.path.join(OUTPUT_FOLDER, f"tag_mappings_{time_prefix}.json"),
        "w",
    ) as f:
        json.dump(genre_mappings, f, indent=4)


def test_gender_test(source_data, eval_lyrics, num_gens_per_source=1, time_prefix=""):
    genre_mappings = {}

    conductor = EvalConductorStub()

    male_verse = ["instrumentals", "male singer", "female singer", "male singer"]
    female_verse = ["instrumentals", "female singer", "male singer", "female singer"]

    for genre, genre_data in source_data.items():
        genre_mappings[genre] = []
        inputs = []

        for style_description in genre_data:
            genre_lyrics = eval_lyrics[genre]
            for _ in range(num_gens_per_source):
                lyrics = random.choice(genre_lyrics)
                for replacements in [male_verse, female_verse]:
                    final_lyrics = replace_placeholders(lyrics, replacements)
                    uid = str(uuid4())
                    genre_mappings[genre].append(
                        dict(
                            s3_id=uid, tags=style_description, lyrics=final_lyrics, mappings=replacements
                        )
                    )

                    inputs.append(
                        json.dumps(
                            dict(
                                id=uid,
                                prompt_text=final_lyrics,
                                model_name=f"{MODEL_NAME}",
                                metadata={
                                    "lang": "English",
                                    "tags": "duet, " + style_description,
                                    "type": "gen",
                                    "is_bot": False,
                                    "forced_infer_config": INFER_CONFIG,
                                },
                            ),
                        ),
                    )

        history = None
        for input in inputs:
            conductor.redirect_and_generate.spawn(
                input, history, time_prefix=time_prefix, task="gender_inline"
            )
            time.sleep(2)

    with open(
        os.path.join(OUTPUT_FOLDER, f"vocalist_gender_mappings_{time_prefix}.json"),
        "w",
    ) as f:
        json.dump(genre_mappings, f, indent=4)


def test_genre_test(source_data, num_gens_per_source=1, time_prefix=""):
    genre_mappings = {}

    conductor = EvalConductorStub()

    for genre in source_data:
        genre_name = genre["name"]
        genre_mappings[genre_name] = []
        inputs = []
        for _ in range(num_gens_per_source):
            style = genre_name
            parent_style = genre["parent_genre"]
            uid = str(uuid4())
            genre_mappings[genre_name].append(dict(s3_id=uid, tags=style, parent_style=parent_style))
            inputs.append(
                json.dumps(
                    dict(
                        id=uid,
                        prompt_text="",
                        model_name=MODEL_NAME,
                        metadata={
                            "lang": "English",
                            "tags": style,
                            "type": "gen",
                            "is_bot": False,
                            "forced_infer_config": INFER_CONFIG,
                        },
                    ),
                ),
            )

        history = None
        for input in inputs:
            conductor.redirect_and_generate.spawn(
                input, history, time_prefix=time_prefix, task="genre_encode"
            )
            time.sleep(2)

    with open(
        os.path.join(OUTPUT_FOLDER, f"subgenre_mappings_{time_prefix}.json"),
        "w",
    ) as f:
        json.dump(genre_mappings, f, indent=4)


@app.local_entrypoint()
def main():
    # Get current timestamp as readable string
    time_label = datetime.now().strftime("%Y_%m_%d-%H_%M_%S")

    if TEST_TYPE == EvalTestType.BASE_GENERATION:
        with open(FEATURE_EVAL_DATA_PATH) as f:
            eval_data = json.load(f)

        test_generation_test(eval_data, num_gens_per_source=GENS_PER_SOURCE, time_prefix=time_label)
    elif TEST_TYPE == EvalTestType.COVER_AND_ARTIST:
        # load evaluation sources
        with open(FEATURE_EVAL_DATA_PATH) as f:
            eval_data = json.load(f)

        if not os.path.exists(OUTPUT_FOLDER):
            os.makedirs(OUTPUT_FOLDER)

        cover_data = eval_data["cover"]
        artist_data = eval_data["artist"]

        test_cover_test(cover_data, num_gens_per_source=GENS_PER_SOURCE, time_prefix=time_label)
        test_artist_test(artist_data, num_gens_per_source=GENS_PER_SOURCE, time_prefix=time_label)
    elif TEST_TYPE == EvalTestType.INLINE_TAG_EVAL:
        with open(FEATURE_EVAL_DATA_PATH) as f:
            eval_data = json.load(f)

        with open(LYRICS_PATH) as f:
            eval_lyrics = json.load(f)

        if not os.path.exists(OUTPUT_FOLDER):
            os.makedirs(OUTPUT_FOLDER)

        test_inline_tag_test(
            eval_data, eval_lyrics, num_gens_per_source=GENS_PER_SOURCE, time_prefix=time_label
        )
    elif TEST_TYPE == EvalTestType.VOCALIST_GENDER_EVAL:
        with open(FEATURE_EVAL_DATA_PATH) as f:
            eval_data = json.load(f)

        with open(LYRICS_PATH) as f:
            eval_lyrics = json.load(f)

        if not os.path.exists(OUTPUT_FOLDER):
            os.makedirs(OUTPUT_FOLDER)

        test_gender_test(
            eval_data, eval_lyrics, num_gens_per_source=GENS_PER_SOURCE, time_prefix=time_label
        )
    elif TEST_TYPE == EvalTestType.GENRE_EVAL:
        with open(FEATURE_EVAL_DATA_PATH) as f:
            eval_data = json.load(f)

        if not os.path.exists(OUTPUT_FOLDER):
            os.makedirs(OUTPUT_FOLDER)

        test_genre_test(eval_data, num_gens_per_source=GENS_PER_SOURCE, time_prefix=time_label)
    else:
        raise ValueError(f"Test type {TEST_TYPE} not recognized")

    time.sleep(3600 * 2)