import datasets
import os, jsonlines, random
from tqdm import tqdm

dataset = datasets.load_dataset('Salesforce/dialogstudio', 'MULTIWOZ2_2', cache_dir='./topic_dia_cache')['train']

dial_p = ''
dial = []
with jsonlines.open(dial_p, 'r') as f:
    for line in f:
        dial.append(line)
# dial = dial[:10]
for i in tqdm(range(len(dial))):
    
    other_topic = random.choice(range(len(dataset)))
    insert = ''
    for ti in dataset[other_topic]['log']:
        insert += 'Human: '+ti['user utterance']+'Assistant: '+ti['system response'] + '\n'
         
    turns = dial[i]['dial'].split('Human ')
    pos = random.choice(range(1, len(turns)))
    turns[pos] = insert + turns[pos]
    dial[i]['dial'] = 'Human '.join(turns)
    
with jsonlines.open(dial_p.split('.')[0] + 'inserted.jsonl', 'w') as f:
    for di in dial:
        f.write(di)