from utils.bct import BlockType

TextBlockType = BlockType(
    name="text",
    is_causal=True,
    description="Used for text input. Could be lyrics, instructions, tags, etc.",
)
MMBertTextBlockType = BlockType(
    name="mmbert_text",
    is_causal=True,
    description="Text embeddings of text from mmBERT. Provides better semantic representation of text.",
)
HootTextBlockType = BlockType(
    name="hoot_text",
    is_causal=True,
    description="""
    Greedy text predictions from Hoot, a CTC-based text recognition model. Used to control when lyrics are sung.
    Can transfer the lyrics timing from the reference track to the generated track.
    """,
)

DittoBlockType = BlockType(
    name="ditto",
    is_causal=True,
    description="""
    A single embedding from ditto, a contrastive text-audio model. Conditions the vibe of the song. 
    Can be used to interpolate between songs.
    """,
)

CausalSemanticBlockType = BlockType(
    name="semantic",
    is_causal=True,
    description="""
    Discrete semantic tokens post quantization. Discrete semantic contains the semantic content of the audio, but misses fine-grained details.
    This is the standard output of the semantic model.
    """,
)
ContinuousSemanticBlockType = BlockType(
    name="continuous_semantic",
    is_causal=True,
    description="""
    Continuous semantic embeddings pre-quantization. Continuous semantic contains the semantic content of the audio, with fine-grained details.
    Some models use this as input.
    """,
)
ArtistBlockType = BlockType(
    name="artist",
    is_causal=True,
    description="Another song from the same artist. Used to create a new song that could be from the same artist, often keeping style and vocals similar.",
)
PlaylistBlockType = BlockType(
    name="playlist",
    is_causal=True,
    description="""
    Another song from the same playlist. Used to create a new song that could be from the same playlist, depends a lot on what songs are in the playlist. 
    Can be unpredictable.
    """,
)
UnderpaintBlockType = BlockType(
    name="underpaint",
    is_causal=True,
    description="""
    Only the vocals of the song. Used to fill in the instrumental of the song. 
    Input vocals, output full mix including vocals/instrumental.
    """,
)
OverpaintBlockType = BlockType(
    name="overpaint",
    is_causal=True,
    description="""
    Only the instrumental of the song. Used to fill in the vocals of the song.
    Input instrumental, output full mix including vocals/instrumental.
    """,
)
VoxBlockType = BlockType(
    name="vox",
    is_causal=True,
    description="A short vocal clip from the singer of the song. Used to specify the singer of the song.",
)
RemixBlockType = BlockType(
    name="remix",
    is_causal=True,
    description="A whole song to be remixed, maintains some stems or audio from the original.",
)
SampleSourceBlockType = BlockType(
    name="sample_source",
    is_causal=True,
    description="Selects small chunks from a source song to use in the generated output.",
)
MashupBlockType = BlockType(
    name="mashup",
    is_causal=True,
    description="A list of two or more songs that will be layered and transformed in the output.",
)
StemBlockType = BlockType(
    name="stem",
    is_causal=True,
    description="""
    The song with some stems removed. Used to add a single stem to the song.
    Input everthing except the stem, output just the stem.
    """,
)
SampleBlockType = BlockType(
    name="sample",
    is_causal=True,
    description="""
    A short clip from the song. Can either be a clip from the full mix or a clip from a single stem.
    Input sample, output a song that includes the sample.
    """,
)
CoverBlockType = BlockType(
    name="cover",
    is_causal=True,
    description="""
    A cover of the song, likely with similar melody.
    Input cover, output a song that the song is covered by the cover. Used to create a new song that is a cover of the original song, often keeping melody and lyrics similar.
    """,
)
PrefixBlockType = BlockType(
    name="prefix",
    is_causal=True,
    description="""
    A prefix of the song. The model will generate the rest of the song after this prefix. 
    Combined with a suffix, this infills the gap between the prefix and suffix.",
    """,
)
SuffixBlockType = BlockType(
    name="suffix",
    is_causal=True,
    description="""
    A suffix of the song. The model will generate the rest of the song before this suffix.
    Combined with a prefix, this infills the gap between the prefix and suffix.
    """,
)

NonCausalSemanticBlockType = BlockType(
    name="non_causal_semantic",
    is_causal=False,
)

DiffusionBlockType = BlockType(
    name="diffusion",
    is_causal=False,
)

InterleavedSemanticBlockType = BlockType(
    name="interleaved_semantic",
    is_causal=True,
    chunk_size=10,
)

# VAE Input Block Types for conditioning (used when doing diffusion)
VAEArtistBlockType = BlockType(
    name="vae_artist",
    is_causal=False,
)
VAEPlaylistBlockType = BlockType(
    name="vae_playlist",
    is_causal=False,
)
VAEUnderpaintBlockType = BlockType(
    name="vae_underpaint",
    is_causal=False,
)
VAEOverpaintBlockType = BlockType(
    name="vae_overpaint",
    is_causal=False,
)
VAEVoxBlockType = BlockType(
    name="vae_vox",
    is_causal=False,
)
VAEStemBlockType = BlockType(
    name="vae_stem",
    is_causal=False,
)
VAESampleBlockType = BlockType(
    name="vae_sample",
    is_causal=False,
)
VAECoverBlockType = BlockType(
    name="vae_cover",
    is_causal=False,
)
VAEPrefixBlockType = BlockType(
    name="vae_prefix",
    is_causal=False,
)
VAESuffixBlockType = BlockType(
    name="vae_suffix",
    is_causal=False,
)

# Text description block for conditioning pairs - replace indicator tokens
# This is paired with conditioning blocks to provide metadata about the content
CondAudioTextBlockType = BlockType(
    name="text_description",
    is_causal=True,
)
CondAudioBlockType = BlockType(
    name="cond_audio",
    is_causal=True,  # TODO (vibertthio): this could also be non-causal
)
VAECondAudioBlockType = BlockType(
    name="vae_cond_audio",
    is_causal=False,
)

# mapping of block type names to ids, must match GPTConfig.block_type_vocab_size
BLOCK_TYPE_NAME_TO_ID = {
    "text": 0,
    "mmbert_text": 1,
    "hoot_text": 2,
    "ditto": 3,
    "semantic": 4,
    "artist": 5,
    "playlist": 6,
    "underpaint": 7,
    "overpaint": 8,
    "vox": 9,
    "remix": 10,
    "sample_source": 11,
    "mashup": 12,
    "stem": 13,
    "sample": 14,
    "cover": 15,
    "prefix": 16,
    "suffix": 17,
    "non_causal_semantic": 18,
    "diffusion": 19,
    "vae_artist": 20,
    "vae_playlist": 21,
    "vae_underpaint": 22,
    "vae_overpaint": 23,
    "vae_vox": 24,
    "vae_stem": 25,
    "vae_sample": 26,
    "vae_cover": 27,
    "vae_prefix": 28,
    "vae_suffix": 29,
    "text_description": 30,
    "cond_audio": 31,
    "vae_cond_audio": 32,
    "continuous_semantic": 33,
}