

import gym
from gym import spaces
import pickle
import copy
import numpy as np
import socket
import torch
import random

USE_SUMMARY_WRITER = False
CONNECT_IN_CONSTRUCTOR = False

if USE_SUMMARY_WRITER:
    from torch.utils.tensorboard import SummaryWriter

class SpeakerRecognitionEnv(gym.Env):
    """Custom Environment that follows gym interface"""

    metadata = {"render.modes": ["human"]}

    def __init__(self, audios):
        #print("audio:", audio)
        super(SpeakerRecognitionEnv, self).__init__()
        # Define action and observation space
        # They must be gym.spaces objects
        # Example when using discrete actions:
        # self.action_space = spaces.Discrete(N_DISCRETE_ACTIONS)
        # Example for using image as input (channel-first; channel-last also works):
        # self.observation_space = spaces.Box(low=0, high=255,
        #                                    shape=(N_CHANNELS, HEIGHT, WIDTH), dtype=np.uint8)

        N_DISCRETE_ACTIONS = 2
        self.action_space = spaces.Discrete(N_DISCRETE_ACTIONS)
        N_CHANNELS = 1
        HEIGHT = 1
        #WIDTH = 1000
        #WIDTH = 41680
        #WIDTH = 50000
        WIDTH = 4000
        self.WIDTH = WIDTH
        self.PORT = 12943
        self.HOST = ''
        self.observation_space = spaces.Box(
            low=-int(1e15) - 1,
            high=int(1e15) + 1,
            #shape=(N_CHANNELS, HEIGHT, WIDTH),
            shape=(WIDTH,),
            #dtype=np.uint8,
            dtype=np.float,
        )
        #self.local_socket = None
        self.original_audios = audios
        self.steps_in_epoch = 100
        self.end_rewards = []
        if USE_SUMMARY_WRITER:
            self.writer = SummaryWriter()
        self.epoch = 0
        self.local_socket = None
        if CONNECT_IN_CONSTRUCTOR:
            self.local_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.local_socket.connect((self.HOST, self.PORT))

            #original_score = self.score_audio(self.original_audio)
            #print("original_score:", original_score)

    def score_audio(self, audio):
        try:
            #print("sending")
            self.local_socket.sendall(pickle.dumps(audio))
            self.local_socket.sendall(b"=========end")
            for _ in range(1000):
                data = self.local_socket.recv(4096)
                if data:
                    break
            #print("data:", data)
            score = pickle.loads(data)
            #print("score:", score)
        except (socket.timeout, ConnectionResetError) as e:
            print(e)
            score = self.score_audio(audio)

        return score

    def get_observation(self):
        length = self.current_audio.shape[0]
        padding_length = int((self.WIDTH / 2) - length)
        padding_audio = np.zeros((padding_length))
        position_encoding = np.zeros((int(self.WIDTH / 2)))
        position_encoding[self.current_position_in_audio] = 1 
        observation = np.concatenate((self.current_audio / 16000, padding_audio, position_encoding))
        #observation = observation[np.newaxis, np.newaxis, :]
        return observation


    def step(self, action):
        if action == 1:
            self.current_audio = np.insert(self.current_audio, self.current_position_in_audio, 0)
       #print("self.current_audio:", self.current_audio)
       #print("self.current_audio:", self.current_audio.shape)
        self.current_position_in_audio += 1
        #observation = self.current_audio
        observation = self.get_observation()
        done = self.current_position_in_audio > self.steps_in_epoch
        if done:
            reward = 10 * self.score_audio(self.current_audio)
            self.end_rewards.append(reward)
            self.epoch += 1
            if USE_SUMMARY_WRITER:
                self.writer.add_scalar('end_reward', reward, self.epoch)
                self.writer.flush()
            print(f'reward: {reward:.5f}  step: {self.current_position_in_audio}  action: {action}')
        else:
            reward = 0
        info = {}
        return observation, reward, done, info

    def reset(self):
        print('in reset')
        if not CONNECT_IN_CONSTRUCTOR:
            if not self.local_socket:
                self.local_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.local_socket.connect((self.HOST, self.PORT))

        self.current_position_in_audio = 0
        #if self.local_socket:
        #    self.local_socket.close()
        random_original_audio = random.choice(self.original_audios)
        self.current_audio = copy.deepcopy(random_original_audio)
        #observation = self.current_audio
        observation = self.get_observation()

        #print("observation:", observation.shape)
        return observation  # reward, done, info can't be included

    #  def render(self, mode='human'):
    #    ...
    def close(self):
        self.local_socket.close()
        self.writer.close()



#class SpeakerRecognitionEnv(gym.Env):
#    metadata = {"render.modes": ["human"]}
#
#    def __init__(self, audio):
#        #print("audio:", audio)
#        pass
#
#    def score_audio(self, audio):
#        pass
#
#    def get_observation(self):
#        pass
#
#
#    def step(self, action):
#        pass
#
#    def reset(self):
#        pass
#
#    #  def render(self, mode='human'):
#    #    ...
#    def close(self):
#        pass
#



#if __name__ == '__main__':
#    print('test')
#    freeze_support()
#    torch.multiprocessing.freeze_support()
#
#    if False:
#        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
#            print("connecting")
#            s.connect((HOST, PORT))
#            # s.settimeout(1)
#            while True:
#                try:
#                    print("sending")
#                    s.sendall(pickle.dumps(audio_list[0]))
#                    # s.sendall(pickle.dumps('jkl'))
#                    s.sendall(b"=========end")
#                    for _ in range(1000):
#                        data = s.recv(4096)
#                        if data:
#                            break
#                    if not data:
#                        continue
#                    print("data:", data)
#                    score = pickle.loads(data)
#                    print("score:", score)
#                except socket.timeout as e:
#                    print(e)
