# #
#  Batch GPU encode modal scripts
# #

import funcy
import gc
import os
import pathlib
import re
import tempfile
import time
import torch

import modal
from modal import enter
import numpy as np
from scipy.special import softmax

from suno_utils.utils.display import suppress_logging
from contextlib import redirect_stderr
from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir
from suno_utils.utils.text import read_jsonl, write_jsonl
from suno_utils.tasks.data_loader import load_audio_mp

from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import InferenceClient


N_MAX_REPLICAS = 1


image = (
    modal.Image.debian_slim()
    .apt_install(
        "ffmpeg", "sox", "libsox-fmt-mp3", "libcairo2", "libcairo2-dev", "texlive-full"
    )
    # .pip_install_private_repos(
    #     "github.com/suno-ai/fairseq.git@9898715cdff75",
    #     git_user="gkucsko",
    #     secrets=[modal.Secret.from_name("glock_pat")],
    # )
    # .pip_install_private_repos(
    #     "github.com/suno-ai/hoot.git@1ad12a3",
    #     git_user="gkucsko",
    #     secrets=[modal.Secret.from_name("glock_pat")],
    # )
    # .pip_install_from_pyproject(
    #     str(pathlib.Path(__file__).parent.parent.parent.parent.parent / "descript-audio-codec/pyproject.toml"),
    #     force_build=True,
    # )
    .pip_install_from_pyproject(
        str(pathlib.Path(__file__).parent.parent.parent.parent / "pyproject.toml"),
        # force_build=True,
    )
    .pip_install(
        "boto3",
        "transformers",
        "tokenizers",
        "encodec",
        "ctc_segmentation",
        "psutil",
        "nnAudio",
        "torch==2.1.1",
        "torchaudio==2.1.1",
        "importlib-resources==5.13.0",
    )
    # .pip_install(
    #    "stable-audio-tools",
    # )
    # .pip_install_private_repos(
    #     "github.com/suno-ai/glockenspiel.git@c35e74c58299#subdirectory=descript-audio-codec&egg=descript-audio-codec",
    #     git_user="gkucsko",
    #     secrets=[modal.Secret.from_name("glock90_pat")],
    #     force_build=False,
    # )
)
stub = modal.App("batch-encode-gpu", image=image)


################## Prompt Engineering Here ##############################
mistral_system_prompt = """
Describe the song with a comma separated list of words. These words should include detailed genre tags, mood, style, time period, tempo, and vocal type. 
If notable, also include details of production, audio gear (e.g. mic type) and instruments (e.g. guitar amp type) used.
If the song has lyrics include the language. Do not include reference to any musicians or musical artists. 
Do not respond with multiple sentences. Do not repeat the title of the song.
Include as much detail as possible but make sure your entire response is a single list of words. 
The title of the video is: 
"""

mistral_example_title_1 = "Party in LA - Miley Cyrus"
mistral_example_response_1 = "Pop, 2010s, upbeat, carefree, modern, mid-tempo, English lyrics, catchy hooks, polished production, digital recording, bright synths, groovy bassline, clean electric guitar, danceable rhythm, youthful energy, fun mood, vibrant, radio-friendly, layered vocals, punchy drums, contemporary pop, light reverb, compressed dynamics, summer vibe, party atmosphere."

mistral_example_title_2 = "Dreams - Fleetwood Mac"
mistral_example_response_2 = "Soft rock, 1970s, mellow, laid-back, melancholic, moderate tempo, English lyrics, smooth vocals, warm production, analog recording, Fender Rhodes, clean electric guitar, light percussion, natural reverb, vintage vibe, introspective mood, subtle harmonies, analog tape warmth, dynamic range, emotional tone, intimate feel, soft sustain, understated bass, acoustic drum kit."

mistral_example_title_3 = "Enter Sandman - Metallica"
mistral_example_response_3 = "Heavy metal, 1990s, dark, aggressive, high energy, fast tempo, English lyrics, distorted guitars, powerful vocals, driving bassline, thunderous drums, analog recording, tight production, deep reverb, punchy dynamics, ominous mood, crunchy guitar riffs, solid-state amps, gritty tone, headbanging rhythm, intense atmosphere, raw power, layered guitars, stadium rock, heavy distortion, suspenseful intro, commanding presence."


# ## Defining the prediction function
#
# * Container lifecycle hook: this lets us load the model only once in each container
#
# Instead of a using `@stub.function()` in the global scope,
# we put the method on a class, and define an `__enter__` method on that class.
# Modal reuses containers for successive calls to the same function, so
# we want to take advantage of this and avoid setting up the same model
# for every function call.

# Notes:
# relative speeds: 5.5s on a5000, 7.5s on A10G, 18.5s on T4
# n_cpu on T4/A10G: 24
# audio load: 75s for 1k files on 24 cores (70h audio) ~3-4k realtime


@stub.cls(
    gpu="A10G",
    timeout=2 * 60 * 60,
    concurrency_limit=N_MAX_REPLICAS,
    secrets=[modal.Secret.from_name("aws-bucket")],
    # cloud="oci",
    memory=150000,  # youtube 5k needs about 100gb
    retries=3,
)
class Worker:
    def __init__(
        self,
        encode_s3_dir,
    ):
        self.encode_s3_dir = encode_s3_dir

    @enter()
    def set_up_models(self):

        self.device = "cuda"  # the device to load the model onto
        self.access_token = "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"
        self.model = AutoModelForCausalLM.from_pretrained(
            "mistralai/Mistral-7B-Instruct-v0.1", token=self.access_token
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            "mistralai/Mistral-7B-Instruct-v0.1", token=self.access_token
        )
        self.model.to(self.device)

    @modal.method()
    def pseudo_label(self, work_item):
        work_item_id = work_item["id"]
        youtube_title = work_item["youtube_title"]

        # construct messages
        messages = [
            {
                "role": "user",
                "content": f"{mistral_system_prompt} {mistral_example_title_1}",
            },
            {"role": "assistant", "content": mistral_example_response_1},
            {"role": "user", "content": f"{mistral_example_title_2}"},
            {"role": "assistant", "content": mistral_example_response_2},
            {"role": "user", "content": f"{mistral_example_title_3}"},
            {"role": "assistant", "content": mistral_example_response_3},
            {"role": "user", "content": f"{youtube_title}"},
        ]

        with torch.no_grad():
            encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt")
            model_inputs = encodeds.to(self.device)
            generated_ids = self.model.generate(
                model_inputs, max_new_tokens=256, do_sample=True
            )
            decoded = self.tokenizer.batch_decode(generated_ids)
            response = decoded[0].split("[/INST]")[-1]
            response = response.strip("</s>")
            response = response.strip('"')
            response = response.strip("[")
            response = response.strip("]")
            tags = response.split(",")
            tags = [tag.strip() for tag in tags]
            tags = [tag.strip('"') for tag in tags]
            tags = [tag.strip(".") for tag in tags]

        new_tags = []
        for tag in tags:
            # check for newlines
            if "\n" in tag:
                sub_tags = tag.split("\n")
                for sub_tag in sub_tags:
                    new_tags.append(sub_tag)
            else:
                new_tags.append(tag)

        # save tags to a text file in comma separated format
        with open(f"{work_item_id}.txt", "w") as f:
            f.write(",".join(new_tags))

        # upload the text file to the s3 bucket
        upload_s3_files(f"{work_item_id}.txt", self.encode_s3_dir)


## Run as eg:
# modal run test_modal.py \


@stub.local_entrypoint()
def main(
    s3_metas_filepath: str = "s3://suno-data/datasets/metadata/chirp_v4/genius_hq_metas_filtered.jsonl",
    output_name: str = "mistral_tags",
    chunksize: int = 500,
    first_only: bool = True,
):

    print("Downloading data...")
    metas = read_from_s3(s3_metas_filepath, read_f=read_jsonl)
    print(len(metas), "data items loaded")

    work_items = list(funcy.chunks(chunksize, metas))
    work_items = list(zip(range(len(work_items)), work_items))
    print(len(work_items), "work items")

    worker = Worker()

    print("Testing inference...")
    t0 = time.time()
    for work_item in work_items[:1]:
        _ = worker.pseudo_label.remote(work_item)
    print(f"{int(round(time.time()-t0))}s for test")

    if first_only:
        print("done with test.")
        return

    print("Running batch inference...")
    t0 = time.time()
    _ = list(worker.embed.map(work_items[1:]))
    print(round((time.time() - t0) / 60 / 60), "h total for batch embed")

    # verify that all is embedded
    done_archives = [fn for fn, _ in list_s3_dir(encode_s3_dir) if fn.endswith(".npz")]
    assert len(list(funcy.chunks(chunksize, range(len(metas))))) == len(done_archives)
    print(f"all {len(done_archives)} parts done.")


# ## Further optimization notes
#
# Every container downloads the model when it starts, which is a bit inefficient.
# In order to improve this, what you could do is to set up a shared volume that gets
# mounted to each container.
# See [shared volumes](/docs/guide/shared-volumes).
#
# In order for Huggingface to use the shared volume, you need to set the value of
# the `TRANSFORMERS_CACHE` environment variable to the path of the shared volume.
# See [secrets](/docs/guide/secrets).