import pickle
# zz = torch.stack([z.repeat(50,1).transpose(0,1)])

# GetWeather
getWeather = ['state', 'city', 'timeRange', 'geographic_poi', 'country', 'facility']

# PlayMusic
playMusic = ['artist', 'album', 'playlist', 'service', 'rating_value', 'best_rating', 'music_item', 'track', 'playlist_owner', 'year', 'genre', 'rating_unit']

# AddToPlaylist
addToPlaylist = ['artist', 'album', 'playlist', 'service', 'rating_value', 'best_rating', 'music_item', 'track', 'playlist_owner', 'year', 'genre', 'rating_unit']

# RateBook
## I don't care

# SearchScreeningEvent
## I don't care

# BookRestaurant
## I don't care

# SearchCreativeWork
## I don't care

generic = ['O', 'B-entity_name', 'B-object_name', 'B-object_type', 'I-entity_name', 'I-object_name', 'I-object_type']

getWeatherOut = generic[:]
for i in getWeather:
    getWeatherOut.append('B-{}'.format(i))
    getWeatherOut.append('I-{}'.format(i))

playMusicOut = generic[:]
for i in playMusic:
    playMusicOut.append('B-{}'.format(i))
    playMusicOut.append('I-{}'.format(i))

addToPlaylistOut = generic[:]
for i in addToPlaylist:
    addToPlaylistOut.append('B-{}'.format(i))
    addToPlaylistOut.append('I-{}'.format(i))

out = {"GetWeather": getWeatherOut,
       "PlayMusic": playMusicOut,
       "AddToPlaylist": addToPlaylist}

pickle.dump(out, open("slot_filters.pkl", "wb"))