import random
from typing import Any, Optional

import jax
import numpy as np
import transformers
from jax import numpy as jnp
from jax.experimental import maps
from mesh_transformer.checkpoint import read_ckpt_lowmem
from mesh_transformer.sampling import nucleaus_sample
from mesh_transformer.transformer_shard import CausalTransformer

from .constants import ModelParams, InferConfig

def default(value: Any, fallback: Any) -> Any:
    # luke prefers making a function that chooses between `value` and `feedback` so i am gonna keep it
    if value is None:
        return fallback

    return value

_cores_per_replica = ModelParams.cores_per_replica
_mesh_shape = (jax.device_count() // _cores_per_replica, _cores_per_replica)
_devices = np.array(jax.devices()).reshape(_mesh_shape)

maps.thread_resources.env = maps.ResourceEnv(maps.Mesh(_devices, ("dp", "mp")), ())

class Inference:
    _NP_ONE = np.ones((1,))

    def __init__(
            self,
            path: Optional[str] = None,
            parameters: Optional[ModelParams] = None,
            config: Optional[InferConfig] = None,
    ):
        path = "checkpoint_slim/" if path is None else path

        self.params = ModelParams() if parameters is None else parameters
        self.params.sampler = nucleaus_sample
        self.config = InferConfig() if config is None else config

        self.model = CausalTransformer(self.params.__dict__)
        self.tokenizer = transformers.GPT2TokenizerFast.from_pretrained("gpt2")
        self.model.state = read_ckpt_lowmem(
            self.model.state, path, self.params.cores_per_replica, load_opt=False
        )

    def generate_tokens(
            self,
            prompt: np.ndarray,
            length: Optional[int] = None,
            top_p: Optional[float] = None,
            temperature: Optional[float] = None,
    ) -> np.ndarray:
        length = default(length, self.config.token_length)
        top_p = default(top_p, self.config.top_p)
        new_temp = random.random() * (self.config.max_temperature - self.config.min_temperature)
        new_temp += self.config.min_temperature
        temperature = default(temperature, new_temp)
        prompt = prompt[:, -2048:]
        prompt = prompt[:, length:]

        source = jnp.array(
            np.pad(
                prompt,
                (
                    (0, 0),
                    (self.params.seq - prompt.shape[1], 0),
                ),
            )
        )
        sampler_options = {
            "top_p": self._NP_ONE * top_p,
            "temp": self._NP_ONE * temperature,
        }

        with jax.experimental.maps.mesh(_devices, ("dp", "mp")):
            out = self.model.generate(
                source, self._NP_ONE * prompt.shape[1], length, sampler_options
            )

        return out[0][0, :, 0]

    def generate(
            self,
            prompt: str,
            length: Optional[int] = None,
            top_p: Optional[float] = None,
            temperature: Optional[float] = None,
    ) -> str:
        inp_tokens = self.tokenizer([prompt], verbose=False, return_tensors="np")
        inp_tokens = inp_tokens["input_ids"][0]
        out_tokens = self.generate_tokens(
            inp_tokens.reshape(1, -1), length, top_p, temperature
        )

        return self.tokenizer.decode(out_tokens)


import time
import numpy as np
import pyaudio

class Microphone:
    def __init__(self):
        self.stream = None
        self.p = None
        self.continue_mic_stream = True

    def start_stream(self, callback, mic_rate=44100, fps=60):
        self.continue_mic_stream = True
        self.p = pyaudio.PyAudio()
        frames_per_buffer = int(mic_rate / fps)
        self.stream = self.p.open(format=pyaudio.paInt16,
                        channels=1,
                        rate=mic_rate,
                        input=True,
                        frames_per_buffer=frames_per_buffer)
        overflows = 0
        prev_ovf_time = time.time()
        while self.continue_mic_stream:
            try:
                y = np.fromstring(self.stream.read(frames_per_buffer, exception_on_overflow=False), dtype=np.int16)
                y = y.astype(np.float32)
                self.stream.read(self.stream.get_read_available(), exception_on_overflow=False)
                callback(y)
            except IOError:
                overflows += 1
                if time.time() > prev_ovf_time + 1:
                    prev_ovf_time = time.time()
                    print('Audio buffer has overflowed {} times'.format(overflows))
            except:
                pass
        try:
            self.stream.stop_stream()
            self.stream.close()
            del self.stream
            self.stream = None
        except:
            pass
        try:
            self.p.terminate()
            del self.p
            self.p = None
        except:
            pass

    def stop_stream(self):
        self.continue_mic_stream = False
        time.sleep(1)
        if self.stream:
            try:
                self.stream.stop_stream()
                self.stream.close()
                del self.stream
                self.stream = None
            except:
                pass
        if self.p:
            try:
                self.p.terminate()
                del self.p
                self.p = None
            except:
                pass


from itertools import permutations
import torch
import time
import pdb
from pathlib import Path
import os
import numpy as np
import scipy.io.wavfile as wf
import scipy
from pathlib import Path
import soundfile as sf
from numpy.linalg import solve
import warnings

class Tester_Separate(object):
    def __init__(self,dataset,tr_loader,dt_loader,test_loader,model,num_ch_utilize,device,num_spks,chunk_time,save_rootDir,ref_ch,cuda_flag,tr_inference_flag, **ISTFT_args):

        self.dataset = dataset
        self.tr_loader = tr_loader
        self.dt_loader = dt_loader
        self.test_loader = test_loader
        self.model = model
        self.device = device
        self.num_spks = num_spks
        self.num_ch_utilize = num_ch_utilize

        #ISTFT params
        self.fs = ISTFT_args['fs']
        self.window = ISTFT_args['window']
        self.nperseg = ISTFT_args['length']
        self.noverlap = ISTFT_args['overlap']

        hann_win = scipy.signal.get_window(self.window, self.nperseg)
        self.scale = np.sqrt(1.0/hann_win.sum()**2)
        self.MaxINT16 = np.iinfo(np.int16).max

        self.chunk_size = int(chunk_time * self.fs)
        self.save_rootDir = save_rootDir
        self.ref_ch = ref_ch
        self.cuda_flag = cuda_flag
        self.tr_inference_flag = tr_inference_flag

    def test(self):
        print('testing...')

        if self.tr_inference_flag:
            # print('-'*85)
            # print(' train dataset inference...')
            # print('-'*85)
            # saveDir = os.path.join(self.save_rootDir+'_{}ch'.format(self.num_ch_utilize), 'train_si284')
            # Path(saveDir).mkdir(exist_ok=True, parents=True)
            # self.inference(self.tr_loader, saveDir)
            # print('train dataset inference finish !')

            print('-'*85)
            print('development dataset inference...')
            print('-'*85)
            saveDir = os.path.join(self.save_rootDir+'_{}ch'.format(self.num_ch_utilize),'cv_dev93')
            Path(saveDir).mkdir(exist_ok=False, parents=True)
            self.inference(self.dt_loader, saveDir)
            print('development dataset inference finish !')

        else:
            print('-'*85)
            print('development dataset inference...')
            print('-'*85)
            saveDir = os.path.join(self.save_rootDir+'_{}ch'.format(self.num_ch_utilize),'cv_dev93')
            Path(saveDir).mkdir(exist_ok=False, parents=True)
            self.inference(self.dt_loader, saveDir)
            print('development dataset inference finish !')

            print('-'*85)
            print('test dataset inference...')
            print('-'*85)
            saveDir = os.path.join(self.save_rootDir+'_{}ch'.format(self.num_ch_utilize), 'test_eval92')
            Path(saveDir).mkdir(exist_ok=False, parents=True)
            self.inference(self.test_loader,saveDir)
            print('test dataset inference finish !')

    def inference(self,data_loader,saveDir):
        """
            split_observe_dict : dictionary, keys : split index
                                split_observe_dict[split_index] : [B,Ch,F,T]
            split_clean_s0_dict :        "
            split_clean_s1_dict :        "

            gap : zero pad size of last split wav,  [B]
            wav_name : file name to test
        """
        for idx, (data) in enumerate(data_loader):
            split_observe_dict, split_clean_s0_dict, split_clean_s1_dict, gap, wav_name = data
            split_len = len(split_observe_dict)

            for split_idx in range(split_len):

                observe = split_observe_dict[str(split_idx)] #[B,Ch,F,T]
                B, Ch, T, F = observe.size()
                if self.cuda_flag:
                    observe = observe.cuda(self.device)

                if split_idx == 0:
                    t_e_clean_s0 = np.empty([B], dtype=np.ndarray)
                    t_e_clean_s1 = np.empty([B], dtype=np.ndarray)

                s0 = split_clean_s0_dict[str(split_idx)][:,self.ref_ch,:,:] # [B,1,T,F]
                s1 = split_clean_s1_dict[str(split_idx)][:,self.ref_ch,:,:]
                clean = torch.stack((s0,s1), dim=1) #[B,Ch,T,F]
                s_clean = torch.unsqueeze(clean, dim=2) #[B,Ch,1,T,F]

                """
                    Apply Source Separation
                """
                e_clean_MISO1_temp = self.MISO1_Inference(observe, ref_ch=self.ref_ch)
                if self.cuda_flag:
                    observe = observe.detach().cpu()
                    for spk_idx in range(self.num_spks):
                        e_clean_MISO1_temp[spk_idx] = e_clean_MISO1_temp[spk_idx].detach().cpu()

                """
                    Source Alignment between Clean reference signal and MISO1 signal
                    calculate magnitude distance between ref mic(ch0) and target signal(reference mic : ch0)
                """
                e_clean_MISO1 = [[] for _ in range(self.num_spks)]
                for spk_idx in range(self.num_spks):
                    e_clean_MISO1[spk_idx] = torch.zeros_like(e_clean_MISO1_temp[spk_idx])
                    if spk_idx == 0:
                        e_clean_MISO1_ref_ch = e_clean_MISO1_temp[spk_idx][:,self.ref_ch,:,:] #[B,T,F]
                    else:
                        e_clean_MISO1_ref_ch = torch.stack((e_clean_MISO1_ref_ch, e_clean_MISO1_temp[spk_idx][:,self.ref_ch,:,:]),dim=1)
                # MISO1_stft_ref_ch : [B,Spks,T,F]

                s_e_clean_MISO1_ref_ch = torch.unsqueeze(e_clean_MISO1_ref_ch, dim=1) #[B,1,Spks,T,F]
                magnitude_e_clean_MISO1_ref_ch = torch.abs(torch.sqrt(s_e_clean_MISO1_ref_ch.real**2 + s_e_clean_MISO1_ref_ch.imag**2))
                magnitude_distance = torch.sum(torch.abs(magnitude_e_clean_MISO1_ref_ch - abs(s_clean)),[3,4])
                perms = clean.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                index_ = torch.unsqueeze(perms, dim=2)
                perms_one_hot = clean.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                batchwise_distance = torch.einsum('bij,pij->bp', [magnitude_distance, perms_one_hot])
                min_distance_idx = torch.argmin(batchwise_distance, dim=1)

                for b_idx in range(B):
                    align_index = torch.argmax(perms_one_hot[min_distance_idx[b_idx]], dim=1)
                    for spk_idx in range(self.num_spks):
                        target_index = align_index[spk_idx]
                        e_clean_MISO1[spk_idx][b_idx,...] = e_clean_MISO1_temp[target_index][b_idx,...]

                for b_idx in range(B):
                    temp_s0 = e_clean_MISO1[0]
                    temp_s0 = temp_s0[b_idx,...] * self.scale
                    if len(temp_s0.shape) == 2:
                        temp_s0 = torch.unsqueeze(temp_s0, dim=0)
                    temp_s0 = torch.permute(temp_s0, [0,2,1])
                    t_split_e_clean_s0 = self.ISTFT(temp_s0)
                    t_split_e_clean_s0 = t_split_e_clean_s0 * self.MaxINT16
                    t_split_e_clean_s0 = t_split_e_clean_s0.astype(np.int16)
                    assert t_split_e_clean_s0.shape[0] == self.chunk_size, ('estimate source 0 wav length does not match chunk size, please check ISTFT function ')

                    temp_s1 = e_clean_MISO1[1]
                    temp_s1 = temp_s1[b_idx,...] * self.scale
                    if len(temp_s1.shape) == 2:
                        temp_s1 = torch.unsqueeze(temp_s1, dim=0)
                    temp_s1 = torch.permute(temp_s1, [0,2,1])
                    t_split_e_clean_s1 = self.ISTFT(temp_s1)
                    t_split_e_clean_s1 = t_split_e_clean_s1 * self.MaxINT16
                    t_split_e_clean_s1 = t_split_e_clean_s1.astype(np.int16)
                    assert t_split_e_clean_s1.shape[0] == self.chunk_size, ('estimate source 1 wav length does not match chunk size, please check ISTFT function ')

                    if split_idx == split_len -1:
                        t_split_e_clean_s0 = t_split_e_clean_s0[:len(t_split_e_clean_s0)-gap[b_idx]]
                        t_split_e_clean_s1 = t_split_e_clean_s1[:len(t_split_e_clean_s1)-gap[b_idx]]

                    if split_idx == 0:
                        t_e_clean_s0[b_idx] = t_split_e_clean_s0
                        t_e_clean_s1[b_idx]=  t_split_e_clean_s1
                    else:
                        t_e_clean_s0[b_idx] = np.append(t_e_clean_s0[b_idx] ,t_split_e_clean_s0,  axis=0)
                        t_e_clean_s1[b_idx] = np.append(t_e_clean_s1[b_idx], t_split_e_clean_s1, axis=0)
            for b_idx in range(B):
                sf.write('{}_0.wav'.format(os.path.join(saveDir,wav_name[b_idx])), t_e_clean_s0[b_idx], self.fs, 'PCM_24')
                sf.write('{}_1.wav'.format(os.path.join(saveDir,wav_name[b_idx])), t_e_clean_s1[b_idx], self.fs, 'PCM_24')
            print('[MISO1] Testing ... ','Save Directory : {}'.format(saveDir), 'process : {:.2f}'.format(idx/len(data_loader)))

    def ISTFT(self,FT_sig):

        '''
        input : [C,F,T]
        output : [C,T]
        '''
        # if FT_sig.shape[1] != self.config['ISTFT']['length']+1:
            # FT_sig = np.transpose(FT_sig,(0,1)) # [C,T,F] -> [C,F,T]

        _, t_sig = scipy.signal.istft(FT_sig,fs=self.fs, window=self.window, nperseg=self.nperseg, noverlap=self.noverlap) #[C,F,T] -> [C,T]

        t_sig = t_sig.T
        return t_sig

    def MISO1_Inference(self,mix_stft,ref_ch=0):
        """
        Input:
            mix_stft : observe STFT, size - [B, Mic, T, F]
        Output:
            MISO1_stft : list of separated source, - [B, reference Mic, T, F]

            1. circular shift the microphone array at run time for the prediction of each microphone signal
               If the microphones are arranged uniformly on a circle, Select the reference microphone by circular shifting the microphone. e.g reference mic q -> [Yq, Yq+1, ..., Yp, Y1, ..., Yq-1]
            2. Using Permutation Invariance Alignmnet method to match between clean target signal and estimated signal
        """
        B, M, T, F = mix_stft.size()

        MISO1_stft = [torch.empty(B,M,T,F, dtype=torch.complex64) for _ in range(self.num_spks)]

        Mic_array = [x for x in range(M)]
        Mic_array = np.roll(Mic_array, -ref_ch)  # [ref_ch, ref_ch+1, ..., 0, 1, ..., ref_ch-1]
        # print('Mic_array : ', Mic_array)

        with torch.no_grad():
            mix_stft_refCh = torch.roll(mix_stft,-ref_ch, dims=1)
            MISO1_refCh = self.model(mix_stft_refCh)

        for spk_idx in range(self.num_spks):
            MISO1_stft[spk_idx][:,ref_ch,...] = MISO1_refCh[:,spk_idx,...]

        # MISO1_spk1[:,ref_ch,...] = MISO1_refCh[:,0,...]
        # MISO1_spk2[:,ref_ch,...] = MISO1_refCh[:,1,...]

        s_MISO1_refCh = torch.unsqueeze(MISO1_refCh, dim=2)
        s_Magnitude_refCh = torch.abs(torch.sqrt(s_MISO1_refCh.real**2 + s_MISO1_refCh.imag**2)) # [B,Spks,1,T,F]

        with torch.no_grad():
            for shiftIdx in Mic_array[1:]:
                # print('shift Micnumber', shiftIdx)

                mix_stft_shift = torch.roll(mix_stft,-shiftIdx, dims=1)
                MISO1_chShift = self.model(mix_stft_shift)

                s_MISO1_chShift = torch.unsqueeze(MISO1_chShift, dim=1) #[B,1,Spks,T,F]
                s_magnitude_chShift = torch.sum(torch.abs(s_Magnitude_refCh - abs(s_MISO1_chShift)),[3,4]) #[B,Spks,Spks,T,F]
                perms = MISO1_chShift.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                index_ = torch.unsqueeze(perms, dim=2)
                perms_one_hot = MISO1_chShift.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                batchwise_distance = torch.einsum('bij,pij->bp', [s_magnitude_chShift, perms_one_hot])
                min_distance_idx = torch.argmin(batchwise_distance,dim=1)

                for batch_idx in range(B):
                    align_index = torch.argmax(perms_one_hot[min_distance_idx[batch_idx]],dim=1)
                    for spk_idx in range(self.num_spks):
                        target_index = align_index[spk_idx]
                        MISO1_stft[spk_idx][:,shiftIdx,...] = MISO1_chShift[batch_idx,target_index,...]

        return MISO1_stft

class Tester_Beamforming(object):
    def __init__(self,dataset,tr_loader, dt_loader,test_loader,model,num_ch_utilize,device,num_spks,chunk_time,save_rootDir,
    ref_ch, cuda_flag, tr_inference_flag, utterance_flag, **ISTFT_args):

        self.dataset = dataset
        self.tr_loader = tr_loader
        self.dt_loader = dt_loader
        self.test_loader = test_loader
        self.model = model
        self.num_ch_utilize = num_ch_utilize
        self.device = device
        self.num_spks = num_spks

        # ISTFT params
        self.fs = ISTFT_args['fs']
        self.window = ISTFT_args['window']
        self.nperseg = ISTFT_args['length']
        self.noverlap = ISTFT_args['overlap']

        hann_win = scipy.signal.get_window(self.window, self.nperseg)
        self.scale = np.sqrt(1.0/hann_win.sum()**2)
        self.MaxINT16 = np.iinfo(np.int16).max

        self.chunk_size = int(chunk_time * self.fs)
        self.save_rootDir = save_rootDir
        self.ref_ch = ref_ch
        self.cuda_flag = cuda_flag
        self.tr_inference_flag =tr_inference_flag
        self.utterance_flag = utterance_flag

    def test(self):
        print('testing...')
        if self.tr_inference_flag:
            print('-'*85)
            print(' train dataset inference...')
            print('-'*85)
            saveDir = os.path.join(self.save_rootDir, 'train_si284')
            Path(saveDir).mkdir(exist_ok=True, parents=True)
            self.inference(self.tr_loader, saveDir)
            print('train dataset inference finish !')

            # print('-'*85)
            # print('development dataset inference ...')
            # print('-'*85)
            # saveDir = os.path.join(self.save_rootDir, 'cv_dev93')
            # Path(saveDir).mkdir(exist_ok=True, parents=True)
            # self.inference(self.dt_loader,saveDir)
            # print('development dataset inference finish !')

        else:
            print('-'*85)
            print('development dataset inference ...')
            print('-'*85)
            saveDir = os.path.join(self.save_rootDir, 'cv_dev93')
            Path(saveDir).mkdir(exist_ok=True, parents=True)
            self.inference(self.dt_loader,saveDir)
            print('development dataset inference finish !')

            print('-'*85)
            print('test dataset inference...')
            print('-'*85)
            saveDir = os.path.join(self.save_rootDir, 'test_eval92')
            Path(saveDir).mkdir(exist_ok=True, parents=True)
            self.inference(self.test_loader, saveDir)
            print('test dataset inference finish !')

    def inference(self, data_loader, saveDir):
        """
            split_observe_dict : dictionary, keys : split index
                                split_observe_dict[split_index] : [B,Ch,F,T]
            split_clean_s0_dict :       "
            split_clean_s1_dict :       "

            gap : zero pad size of last split wav, [B]
            wav_name : file name to test
        """

        for idx, (data) in enumerate(data_loader):
            split_observe_dict, split_clean_s0_dict, split_clean_s1_dict, gap, wav_name = data
            split_len = len(split_observe_dict)

            if self.utterance_flag:

                for split_idx in range(split_len):

                    observe = split_observe_dict[str(split_idx)] #[B,Ch,F,T]
                    if self.cuda_flag:
                        observe = observe.cuda(self.device)
                    B, Ch, T, F = observe.size()

                    if split_idx == 0:
                        t_e_clean= np.empty([self.num_spks, B], dtype=np.ndarray)
                        t_observe = np.empty([B], dtype=np.ndarray)

                    s0 = split_clean_s0_dict[str(split_idx)][:,self.ref_ch,:,:] # [B,1,T,F]
                    s1 = split_clean_s1_dict[str(split_idx)][:,self.ref_ch,:,:]
                    clean = torch.stack((s0,s1), dim=1) #[B,Ch,T,F]
                    s_clean = torch.unsqueeze(clean, dim=2) #[B,Ch,1,T,F]

                    """
                        Apply Source Separation
                    """
                    e_clean_MISO1_temp = self.MISO1_Inference(observe, ref_ch=self.ref_ch)
                    if self.cuda_flag:
                        observe = observe.detach().cpu()
                        for spk_idx in range(self.num_spks):
                            e_clean_MISO1_temp[spk_idx] = e_clean_MISO1_temp[spk_idx].detach().cpu()

                    """
                        Source Alignment between Clean reference signal and MISO1 signal
                        calculate magnitude distance between ref mic(ch0) and target signal(reference mic : ch0)
                    """
                    e_clean_MISO1 = [[] for _ in range(self.num_spks)]
                    for spk_idx in range(self.num_spks):
                        e_clean_MISO1[spk_idx] = torch.zeros_like(e_clean_MISO1_temp[spk_idx])
                        if spk_idx == 0:
                            e_clean_MISO1_ref_ch = e_clean_MISO1_temp[spk_idx][:,self.ref_ch,:,:] #[B,T,F]
                        else:
                            e_clean_MISO1_ref_ch = torch.stack((e_clean_MISO1_ref_ch, e_clean_MISO1_temp[spk_idx][:,self.ref_ch,:,:]),dim=1)
                    # MISO1_stft_ref_ch : [B,Spks,T,F]

                    s_e_clean_MISO1_ref_ch = torch.unsqueeze(e_clean_MISO1_ref_ch, dim=1) #[B,1,Spks,T,F]
                    magnitude_e_clean_MISO1_ref_ch = torch.abs(torch.sqrt(s_e_clean_MISO1_ref_ch.real**2 + s_e_clean_MISO1_ref_ch.imag**2))
                    magnitude_distance = torch.sum(torch.abs(magnitude_e_clean_MISO1_ref_ch - abs(s_clean)),[3,4])
                    perms = clean.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                    index_ = torch.unsqueeze(perms, dim=2)
                    perms_one_hot = clean.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                    batchwise_distance = torch.einsum('bij,pij->bp', [magnitude_distance, perms_one_hot])
                    min_distance_idx = torch.argmin(batchwise_distance, dim=1)

                    for b_idx in range(B):
                        align_index = torch.argmax(perms_one_hot[min_distance_idx[b_idx]], dim=1)
                        for spk_idx in range(self.num_spks):
                            target_index = align_index[spk_idx]
                            e_clean_MISO1[spk_idx][b_idx,...] = e_clean_MISO1_temp[target_index][b_idx,...]

                    for b_idx in range(B):
                        for spk_idx in range(self.num_spks):
                            e_clean_MISO1_spk = e_clean_MISO1[spk_idx][b_idx,...] * self.scale
                            if len(e_clean_MISO1_spk.shape) == 2:
                                e_clean_MISO1_spk = torch.unsqueeze(e_clean_MISO1_spk, dim=0)
                            e_clean_MISO1_spk = torch.permute(e_clean_MISO1_spk, [0,2,1]) # [B,T,F] -> [B,F,T]
                            t_split_e_clean = self.ISTFT(e_clean_MISO1_spk)
                            assert t_split_e_clean.shape[1] == self.chunk_size, ('estimate source{} length : {} does not match chnuk size : {}'.format(spk_idx, t_split_e_clean.shape[1], self.chunk_size))

                            if split_idx == split_len -1:
                                t_split_e_clean = t_split_e_clean[:,:t_split_e_clean.shape[1] - gap[b_idx].item()]

                            if split_idx == 0:
                                t_e_clean[spk_idx,b_idx] = t_split_e_clean
                            else:
                                t_e_clean[spk_idx,b_idx] = np.append(t_e_clean[spk_idx,b_idx], t_split_e_clean, axis=1)

                        if isinstance(observe, np.ndarray): observe = torch.from_numpy(observe)
                        t_split_observe = observe[b_idx,...]* self.scale
                        t_split_observe = torch.permute(t_split_observe, [0,2,1]) # [B,T,F] -> [B,F,T]
                        t_split_observe = self.ISTFT(t_split_observe)
                        assert t_split_observe.shape[1] == self.chunk_size, ('input length : {} does not match chnuk size : {}'.format(t_split_observe.shape[1], self.chunk_size))

                        if split_idx == split_len -1:
                            t_split_observe = t_split_observe[:,:t_split_observe.shape[1]-gap[b_idx].item()]
                        if split_idx == 0:
                            t_observe[b_idx] = t_split_observe
                        else:
                            t_observe[b_idx] = np.append(t_observe[b_idx], t_split_observe, axis=1)

                for b_idx in range(B):
                    stft_observe_utterance = self.STFT(t_observe[b_idx])
                    stft_observe_utterance = stft_observe_utterance / self.scale
                    stft_observe_utterance = torch.unsqueeze(torch.permute(stft_observe_utterance,[1,0,2]),dim=0) #[1,F,C,T]
                    e_clean_Beamformer = [[] for _ in range(self.num_spks)]

                    for spk_idx in range(self.num_spks):

                        clean_stft_utterance = self.STFT(t_e_clean[spk_idx,b_idx])
                        clean_stft_utterance = clean_stft_utterance / self.scale
                        # # Test
                        # temp = self.ISTFT(clean_stft_utterance * self.scale)
                        # sf.write('clean.wav', temp.T, self.fs, 'PCM_24')
                        # temp = self.ISTFT(torch.squeeze(torch.permute(stft_observe_utterance * self.scale, [0,2,1,3])))
                        # sf.write('mix.wav', temp.T, self.fs, 'PCM_24')

                        clean_stft_utterance = torch.unsqueeze(torch.permute(clean_stft_utterance,[1,0,2]), dim=0)
                        e_clean_Beamformer[spk_idx] = torch.squeeze(self.Apply_Beamforming(clean_stft_utterance, stft_observe_utterance),dim=1)

                        t_utterance_e_clean = self.ISTFT(torch.permute(e_clean_Beamformer[spk_idx][b_idx,...],[1,0]) * self.scale)
                        t_utterance_e_clean = t_utterance_e_clean * self.MaxINT16
                        t_utterance_e_clean = t_utterance_e_clean.astype(np.int16)
                        sf.write('{}_{}.wav'.format(os.path.join(saveDir,wav_name[b_idx]),spk_idx) , t_utterance_e_clean.T, self.fs, 'PCM_24')
                        print('[Beamforming] Testing ... ','Save Directory : {}'.format(saveDir), 'file : {}'.format(wav_name[b_idx]) , 'process : {:.2f}'.format(idx/len(data_loader)))

            else:
                for split_idx in range(split_len):

                    observe = split_observe_dict[str(split_idx)] #[B,Ch,F,T]
                    if self.cuda_flag:
                        observe = observe.cuda(self.device)
                    B, Ch, T, F = observe.size()

                    if split_idx == 0:
                        t_e_clean_s0 = np.empty([B], dtype=np.ndarray)
                        t_e_clean_s1 = np.empty([B], dtype=np.ndarray)

                    s0 = split_clean_s0_dict[str(split_idx)][:,self.ref_ch,:,:] # [B,1,T,F]
                    s1 = split_clean_s1_dict[str(split_idx)][:,self.ref_ch,:,:]
                    clean = torch.stack((s0,s1), dim=1) #[B,Ch,T,F]
                    s_clean = torch.unsqueeze(clean, dim=2) #[B,Ch,1,T,F]

                    """
                        Apply Source Separation
                    """
                    e_clean_MISO1_temp = self.MISO1_Inference(observe, ref_ch=self.ref_ch)
                    if self.cuda_flag:
                        observe = observe.detach().cpu()
                        for spk_idx in range(self.num_spks):
                            e_clean_MISO1_temp[spk_idx] = e_clean_MISO1_temp[spk_idx].detach().cpu()

                    """
                        Source Alignment between Clean reference signal and MISO1 signal
                        calculate magnitude distance between ref mic(ch0) and target signal(reference mic : ch0)
                    """
                    e_clean_MISO1 = [[] for _ in range(self.num_spks)]
                    for spk_idx in range(self.num_spks):
                        e_clean_MISO1[spk_idx] = torch.zeros_like(e_clean_MISO1_temp[spk_idx])
                        if spk_idx == 0:
                            e_clean_MISO1_ref_ch = e_clean_MISO1_temp[spk_idx][:,self.ref_ch,:,:] #[B,T,F]
                        else:
                            e_clean_MISO1_ref_ch = torch.stack((e_clean_MISO1_ref_ch, e_clean_MISO1_temp[spk_idx][:,self.ref_ch,:,:]),dim=1)
                    # MISO1_stft_ref_ch : [B,Spks,T,F]

                    s_e_clean_MISO1_ref_ch = torch.unsqueeze(e_clean_MISO1_ref_ch, dim=1) #[B,1,Spks,T,F]
                    magnitude_e_clean_MISO1_ref_ch = torch.abs(torch.sqrt(s_e_clean_MISO1_ref_ch.real**2 + s_e_clean_MISO1_ref_ch.imag**2))
                    magnitude_distance = torch.sum(torch.abs(magnitude_e_clean_MISO1_ref_ch - abs(s_clean)),[3,4])
                    perms = clean.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                    index_ = torch.unsqueeze(perms, dim=2)
                    perms_one_hot = clean.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                    batchwise_distance = torch.einsum('bij,pij->bp', [magnitude_distance, perms_one_hot])
                    min_distance_idx = torch.argmin(batchwise_distance, dim=1)

                    for b_idx in range(B):
                        align_index = torch.argmax(perms_one_hot[min_distance_idx[b_idx]], dim=1)
                        for spk_idx in range(self.num_spks):
                            target_index = align_index[spk_idx]
                            e_clean_MISO1[spk_idx][b_idx,...] = e_clean_MISO1_temp[target_index][b_idx,...]

                        """
                            Apply MVDR Beamforming
                        """

                        e_clean_Beamformer = [[] for _ in range(self.num_spks)]
                        observe = torch.permute(observe, [0,3,1,2]).numpy()
                        for spk_idx in range(self.num_spks):
                            source = torch.permute(e_clean_MISO1[spk_idx], [0,3,1,2]).numpy()
                            e_clean_Beamformer[spk_idx] = torch.squeeze(self.Apply_Beamforming(source,observe),dim=1)

                        for b_idx in range(B):
                            t_split_e_clean_s0 = self.ISTFT(torch.permute(e_clean_Beamformer[0][b_idx,...],[1,0]) * self.scale)
                            t_split_e_clean_s0 = t_split_e_clean_s0 * self.MaxInt16
                            t_split_e_clean_s0 = t_split_e_clean_s0.astype(np.int16)
                            assert t_split_e_clean_s0.shape[0] == self.chunk_size, ('estimate source 0 wav length does not match chunk size, please check ISTFT function ')

                            t_split_e_clean_s1 = self.ISTFT(torch.permute(e_clean_Beamformer[1][b_idx,...],[1,0]) * self.scale)
                            t_split_e_clean_s1 = t_split_e_clean_s1 * self.MaxInt16
                            t_split_e_clean_s1 = t_split_e_clean_s1.astype(np.int16)
                            assert t_split_e_clean_s1.shape[0] == self.chunk_size, ('estimate source 1 wav length does not match chunk size, please check ISTFT function ')

                            if split_idx == split_len -1:
                                t_split_e_clean_s0 = t_split_e_clean_s0[:len(t_split_e_clean_s0)-gap[b_idx]]
                                t_split_e_clean_s1 = t_split_e_clean_s1[:len(t_split_e_clean_s1)-gap[b_idx]]

                            if split_idx == 0:
                                t_e_clean_s0[b_idx] = t_split_e_clean_s0
                                t_e_clean_s1[b_idx]=  t_split_e_clean_s1
                            else:
                                t_e_clean_s0[b_idx] = np.append(t_e_clean_s0[b_idx] ,t_split_e_clean_s0)
                                t_e_clean_s1[b_idx] = np.append(t_e_clean_s1[b_idx], t_split_e_clean_s1)

                for b_idx in range(B):
                    sf.write('{}_0.wav'.format(os.path.join(saveDir,wav_name[b_idx])), t_e_clean_s0[b_idx].T, self.fs, 'PCM_24')
                    sf.write('{}_1.wav'.format(os.path.join(saveDir,wav_name[b_idx])), t_e_clean_s1[b_idx].T, self.fs, 'PCM_24')
                print('[Beamforming] Testing ... ','Save Directory : {}'.format(saveDir), 'process : {:.2f}'.format(idx/len(data_loader)))

    def ISTFT(self,FT_sig):

        '''
        input : [F,T]
        output : [C, T]
        '''
        # if FT_sig.shape[1] != self.config['ISTFT']['length']+1:
            # FT_sig = np.transpose(FT_sig,(0,1)) # [C,T,F] -> [C,F,T]

        _, t_sig = scipy.signal.istft(FT_sig,fs=self.fs, window=self.window, nperseg=self.nperseg, noverlap=self.noverlap) #[C,F,T] -> [C,T]

        return t_sig

    def STFT(self,time_sig):
        '''
        input : [T,Nch]
        output : [Nch,F,T]
        '''
        if time_sig.shape[1] > time_sig.shape[0]:
            time_sig = time_sig.T
            warnings.warn(f'the STFT input dimension warning, input size :[Nch, T] ---> [T,Nch]')
        num_ch = time_sig.shape[1]
        for num_ch in range(num_ch):
            # scipy.signal.stft : output : [F range, T range, FxT components]
            _,_,stft_ch = scipy.signal.stft(time_sig[:,num_ch],fs=self.fs,window=self.window,nperseg=self.nperseg,noverlap=self.noverlap)
            # output : [FxT]
            stft_ch = np.expand_dims(stft_ch,axis=0)
            if num_ch == 0:
                stft_chcat = stft_ch
            else:
                stft_chcat = np.append(stft_chcat,stft_ch,axis=0)
        if isinstance(stft_chcat, np.ndarray): stft_chcat = torch.from_numpy(stft_chcat)

        return stft_chcat

    def MISO1_Inference(self,mix_stft,ref_ch=0):
        """
        Input:
            mix_stft : observe STFT, size - [B, Mic, T, F]
        Output:
            MISO1_stft : list of separated source, size - [B, reference Mic, T, F]

            1. circular shift the microphone array at run time for the prediction of each microphone signal
               If the microphones are arranged uniformly on a circle, Select the reference microphone by circular shifting the microphone. e.g reference mic q -> [Yq, Yq+1, ..., Yp, Y1, ..., Yq-1]
            2. Using Permutation Invariance Alignmnet method to match between clean target signal and estimated signal
        """
        B, M, T, F = mix_stft.size()

        MISO1_stft = [torch.empty(B,M,T,F, dtype=torch.complex64) for _ in range(self.num_spks)]

        Mic_array = [x for x in range(M)]
        Mic_array = np.roll(Mic_array, -ref_ch)  # [ref_ch, ref_ch+1, ..., 0, 1, ..., ref_ch-1]
        # print('Mic_array : ', Mic_array)

        with torch.no_grad():
            mix_stft_refCh = torch.roll(mix_stft,-ref_ch, dims=1)
            MISO1_refCh = self.model(mix_stft_refCh)

        for spk_idx in range(self.num_spks):
            MISO1_stft[spk_idx][:,ref_ch,...] = MISO1_refCh[:,spk_idx,...]

        # MISO1_spk1[:,ref_ch,...] = MISO1_refCh[:,0,...]
        # MISO1_spk2[:,ref_ch,...] = MISO1_refCh[:,1,...]

        s_MISO1_refCh = torch.unsqueeze(MISO1_refCh, dim=2)
        s_Magnitude_refCh = torch.abs(torch.sqrt(s_MISO1_refCh.real**2 + s_MISO1_refCh.imag**2)) # [B,Spks,1,T,F]

        with torch.no_grad():
            for shiftIdx in Mic_array[1:]:
                # print('shift Micnumber', shiftIdx)

                mix_stft_shift = torch.roll(mix_stft,-shiftIdx, dims=1)
                MISO1_chShift = self.model(mix_stft_shift)

                s_MISO1_chShift = torch.unsqueeze(MISO1_chShift, dim=1) #[B,1,Spks,T,F]
                s_magnitude_chShift = torch.sum(torch.abs(s_Magnitude_refCh - abs(s_MISO1_chShift)),[3,4]) #[B,Spks,Spks,T,F]
                perms = MISO1_chShift.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                index_ = torch.unsqueeze(perms, dim=2)
                perms_one_hot = MISO1_chShift.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                batchwise_distance = torch.einsum('bij,pij->bp', [s_magnitude_chShift, perms_one_hot])
                min_distance_idx = torch.argmin(batchwise_distance,dim=1)

                for batch_idx in range(B):
                    align_index = torch.argmax(perms_one_hot[min_distance_idx[batch_idx]],dim=1)
                    for spk_idx in range(self.num_spks):
                        target_index = align_index[spk_idx]
                        MISO1_stft[spk_idx][:,shiftIdx,...] = MISO1_chShift[batch_idx,target_index,...]

        return MISO1_stft

    def Apply_Beamforming(self, source_stft, mix_stft, epsi=1e-6):
        """
        Input :
            mix_stft : observe STFT, size - [B, F, Ch, T], np.ndarray
            source_stft : estimated source STFT, size - [B, F, Ch, T], np.ndarray
        Output :
            Beamform_stft : MVDR Beamforming output, size - [B, 1, T, F], torch.Tensor

            1. estimate target steering using EigenValue decomposition
            2. get source, noise Spatial Covariance Matrix,  S = 1/T * xx_h
            3. MVDR Beamformer
        """
        B, F, M, T = source_stft.shape

        # Apply small Diagonal matrix to prevent matrix inversion error
        eye = np.eye(M)
        eye = eye.reshape(1,1,M,M)
        delta = epsi * np.tile(eye,[B,F,1,1])

        ''' Source '''
        source_SCM = self.get_spatial_covariance_matrix(source_stft,normalize=True) # target covariance matrix, size : [B,F,C,C]
        source_SCM = 0.5 * (source_SCM + np.conj(source_SCM.swapaxes(-1,-2))) # verify hermitian symmetric

        ''' Noise Spatial Covariance '''
        noise_signal = mix_stft - source_stft
        # s1_noise_signal = mix_stft  #MPDR
        noise_SCM = self.get_spatial_covariance_matrix(noise_signal,normalize = True) # noise covariance matrix, size : [B,F,C,C]
        # s1_SCMn = self.condition_covariance(s1_SCMn, 1e-6)
        # s1_SCMn /= np.trace(s1_SCMn, axis1=-2, axis2= -1)[...,None, None]
        noise_SCM = 0.5 * (noise_SCM + np.conj(noise_SCM.swapaxes(-1,-2))) # verify hermitian symmetric

        ''' Get Steering vector : Eigen-decomposition '''
        shape = source_SCM.shape
        source_steering = np.empty(shape[:-1], dtype=np.complex)

        # s1_SCMs += delta
        source_SCM = np.reshape(source_SCM, (-1,) + shape[-2:])
        eigenvals, eigenvecs = np.linalg.eigh(source_SCM)
        # Find max eigenvals
        vals = np.argmax(eigenvals, axis=-1)
        # Select eigenvec for max eigenval
        source_steering = np.array([eigenvecs[i,:,vals[i]] for i in range(eigenvals.shape[0])])
        # s1_steering = np.array([eigenvecs[i,:,vals[i]] * np.sqrt(Mic/np.linalg.norm(eigenvecs[i,:,vals[i]])) for i in range(eigenvals.shape[0])]) # [B*F,Ch,Ch]
        source_steering = np.reshape(source_steering, shape[:-1]) # [B,F,Ch]
        source_SCM = np.reshape(source_SCM, shape)

        ''' steering normalize with respect to the reference microphone '''
        # ver 1
        source_steering = source_steering / np.expand_dims(source_steering[:,:,0], axis=2)
        for b_idx in range(0,B):
            for f_idx in range(0,F):
                # s1_steering[b_idx,f_idx,:] = s1_steering[b_idx,f_idx,:] / s1_steering[b_idx,f_idx,0]
                source_steering[b_idx,f_idx,:] = source_steering[b_idx,f_idx,:] * np.sqrt(M/(np.linalg.norm(source_steering[b_idx,f_idx,:])))

        # ver 2
        # s1_steering = self.normalize(s1_steering)

        source_steering = self.PhaseCorrection(source_steering)
        beamformer = self.get_mvdr_beamformer(source_steering, noise_SCM, delta)
        # s1_beamformer = self.blind_analytic_normalization(s1_beamformer,s1_SCMn)
        source_bf = self.apply_beamformer(beamformer,mix_stft)

        if isinstance(source_bf, np.ndarray): source_bf = torch.from_numpy(source_bf)
        source_bf = torch.permute(source_bf, [0,2,1])

        return source_bf

    def get_spatial_covariance_matrix(self,observation,normalize):
        '''
        Input :
            observation : complex
                            size : [B,F,C,T]
        Return :
                R       : double
                            size : [B,F,C,C]
        '''
        B,F,C,T = observation.shape
        R = np.einsum('...dt,...et-> ...de', observation, observation.conj())
        if normalize:
            normalization = np.sum(np.ones((B,F,1,T)),axis=-1, keepdims=True)
            R /= normalization
        return R

    def PhaseCorrection(self,W): #Matlab과 동일
        """
        Phase correction to reduce distortions due to phase inconsistencies.
        Input:
                W : steering vector
                    size : [B,F,Ch]
        """
        w = W.copy()
        B, F, Ch = w.shape
        for b_idx in range(0,B):
            for f in range(1, F):
                w[b_idx,f, :] *= np.exp(-1j*np.angle(
                    np.sum(w[b_idx,f, :] * w[b_idx,f-1, :].conj(), axis=-1, keepdims=True)))
        return w

    def condition_covariance(self,x,gamma):
        """see https://stt.msu.edu/users/mauryaas/Ashwini_JPEN.pdf (2.3)"""
        B,F,_,_ = x.shape
        for b_idx in range(0,B):
            scale = gamma * np.trace(x[b_idx,...]) / x[b_idx,...].shape[-1]
            scaled_eye = np.eye(x.shape[-1]) * scale
            x[b_idx,...] = (x[b_idx,...]+scaled_eye) / (1+gamma)
        return x

    def normalize(self,vector):
        B,F,Ch = vector.shape
        for b_idx in range(0,B):
            for ii in range(0,F):
                weight = np.matmul(np.conjugate(vector[b_idx,ii,:]).reshape(1,-1), vector[b_idx,ii,:])
                vector[b_idx,ii,:] = (vector[b_idx,ii,:] / weight)
        return vector

    def blind_analytic_normalization(self,vector, noise_psd_matrix, eps=0):
        """Reduces distortions in beamformed ouptput.

        :param vector: Beamforming vector
            with shape (..., sensors)
        :param noise_psd_matrix:
            with shape (..., sensors, sensors)
        :return: Scaled Deamforming vector
            with shape (..., sensors)
        """
        nominator = np.einsum(
            '...a,...ab,...bc,...c->...',
            vector.conj(), noise_psd_matrix, noise_psd_matrix, vector
        )
        nominator = np.abs(np.sqrt(nominator))

        denominator = np.einsum(
            '...a,...ab,...b->...', vector.conj(), noise_psd_matrix, vector
        )
        denominator = np.abs(denominator)

        normalization = nominator / (denominator + eps)
        return vector * normalization[..., np.newaxis]

    def get_mvdr_beamformer(self, steering_vector, R_noise, delta):
        """
        Returns the MVDR beamformers vector

        Input :
            steering_vector : Acoustic transfer function vector
                                shape : [B, F, Ch]
                R_noise     : Noise spatial covariance matrix
                                shape : [B, F, Ch, Ch]
        """
        R_noise += delta
        numer = solve(R_noise, steering_vector)
        denom = np.einsum('...d,...d->...', steering_vector.conj(), numer)
        beamformer = numer / np.expand_dims(denom, axis=-1)
        return beamformer

    def apply_beamformer(self, beamformer, mixture):
        return np.einsum('...a,...at->...t',beamformer.conj(), mixture)

class Tester_Enhance(object):
    def __init__(self, dataset, enhance_mode, dt_loader, test_loader, model_sep, model, num_ch_utilize,
                device, num_spks, chunk_time, save_rootDir, ref_ch, cuda_flag, **ISTFT_args):

        self.dataset = dataset
        self.dt_loader = dt_loader
        self.test_loader = test_loader
        self.model_sep = model_sep
        self.model = model
        self.num_ch_utilize = num_ch_utilize
        self.device = device
        self.num_spks = num_spks
        self.enhance_mode= enhance_mode

        #ISTFT params
        self.fs = ISTFT_args['fs']
        self.window = ISTFT_args['window']
        self.nperseg = ISTFT_args['length']
        self.noverlap = ISTFT_args['overlap']

        hann_win = scipy.signal.get_window(self.window, self.nperseg)
        self.scale = np.sqrt(1.0/hann_win.sum()**2)
        self.MaxInt16 = np.iinfo(np.int16).max

        self.chunk_size = int(chunk_time * self.fs)
        self.save_rootDir = save_rootDir
        self.ref_ch = ref_ch
        self.cuda_flag = cuda_flag

    def test(self):
        print('testing...')

        print('-'*85)
        print('development dataset inference ...')
        print('-'*85)
        saveDir = os.path.join(self.save_rootDir, 'cv_dev93')
        Path(saveDir).mkdir(exist_ok=True, parents=True)
        self.inference(self.dt_loader,saveDir)
        print('development dataset inference finish !')

        print('-'*85)
        print('test dataset inference...')
        print('-'*85)
        saveDir = os.path.join(self.save_rootDir, 'test_eval92')
        Path(saveDir).mkdir(exist_ok=True, parents=True)
        self.inference(self.test_loader, saveDir)
        print('test dataset inference finish !')

    def inference(self, data_loader, saveDir):
        """
            split_observe_dict : dictionary, keys : split index
            split_observe_dict[split_index] : [B,Ch,F,T]
            split_clean_s0_dict :       "
            split_clean_s1_dict :       "

            gap : zero pad size of last split wav, [B]
            wav_name : file name to test
        """

        for idx, (data) in enumerate(data_loader):
            split_observe_dict, split_clean_s0_dict, split_clean_s1_dict, gap, wav_name = data
            split_len = len(split_observe_dict)

            ### self.utterance_flag 추가하기

            for split_idx in range(split_len):

                observe = split_observe_dict[str(split_idx)]
                if self.cuda_flag:
                    observe = observe.cuda(self.device)

                B, Ch, T, F = observe.size()

                """
                    Apply Source Separation
                """
                e_clean_MISO1_temp = self.MISO1_Inference(observe, ref_ch = self.ref_ch)
                if self.cuda_flag:
                    observe = observe.detach().cpu()
                    for spk_idx in range(self.num_spks):
                        e_clean_MISO1_temp[spk_idx] = e_clean_MISO1_temp[spk_idx].detach().cpu()

                if split_idx == 0:
                    t_e_clean_s0 = np.empty([B], dtype=np.ndarray)
                    t_e_clean_s1 = np.empty([B], dtype=np.ndarray)

                """
                    Source Alignment between Clean reference signal and MISO1 signal
                    calculate magnitude distance between ref mic(ch0) and target signal(reference mic : ch0)
                """

                s0 = split_clean_s0_dict[str(split_idx)][:, self.ref_ch, :,:] #[B,1,T,F]
                s1 = split_clean_s1_dict[str(split_idx)][:, self.ref_ch, :,:]
                clean = torch.stack((s0,s1), dim=1)
                s_clean = torch.unsqueeze(clean, dim=2)

                e_clean_MISO1 = [[] for _ in range(self.num_spks)]
                for spk_idx in range(self.num_spks):
                    e_clean_MISO1[spk_idx] = torch.zeros_like(e_clean_MISO1_temp[spk_idx])
                    if spk_idx == 0:
                        e_clean_MISO1_ref_ch = e_clean_MISO1_temp[spk_idx][:, self.ref_ch, :,:] # [B,T,F]
                    else:
                        e_clean_MISO1_ref_ch = torch.stack((e_clean_MISO1_ref_ch, e_clean_MISO1_temp[spk_idx][:,self.ref_ch, :,:]), dim=1)

                s_e_clean_MISO1_ref_ch = torch.unsqueeze(e_clean_MISO1_ref_ch, dim=1)
                magnitude_e_clean_MISO1_ref_ch = torch.abs(torch.sqrt(s_e_clean_MISO1_ref_ch.real**2+s_e_clean_MISO1_ref_ch.imag**2))
                magnitude_distance = torch.sum(torch.abs(magnitude_e_clean_MISO1_ref_ch - abs(s_clean)),[3,4])
                perms = clean.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                index_ = torch.unsqueeze(perms, dim=2)
                perms_one_hot = clean.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                batchwise_distance = torch.einsum('bij,pij->bp', [magnitude_distance, perms_one_hot])
                min_distance_idx = torch.argmin(batchwise_distance, dim=1)

                for b_idx in range(B):
                    align_index = torch.argmax(perms_one_hot[min_distance_idx[b_idx]], dim=1)
                    for spk_idx in range(self.num_spks):
                        target_index = align_index[spk_idx]
                        e_clean_MISO1[spk_idx][b_idx,...] = e_clean_MISO1_temp[target_index][b_idx,...]

                    """
                        Apply MVDR Beamforming
                    """
                    e_clean_Beamformer = [[] for _ in range(self.num_spks)]
                    observe_bf = torch.permute(observe, [0,3,1,2]).numpy()
                    for spk_idx in range(self.num_spks):
                        source = torch.permute(e_clean_MISO1[spk_idx], [0,3,1,2]).numpy()
                        e_clean_Beamformer[spk_idx] = torch.unsqueeze(self.Apply_Beamforming(source, observe_bf),dim=1)
                    """
                        Apply Enhancement
                    """
                    if self.cuda_flag:
                        observe = observe.cuda(self.device)
                        for spk_idx in range(self.num_spks):
                            e_clean_Beamformer[spk_idx] = e_clean_Beamformer[spk_idx].cuda(self.device) # [B,1,T,F]
                            e_clean_MISO1[spk_idx] = e_clean_MISO1[spk_idx].cuda(self.device) # [B, Mic, T, F]

                    e_clean_Enhance = [[] for _ in range(self.num_spks)]
                    if self.enhance_mode == 'MISO3':
                        for spk_idx in range(self.num_spks):
                            MISO1_stft = torch.unsqueeze(e_clean_MISO1[spk_idx][:,self.ref_ch,...], dim=1)
                            bf_stft=  e_clean_Beamformer[spk_idx]
                            e_clean_Enhance[spk_idx] = torch.squeeze(self.MISO3_inference(observe,bf_stft,MISO1_stft), dim=1).detach().cpu()
                    else:
                        MISO1_stft_spk1 = torch.unsqueeze(e_clean_MISO1[0][:,self.ref_ch,...], dim=1)
                        MISO1_stft_spk2 = torch.unsqueeze(e_clean_MISO1[1][:,self.ref_ch,...], dim=1)
                        bf_stft_spk1 = e_clean_Beamformer[0]
                        bf_stft_spk2 = e_clean_Beamformer[1]
                        e_clean_Enhance_temp = self.MISO2_inference(observe, bf_stft_spk1, bf_stft_spk2, MISO1_stft_spk1, MISO1_stft_spk2)
                        for spk_idx in range(self.num_spks):
                            e_clean_Enhance[spk_idx] = torch.squeeze(e_clean_Enhance_temp[:,spk_idx,...]).detach().cpu()

                    for b_idx in range(B):
                        t_split_e_clean_s0 = self.ISTFT(torch.permute(e_clean_Enhance[0][b_idx,...],[1,0]) * self.scale)
                        t_split_e_clean_s0 = t_split_e_clean_s0 * self.MaxInt16
                        t_split_e_clean_s0 = t_split_e_clean_s0.astype(np.int16)
                        assert t_split_e_clean_s0.shape[0] == self.chunk_size, ('estimate source 0 wav length does not match chunk size, please check ISTFT function ')

                        t_split_e_clean_s1 = self.ISTFT(torch.permute(e_clean_Enhance[1][b_idx,...],[1,0]) * self.scale)
                        t_split_e_clean_s1 = t_split_e_clean_s1 * self.MaxInt16
                        t_split_e_clean_s1 = t_split_e_clean_s1.astype(np.int16)
                        assert t_split_e_clean_s1.shape[0] == self.chunk_size, ('estimate source 1 wav length does not match chunk size, please check ISTFT function ')

                        if split_idx == split_len -1:
                            t_split_e_clean_s0 = t_split_e_clean_s0[:len(t_split_e_clean_s0)-gap[b_idx]]
                            t_split_e_clean_s1 = t_split_e_clean_s1[:len(t_split_e_clean_s1)-gap[b_idx]]

                        if split_idx == 0:
                            t_e_clean_s0[b_idx] = t_split_e_clean_s0
                            t_e_clean_s1[b_idx]=  t_split_e_clean_s1
                        else:
                            t_e_clean_s0[b_idx] = np.append(t_e_clean_s0[b_idx] ,t_split_e_clean_s0)
                            t_e_clean_s1[b_idx] = np.append(t_e_clean_s1[b_idx], t_split_e_clean_s1)

            for b_idx in range(B):
                sf.write('{}_0.wav'.format(os.path.join(saveDir,wav_name[b_idx])), t_e_clean_s0[b_idx].T, self.fs, 'PCM_24')
                sf.write('{}_1.wav'.format(os.path.join(saveDir,wav_name[b_idx])), t_e_clean_s1[b_idx].T, self.fs, 'PCM_24')
            print('[Enhance {}] Testing ... '.format(self.enhance_mode),'Save Directory : {}'.format(saveDir), 'process : {:.2f}'.format(idx/len(data_loader)))

    def ISTFT(self,FT_sig):

        '''
        input : [F,T]
        output : [C, T]
        '''
        # if FT_sig.shape[1] != self.config['ISTFT']['length']+1:
            # FT_sig = np.transpose(FT_sig,(0,1)) # [C,T,F] -> [C,F,T]

        _, t_sig = scipy.signal.istft(FT_sig,fs=self.fs, window=self.window, nperseg=self.nperseg, noverlap=self.noverlap) #[C,F,T] -> [C,T]

        return t_sig

    def STFT(self,time_sig):
        '''
        input : [T,Nch]
        output : [Nch,F,T]
        '''
        if time_sig.shape[1] > time_sig.shape[0]:
            time_sig = time_sig.T
            warnings.warn(f'the STFT input dimension warning, input size :[Nch, T] ---> [T,Nch]')
        num_ch = time_sig.shape[1]
        for num_ch in range(num_ch):
            # scipy.signal.stft : output : [F range, T range, FxT components]
            _,_,stft_ch = scipy.signal.stft(time_sig[:,num_ch],fs=self.fs,window=self.window,nperseg=self.nperseg,noverlap=self.noverlap)
            # output : [FxT]
            stft_ch = np.expand_dims(stft_ch,axis=0)
            if num_ch == 0:
                stft_chcat = stft_ch
            else:
                stft_chcat = np.append(stft_chcat,stft_ch,axis=0)
        if isinstance(stft_chcat, np.ndarray): stft_chcat = torch.from_numpy(stft_chcat)

        return stft_chcat

    def MISO1_Inference(self,mix_stft,ref_ch=0):
        """
        Input:
            mix_stft : observe STFT, size - [B, Mic, T, F]
        Output:
            MISO1_stft : list of separated source, - [B, reference Mic, T, F]

            1. circular shift the microphone array at run time for the prediction of each microphone signal
               If the microphones are arranged uniformly on a circle, Select the reference microphone by circular shifting the microphone. e.g reference mic q -> [Yq, Yq+1, ..., Yp, Y1, ..., Yq-1]
            2. Using Permutation Invariance Alignmnet method to match between clean target signal and estimated signal
        """
        B, M, T, F = mix_stft.size()

        MISO1_stft = [torch.empty(B,M,T,F, dtype=torch.complex64) for _ in range(self.num_spks)]

        Mic_array = [x for x in range(M)]
        Mic_array = np.roll(Mic_array, -ref_ch)  # [ref_ch, ref_ch+1, ..., 0, 1, ..., ref_ch-1]
        # print('Mic_array : ', Mic_array)

        with torch.no_grad():
            mix_stft_refCh = torch.roll(mix_stft,-ref_ch, dims=1)
            MISO1_refCh = self.model_sep(mix_stft_refCh)

        for spk_idx in range(self.num_spks):
            MISO1_stft[spk_idx][:,ref_ch,...] = MISO1_refCh[:,spk_idx,...]

        # MISO1_spk1[:,ref_ch,...] = MISO1_refCh[:,0,...]
        # MISO1_spk2[:,ref_ch,...] = MISO1_refCh[:,1,...]

        s_MISO1_refCh = torch.unsqueeze(MISO1_refCh, dim=2)
        s_Magnitude_refCh = torch.abs(torch.sqrt(s_MISO1_refCh.real**2 + s_MISO1_refCh.imag**2)) # [B,Spks,1,T,F]

        with torch.no_grad():
            for shiftIdx in Mic_array[1:]:
                # print('shift Micnumber', shiftIdx)

                mix_stft_shift = torch.roll(mix_stft,-shiftIdx, dims=1)
                MISO1_chShift = self.model_sep(mix_stft_shift)

                s_MISO1_chShift = torch.unsqueeze(MISO1_chShift, dim=1) #[B,1,Spks,T,F]
                s_magnitude_chShift = torch.sum(torch.abs(s_Magnitude_refCh - abs(s_MISO1_chShift)),[3,4]) #[B,Spks,Spks,T,F]
                perms = MISO1_chShift.new_tensor(list(permutations(range(self.num_spks))), dtype=torch.long)
                index_ = torch.unsqueeze(perms, dim=2)
                perms_one_hot = MISO1_chShift.new_zeros((*perms.size(), self.num_spks), dtype=torch.float).scatter_(2,index_,1)
                batchwise_distance = torch.einsum('bij,pij->bp', [s_magnitude_chShift, perms_one_hot])
                min_distance_idx = torch.argmin(batchwise_distance,dim=1)

                for batch_idx in range(B):
                    align_index = torch.argmax(perms_one_hot[min_distance_idx[batch_idx]],dim=1)
                    for spk_idx in range(self.num_spks):
                        target_index = align_index[spk_idx]
                        MISO1_stft[spk_idx][:,shiftIdx,...] = MISO1_chShift[batch_idx,target_index,...]

        return MISO1_stft

    def Apply_Beamforming(self, source_stft, mix_stft, epsi=1e-6):
        """
        Input :
            mix_stft : observe STFT, size - [B, F, Ch, T], np.ndarray
            source_stft : estimated source STFT, size - [B, F, Ch, T], np.ndarray
        Output :
            Beamform_stft : MVDR Beamforming output, size - [B, 1, T, F], torch.Tensor

            1. estimate target steering using EigenValue decomposition
            2. get source, noise Spatial Covariance Matrix,  S = 1/T * xx_h
            3. MVDR Beamformer
        """
        B, F, M, T = source_stft.shape

        # Apply small Diagonal matrix to prevent matrix inversion error
        eye = np.eye(M)
        eye = eye.reshape(1,1,M,M)
        delta = epsi * np.tile(eye,[B,F,1,1])

        ''' Source '''
        source_SCM = self.get_spatial_covariance_matrix(source_stft,normalize=True) # target covariance matrix, size : [B,F,C,C]
        source_SCM = 0.5 * (source_SCM + np.conj(source_SCM.swapaxes(-1,-2))) # verify hermitian symmetric

        ''' Noise Spatial Covariance '''
        noise_signal = mix_stft - source_stft
        # s1_noise_signal = mix_stft  #MPDR
        noise_SCM = self.get_spatial_covariance_matrix(noise_signal,normalize = True) # noise covariance matrix, size : [B,F,C,C]
        # s1_SCMn = self.condition_covariance(s1_SCMn, 1e-6)
        # s1_SCMn /= np.trace(s1_SCMn, axis1=-2, axis2= -1)[...,None, None]
        noise_SCM = 0.5 * (noise_SCM + np.conj(noise_SCM.swapaxes(-1,-2))) # verify hermitian symmetric

        ''' Get Steering vector : Eigen-decomposition '''
        shape = source_SCM.shape
        source_steering = np.empty(shape[:-1], dtype=np.complex)

        # s1_SCMs += delta
        source_SCM = np.reshape(source_SCM, (-1,) + shape[-2:])
        eigenvals, eigenvecs = np.linalg.eigh(source_SCM)
        # Find max eigenvals
        vals = np.argmax(eigenvals, axis=-1)
        # Select eigenvec for max eigenval
        source_steering = np.array([eigenvecs[i,:,vals[i]] for i in range(eigenvals.shape[0])])
        # s1_steering = np.array([eigenvecs[i,:,vals[i]] * np.sqrt(Mic/np.linalg.norm(eigenvecs[i,:,vals[i]])) for i in range(eigenvals.shape[0])]) # [B*F,Ch,Ch]
        source_steering = np.reshape(source_steering, shape[:-1]) # [B,F,Ch]
        source_SCM = np.reshape(source_SCM, shape)

        ''' steering normalize with respect to the reference microphone '''
        # ver 1
        source_steering = source_steering / np.expand_dims(source_steering[:,:,0], axis=2)
        for b_idx in range(0,B):
            for f_idx in range(0,F):
                # s1_steering[b_idx,f_idx,:] = s1_steering[b_idx,f_idx,:] / s1_steering[b_idx,f_idx,0]
                source_steering[b_idx,f_idx,:] = source_steering[b_idx,f_idx,:] * np.sqrt(M/(np.linalg.norm(source_steering[b_idx,f_idx,:])))

        # ver 2
        # s1_steering = self.normalize(s1_steering)

        source_steering = self.PhaseCorrection(source_steering)
        beamformer = self.get_mvdr_beamformer(source_steering, noise_SCM, delta)
        # s1_beamformer = self.blind_analytic_normalization(s1_beamformer,s1_SCMn)
        source_bf = self.apply_beamformer(beamformer,mix_stft)

        if isinstance(source_bf, np.ndarray): source_bf = torch.from_numpy(source_bf)
        source_bf = torch.permute(source_bf, [0,2,1])

        return source_bf

    def get_spatial_covariance_matrix(self,observation,normalize):
        '''
        Input :
            observation : complex
                            size : [B,F,C,T]
        Return :
                R       : double
                            size : [B,F,C,C]
        '''
        B,F,C,T = observation.shape
        R = np.einsum('...dt,...et-> ...de', observation, observation.conj())
        if normalize:
            normalization = np.sum(np.ones((B,F,1,T)),axis=-1, keepdims=True)
            R /= normalization
        return R

    def PhaseCorrection(self,W): #Matlab과 동일
        """
        Phase correction to reduce distortions due to phase inconsistencies.
        Input:
                W : steering vector
                    size : [B,F,Ch]
        """
        w = W.copy()
        B, F, Ch = w.shape
        for b_idx in range(0,B):
            for f in range(1, F):
                w[b_idx,f, :] *= np.exp(-1j*np.angle(
                    np.sum(w[b_idx,f, :] * w[b_idx,f-1, :].conj(), axis=-1, keepdims=True)))
        return w

    def condition_covariance(self,x,gamma):
        """see https://stt.msu.edu/users/mauryaas/Ashwini_JPEN.pdf (2.3)"""
        B,F,_,_ = x.shape
        for b_idx in range(0,B):
            scale = gamma * np.trace(x[b_idx,...]) / x[b_idx,...].shape[-1]
            scaled_eye = np.eye(x.shape[-1]) * scale
            x[b_idx,...] = (x[b_idx,...]+scaled_eye) / (1+gamma)
        return x

    def normalize(self,vector):
        B,F,Ch = vector.shape
        for b_idx in range(0,B):
            for ii in range(0,F):
                weight = np.matmul(np.conjugate(vector[b_idx,ii,:]).reshape(1,-1), vector[b_idx,ii,:])
                vector[b_idx,ii,:] = (vector[b_idx,ii,:] / weight)
        return vector

    def blind_analytic_normalization(self,vector, noise_psd_matrix, eps=0):
        """Reduces distortions in beamformed ouptput.

        :param vector: Beamforming vector
            with shape (..., sensors)
        :param noise_psd_matrix:
            with shape (..., sensors, sensors)
        :return: Scaled Deamforming vector
            with shape (..., sensors)
        """
        nominator = np.einsum(
            '...a,...ab,...bc,...c->...',
            vector.conj(), noise_psd_matrix, noise_psd_matrix, vector
        )
        nominator = np.abs(np.sqrt(nominator))

        denominator = np.einsum(
            '...a,...ab,...b->...', vector.conj(), noise_psd_matrix, vector
        )
        denominator = np.abs(denominator)

        normalization = nominator / (denominator + eps)
        return vector * normalization[..., np.newaxis]

    def get_mvdr_beamformer(self, steering_vector, R_noise, delta):
        """
        Returns the MVDR beamformers vector

        Input :
            steering_vector : Acoustic transfer function vector
                                shape : [B, F, Ch]
                R_noise     : Noise spatial covariance matrix
                                shape : [B, F, Ch, Ch]
        """
        R_noise += delta
        numer = solve(R_noise, steering_vector)
        denom = np.einsum('...d,...d->...', steering_vector.conj(), numer)
        beamformer = numer / np.expand_dims(denom, axis=-1)
        return beamformer

    def apply_beamformer(self, beamformer, mixture):
        return np.einsum('...a,...at->...t',beamformer.conj(), mixture)

    def MISO3_inference(self, mix_stft, bf_stft, MISO1_stft):
        """
        Input:
            mix_stft : observe STFT, size - [B, Mic, T, F]
            bf_stft : beamformer ouput STFT, size - [B, 1, T, F]
            MISO1_stft : reference mic MISO1 output STFT, size - [B, reference Mic, T, F]
        Output:
            MISO3_stft : enhanced source, size - [B, reference Mic, T, F]
        """

        with torch.no_grad():
            MISO3_stft = self.model(mix_stft, bf_stft, MISO1_stft)

        return MISO3_stft

    def MISO2_inference(self, mix_stft ,bf_stft_spk1, bf_stft_spk2, MISO1_stft_spk1, MISO1_stft_spk2):
        """
        Input:
            mix_stft : observe STFT, size - [B, Mic, T, F]
            bf_stft_spk1, _spk2 : beamformer ouput STFT, size - [B, 1, T, F]
            MISO1_stft_spk1, _spk2 : reference mic MISO1 output STFT, size - [B, reference Mic, T, F]
        Output:
            MISO3_stft : enhanced source, size - [B, Spks, T, F]
        """
        with torch.no_grad():
            MISO2_stft = self.model(mix_stft, bf_stft_spk1, bf_stft_spk2, MISO1_stft_spk1, MISO1_stft_spk2)
        return MISO2_stft


import os
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import random
import shutil
from random import sample
from datetime import datetime
import cv2
import matplotlib.font_manager as fm
from matplotlib.collections import QuadMesh

# Model
import tensorflow as tf
from tensorflow.keras.applications import InceptionV3
from tensorflow.keras.models import Model
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from tensorflow.keras.optimizers import Adam

#Image preprocessing
from tensorflow import keras
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.inception_v3 import preprocess_input
from sklearn.utils import class_weight
from sklearn.metrics import classification_report, confusion_matrix
_ROOT = os.getcwd()
# dictionary of paths
all_paths_dic = {'summary_file_path':os.path.join('train','train.csv'),
         'data_train_path':os.path.join('train','train'),
         'test_train_path':get_data('test'),
         'glaucoma_train_path':os.path.join(os.getcwd(),'glaucoma_data','train'),
         'glaucoma_test_path':os.path.join(os.getcwd(),'glaucoma_data','test'),
         'diabetic_ret_train_path':os.path.join(os.getcwd(),'diabetic_ret_data','train'),
         'diabetic_ret_test_path':os.path.join(os.getcwd(),'diabetic_ret_data','test')}

##################################################
#
#             UTILITY FUNCTIONS
#
##################################################

# Access Utility Function
def get_data(path):
    return os.path.join(_ROOT,'Data', path)

# Utility function to split dataset
def split_train_test(df_orig,train_split = 0.7,paths_dic=all_paths_dic):
  # Method to create train / split and copy all files accordingly to respective folder
  # Args:
    # train_split -> float [0,1] representing train split
    # paths_dic -> dictionary representing all paths
    # df_orig -> original data frame
  # Returns:
    # ---

  def copy_files(current_dir,dist_dir,file_names):
    # helper function to copy files from one directory to another
    for file in file_names:
      src = os.path.join(current_dir,file)
      dst = os.path.join(os.getcwd(),dist_dir,file)
      shutil.copy(src, dst)

  # Getting inx for each positive case
  glaucoma_idx     = (df_orig["glaucoma"]==1).values
  diabetic_ret_idx = (df_orig["diabetic retinopathy"]==1).values
  normal_idx       =  (df_orig["normal"]==1).values

  # Getting train size
  train_size_glaucoma     = int(glaucoma_idx.sum() * train_split)
  train_size_diabetic_ret = int(diabetic_ret_idx.sum() * train_split)
  train_size_normal      = int(normal_idx.sum() * train_split)

  random.seed(1)

  # glaucoma file names (train / test)
  glaucoma_train     = random.sample(df_orig.loc[glaucoma_idx,"filename"].values.tolist(),train_size_glaucoma)
  glaucoma_test      = [file for file in df_orig.loc[glaucoma_idx,"filename"].values.tolist() if file not in glaucoma_train]
  # diabetic ret file names (train / test)
  diabetic_ret_train = random.sample(df_orig.loc[diabetic_ret_idx,"filename"].values.tolist(),train_size_diabetic_ret)
  diabetic_ret_test      = [file for file in df_orig.loc[diabetic_ret_idx,"filename"].values.tolist() if file not in diabetic_ret_train]
  # neither file names (train / test)
  normal_train      = random.sample(df_orig.loc[normal_idx,"filename"].values.tolist(),train_size_normal)
  normal_test       = [file for file in df_orig.loc[normal_idx,"filename"].values.tolist() if file not in normal_train]

  # creating all directories
  # creating training directories
  try:
    os.makedirs(os.path.join(all_paths_dic['glaucoma_train_path'],'positive'))
  except:
    print("directory exists!")
  try:
    os.makedirs(os.path.join(all_paths_dic['glaucoma_train_path'],'negative'))
  except:
    print("directory exists!")
  try:
    os.makedirs(os.path.join(all_paths_dic['diabetic_ret_train_path'],'positive'))
  except:
    print("directory exists!")
  try:
    os.makedirs(os.path.join(all_paths_dic['diabetic_ret_train_path'],'negative'))
  except:
    print("directory exists!")

  #creating testing directories
  try:
    os.makedirs(os.path.join(all_paths_dic['glaucoma_test_path'],'positive'))
  except:
    print("directory exists!")
  try:
    os.makedirs(os.path.join(all_paths_dic['glaucoma_test_path'],'negative'))
  except:
    print("directory exists!")
  try:
    os.makedirs(os.path.join(all_paths_dic['diabetic_ret_test_path'],'positive'))
  except:
    print("directory exists!")
  try:
    os.makedirs(os.path.join(all_paths_dic['diabetic_ret_test_path'],'negative'))
  except:
    print("directory exists!")

  # copying files from main directory to glaucome train / test
  main_directory = get_data(all_paths_dic['data_train_path'])
  copy_files(main_directory,os.path.join(all_paths_dic['glaucoma_train_path'],'positive'),glaucoma_train)
  copy_files(main_directory,os.path.join(all_paths_dic['glaucoma_train_path'],'negative'),normal_train)

  copy_files(main_directory,os.path.join(all_paths_dic['glaucoma_test_path'],'positive'),glaucoma_test)
  copy_files(main_directory,os.path.join(all_paths_dic['glaucoma_test_path'],'negative'),normal_test)

  # copying files from main directory to diabetic_ret train / test
  copy_files(main_directory,os.path.join(all_paths_dic['diabetic_ret_train_path'],'positive'),diabetic_ret_train)
  copy_files(main_directory,os.path.join(all_paths_dic['diabetic_ret_train_path'],'negative'),normal_train)

  copy_files(main_directory,os.path.join(all_paths_dic['diabetic_ret_test_path'],'positive'),diabetic_ret_test)
  copy_files(main_directory,os.path.join(all_paths_dic['diabetic_ret_test_path'],'negative'),normal_test)

  # calculating size of each sample:
  glaucoma_train_pos = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['glaucoma_train_path'],'positive'))))
  glaucoma_train_neg = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['glaucoma_train_path'],'negative'))))
  glaucoma_test_pos = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['glaucoma_test_path'],'positive'))))
  glaucoma_test_neg = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['glaucoma_test_path'],'negative'))))

  diabetic_ret_train_pos = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['diabetic_ret_train_path'],'positive'))))
  diabetic_ret_train_neg = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['diabetic_ret_train_path'],'negative'))))
  diabetic_ret_test_pos = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['diabetic_ret_test_path'],'positive'))))
  diabetic_ret_test_neg = len(os.listdir(os.path.join(os.getcwd(),os.path.join(all_paths_dic['diabetic_ret_test_path'],'negative'))))

  # Printing results:
  print("\n\nGlaucoma Retinopathy Summary:")
  print("\nGlaucoma positive cases in train dataset: ",glaucoma_train_pos)
  print("Glaucoma positive cases in test dataset: ",glaucoma_test_pos)
  print("Glaucoma negative cases in train dataset: ",glaucoma_train_neg)
  print("Glaucoma negative cases in test dataset: ",glaucoma_test_neg)

  print("\n\nDiabetic Retinopathy Summary:")
  print("\nDiabetic retinopathy positive cases in train dataset: ",diabetic_ret_train_pos)
  print("Diabetic retinopathy positive cases in test dataset: ",diabetic_ret_test_pos)
  print("Diabetic retinopathy negative cases in train dataset: ",diabetic_ret_train_neg)
  print("Diabetic retinopathy negative cases in test dataset: ",diabetic_ret_test_neg)

def get_data_summary(path_dic):
  # Function to get data summary
  # Args:
      # path: path to train.csv summary file
  # returns:
      # summary: pd.DataFrame() describing the instances of interest
  path = path_dic['summary_file_path']
  summary_path = get_data(path)
  df = pd.read_csv(summary_path,usecols=['filename','diabetic retinopathy','glaucoma','normal'])
  # useful idx
  glaucoma_and_diabetic_retinopathy_idx = (df['diabetic retinopathy'].values==1) & (df['glaucoma'].values==1)
  normal_idx  = df['normal'].values==1

  glaucoma_num = df['glaucoma'].values.sum()
  diabetic_retinopathy_num = df['diabetic retinopathy'].values.sum()
  glaucoma_and_diabetic_retinopathy_num = glaucoma_and_diabetic_retinopathy_idx.sum()
  background = df[normal_idx].shape[0]

  summary = pd.DataFrame([["glaucoma",'diabetic retinopathy','glaucoma & diabetic retinopathy','normal'],
                          [glaucoma_num,diabetic_retinopathy_num,glaucoma_and_diabetic_retinopathy_num,background]]).T

  summary.columns = ["instance name","number of cases"]
  summary.set_index("instance name",inplace=True)

  plt.figure(figsize=(15, 10))
  # Plotting
  # Defining the values for x-axis, y-axis
  # and from which datafarme the values are to be picked
  plots = sns.barplot(x=summary.index, y="number of cases", data=summary)

  # Iterrating over the bars one-by-one
  for bar in plots.patches:
      plots.annotate(format(bar.get_height(), '.2f'),
                      (bar.get_x() + bar.get_width() / 2,
                      bar.get_height()), ha='center', va='center',
                      size=15, xytext=(0, 8),
                      textcoords='offset points')

  # Setting the label for x-axis
  plt.xlabel("Pathologies", size=14)

  # Setting the label for y-axis
  plt.ylabel("Number of Cases", size=14)

  # Setting the title for the graph
  plt.title("Data Summary")

  return [summary,plt,df]

# Method to plot sample of data from training set
def plot_sampleData(data_path):
  # Method to plot samples of retinal images
  # Args:
    # data_path -> str representing path to the directory containing the images
  # Retuns:
    # --- plots samples
  nrows = 1
  ncols = 2
  pic_index = 0
  fig = plt.gcf()
  fig.set_size_inches(ncols*6, nrows*4)
  sample_images = []
  directories = os.listdir(data_path)
  for i,dir in enumerate(directories):
    img_lst = os.listdir(os.path.join(data_path,dir))
    image_file = sample( img_lst,1)[0]
    image_path = os.path.join(data_path,dir,image_file)
    sp = plt.subplot(nrows, ncols, i + 1)
    sp.axis('Off') # Don't show axes (or gridlines)
    img = mpimg.imread(image_path)
    plt.imshow(img)
    plt.title(directories[i])
  plt.show()
  print("image size is:" + str(img.shape[0]) + "x" + str(img.shape[1]) + " pixels")

##################################################
#
#                       MODEL
#
##################################################

# Building Model
def get_model(params,print_summary=False):
  # Mathod to build model
  # Args:
    # params -> dict holding model params
    # print_summary -> bool (default is False). if true, prints model summary design
  # Returns:
    #  Model object from tensorflow.keras.models

  # get params:
  number_of_fixed_layers = params['number_of_fixed_layers']
  l2_penalty = params['l2_penalty']
  pre_trained_model   = InceptionV3(weights='imagenet',include_top=True)
  # Freezing all layers until layer conv2D_181
  for layer in pre_trained_model.layers[:number_of_fixed_layers]:
    layer.trainable=False
  last_layer  = pre_trained_model.get_layer('avg_pool')
  last_output = last_layer.output
  x = layers.Dense(1024,activation='relu')(last_output)
  x = layers.Dense(256,activation='relu')(x)
  x = layers.Dense(2,activation='softmax',kernel_regularizer=tf.keras.regularizers.l2(l2=l2_penalty))(x)
  model=Model(pre_trained_model.input,x)
  return model

def get_callbacks(checkpoint_dir,logs_dir):
  # Method to get callbacks
  # Args:
    # checkpoint_dir -> str representing checkpoint dir
    # logs_dir -> str representing log dir
  # Returns:
    # list of callbacks

  # Tensorboard callbacks
  tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logs_dir)

  # checkpoint callback
  model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath                = checkpoint_dir,
    save_freq               = 'epoch',
    save_weights_only       = True,
    monitor                 = 'val_accuracy',
    save_best_only          = True)
  return [tensorboard_callback , model_checkpoint_callback]

##################################################
#
#              PRE PROCESSING
#
##################################################
def get_generators(params,train_path,test_path,FinalTest_path):
  # method for pre-prcessing and providing true ground y_labels
  # Args:
    # params -> dict describing pre-processing parameters
  # returns:
    # generators -> dictionary of train and test generators as well as true y_label

  #get params:
  rotation_range = params['rotation_range']
  zca_whitening = params['zca_whitening']
  zca_epsilon = params['zca_epsilon']
  width_shift_range = params['width_shift_range']
  height_shift_range = params['height_shift_range']
  shear_range = params['shear_range']
  zoom_range = params['zoom_range']
  horizontal_flip = params['horizontal_flip']
  fill_mode = params['fill_mode']
  brightness_range = params['brightness_range']
  batch_size = params['batch_size']
  target_size = params['target_size']

  train_datagen=ImageDataGenerator(preprocessing_function = preprocess_input,
                                  rotation_range          = rotation_range,
                                  zca_whitening           = zca_whitening,
                                  zca_epsilon             = zca_epsilon,
                                  width_shift_range       =width_shift_range,
                                  height_shift_range      =height_shift_range,
                                  shear_range             =shear_range,
                                  zoom_range              =zoom_range,
                                  horizontal_flip         =horizontal_flip,
                                  fill_mode               =fill_mode,
                                  brightness_range        = brightness_range)

  test_datagen  = ImageDataGenerator(preprocessing_function = preprocess_input)
  TEST_DATA_GEN = ImageDataGenerator(preprocessing_function = preprocess_input)

  generator_train      = train_datagen.flow_from_directory(train_path,target_size = target_size,batch_size=batch_size)
  generator_test       = test_datagen.flow_from_directory(test_path,target_size = target_size,batch_size=batch_size)
  label_generator      = test_datagen.flow_from_directory(test_path,target_size = target_size,batch_size=batch_size,shuffle=False)
  TEST_GEN_unlabeled   = TEST_DATA_GEN.flow_from_directory(FinalTest_path
                                                           ,target_size = target_size,batch_size=batch_size,shuffle=False)

  generator = {'generator_train': generator_train,'generator_test':generator_test,'label_generator':label_generator, 'TEST_GEN_unlabeled':TEST_GEN_unlabeled}
  return generator

##################################################
#
#              Grad CAM
#
##################################################

def img_is_color(img):

    if len(img.shape) == 3:
        # Check the color channels to see if they're all the same.
        c1, c2, c3 = img[:, : , 0], img[:, :, 1], img[:, :, 2]
        if (c1 == c2).all() and (c2 == c3).all():
            return True

    return False
def show_image_list(list_images, list_titles=None, list_cmaps=None, grid=True, num_cols=2, figsize=(15, 7), title_fontsize=30):
    '''
    Shows a grid of images, where each image is a Numpy array. The images can be either
    RGB or grayscale.

    Parameters:
    ----------
    images: list
        List of the images to be displayed.
    list_titles: list or None
        Optional list of titles to be shown for each image.
    list_cmaps: list or None
        Optional list of cmap values for each image. If None, then cmap will be
        automatically inferred.
    grid: boolean
        If True, show a grid over each image
    num_cols: int
        Number of columns to show.
    figsize: tuple of width, height
        Value to be passed to pyplot.figure()
    title_fontsize: int
        Value to be passed to set_title().
    '''

    assert isinstance(list_images, list)
    assert len(list_images) > 0
    assert isinstance(list_images[0], np.ndarray)

    if list_titles is not None:
        assert isinstance(list_titles, list)
        assert len(list_images) == len(list_titles), '%d imgs != %d titles' % (len(list_images), len(list_titles))

    if list_cmaps is not None:
        assert isinstance(list_cmaps, list)
        assert len(list_images) == len(list_cmaps), '%d imgs != %d cmaps' % (len(list_images), len(list_cmaps))

    num_images  = len(list_images)
    num_cols    = min(num_images, num_cols)
    num_rows    = int(num_images / num_cols) + (1 if num_images % num_cols != 0 else 0)

    # Create a grid of subplots.
    fig, axes = plt.subplots(num_rows, num_cols, figsize=figsize)

    # Create list of axes for easy iteration.
    list_axes = list(axes.flat)

    for i in range(num_images):

        img    = list_images[i]
        title  = list_titles[i] if list_titles is not None else 'Image %d' % (i)
        cmap   = list_cmaps[i] if list_cmaps is not None else (None if img_is_color(img) else 'gray')

        list_axes[i].imshow(img, cmap=cmap)
        list_axes[i].set_title(title, fontsize=title_fontsize)
        list_axes[i].grid(grid)

    for i in range(num_images, len(list_axes)):
        list_axes[i].set_visible(False)

    fig.tight_layout()
    _ = plt.show()

###############################################################################
#    gradCam
###############################################################################

def gradCAM(model=None,test_path=None,image_num = 0, intensity=0.5, res=250):

    datagen = ImageDataGenerator(preprocessing_function=preprocess_input)
    generator      = datagen.flow_from_directory(test_path,target_size = (299,299),batch_size=250,shuffle=False)
    images,_ = next(generator)
    images,_ = next(generator)

    # Path to original image
    path_to_image = test_path + "/"+generator.filenames[image_num]

    x = images[image_num]
    x=np.expand_dims(x, axis=0)
    preds = model.predict(x)

    with tf.GradientTape() as tape:
        last_conv_layer = model.get_layer('mixed10')
        iterate = tf.keras.models.Model([model.inputs], [model.output, last_conv_layer.output])
        model_out, last_conv_layer = iterate(x)
        class_out = model_out[:, np.argmax(model_out[0])]
        grads = tape.gradient(class_out, last_conv_layer)
        pooled_grads = K.mean(grads, axis=(0, 1, 2))

    heatmap = tf.reduce_mean(tf.multiply(pooled_grads, last_conv_layer), axis=-1)
    heatmap = np.maximum(heatmap, 0)
    heatmap /= np.max(heatmap)
    heatmap = heatmap.reshape((8,8))

    orig = cv2.imread(path_to_image)
    orig = cv2.cvtColor(orig, cv2.COLOR_BGR2RGB)

    heatmap = cv2.resize(heatmap, (orig.shape[1], orig.shape[0]))
    heatmap = cv2.applyColorMap(np.uint8(255*heatmap), cv2.COLORMAP_JET)

    img = heatmap * intensity + orig
    show_image_list([orig.astype('uint8'),img.astype('uint8')],list_titles=["Original","gradCam Applied"])

##################################################
#
#              Semi-Supervised Split
#
##################################################

def semi_supervised_split(clf,src_generator,model='glaucoma'):
  # Method to copy files based on predicted earlier models
  # Args:
  #   clf: classifier
  #   src_generator: Source generator flowing data from test directory
  #   model : str can be either glaucoma or retinopathy

  y_proba = clf.predict(src_generator)
  y_pred  = y_proba.argmax(axis=1)

  src_path_base     = os.path.join(all_paths_dic['test_train_path'],'test')
  dst_path_base ='/content/glaucoma_data/train' if model=='glaucoma' else '/content/diabetic_ret_data/train'

  file_list = os.listdir(src_path_base)

  for idx,label in enumerate(y_pred):
    file_name = file_list[idx]
    src = os.path.join(src_path_base,file_name)
    if label==1:
      dst = os.path.join(dst_path_base,'positive')
    else:
      dst = os.path.join(dst_path_base,'negative')
    shutil.copy(src,dst)


from styx_msgs.msg import TrafficLight

import numpy as np
import tensorflow as tf

import rospy

class TLClassifier(object):
    def __init__(self, is_site=False):
        if is_site:
            self.graph = self.load_inference_graph(r'./frozen_inference_graph_real.pb')
        else:
            self.graph = self.load_inference_graph(r'./frozen_inference_graph.pb')

        self.sess = tf.Session(graph=self.graph)

    def load_inference_graph(self, path):
        graph = tf.Graph()
        with graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(path, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        return graph
    def get_classification(self, image):
        """Determines the color of the traffic light in the image

        Args:
            image (cv::Mat): image containing the traffic light

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        #TODO implement light color prediction
        np_image = self.load_image_into_numpy_array(image)
        image_expanded = np.expand_dims(np_image, axis=0)
        with self.graph.as_default():
            #with tf.Session(graph=self.graph) as sess:
            image_tensor = self.graph.get_tensor_by_name('image_tensor:0')
            detect_boxes = self.graph.get_tensor_by_name('detection_boxes:0')
            detect_scores = self.graph.get_tensor_by_name('detection_scores:0')
            detect_classes = self.graph.get_tensor_by_name('detection_classes:0')
            num_detections = self.graph.get_tensor_by_name('num_detections:0')
            (boxes, scores, classes, num) = self.sess.run(
                [detect_boxes, detect_scores, detect_classes, num_detections],
                feed_dict={image_tensor: image_expanded})

        return self.transform_idx(classes[0][0])

    def transform_idx(self,idx):
        if idx==1:
            return TrafficLight.GREEN
        if idx==2:
            return TrafficLight.RED
        if idx ==3:
            return TrafficLight.YELLOW
        if idx ==4:
	    return TrafficLight.UNKNOWN
	return TrafficLight.UNKNOWN

    def load_image_into_numpy_array(self,image):

        np_arr = np.asarray(image[:,:])
        return np_arr


# -*- coding: utf-8 -*-
"""VGG16model.py

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1aLaIGA-XgtIbU2t_ZgepT3gxVY7RIgym
"""

import cv2
from PIL import Image
import matplotlib.pyplot as plt
import sklearn
import imutils
import os
import numpy as np
from IPython.display import Image, display
from keras.preprocessing.image import img_to_array
from keras.preprocessing.image import load_img
from keras.utils import  np_utils
import keras
from keras.models import Sequential
from keras.layers import Dense,Flatten,Dropout
from keras.layers import Conv2D,MaxPool2D,ZeroPadding2D,Convolution2D,MaxPooling2D
from keras import backend as K
from keras.layers import BatchNormalization
from keras.optimizers import SGD,RMSprop
from keras.models import load_model

#insert image addess
img_address=""

data=[]
img_read=plt.imread(img_address)
img_resize=cv2.resize(img_read,(50,50))
img_ch = cv2.cvtColor(img_resize, cv2.COLOR_BGR2GRAY)
img_array=img_to_array(img_ch)
data.append(img_array)

image_data=np.array(data)

model = load_model('vgg_model.h5')
y=model.predict(img)
#0 for Normal Lungs
#1 for infected lungs
print(y)


import cv2
import numpy as np
import time

# https://raw.githubusercontent.com/opencv/opencv/master/samples/dnn/face_detector/deploy.prototxt
prototxt_path = "weights/deploy.prototxt.txt"
# https://raw.githubusercontent.com/opencv/opencv_3rdparty/dnn_samples_face_detector_20180205_fp16/res10_300x300_ssd_iter_140000_fp16.caffemodel
model_path = "weights/res10_300x300_ssd_iter_140000_fp16.caffemodel"

# load Caffe model
model = cv2.dnn.readNetFromCaffe(prototxt_path, model_path)

cap = cv2.VideoCapture(0)
while True:
    start = time.time()
    _, image = cap.read()
    # get width and height of the image
    h, w = image.shape[:2]
    kernel_width = (w // 7) | 1
    kernel_height = (h // 7) | 1
    # preprocess the image: resize and performs mean subtraction
    blob = cv2.dnn.blobFromImage(image, 1.0, (300, 300), (104.0, 177.0, 123.0))
    # set the image into the input of the neural network
    model.setInput(blob)
    # perform inference and get the result
    output = np.squeeze(model.forward())
    for i in range(0, output.shape[0]):
        confidence = output[i, 2]
        # get the confidence
        # if confidence is above 40%, then blur the bounding box (face)
        if confidence > 0.4:
            # get the surrounding box cordinates and upscale them to original image
            box = output[i, 3:7] * np.array([w, h, w, h])
            # convert to integers
            start_x, start_y, end_x, end_y = box.astype(np.int)
            # get the face image
            face = image[start_y: end_y, start_x: end_x]
            # apply gaussian blur to this face
            face = cv2.GaussianBlur(face, (kernel_width, kernel_height), 0)
            # put the blurred face into the original image
            image[start_y: end_y, start_x: end_x] = face
    cv2.imshow("image", image)
    if cv2.waitKey(1) == ord("q"):
        break
    time_elapsed = time.time() - start
    fps = 1 / time_elapsed
    print("FPS:", fps)

cv2.destroyAllWindows()
cap.release()