'''Audio Handler

2021-09-14 <NAME>
'''

import io
import sys
sys.path.insert(0, '/home/jkang/project/serve')
sys.path.insert(0, '../')
from scipy.io import wavfile

import torch
from ts.torch_handler.base_handler import BaseHandler
from resemblyzer import preprocess_wav, VoiceEncoder


class AudioHandler(BaseHandler):
    def __init__(self):
        self._context = None
        self.initialized = False
        # self.device = None
        self.device = torch.device('cuda:0')

    def initialize(self, context):
        properties = context.system_properties
        if self.device is None:
            properties = context.system_properties
            self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")
        self.encoder = VoiceEncoder(device=self.device)

    def inference(self, data):
        return self.encoder.embed_utterance(data)

    def handle(self, file, context):
        # Preprocess
        model_input = self.preprocess_file(file)
        model_output = self.inference(model_input)
        final_output = self.postprocess(model_output)
        return final_output

    def preprocess_file(self, file):
        raw_data = file[0].get('body')
        rate, data = wavfile.read(io.BytesIO(raw_data))
        print(f'sample rate: {rate}')
        #print('>>data_shape:', data.shape)  # (N,)
        return preprocess_wav(data)

    def preprocess_stream(self, stream):
        return NotImplementedError

    def postprocess(self, output_data):
        # output_data: float32
        return [output_data.tolist()]


if __name__ == '__main__':
    audiohandler = AudioHandler()
    audiohandler.initialize(None)
    out = audiohandler.handle('serving/SA1.wav', None)
    print(out)
    print('length:', len(out))