import math
import os
import random
import re

import numpy as np
import torch
from tokenizers import AddedToken
from transformers import PreTrainedTokenizerFast

# to avoid: "The current process just got forked, after parallelism has already been used"
os.environ["TOKENIZERS_PARALLELISM"] = "False"


def detect_chorus_sections(text_aligned):
    """Detect chorus sections from aligned lyrics based on section tags in brackets.

    Args:
        text_aligned: List of [start_time, end_time, text] tuples

    Returns:
        List of (start_time, end_time) tuples for detected chorus sections
    """
    if not text_aligned:
        return []

    chorus_sections = []
    current_chorus_start = None

    for start_time, end_time, text in text_aligned:
        # Convert text to lowercase for case-insensitive matching
        text_lower = text.lower().strip()

        # Check if line contains [chorus] or similar chorus indicators
        if "[chorus]" in text_lower:
            if current_chorus_start is None:
                current_chorus_start = start_time
        # Check for section end markers or new sections
        elif "[" in text_lower:
            if current_chorus_start is not None:
                # End the current chorus section
                chorus_sections.append((current_chorus_start, start_time))
                current_chorus_start = None

    # Handle case where song ends during chorus
    if current_chorus_start is not None and text_aligned:
        last_end_time = text_aligned[-1][1]
        chorus_sections.append((current_chorus_start, last_end_time))

    return chorus_sections


def get_chorus_section_offset(data_meta):
    """Get the start time of the first chorus section from aligned lyrics.

    Args:
        data_meta: Metadata dictionary containing text_aligned

    Returns:
        float or None: Start time in seconds of first chorus section, or None if no chorus found
    """
    if not data_meta.get("text_aligned"):
        return None

    chorus_sections = detect_chorus_sections(data_meta["text_aligned"])

    if chorus_sections:
        return max(0, chorus_sections[0][0] - 2)  # Return start time of first chorus, 2 seconds before

    return None


def calculate_activity_percentage(loudness_array, threshold=0.01):
    """Calculate percentage of time that loudness exceeds threshold.

    Args:
        loudness_array: Array of loudness values
        threshold: Minimum loudness to consider "active"

    Returns:
        float: Percentage of active time (0-1)
    """
    if loudness_array is None or len(loudness_array) == 0:
        return 0.0

    active_frames = np.sum(loudness_array > threshold)
    return active_frames / len(loudness_array)


def generate_span_activity_tags(loudness_seq, num_spans=None, rate_hz=25):
    """Generate span-based activity tags from loudness data.

    Args:
        loudness_seq: Array of loudness values at target rate (e.g., 25Hz)
        num_spans: Number of spans to generate (default: random 1-3)
        rate_hz: Feature rate in Hz (default 25), used to determine minimum duration

    Returns:
        list: Activity tags in format "activity[start:end]:percentage%"
    """
    # Minimum 2 seconds of audio required
    min_frames = int(2 * rate_hz)
    if loudness_seq is None or len(loudness_seq) <= min_frames:
        return []

    if num_spans is None:
        num_spans = random.randint(1, 3)

    tags = []
    loudness_length = len(loudness_seq)

    for _ in range(num_spans):
        # Generate random span within the loudness array
        # Minimum span: 0.4 seconds (10 frames at 25Hz)
        min_span_frames = max(1, int(0.4 * rate_hz))
        span_length = random.randint(min_span_frames, loudness_length // 2)
        start_frame = random.randint(0, max(0, loudness_length - span_length))
        end_frame = start_frame + span_length

        # Ensure valid bounds
        end_frame = min(end_frame, loudness_length)

        # Calculate activity for this span
        span_loudness = loudness_seq[start_frame:end_frame]
        if len(span_loudness) > 0:
            span_activity_pct = calculate_activity_percentage(span_loudness) * 100
            rounded_span_activity = round(span_activity_pct / 10) * 10
            tags.append(f"activity[{start_frame}:{end_frame}]:{int(rounded_span_activity)}%")

    return tags


def generate_activity_tags(loudness_seq, rate_hz=25):
    """Generate all activity-related tags from loudness data.

    Args:
        loudness_seq: Array of loudness values at target rate (e.g., 25Hz)
        rate_hz: Feature rate in Hz (default 25), used to determine minimum duration

    Returns:
        list: All activity tags (basic + threshold + span-based)
    """
    if loudness_seq is None:
        return []

    tags = []

    # Calculate overall activity percentage
    activity_percentage = calculate_activity_percentage(loudness_seq)

    # Basic activity tag
    rounded_activity = round(activity_percentage * 100 / 10) * 10
    tags.append(f"activity:{int(rounded_activity)}%")

    # Threshold-based activity tag
    activity_pct = activity_percentage * 100
    random_threshold = random.uniform(0, activity_pct)
    rounded_threshold = round(random_threshold / 10) * 10
    tags.append(f"activity>{int(rounded_threshold)}%")

    # Span-based activity tags
    span_tags = generate_span_activity_tags(loudness_seq, rate_hz=rate_hz)
    tags.extend(span_tags)

    return tags


def calculate_spectral_metric_percentage(array, threshold):
    """Calculate percentage of frames where metric exceeds threshold.

    Similar to calculate_activity_percentage but for generic spectral metrics.

    Args:
        array: Array of metric values
        threshold: Minimum value to consider "above threshold"

    Returns:
        float: Percentage of frames above threshold (0-1)
    """
    if array is None or len(array) == 0:
        return 0.0

    above_threshold = np.sum(array > threshold)
    return above_threshold / len(array)


def generate_loudness_contour_tags(loudness_normalized, is_warped=False):
    """Generate single full-length span tag with ALL frame values for loudness contour.

    Example output:
        loudness_contour:[45,52,68,71,...]  (normal)
        loudness_warped_contour:[45,52,68,71,...]  (when warped)

    Args:
        loudness_normalized: Array of normalized loudness values [0,1]
        is_warped: If True, use "_warped_contour" suffix

    Returns:
        list: Single-element list with full-length value tag in 0-100 format
    """
    if loudness_normalized is None or len(loudness_normalized) == 0:
        return []

    # Scale normalized [0,1] values to integer [0,100] range
    scaled_values = [int(round(v * 100)) for v in loudness_normalized]
    values_str = ",".join([str(v) for v in scaled_values])

    tag_name = "loudness_warped_contour" if is_warped else "loudness_contour"
    return [f"{tag_name}:[{values_str}]"]


def generate_spectral_centroid_span_with_values(centroid_normalized, is_warped=False):
    """Generate single full-length span tag with ALL frame values for spectral centroid.

    Example output:
        spectral_centroid_contour:[53,57,60,58,62,...]  (normal)
        spectral_centroid_warped_contour:[53,57,60,58,62,...]  (when warped)

    Args:
        centroid_normalized: Array of normalized spectral centroid values [0,1]
        is_warped: If True, use "_warped_contour" suffix

    Returns:
        list: Single-element list with full-length value tag in 0-100 format
    """
    if centroid_normalized is None or len(centroid_normalized) == 0:
        return []

    # Scale normalized [0,1] values to integer [0,100] range
    scaled_values = [int(round(v * 100)) for v in centroid_normalized]
    values_str = ",".join([str(v) for v in scaled_values])

    tag_name = "spectral_centroid_warped_contour" if is_warped else "spectral_centroid_contour"
    return [f"{tag_name}:[{values_str}]"]


def generate_spectral_complexity_span_with_values(complexity_normalized, is_warped=False):
    """Generate single full-length span tag with ALL frame values for spectral complexity.

    Example output:
        spectral_complexity_contour:[68,71,73,70,75,...]  (normal)
        spectral_complexity_warped_contour:[68,71,73,70,75,...]  (when warped)

    Args:
        complexity_normalized: Array of normalized spectral complexity values [0,1]
        is_warped: If True, use "_warped_contour" suffix

    Returns:
        list: Single-element list with full-length value tag in 0-100 format
    """
    if complexity_normalized is None or len(complexity_normalized) == 0:
        return []

    # Scale normalized [0,1] values to integer [0,100] range
    scaled_values = [int(round(v * 100)) for v in complexity_normalized]
    values_str = ",".join([str(v) for v in scaled_values])

    tag_name = "spectral_complexity_warped_contour" if is_warped else "spectral_complexity_contour"
    return [f"{tag_name}:[{values_str}]"]


def generate_spectral_centroid_tags(centroid_normalized, is_warped=False):
    """Generate spectral centroid tag from normalized centroid data.

    Args:
        centroid_normalized: Array of normalized spectral centroid values [0,1]
        is_warped: If True, use "_warped_contour" suffix

    Returns:
        list: Single centroid value list tag in 0-100 format
    """
    if centroid_normalized is None or len(centroid_normalized) == 0:
        return []

    return generate_spectral_centroid_span_with_values(centroid_normalized, is_warped=is_warped)


def generate_spectral_complexity_tags(complexity_normalized, is_warped=False):
    """Generate spectral complexity tag from normalized complexity data.

    Args:
        complexity_normalized: Array of normalized spectral complexity values [0,1]
        is_warped: If True, use "_warped_contour" suffix

    Returns:
        list: Single complexity value list tag in 0-100 format
    """
    if complexity_normalized is None or len(complexity_normalized) == 0:
        return []

    return generate_spectral_complexity_span_with_values(complexity_normalized, is_warped=is_warped)


global tokenizer
g_tokenizer = None

MAX_TAG_LEN = 256  # characters, not tokens
MAX_TOT_TAGS_LEN = 1024  # characters, not tokens
MAX_N_TAGS = 10  # to avoid overfitting to an artist


def load_tokenizer(tokenizer_fp=None):
    global g_tokenizer
    if g_tokenizer is not None:
        return g_tokenizer
    assert os.path.exists(tokenizer_fp)
    g_tokenizer = PreTrainedTokenizerFast(
        tokenizer_file=tokenizer_fp,
        unk_token="[UNK]",
        pad_token="[PAD]",
        clean_up_tokenization_spaces=True,
    )
    g_tokenizer.add_special_tokens({"additional_special_tokens": [AddedToken("\n")]})
    return g_tokenizer


def _space_repl(m):
    s = m.group()
    n_newline = s.count("\n")
    if n_newline >= 2:
        return "\n\n"
    elif n_newline == 1:
        return "\n"
    return " "


def _simplify_whitespace(text, retain_newlines=True):
    """simplify while respecting up to 2 newlines"""
    if retain_newlines:
        text = re.sub(r"\s+", _space_repl, text).strip()
    else:
        text = re.sub(r"\s+", " ", text).strip()
    return text


def tokenize_batch(
    text_list,
    max_tokens=None,
    pad_token_id=0,
    retain_newlines=True,
    tokenizer_fp=None,
):
    tokenizer = load_tokenizer(tokenizer_fp)
    text_list = [_simplify_whitespace(s, retain_newlines=retain_newlines) for s in text_list]
    text_enc = tokenizer(
        text_list,
        add_special_tokens=False,
        truncation=max_tokens is not None,
        max_length=max_tokens,
        padding="longest",
        return_tensors="pt",
    )["input_ids"].type(torch.long)
    text_enc[text_enc == tokenizer.pad_token_id] = pad_token_id
    return text_enc


CASE_AUGMENT_FUNCS = [
    str.upper,
    str.lower,
    str.capitalize,
    str.title,
]


def _augment_tag(s):
    # case augment
    if random.random() >= 0.8:
        s = random.choice(CASE_AUGMENT_FUNCS)(s)
    # other misc formatting
    if random.random() >= 0.5:
        s = s.replace("-", " ").strip()
    return s


def get_control_tags(
    sample_duration_s,
    sample_duration_toks,
    sample_offset_toks=0,
    sample_vocal_start_s=None,
    hook_only=False,
    start_from_chorus=False,
    do_augment=True,
    max_possible_duration_s: int = 60 * 40,
    critical_control_tags: list[str] | None = None,
    semantic_rate_hz: int = 25,
    audio_sample_start_times_s=None,  # Audio sample timings for conditioning (list)
    audio_sample_sources=None,  # Source types for each audio sample ("vocal", "drum", "full_mix", etc.)
    loudness_25hz=None,  # Loudness at 25Hz for activity tag calculations
    loudness_seq=None,  # Loudness contour with randomized rate/smoothing/warping
    spectral_centroid_seq=None,  # Spectral centroid sequence with randomized rate/smoothing/warping
    spectral_complexity_seq=None,  # Spectral complexity sequence with randomized rate/smoothing/warping
    contour_rate_hz: int
    | None = None,  # Rate for loudness_seq and spectral contour features (randomized 0.2-1Hz, for debugging)
    contour_is_warped: bool = False,  # Whether time-warping was applied to contour features
    vocal_pitch_hz_min=None,  # Minimum vocal pitch in Hz for pitch range control tags
    vocal_pitch_hz_max=None,  # Maximum vocal pitch in Hz for pitch range control tags
):
    control_tags = []
    control_tags.append(f"duration:{int(round(sample_duration_s))}")
    control_tags.append(f"duration_toks:{sample_duration_toks}")

    increment_s = 10
    min_durations = [
        dur for dur in range(0, max_possible_duration_s, increment_s) if dur <= sample_duration_s
    ]
    max_durations = [
        dur for dur in range(0, max_possible_duration_s, increment_s) if dur >= sample_duration_s
    ]
    if len(min_durations) > 0:
        control_tags.append(f"min_duration:{int(random.choice(min_durations))}")
    if len(max_durations) > 0:
        control_tags.append(f"max_duration:{int(random.choice(max_durations))}")

    if sample_vocal_start_s is not None:
        if sample_vocal_start_s <= 5:
            control_tags.append("vocals:early")
        if sample_vocal_start_s <= 15:
            control_tags.append("vocals:normal")
        if 10 <= sample_vocal_start_s <= 20:
            control_tags.append("vocals:intro")

    if critical_control_tags is None:
        critical_control_tags = []
    if sample_offset_toks > 0:
        critical_control_tags.append(f"start_offset:{sample_offset_toks}")
    if hook_only:
        critical_control_tags.append("hook_offset:random")
    if start_from_chorus:
        critical_control_tags.append("start_from_chorus")

    # Add audio sample timings and source tags if available
    if audio_sample_start_times_s is not None and len(audio_sample_start_times_s) > 0:
        for i, start_time in enumerate(audio_sample_start_times_s):
            start_time_s = int(round(start_time))
            start_time_toks = int(round(start_time * semantic_rate_hz))
            # Always use zero-based indexing for consistency
            control_tags.append(f"audio_sample_time_{i}:{start_time_s}")
            control_tags.append(f"audio_sample_start_toks_{i}:{start_time_toks}")

            # Add source type tag if available
            if audio_sample_sources is not None and i < len(audio_sample_sources):
                source_type = audio_sample_sources[i]
                control_tags.append(f"audio_sample_{source_type}_{i}")

    # Add activity-based tags from loudness_25hz data (precise alignment at 25Hz)
    activity_tags = generate_activity_tags(loudness_25hz, rate_hz=semantic_rate_hz)
    control_tags.extend(activity_tags)

    # Add loudness contour tags (with randomized rate/smoothing/warping)
    loudness_contour_tags = generate_loudness_contour_tags(loudness_seq, is_warped=contour_is_warped)
    control_tags.extend(loudness_contour_tags)

    # Add spectral centroid tags (brightness control)
    centroid_tags = generate_spectral_centroid_tags(spectral_centroid_seq, is_warped=contour_is_warped)
    control_tags.extend(centroid_tags)

    # Add spectral complexity tags (harmonic richness control)
    complexity_tags = generate_spectral_complexity_tags(
        spectral_complexity_seq, is_warped=contour_is_warped
    )
    control_tags.extend(complexity_tags)

    # Add vocal pitch range control tags if available
    if vocal_pitch_hz_min is not None:
        control_tags.append(f"vocal_pitch_hz_min:{int(round(vocal_pitch_hz_min))}")
    if vocal_pitch_hz_max is not None:
        control_tags.append(f"vocal_pitch_hz_max:{int(round(vocal_pitch_hz_max))}")

    if do_augment:
        if random.random() >= 0.5:
            random.shuffle(control_tags)
            control_tags = control_tags[: random.randint(0, len(control_tags))]
    if critical_control_tags is not None:
        control_tags += critical_control_tags
        random.shuffle(control_tags)
    if len(control_tags) == 0:
        return None
    return "{" + ";".join(control_tags) + "}"


def augment_tags(tags):
    random.shuffle(tags)
    if random.random() <= 0.5:
        tags = tags[: random.randint(0, len(tags))]
        tags = [_augment_tag(tag) for tag in tags]
    return tags


def clean_tags(tags):
    return [
        clean_tag[:MAX_TAG_LEN]
        for tag in tags
        if len(clean_tag := _simplify_whitespace(tag, retain_newlines=False)) > 0
    ]


def clean_inline_tags(m):
    tags = m.group(2).split(";")
    tags = clean_tags(tags)
    ts = ";".join(tags)[:MAX_TOT_TAGS_LEN]
    if len(ts) > 0:
        return f"[{m.group(1)}: {ts}]"
    return f"[{m.group(1)}]"


def augment_inline_tags(m):
    tags = m.group(2).split(";")
    tags = augment_tags(tags)
    merge_char = random.choice([", ", " ", "; ", ",", ";", ". "])
    ts = merge_char.join(tags)
    if len(ts) > 0:
        return f"[{m.group(1)}: {ts}]"
    return f"[{m.group(1)}]"


def randomize_lyrics(lyrics, max_ngram=5):
    # remove some basics and split into words
    lyrics = lyrics.lower()
    lyrics = re.sub(r"\[.*?\]", " ", lyrics)
    lyrics = re.sub(r"[\s+\,\;\(\)]", " ", lyrics).strip()
    words = lyrics.split()

    # Generate fragments with variable length, biased toward single words
    n_grams = random.randint(1, max_ngram)
    fragments = [
        " ".join(words[n * n_grams : (n + 1) * n_grams]) for n in range(math.ceil(len(words) / n_grams))
    ]

    # dedupe and shuffle
    fragments = list(set(fragments))
    random.shuffle(fragments)

    return "; ".join(fragments)


def format_timestamped_lyrics(text_lines, augment_probability=0.4):
    """
    Format lyrics with timestamps based on text_lines data.

    Examples:
        # all_timestamps strategy:
        "[0.5-2.1] Verse one starts here"
        "[2.1-4.3] And continues with more lyrics"

        # start_end_only strategy:
        "[0.5] Verse one starts here"
        "And continues with more lyrics"
        "Final line here [4.3]"

        # random_dropout strategy:
        "[0.5] Verse one starts here"
        "And continues with more lyrics"
        "[3.2-4.3] Final line here"

        # start_times_only strategy:
        "[0.5] Verse one starts here"
        "[2.1] And continues with more lyrics"
        "[3.2] Final line here"
    """
    if not text_lines or random.random() >= augment_probability:
        # Return regular text without timestamps
        if not text_lines:
            return ""
        return "\n".join([line[2] for line in text_lines])

    # Decide on timestamp dropout strategy
    dropout_strategy = random.choice(
        [
            "all_timestamps",  # Include all timestamps
            "start_end_only",  # Only first start and last end
            "per_line_random",  # Randomly decide per line independently
        ]
    )

    def format_line_with_timestamps(text, start_s=None, end_s=None):
        """Helper function to format a line with optional timestamps."""
        # Handle multiple trailing newlines
        trailing_newlines = ""
        clean_text = text
        while clean_text.endswith("\n"):
            trailing_newlines += "\n"
            clean_text = clean_text[:-1]

        if start_s is not None and end_s is not None:
            formatted = f"[{start_s:.1f}] {clean_text} [{end_s:.1f}]"
        elif start_s is not None:
            formatted = f"[{start_s:.1f}] {clean_text}"
        elif end_s is not None:
            formatted = f"{clean_text} [{end_s:.1f}]"
        else:
            formatted = clean_text

        return formatted + trailing_newlines

    formatted_lines = []
    for i, line in enumerate(text_lines):
        start_s, end_s, text = line
        if dropout_strategy == "all_timestamps":
            formatted_lines.append(format_line_with_timestamps(text, start_s, end_s))
        elif dropout_strategy == "start_end_only":
            if i == 0:
                formatted_lines.append(format_line_with_timestamps(text, start_s))
            elif i == len(text_lines) - 1:
                formatted_lines.append(format_line_with_timestamps(text, end_s=end_s))
            else:
                formatted_lines.append(text)
        elif dropout_strategy == "per_line_random":
            rand_choice = random.random()
            if rand_choice < 0.25:  # 25% no timestamp
                formatted_lines.append(text)
            elif rand_choice < 0.5:  # 25% start only
                formatted_lines.append(format_line_with_timestamps(text, start_s))
            elif rand_choice < 0.75:  # 25% end only
                formatted_lines.append(format_line_with_timestamps(text, end_s=end_s))
            else:  # 25% start-end range
                formatted_lines.append(format_line_with_timestamps(text, start_s, end_s))

    return "\n".join(formatted_lines)


def add_vocal_suffixes(vocal_tags):
    """Add vocal suffixes to vocal tags with random probability."""
    if not vocal_tags:
        return vocal_tags

    vocal_suffixes = ["vocal", "voice", "vox"]
    processed_vocal_tags = []

    for tag in vocal_tags:
        if random.random() <= 0.7:  # 70% chance to add suffix
            suffix = random.choice(vocal_suffixes)
            processed_vocal_tags.append(f"{tag} {suffix}")
        else:
            processed_vocal_tags.append(tag)

    return processed_vocal_tags


def add_tags_element(text_elements, tags, separator, prefix=""):
    """Add a tag element to text_elements and return new list.

    Args:
        text_elements: List to add the tag element to
        tags: List of tags to process
        separator: Character(s) to join tags with
        prefix: Optional prefix for the tag element (e.g., "vocal:")

    Returns:
        New text_elements list with tag element added (if valid)
    """
    if not tags:
        return text_elements

    tags_str = f"{separator.join(tags[:MAX_N_TAGS])}"[:MAX_TOT_TAGS_LEN]
    if len(tags_str) > 0:
        element = f"[{prefix}{tags_str}]" if prefix else f"[{tags_str}]"
        return text_elements + [element]

    return text_elements


def build_text(
    tags,
    text,
    sample_duration_s,
    sample_duration_toks,
    sample_vocal_start_s=None,
    inference=False,
    suppress_text=False,
    enable_control_tags=True,
    hook_only=False,
    start_from_chorus=False,
    sample_offset_toks=0,
    max_possible_duration_s=60 * 40,
    critical_control_tags=None,
    audio_sample_start_times_s=None,  # Audio sample timings for text conditioning (list)
    audio_sample_sources=None,  # Source types for each audio sample ("vocal", "drum", "full_mix", etc.)
    semantic_rate_hz=25,  # Default semantic rate for control tag calculations
    loudness_25hz=None,  # Loudness at 25Hz for activity tag calculations
    loudness_seq=None,  # Loudness contour with randomized rate/smoothing/warping
    spectral_centroid_seq=None,  # Spectral centroid sequence with randomized rate/smoothing/warping
    spectral_complexity_seq=None,  # Spectral complexity sequence with randomized rate/smoothing/warping
    contour_rate_hz=None,  # Rate for loudness_seq and spectral features (randomized 0.2-1Hz, for debugging)
    contour_is_warped=False,  # Whether time-warping was applied to contour features
    vocal_tags=None,  # Vocal characteristic tags for voice conditioning
    vocal_pitch_hz_min=None,  # Minimum vocal pitch in Hz for pitch range control tags
    vocal_pitch_hz_max=None,  # Maximum vocal pitch in Hz for pitch range control tags
):
    if suppress_text:
        return ""
    text_elements = []

    # Process tags and vocal tags with shared logic
    tags = clean_tags(tags) if tags else []
    vocal_tags = clean_tags(vocal_tags) if vocal_tags else []

    if not inference:
        # Training mode: augment and use random joining
        tags = augment_tags(tags) if tags else []
        vocal_tags = augment_tags(vocal_tags) if vocal_tags else []

        merge_char = random.choice([", ", " ", "; ", ",", ";", ". "])

        # Two paths for vocal tags processing
        if random.random() <= 0.5:
            # Path 1: Add vocal suffixes and merge into regular tags
            vocal_tags_with_suffixes = add_vocal_suffixes(vocal_tags)
            combined_tags = tags + vocal_tags_with_suffixes
            random.shuffle(combined_tags)  # Randomize the order of combined tags
            text_elements = add_tags_element(text_elements, combined_tags, merge_char)
        else:
            # Path 2: Separate processing - regular tags + [vocal:...] element
            text_elements = add_tags_element(text_elements, tags, merge_char)
            text_elements = add_tags_element(text_elements, vocal_tags, merge_char, prefix="vocal:")
    else:
        # Inference mode: keep the non-inference behavior
        text_elements = add_tags_element(text_elements, tags, ",")
        text_elements = add_tags_element(text_elements, vocal_tags, ",", prefix="vocal:")

    # get lyrics
    if len(text) > 0 and not inference:
        text = re.sub(r"\[(.*?)\:(.*?)\]", clean_inline_tags, text)
        # augment tags inside text:
        text = re.sub(r"\[(.*?)\:(.*?)\]", augment_inline_tags, text)
        if random.random() >= 0.95:
            text = text.lower()
        if random.random() >= 0.95:
            text = re.sub(r"\n+", " ", text)
    if len(text) > 0:
        text_elements.append(text.strip())
    # get control tags
    for n in range(len(text_elements)):
        text_elements[n] = text_elements[n].replace("{", "").replace("}", "")
    if (inference or random.random() >= 0.1) and enable_control_tags:
        control_tags = get_control_tags(
            sample_duration_s,
            sample_duration_toks,
            sample_vocal_start_s=sample_vocal_start_s,
            hook_only=hook_only,
            start_from_chorus=start_from_chorus,
            sample_offset_toks=sample_offset_toks,
            do_augment=not inference,
            max_possible_duration_s=max_possible_duration_s,
            critical_control_tags=critical_control_tags,
            semantic_rate_hz=semantic_rate_hz,
            audio_sample_start_times_s=audio_sample_start_times_s,
            audio_sample_sources=audio_sample_sources,
            loudness_25hz=loudness_25hz,
            loudness_seq=loudness_seq,
            spectral_centroid_seq=spectral_centroid_seq,
            spectral_complexity_seq=spectral_complexity_seq,
            contour_rate_hz=contour_rate_hz,
            contour_is_warped=contour_is_warped,
            vocal_pitch_hz_min=vocal_pitch_hz_min,
            vocal_pitch_hz_max=vocal_pitch_hz_max,
        )
        if control_tags is not None:
            text_elements = [control_tags] + text_elements

    if inference or random.random() >= 0.5:
        text = "\n\n".join(text_elements)
    else:
        text = ""
        for t in text_elements:
            text += random.choice([" ", "\n", "\n\n"]) + t
    text = text.strip()
    return text