from xtuner.utils import DEFAULT_SPEECH_TOKEN

# SYSTEM_LS = ('Transcribe the audio into English.\n')
def librispeech_map_fn(example):
    return {
        'conversation': [{
            'system': '',
            'input': f'{DEFAULT_SPEECH_TOKEN}\nTranscribe the audio into English.\n',
            'output': example['text'].lower(),
        }]
    }

def librispeech_test_map_fn(example):
    return {
        'conversation': [{
            'system': '',
            'input': f'{DEFAULT_SPEECH_TOKEN}\nTranscribe the audio into English.\n',
            'output': '',
        }]
    }

def para_map_fn(example):
    return {
        'conversation': [{
            'system': '',
            'input': DEFAULT_SPEECH_TOKEN,
            'output': example['target'],
        }]
    }

def para_map_with_tag_fn(example):
    class_type = example['class']
    if example['class'] == 'accent':
        class_info = example['accent']
    elif example['class'] == 'emotion':
        class_info = example['emotion']
    elif example['class'] == 'age':
        if "age" in example and  "18" in example['age']:
            class_info = 'neutral'
            class_type = "emotion"
        else:
            class_info = 'ten years old'
    elif example['class'] == 'environment':
        class_info = example['background']
    return {
        'conversation': [{
            'system': '',
            'input': f"{DEFAULT_SPEECH_TOKEN}|{class_type}:{class_info}.",
            'output': example['target'],
        }]
    }

def para_test_map_fn(example):
    return {
        'conversation': [{
            'system': '',
            'input': DEFAULT_SPEECH_TOKEN,
            'output': '',
        }]
    }

def para_test_map_with_tag_fn(example):
    class_type = example['class']
    if example['class'] == 'accent':
        class_info = example['accent']
    elif example['class'] == 'emotion':
        class_info = example['emotion']
    elif example['class'] == 'age':
        if "age" in example and  "18" in example['age']:
            class_info = 'neutral'
            class_type = "emotion"
        else:
            class_info = 'ten years old'
    elif example['class'] == 'environment':
        class_info = example['background']
    return {
        'conversation': [{
            'system': '',
            'input': f"{DEFAULT_SPEECH_TOKEN}|{class_type}:{class_info}.",
            'output': '',
        }]
    }

def para_text_map_fn(example):
    return {
        'conversation': [{
            'system': '',
            'input': example['transcript'],
            'output': example['target'],
        }]
    }

def para_text_test_map_fn(example):
    return {
        'conversation': [{
            'system': '',
            'input': example['transcript'],
            'output': '',
        }]
    }

def para_text_map_with_tag_fn(example):
    class_type = example['class']
    if example['class'] == 'accent':
        class_info = example['accent']
    elif example['class'] == 'emotion':
        class_info = example['emotion']
    elif example['class'] == 'age':
        if "age" in example and  "18" in example['age']:
            class_info = 'neutral'
            class_type = "emotion"
        else:
            class_info = 'ten years old'
    elif example['class'] == 'environment':
        class_info = example['background']
    return {
        'conversation': [{
            'system': '',
            'input': f"{example['transcript']}|{class_type}:{class_info}.",
            'output': example['target'],
        }]
    }

def para_text_test_with_tag_map_fn(example):
    class_type = example['class']
    if example['class'] == 'accent':
        class_info = example['accent']
    elif example['class'] == 'emotion':
        class_info = example['emotion']
    elif example['class'] == 'age':
        if "age" in example and  "18" in example['age']:
            class_info = 'neutral'
            class_type = "emotion"
        else:
            class_info = 'ten years old'
    elif example['class'] == 'environment':
        class_info = example['background']
    return {
        'conversation': [{
            'system': '',
            'input': f"{example['transcript']}|{class_type}:{class_info}.",
            'output': '',
        }]
    }
