#!/usr/bin/env python3

import gym
from stable_baselines3 import HER, SAC, DQN, PPO
from speaker_recognition_env import SpeakerRecognitionEnv
import os
from scipy.io.wavfile import read
from stable_baselines3.common.vec_env import SubprocVecEnv
import time





test_dir = "../data/test-set"
adversarial_dir = "../adversarial-audio"
illegal_dir = "../data/illegal-set"


spk_id = '1580'
audio_list = []
spk_dir = os.path.join(test_dir, spk_id)
audio_iter = os.listdir(spk_dir)
for i, audio_name in enumerate(audio_iter):
    path = os.path.join(spk_dir, audio_name)
    _, audio = read(path)
    audio_list.append(audio)


#print(type(audio_list[0]))
#env = gym.make('SpeakerRecognitionEnv')

if True:
    env = SpeakerRecognitionEnv(audio_list[0][:1000])
else:
    env_list = []
    for i in range(10):
        env = SpeakerRecognitionEnv(audio_list[0][:1000])
        env_list.append(env)
        time.sleep(0.3)

    time.sleep(9)
    env = SubprocVecEnv(env_list)

model = PPO("MlpPolicy", env, verbose=1, gamma=1)

model.learn(total_timesteps=int(1e12), log_interval=4)
