import torch
import torch.nn.functional

from transformers import BertTokenizer
from PIL import Image
from avatar_sgg.captioning.catr.hubconf import v3
from avatar_sgg.captioning.catr.datasets import coco
from avatar_sgg.captioning.catr.configuration import Config
from avatar_sgg.config.util import get_config
import numpy as np
import os

class CATRInference():
    def __init__(self):
        self.config = Config()
        # set local file to True if you have connection issues...
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', local_files_only=False)
        self.max_length = self.config.max_position_embeddings
        self.start_token = self.tokenizer.convert_tokens_to_ids(self.tokenizer._cls_token)
        self.end_token = self.tokenizer.convert_tokens_to_ids(self.tokenizer._sep_token)
        self.model = v3(pretrained=True)#torch.hub.load('saahiluppal/catr', 'v3', pretrained=True)
        self.model.eval()
        catr_config = get_config()["captioning"]["catr"]
        self.cuda_device = catr_config["cuda_device"]
        self.beam_size = catr_config["beam_size"]

        if type(self.cuda_device) is str and self.cuda_device.startswith("cuda"):
            print("Use CATR Model with GPU", self.cuda_device)
            self.model.cuda(self.cuda_device)
        else:
            print("Use CATR Model with CPU")

    def create_caption_and_mask(self):

        self.caption_template = torch.zeros((1, self.max_length), dtype=torch.long)
        self.mask_template = torch.ones((1, self.max_length), dtype=torch.bool)

        self.caption_template[:, 0] = self.start_token
        self.mask_template[:, 0] = False

        return self.caption_template, self.mask_template
    @torch.no_grad()
    def infer_beam(self, image_path):
        """
        Beam Search still broken. Does not deliver better bleu score than greedy search
        :param image_path:
        :return:
        """
        image = Image.open(image_path)
        image = coco.val_transform(image)
        image = image.unsqueeze(0)
        beam_size = self.beam_size
        caption, cap_mask = self.create_caption_and_mask()

        if self.cuda_device.startswith("cuda"):
            image = image.cuda(self.cuda_device)
            caption = caption.cuda(self.cuda_device)
            cap_mask = cap_mask.cuda(self.cuda_device)
        src, mask, pos = self.model.init_sample(image)

        predictions = self.model.infer(src, mask, pos, caption, cap_mask) #self.model(image, caption, cap_mask)
        predictions = torch.nn.functional.log_softmax(predictions[:, 0, :], dim=-1)#predictions[:, 0, :]#torch.nn.functional.log_softmax(predictions[:, 0, :])
        previous_log_prob, candidate_indices = torch.topk(predictions, beam_size)
        preds = {i: np.zeros(self.max_length, dtype=int) for i in range(beam_size)}
        for i in range(beam_size):
            preds[i][0] = candidate_indices[0][i]
        # Copy entries a number of time equal to the beam size (the number of alternative paths)
        # 1 means the dimensions stay untouched
        #image = image.repeat(beam_size, 1, 1, 1)
        caption = caption.repeat(beam_size, 1)
        cap_mask = cap_mask.repeat(beam_size, 1)
        src = src.repeat(beam_size,1 ,1 ,1)
        mask = mask.repeat(beam_size,1 ,1)
        pos[0] = pos[0].repeat(beam_size,1, 1,1)
        candidates = []
        caption[:, 1] = candidate_indices
        cap_mask[:, 1] = False
        for step in range(1, self.max_length - 1):
            predictions = self.model.infer(src, mask, pos, caption, cap_mask) #self.model(image, caption, cap_mask)
            predictions = torch.nn.functional.log_softmax(predictions[:, step, :], dim=-1)#predictions[:, step, :]
            candidates_log_prob, candidate_indices = torch.topk(predictions, beam_size)
            candidates_log_prob = torch.reshape(candidates_log_prob + previous_log_prob, (-1,))
            candidate_indices = torch.reshape(candidate_indices, (-1,))
            current_top_candidates, current_top_candidates_idx = torch.topk(candidates_log_prob, k=beam_size)

            # Do the mapping best candidate and "source" of the best candidates
            k_idx = torch.index_select(candidate_indices, dim=0, index=current_top_candidates_idx)
            prev_idx = torch.floor(current_top_candidates_idx / beam_size).to(torch.int32)

            previous_log_prob = torch.unsqueeze(current_top_candidates, dim=1)
            np_prev_idx = prev_idx.cpu().numpy()
            # Overwrite the previous predictions due to the new best candidates
            temp = caption.clone()
            for i in range(prev_idx.shape[0]):
                temp[i][:step + 1] = caption[np_prev_idx[i]][:step + 1]
            caption = temp
            preds = {i: preds[np_prev_idx[i]].copy() for i in range(prev_idx.shape[0])}
            caption[:, step + 1] = k_idx
            cap_mask[:, step + 1] = False
            stop_idx = []
            for i in range(k_idx.shape[0]):
                preds[i][step] = k_idx[i]
                if k_idx[i] == self.end_token:
                    stop_idx.append(i)

            # remove all finished captions and adjust all tensors accordingly...
            if len(stop_idx):
                for i in reversed(sorted(stop_idx)):
                    candidate = preds.pop(i)
                    loss = current_top_candidates[i]
                    length = np.where(candidate == self.end_token)[0]
                    normalized_loss = loss / float(length)
                    candidates.append((candidate, normalized_loss))
                beam_size = beam_size - len(stop_idx)
                if beam_size > 0:
                    left_idx = torch.LongTensor([i for i in range(k_idx.shape[0]) if i not in stop_idx])

                    if self.cuda_device.startswith("cuda"):
                        left_idx = left_idx.cuda(self.cuda_device)
                    # current_top_candidates = torch.IntTensor(
                    #     [current_top_candidates[i] for i in range(current_top_candidates.shape[0]) if
                    #      i not in stop_idx])
                    caption = torch.index_select(caption, dim=0, index=left_idx)
                    cap_mask = torch.index_select(cap_mask, dim=0, index=left_idx)
                    #image = torch.index_select(image, dim=0, index=left_idx)
                    previous_log_prob = torch.index_select(previous_log_prob, dim=0, index=left_idx)
                    src = torch.index_select(src, dim=0, index=left_idx)
                    mask = torch.index_select(mask, dim=0, index=left_idx)
                    pos[0] = torch.index_select(pos[0], dim=0, index=left_idx)
                    # now that the finished sentences have been removed, we need to update the predictions dict accordingly
                    for i, key in enumerate(sorted(preds.keys())):
                        preds[i] = preds.pop(key)
                else:
                    break  # No sequences unfinished

        if len(candidates) > 0:
            result, _ = max(candidates, key=lambda c: c[1])
        else:
            result = preds[0]

        output = self.tokenizer.decode(result, skip_special_tokens=True)
        return output

    @torch.no_grad()
    def infer(self, image_path):
        image = Image.open(image_path)
        image = coco.val_transform(image)
        image = image.unsqueeze(0)

        caption, cap_mask = self.create_caption_and_mask()
        if self.cuda_device.startswith("cuda"):
            image = image.cuda(self.cuda_device)
            caption = caption.cuda(self.cuda_device)
            cap_mask = cap_mask.cuda(self.cuda_device)

        src, mask, pos = self.model.init_sample(image)
        #model.eval()
        for i in range(self.max_length - 1):
            predictions = self.model.infer(src, mask, pos, caption, cap_mask)
            predictions = predictions[:, i, :]
            predicted_id = torch.argmax(predictions, axis=-1)
            if predicted_id[0] == self.end_token:
                #return caption
                break
            caption[:, i+1] = predicted_id[0]
            cap_mask[:, i+1] = False
        output = self.tokenizer.decode(caption[0].tolist(), skip_special_tokens=True)
        return output

if __name__ == "__main__":
    ade20k = get_config()["ade20k"]["root_dir"]
    image_path = os.path.join(ade20k, "images/training/u/utility_room/ADE_train_00019432.jpg")
    catr = CATRInference()
    output = catr.infer(image_path)
    #result = catr.tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
    #result = tokenizer.decode(output[0], skip_special_tokens=True)
    print(output)


import numpy as np
import matplotlib.pyplot as plt
from ImageD11 import columnfile, parameters, grain
from s3dxrd.utils import peak_mapper, reconstruct_grainshapes
from xfab import tools

class RawMeasurements(object):
    """Container of semi-raw diffraction data and grain properties.

    Performs diffraction peak mapping and grain topology reconstruction
    working of ImageD11 file formats. These files are already proccessed in
    the sense that the grains have been indexed and detector tilt has
    been calibrated for etc. In scanning-3DXRD each layer (x-y) is scanned
    independently of z translations. Thus the input is a stacks of
    measurements, one file per sample z-translation. The coordinate system
    and experimental setup is described in:

        Henningsson, N. A., Hall, S. A., Wright, J. P. & Hektor, J.
        (2020). J. Appl. Cryst. 53, 314-325.
        https://doi.org/10.1107/S1600576720001016

    Attributes:
        zpos (:obj:`list` of :obj:`float`): Sample z-translations corresponding to grain_slices
            indices and peak_stack indices.
        grain_slices (:obj:`list` of :obj:`lists` of :obj:`grain`): Loaded ImageD11 grain '.map'
            or '.ubi' files. Each z-translation has a list of grains present in
            the slice. The map files are loaded via the ImageD11 grain module.
        peak_stack (:obj:`list` of :obj:`columnfile`): Loaded ImageD11 peak files. Each
            instance of the peak stack list contains a columnfile describing
            diffraction peaks originating from correpsonding z-translation. The
            peak files are loaded via the ImageD11 columnfile module.
        params (:obj:`parameters`): Loaded ImageD11 parameter file. This object
            gives access to meta parameters of the experiment, such as detector tilt,
            sample to detector distance, and crystal cell parameters. The parameter
            file is loaded via the ImageD11 parameters module.
        omegastep (:obj:`float`): Rotation stage readout intervall in degrees. Each recorded
            frame in the experiment is was integrated over this intervall.
        ystep (:obj:`float`): y-translation stepsize between measurements. Normally the same
            as the beamsize.
        number_y_scans (:obj:`float`): The number of scanned y-positions made in each z-layer.
        ymin (:obj:`float`): Minium y-coordinate seen by the detector. These values differ from
            what is in the orginal columnfile in the case when the columnfile do not have
            a symmetric scanning intervall over y=0.
        ymax (:obj:`float`): Maximum y-coordinate seen by the detector. These values differ from
            what is in the orginal columnfile in the case when the columnfile do not have
            a symmetric scanning intervall over y=0.
        grain_topology_mask (:obj:`list` of :obj:`lists` of :obj:`numpy arrays`): Grain topology in binarized form.
            Each instance of the list contains all grain shapes at the corresponding z-slice, s.t
            ```grain_topology_mask[i][j]``` gives the 2d numpy binarized array that corresponds to grain number
            j in z-slice number i.
    """

    def __init__(self,
                flt_paths,
                zpos,
                param_path,
                ubi_paths,
                omegastep,
                ymin,
                ymax,
                ystep):
        """Container of semi-raw diffraction data and grain properties.

        Args:
            zpos (list of float): Sample z-translations corresponding to flt_paths.
            ubi_paths (str): Absolute paths to ID11 grain files.
            flt_paths (list of str): Absolute paths to ID11 column files.
            param_path (str): Absolute path to ID11 parameter file.
            omegastep (float): Rotation stage readout intervall in degrees.
            ystep (float): stepsize of sample y-translation between measurements
            ymin (float): Minium y-coordinate seen by the detector,
                as given in ID11 column file.
            ymax (float): Maximum y-coordinate seen by the detector,
                as given in ID11 column file.
        """
        self.zpos = zpos
        self.grain_slices = [ grain.read_grain_file( ubip ) for ubip in ubi_paths ]

        for grs in self.grain_slices:
            for gr in grs:
                gr.u = tools.ubi_to_u( gr.ubi )

        self.peak_stack = [columnfile.columnfile(flt) for flt in flt_paths]
        self.tot_nbr_peaks = sum( [peaks.nrows for peaks in self.peak_stack] )
        self.params = parameters.read_par_file( param_path )
        self.omegastep = omegastep
        self.ystep = ystep
        self.number_y_scans = np.round((abs(ymax-ymin)/ystep)).astype(int)+1
        for p in self.peak_stack:
            p.dty[:] = p.dty - ymin - self.ystep*(self.number_y_scans//2)
        self.ymin = -ystep*(self.number_y_scans//2)
        self.ymax =  ystep*(self.number_y_scans//2)
        self.grain_topology_mask = None

        self.params.parameters['ystep'] = self.ystep
        self.params.parameters['ymin'] = self.ymin
        self.params.parameters['t_x'] = None
        self.params.parameters['t_y'] = None
        self.params.parameters['t_z'] = None

    def map_peaks(self, hkltol, nmedian ):
        """Perform Grain centroid refinement and diffraction peak to grain mappings.

        This method refines the semi-raw diffraction data contained by the input
        peak files. The grains of self.grain_slices will be given masks and refined
        average orientation and strain. The columnfiles of the peak_stack will be
        recomputed based on the centroid of the claimer grain.

        Note:
            This function mutates the attributes of the class and has no return.

        Args:
            hkltol (float): Tolerance on Miller indices (hkl) to match a peak.
            nmedian (float): Outlier threshold. If any of the diffraction angles
                (2*theta, omega, eta) are more than nmedian times deviating from the
                median deviation (between model and measurement) the peak is considered
                an outlier and removed.
        """
        for gs, flt, dtz in zip( self.grain_slices, self.peak_stack, self.zpos ):
            peak_mapper.map_peaks(flt, gs, self.params, self.omegastep, \
                                        hkltol, nmedian, self.ymin, self.ystep, self.number_y_scans)

    def reconstruct_grain_topology(self, rcut):
        """Perform Filtered Back Projection to approximate grain shapes.

        The recorded intensity of the diffraction peaks are used to create
        a grain topology map. Each grain is reconstructed on a pixelated grid,
        and thresholded to a binary 1 - 0, grain or no grain, map.

        Note:
            This function mutates the attributes of the class and has no return.

        Args:
            rcut (float): Threshold for segementing out grain. Reconstructed
            relative intensity of less than rcut is considered void, and intensity
            higher than rcut is considered grain.
        """
        self.grain_topology_mask=[]
        for gs, flt, dtz in zip( self.grain_slices, self.peak_stack, self.zpos ):
            grain_topology_mask = reconstruct_grainshapes.FBP_slice(gs, flt, rcut, self.ymin, \
                                        self.ystep, self.number_y_scans)
            self.grain_topology_mask.append( grain_topology_mask )


#cell 0
import math
import types
import numpy as np
import pandas as pd
from pyqstrat.pq_utils import ReasonCode

class MarketOrder:
    def __init__(self, contract, timestamp, qty, reason_code = ReasonCode.NONE, properties = None, status = 'open'):
        '''
        Args:
            contract (:obj:`Contract`):
            timestamp (:obj:`np.datetime64`): Time the order was placed
            qty (float):  Number of contracts or shares.  Use a negative quantity for sell orders
            reason_code (str, optional): The reason this order was created.
                Prefer a predefined constant from the ReasonCode class if it matches your reason for creating this order.
                Default None
            properties (:obj:`SimpleNamespace`, optional): Any order specific data we want to store.  Default None
            status (str, optional): Status of the order, "open", "filled", etc. (default "open")
        '''
        self.contract = contract
        self.timestamp = timestamp
        if not np.isfinite(qty) or math.isclose(qty, 0): raise Exception(f'order qty must be finite and nonzero: {qty}')
        self.qty = qty
        self.reason_code = reason_code
        self.status = status
        if properties is None: properties = types.SimpleNamespace()
        self.properties = properties

    def __repr__(self):
        timestamp = pd.Timestamp(self.timestamp).to_pydatetime()
        return f'{self.contract.symbol} {timestamp:%Y-%m-%d %H:%M:%S} qty: {self.qty}' + (
            '' if self.reason_code == ReasonCode.NONE else f' {self.reason_code}') + (
            '' if not self.properties.__dict__ else f' {self.properties}') + (
            f' {self.status}')

class LimitOrder:
    def __init__(self, contract, timestamp, qty, limit_price, reason_code = ReasonCode.NONE, properties = None, status = 'open'):
        '''
        Args:
            contract (:obj:`Contract`):
            timestamp (:obj:`np.datetime64`): Time the order was placed
            qty (float):  Number of contracts or shares.  Use a negative quantity for sell orders
            limit_price (float): Limit price (float)
            reason_code (str, optional): The reason this order was created.
                Prefer a predefined constant from the ReasonCode class if it matches your reason for creating this order.
                Default None
            properties (:obj:`SimpleNamespace`, optional): Any order specific data we want to store.  Default None
            status (str, optional): Status of the order, "open", "filled", etc. (default "open")
        '''
        self.contract = contract
        self.timestamp = timestamp
        if not np.isfinite(qty) or math.isclose(qty, 0): raise Exception(f'order qty must be finite and nonzero: {qty}')
        self.qty = qty
        self.reason_code = reason_code
        self.limit_price = limit_price
        if properties is None: properties = types.SimpleNamespace()
        self.properties = properties
        self.properties.limit_price = self.limit_price
        self.status = status

    def __repr__(self):
        timestamp = pd.Timestamp(self.timestamp).to_pydatetime()
        return f'{self.contract.symbol} {timestamp:%Y-%m-%d %H:%M:%S} qty: {self.qty} lmt_prc: {self.limit_price}' + (
            '' if self.reason_code == ReasonCode.NONE else f' {self.reason_code}') + (
            '' if not self.properties.__dict__ else f' {self.properties}') + (
            f' {self.status}')

class RollOrder:
    '''A roll order is used to roll a future from one series to the next.  It represents a sell of one future and the buying of another future.'''
    def __init__(self, contract, timestamp, close_qty, reopen_qty, reason_code = ReasonCode.ROLL_FUTURE, properties = None,
                 status = 'open'):
        '''
        Args:
            contract (:obj:`Contract`):
            timestamp (:obj:`np.datetime64`): Time the order was placed
            close_qty (float): Quantity of the future you are rolling
            reopen_qty (float): Quantity of the future you are rolling to
            reason_code (str, optional): The reason this order was created.
                Prefer a predefined constant from the ReasonCode class if it matches your reason for creating this order.
                Default None
            properties (:obj:`SimpleNamespace`, optional): Any order specific data we want to store.  Default None
            status (str, optional): Status of the order, "open", "filled", etc. (default "open")
        '''
        self.contract = contract
        self.timestamp = timestamp
        if not np.isfinite(close_qty) or math.isclose(close_qty, 0) or not np.isfinite(reopen_qty) or math.isclose(reopen_qty, 0):
            raise Exception(f'order quantities must be non-zero and finite: {close_qty} {reopen_qty}')
        self.close_qty = close_qty
        self.reopen_qty = reopen_qty
        self.reason_code = reason_code
        self.qty = close_qty # For display purposes when we print varying order types
        if properties is None: properties = types.SimpleNamespace()
        self.properties = properties
        self.properties.close_qty = self.close_qty
        self.properties.reopen_qty = self.reopen_qty
        self.status = status

    def __repr__(self):
        timestamp = pd.Timestamp(self.timestamp).to_pydatetime()
        return f'{self.contract.symbol} {timestamp:%Y-%m-%d %H:%M:%S} close_qty: {self.close_qty} reopen_qty: {self.reopen_qty}' + (
            '' if self.reason_code == ReasonCode.NONE else f' {self.reason_code}') + '' if not self.properties.__dict__ else f' {self.properties}' + (
            f' {self.status}')

class StopLimitOrder:
    '''Used for stop loss or stop limit orders.  The order is triggered when price goes above or below trigger price, depending on whether this is a short
      or long order.  Becomes either a market or limit order at that point, depending on whether you set the limit price or not.
    '''
    def __init__(self, contract, timestamp, qty, trigger_price, limit_price = np.nan, reason_code = ReasonCode.NONE,
                 properties = None, status = 'open'):
        '''
        Args:
            contract (:obj:`Contract`):
            timestamp (:obj:`np.datetime64`): Time the order was placed
            qty (float):  Number of contracts or shares.  Use a negative quantity for sell orders
            trigger_price (float): Order becomes a market or limit order if price crosses trigger_price.
            limit_price (float, optional): If not set (default), order becomes a market order when price crosses trigger price.
                Otherwise it becomes a limit order.  Default np.nan
            reason_code (str, optional): The reason this order was created.
                Prefer a predefined constant from the ReasonCode class if it matches your reason for creating this order.
                Default None
            properties (:obj:`SimpleNamespace`, optional): Any order specific data we want to store.  Default None
            status (str, optional): Status of the order, "open", "filled", etc. (default "open")
        '''
        self.contract = contract
        self.timestamp = timestamp
        if not np.isfinite(qty) or math.isclose(qty, 0): raise Exception(f'order qty must be finite and nonzero: {qty}')
        self.qty = qty
        self.trigger_price = trigger_price
        self.limit_price = limit_price
        self.reason_code = reason_code
        self.triggered = False
        if properties is None: properties = types.SimpleNamespace()
        self.properties = properties
        self.properties.trigger_price = trigger_price
        self.properties.limit_price = limit_price
        self.status =  status

    def __repr__(self):
        timestamp = pd.Timestamp(self.timestamp).to_pydatetime()
        return f'{self.contract.symbol} {timestamp:%Y-%m-%d %H:%M:%S} qty: {self.qty} trigger_prc: {self.trigger_price} limit_prc: {self.limit_price}' + (
            '' if self.reason_code == ReasonCode.NONE else f' {self.reason_code}') + ('' if not self.properties.__dict__ else f' {self.properties}') + (
            f' {self.status}')

if __name__ == "__main__":
    import doctest
    doctest.testmod(optionflags = doctest.NORMALIZE_WHITESPACE)


import numpy as np
from mvnc import mvncapi
import matplotlib as plt
import tensorflow as tf
import cv2

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
image_index = 7777 # You may select anything up to 60,000
print(y_train[image_index]) # The label is 8
plt.imshow(x_train[image_index], cmap='Greys')
plt.show()

# Get a list of available device identifiers
device_list = mvncapi.enumerate_devices()
# Initialize a Device
device = mvncapi.Device(device_list[0])
# Initialize the device and open communication
device.open()

# Reading the file from disk
GRAPH_FILEPATH = './graph'
with open(GRAPH_FILEPATH, mode='rb') as f:
    graph_buffer = f.read()

# Initialize a Graph object
graph = mvncapi.Graph('graph1')

# Initialize the input/output queue objects
input_fifo, output_fifo = graph.allocate_with_fifos(device, graph_buffer)

# With your webcam plugged in, this will fetch a reference to it
cam = cv2.VideoCapture(0)
# Set the capture size to 28x28
# Note, this will only force the capture to shrink up untill a certain point.
# For me, this was around 150x120. But it's better than nothing
cam.set(3,28)
cam.set(4,28)

while(True):
    # Capture frame-by-frame
    ret, frame = cam.read()

    # Convert the fame to gray-scale
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # Do the final resizing to 28x28
    resize = cv2.resize(gray, (28, 28))
    # Convert to float32, just like before
    tensor = resize.astype(np.float32)
    # Force the values to be between 0 and 1
    tensor /= 255
    # Display the resulting frame
    #cv2.imshow('frame', tensor)

    graph.queue_inference_with_fifo_elem(input_fifo, output_fifo, tensor, 'user object')

    output, user_obj = output_fifo.read_elem()

    # Print the prediction just like before
    print(np.argmax(output))

    # Press q to kill. Or just hit the stop button on Pycharm
    if cv2.waitKey(1) & 0xFF == ord('q'):
        cam.release()
        cv2.destroyAllWindows()
        break

#boiler plate
input_fifo.destroy()
output_fifo.destroy()
graph.destroy()
device.close()
device.destroy()


"""
Like image_sample.py, but use a noisy image classifier to guide the sampling
process towards more realistic images.
"""

import argparse
import os

import numpy as np
import torch as th
import torch.distributed as dist
import torch.nn.functional as F

from guided_diffusion import dist_util, logger
from guided_diffusion.script_util import (
    model_defaults,
    diffusion_defaults,
    classifier_defaults,
    create_model,
    create_gaussian_diffusion,
    create_classifier,
    add_dict_to_argparser,
    args_to_dict,
)
from guided_diffusion.sample_util import save_samples
from tqdm import tqdm

def main():
    args = create_argparser().parse_args()

    visible_gpus_list = []
    if args.gpus:
        visible_gpus_list = [str(gpu_id) for gpu_id in args.gpus.split(",")]
    dist_util.setup_dist(visible_gpu_list=visible_gpus_list, local_rank=args.local_rank)

    logger.configure(dir=os.path.join(args.log_root, args.save_name))
    logger.log(args)

    logger.log("creating model and diffusion...")
    model = create_model(
        **args_to_dict(args, model_defaults().keys())
    )

    diffusion = create_gaussian_diffusion(
        **args_to_dict(args, diffusion_defaults().keys())
    )
    if args.model_path:
        logger.log("loading model from {}".format(args.model_path))
        pretrained_model_dict = dist_util.load_state_dict(args.model_path, map_location="cpu")

        model.load_state_dict(
            pretrained_model_dict,
            strict=True
        )

    for k, v in model.__dict__.items():
        logger.log(k, v)

    model.to(dist_util.dev())
    if args.use_fp16:
        model.convert_to_fp16()
    model.eval()

    classifier = create_classifier(
        **args_to_dict(args, classifier_defaults().keys())
    )

    if args.classifier_path:
        logger.log("loading classifier from {}".format(args.classifier_path))
        pretrained_classifier_dict = dist_util.load_state_dict(args.classifier_path, map_location="cpu")

        classifier.load_state_dict(
            dist_util.load_state_dict(args.classifier_path, map_location="cpu"),
            strict=True
        )

    for k, v in classifier.__dict__.items():
        logger.log(k, v)
    classifier.to(dist_util.dev())
    if args.classifier_use_fp16:
        classifier.convert_to_fp16()
    classifier.eval()

    def cond_fn(x, t, y=None, prior_variance=1.0, t_range_start=0, t_range_end=1000):
        assert y is not None
        with th.enable_grad():
            x_in = x.detach().requires_grad_(True)
            output = classifier(x_in, t)

            logits = output

            log_probs = F.log_softmax(logits, dim=-1)  # (B, C)
            selected = log_probs[range(len(logits)), y.view(-1)]  # (B, )
            cond_grad = th.autograd.grad(selected.sum(), x_in)[0]

            cond_grad = cond_grad * args.classifier_scale

            if args.use_cond_range_scale:
                if t[0] >= 0 and t[0] < 500:
                    cond_grad = cond_grad * 3

        with th.no_grad():

            if args.use_probability_scale and (t[0] >= t_range_start and t[0] < t_range_end):
                probability = F.softmax(logits, dim=-1)[range(len(logits)), y.view(-1)]
                probability_scale = 1.0 / (1.0 - probability)

                original_grad_norm = th.norm(cond_grad, p=2, dim=(1, 2, 3), dtype=th.float32).detach()
                cond_grad = cond_grad * probability_scale.reshape(-1, 1, 1, 1).repeat(1, *cond_grad[0].shape)
                updated_grad_norm = th.norm(cond_grad, p=2, dim=(1, 2, 3), dtype=th.float32).detach()

                logger.log(
                    't = ', t[0].detach(), '\n',
                    '\t\t mean std median', '\n',
                    '\t\t probability = ', probability.mean(-1).detach(), probability.std(-1).detach(), probability.median(-1).values, '\n',
                    '\t\t probability_scale = ', probability_scale.mean(-1).detach(), probability_scale.std(-1).detach(), probability_scale.median(-1).values, '\n',
                    '\t\t grad_norm =', original_grad_norm.mean(-1).detach(), original_grad_norm.std(-1).detach(), original_grad_norm.median(-1).values, '\n',
                    '\t\t updated_grad_norm = ', updated_grad_norm.mean(-1).detach(), updated_grad_norm.std(-1).detach(), updated_grad_norm.median(-1).values, '\n',
                    '\n'
                )
                return cond_grad

            if args.use_entropy_scale and (t[0] >= t_range_start and t[0] < t_range_end):

                probs = F.softmax(logits, dim=-1)  # (B, C)
                entropy = (-log_probs * probs).sum(dim=-1)  # (B, )
                if args.use_normalized_entropy_scale:
                    entropy_scale = 1.0 / (entropy / np.log(args.classifier_out_channels))  # (B, )
                else:
                    entropy_scale = 1.0 / entropy

                updated_selected = args.classifier_scale * selected * entropy_scale  # (B, )
                original_grad_norm = th.norm(cond_grad, p=2, dim=(1, 2, 3), dtype=th.float32).detach()
                cond_grad = cond_grad * entropy_scale.reshape(-1, 1, 1, 1).repeat(1, *cond_grad[0].shape)
                updated_grad_norm = th.norm(cond_grad, p=2, dim=(1, 2, 3), dtype=th.float32).detach()

                logger.log(
                    't = ', t[0].detach(), '\n',
                    '\t\t mean std median', '\n',
                    '\t\t original_selected = ', selected.mean(-1).detach(), selected.std(-1).detach(), selected.median(-1).values, '\n',
                    '\t\t entropy = ', entropy.mean(-1).detach(), entropy.std(-1).detach(), entropy.median(-1).values, '\n',
                    '\t\t entropy_scale = ', entropy_scale.mean(-1).detach(), entropy_scale.std(-1).detach(), entropy_scale.median(-1).values, '\n',
                    '\t\t grad_norm =', original_grad_norm.mean(-1).detach(), original_grad_norm.std(-1).detach(), original_grad_norm.median(-1).values, '\n',
                    '\t\t updated_grad_norm = ',
                    updated_grad_norm.mean(-1).detach(), updated_grad_norm.std(-1).detach(), updated_grad_norm.median(-1).values, '\n',
                    '\n'
                )
                return cond_grad

            original_grad_norm = th.norm(cond_grad, p=2, dim=(1, 2, 3), dtype=th.float32).detach()
            logger.log(
                't = ', t[0].detach(), '\n',
                '\t\t mean std median', '\n',
                '\t\t grad_norm =', original_grad_norm.mean(-1).detach(), original_grad_norm.std(-1).detach(), original_grad_norm.median(-1).values, '\n'
                                                                                '\t\t logits = ', selected.mean(-1).detach(),
                selected.std(-1).detach(), selected.median(-1).values, '\n',
                '\n'
            )

            return cond_grad

    def model_fn(x, t, y=None, t_range_start=0, t_range_end=1000):
        assert y is not None
        return model(x, t, y if args.class_cond else None)

    logger.log("sampling...")
    all_images = []
    all_labels = []
    while len(all_images) * args.batch_size < args.num_samples:
        model_kwargs = {}
        if args.selected_class == -1:
            classes = th.randint(
                low=0, high=args.classifier_out_channels, size=(args.batch_size,), device=dist_util.dev()
            )
        else:
            classes = th.randint(
                low=args.selected_class, high=args.selected_class + 1, size=(args.batch_size,), device=dist_util.dev()
            )

        model_kwargs["y"] = classes
        model_kwargs['t_range_start'] = args.t_range_start
        model_kwargs['t_range_end'] = args.t_range_end

        sample_fn = (
            diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
        )
        sample = sample_fn(
            model_fn,
            (args.batch_size, 3, args.image_size, args.image_size),
            clip_denoised=args.clip_denoised,
            model_kwargs=model_kwargs,
            cond_fn=cond_fn,
            device=dist_util.dev(),
        )
        sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
        sample = sample.permute(0, 2, 3, 1)
        sample = sample.contiguous()

        gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
        dist.all_gather(gathered_samples, sample)  # gather not supported with NCCL
        all_images.extend([sample.cpu().numpy() for sample in gathered_samples])
        gathered_labels = [th.zeros_like(classes) for _ in range(dist.get_world_size())]
        dist.all_gather(gathered_labels, classes)
        all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
        logger.log(f"created {len(all_images) * args.batch_size} / {args.num_samples} samples")

    arr = np.concatenate(all_images, axis=0)
    arr = arr[: args.num_samples]
    label_arr = np.concatenate(all_labels, axis=0)
    label_arr = label_arr[: args.num_samples]
    if dist.get_rank() == 0:
        shape_str = "x".join([str(x) for x in arr.shape])

        if args.selected_class == -1:
            out_path = os.path.join(logger.get_dir(), "scale{}_steps{}_samples_{}.npz".format(args.classifier_scale, args.timestep_respacing, shape_str))
        else:
            out_path = os.path.join(logger.get_dir(), "scale{}_steps{}_class{}_samples_{}.npz".format(args.classifier_scale, args.timestep_respacing, args.selected_class, shape_str))

        logger.log(f"saving to {out_path}")
        np.savez(out_path, arr, label_arr)

        # sample_dir = "images_" + "scale{}_steps{}_sample{}".format(args.classifier_scale, args.timestep_respacing, args.num_samples)
        # save_samples(arr, label_arr, args.classifier_out_channels, sample_dir=sample_dir)

    dist.barrier()
    logger.log("sampling complete")

def create_argparser():
    defaults = dict(
        clip_denoised=True,
        num_samples=10000,
        batch_size=16,
        use_ddim=False,
        model_path="",
        classifier_path="",
        classifier_scale=1.0,
        log_root="",
        save_name="",
        gpus="",
        t_range_start=0,
        t_range_end=1000,

        use_entropy_scale=False,

        expected_classifier_gradient_value=-1.0,

        selected_class=-1,
        use_cond_range_scale=False,
        use_probability_scale=False,
        use_normalized_entropy_scale=True
    )
    defaults.update(diffusion_defaults())
    defaults.update(model_defaults())
    defaults.update(classifier_defaults())
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, default=0)
    add_dict_to_argparser(parser, defaults)
    return parser

if __name__ == "__main__":
    main()


import string
from seesaw.multiscale_index import MultiscaleIndex
import torch.nn as nn
import torch.nn.functional as F
import clip
import numpy as np
import pandas as pd
from clip.model import build_model
from .clip_module import CLIPFineTunedModel, CLIPTx, MappedDataset
import torch
from .multiscale_index import add_iou_score
from ray.data.extensions import TensorArray
import torch.optim
from .clip_module import configure_optimizer

def join_vecs2annotations(db: MultiscaleIndex, dbidx, annotations):
    patch_box_df = db.get_data(dbidx)
    roi_box_df = pd.DataFrame.from_records([b.dict() for b in annotations])

    dfvec = add_iou_score(patch_box_df, roi_box_df)
    dfvec = dfvec.assign(
        descriptions=dfvec.best_box_idx.map(lambda idx: annotations[idx].description)
    )

    dfbox = add_iou_score(roi_box_df, patch_box_df)

    matched_vecs = np.stack(
        [dfvec.vectors.iloc[i].copy() for i in dfbox.best_box_idx.values]
    )
    dfbox = dfbox.assign(
        descriptions=dfbox.description, vectors=TensorArray(matched_vecs)
    )

    return dfvec, dfbox

def deduplicate_strings(string_list):
    s2id = {}
    sids = []
    for s in string_list:
        if s not in s2id:  # add new string to dict
            s2id[s] = len(s2id)

        sids.append(s2id[s])

    string_ids = np.array(sids)

    reverse_dict = {num: strng for (strng, num) in s2id.items()}
    id2string = np.array([reverse_dict[i] for i in range(len(s2id))])

    return {"strings": id2string, "indices": string_ids}

class LinearScorer(nn.Module):
    def __init__(self, init_weight: torch.Tensor):
        super().__init__()
        self.weight = nn.Parameter(data=init_weight, requires_grad=True)
        self.bias = nn.Parameter(
            torch.tensor(0.0, device=init_weight.device), requires_grad=True
        )
        self.logit_scale = nn.Parameter(
            torch.tensor(0.0, device=init_weight.device), requires_grad=True
        )

    def get_vec(self):
        return self.weight.detach().cpu().numpy()

    def forward(self, X: torch.Tensor):
        return (X @ self.weight.t()) * self.logit_scale.exp() + self.bias

class DynamicLinear(nn.Module):
    def __init__(self, device):
        super().__init__()
        self.device = device
        self.scorers = nn.ModuleDict()

    def add_scorer(self, name: str, init_weight: torch.Tensor):
        assert name not in self.scorers
        self.scorers.add_module(
            name, LinearScorer(init_weight=init_weight).to(self.device)
        )

    def get_vec(self, name):
        return self.scorers[name].get_vec()

    def forward(self, vecs: torch.Tensor):
        assert len(vecs.shape) == 2
        scores = []
        for (_, scorer) in self.scorers.items():
            scores.append(scorer(vecs))

        return torch.stack(scores).t()

### the way we pick vectors for training right now should be revisited
import transformers

def hinge(x, margin):  # 0 loss if x >= margin.
    return F.relu(-x + margin)

def rank_loss(ranking_scores, marked_accepted, margin):
    if 0 < marked_accepted.sum() < marked_accepted.shape[0]:
        pos_scores = ranking_scores[marked_accepted]
        neg_scores = ranking_scores[~marked_accepted]
        image_rank_losses = hinge(
            pos_scores.reshape(-1, 1) - neg_scores.reshape(1, -1), margin
        )
        return image_rank_losses.reshape(-1).mean()
    else:
        return None

class OnlineModel:
    def __init__(self, state_dict, config):
        if not torch.cuda.is_available():
            if config["device"].startswith("cuda"):
                print("Warning: no GPU available, using cpu instead")
                config = {**config, "device": "cpu"}  # overrule gpu if not available

        self.original_weights = {k: v.float() for (k, v) in state_dict.items()}

        self.device = config["device"]
        self.model = build_model(self.original_weights).float().to(self.device)
        self.linear_scorer = None
        self.config = config
        self._reset_model()
        self.mode = self.config["mode"]
        self.losses = []
        self._cache = {}

    def _reset_model(self):  # resets model and optimizers
        print("resetting model state")
        layers = ["text_projection"]
        reset_weights = {
            k: v.to(self.device)
            for (k, v) in self.original_weights.items()
            if k in layers
        }
        self.model.load_state_dict(reset_weights, strict=False)
        self.linear_scorer = DynamicLinear(self.device)
        print("done resetting model")

    def _encode_string(self, tokenized_strings):
        text_features = self.model.encode_text(tokenized_strings.to(self.device))
        text_features = F.normalize(text_features)
        return text_features

    def encode_string(self, string) -> np.ndarray:
        self.model.eval()
        with torch.no_grad():
            tokens = clip.tokenize([string])
            vecs = self._encode_string(tokens)
            return vecs.cpu().numpy()

    def compute_up_to(self, strings, layer) -> torch.Tensor:
        assert layer in ["text_projection", "full"]

        non_cached = []
        for s in strings:
            if s not in self._cache:
                non_cached.append(s)

        def closure(self, tokens):  # pass self.model
            # taken from model.encode_text
            x = self.token_embedding(tokens).type(
                self.dtype
            )  # [batch_size, n_ctx, d_model]

            x = x + self.positional_embedding.type(self.dtype)
            x = x.permute(1, 0, 2)  # NLD -> LND
            x = self.transformer(x)
            x = x.permute(1, 0, 2)  # LND -> NLD
            x = self.ln_final(x).type(self.dtype)

            # x.shape = [batch_size, n_ctx, transformer.width]
            # take features from the eot embedding (eot_token is the highest number in each sequence)
            x = x[torch.arange(x.shape[0]), tokens.argmax(dim=-1)]
            if layer == "text_projection":
                return x
            elif layer == "full":
                x = x @ self.text_projection
                return F.normalize(x)
            else:
                assert False

        if len(non_cached) > 0:
            tokens = clip.tokenize(non_cached).to(self.device)
            new_vecs = closure(self.model, tokens)
            for (s, v) in zip(non_cached, new_vecs):
                self._cache[s] = v

        ans = []
        for s in strings:
            ans.append(self._cache[s])

        return torch.stack(ans)

    def compute_from(self, x, layer) -> torch.Tensor:
        assert layer in ["text_projection", "full"]

        def closure(self, x):
            if layer == "text_projection":
                x = x @ self.text_projection
                return F.normalize(x)
            elif layer == "full":
                return x
            else:
                assert False

        return closure(self.model, x)

    def score_vecs(self, imagevecs):
        with torch.no_grad():
            self.model.eval()
            X = torch.from_numpy(imagevecs).to(self.device)
            logits = self.linear_scorer(X)
            if logits.shape[1] > 1:
                logits = logits.softmax(
                    dim=-1
                )  # want the score to stay within 0 to 1 range for visualization
            return logits[:, 0].cpu().numpy()

    def get_lookup_vec(self, str):
        return self.linear_scorer.get_vec(str)

    ### option 1: discount common mistakes via a softmax.
    ### assume only negatives get text description (including near misses)

    def prepare_annotated_pairs(
        self, imagevecs, marked_accepted, annotations, target_string
    ):
        """returns N image vecs, M unique (non empty) string vecs, and the string id of the annotation for each vec"""

        has_description = annotations != ""
        assert target_string != ""

        imagevecs = imagevecs[has_description]
        annotations = annotations[has_description]

        d = deduplicate_strings([target_string] + list(annotations))
        strings = d["strings"]
        string_ids = torch.from_numpy(d["indices"][1:]).to(self.device)
        imagevecs = torch.from_numpy(imagevecs).to(self.device)

        with torch.no_grad():
            self.model.eval()
            last_layer = "full" if self.mode == "linear" else "text_projection"
            stringvecs = self.compute_up_to(strings, last_layer)

        assert imagevecs.device == stringvecs.device
        assert stringvecs.shape[0] == strings.shape[0]
        return {
            "imagevecs": imagevecs,
            "stringvecs": stringvecs,
            "target": string_ids,
            "unique_strings": strings,
        }

    def update(self, imagevecs, marked_accepted, annotations, target_string):
        assert imagevecs.shape[0] == marked_accepted.shape[0]
        assert annotations.shape[0] == marked_accepted.shape[0]

        all_imagevecs = torch.from_numpy(imagevecs).to(self.device)
        r = self.prepare_annotated_pairs(
            imagevecs, marked_accepted, annotations, target_string
        )

        self._reset_model()
        self.model.train()
        self.linear_scorer.train()

        if self.mode == "linear":
            return self._update_linear(r, all_imagevecs, marked_accepted)
        elif self.mode == "finetune":
            return self._update_finetune(r, all_imagevecs, marked_accepted)
        else:
            assert False

    def _update_linear(self, description_data, all_imagevecs, marked_accepted):
        ## want mask for all vectors where the user has provided a better description
        ## ie. exclude empty descriptions '', and also exclude cases where the description
        ## is identical to the search query used as reference (we would be penalizing it)

        for (lookupstr, vec) in zip(
            description_data["unique_strings"], description_data["stringvecs"]
        ):
            assert lookupstr != ""
            self.linear_scorer.add_scorer(lookupstr, vec)

        # 2 groups: vectors modified more slowly
        # temps and biases more aggressively:
        pgs = [
            dict(
                params=[
                    p
                    for (name, p) in self.linear_scorer.named_parameters()
                    if name.endswith("weight")
                ],
                lr=0.001,
                weight_decay=0.0,
            ),
            dict(
                params=[
                    p
                    for (name, p) in self.linear_scorer.named_parameters()
                    if not name.endswith("weight")
                ],
                lr=0.002,
                weight_decay=0.0,
            ),
        ]

        opt = transformers.AdamW(pgs)
        lr_scheduler = transformers.get_constant_schedule_with_warmup(
            opt, num_warmup_steps=self.config["num_warmup_steps"]
        )

        def _description_loss(score_all_pairs, description_data):
            if len(description_data["unique_strings"]) > 1:
                assert score_all_pairs.shape[1] > 1
                return F.cross_entropy(
                    score_all_pairs, description_data["target"]
                ).mean()
            else:
                print("no textual annotationswith wich to make a loss")
                return None

        def training_step():
            self.model.train()

            score_all_pairs = self.linear_scorer(description_data["imagevecs"])
            label_rank_loss = _description_loss(score_all_pairs, description_data)

            ranking_scores = self.linear_scorer(all_imagevecs).log_softmax(dim=-1)[:, 0]
            image_rank_loss = rank_loss(
                ranking_scores, marked_accepted, margin=self.config["rank_margin"]
            )

            if label_rank_loss is None and image_rank_loss is None:
                return None

            loss1 = 0.0 if label_rank_loss is None else label_rank_loss
            loss2 = 0.0 if image_rank_loss is None else image_rank_loss

            loss = (1.0 - self.config["image_loss_weight"]) * loss1 + self.config[
                "image_loss_weight"
            ] * loss2
            return (loss, loss1, loss2)

        for _ in range(self.config["rounds"] + self.config["num_warmup_steps"]):
            opt.zero_grad()
            loss = training_step()
            if loss is None:
                print("no loss yet: no qualified labels")
                break
            (loss, l1, l2) = loss
            loss.backward()
            print(f"loss:{loss:.02f} label_loss: {l1:.02f} rank_loss: {l2:.02f}")
            opt.step()
            lr_scheduler.step()

        return []

    def _update_finetune(self, description_data, all_imagevecs, marked_accepted):
        r = configure_optimizer(self.model, self.config)
        opt: torch.optim.Optimizer = r["optimizer"]
        lr_scheduler = r["lr_scheduler"]

        d = description_data

        def _compute_label_loss(scores, target):
            if scores.shape[0] > 0 and scores.shape[1] > 1:
                # return hinge(scores[torch.arange(scores.shape[0]), target] - scores[:,0]) # needs to handle special case where target is 0
                return F.multi_margin_loss(
                    scores, target, margin=self.config["label_margin"]
                )
            else:
                return None

        def opt_closure():
            self.model.train()
            text_features = self.compute_from(d["stringvecs"], "text_projection")
            scores = d["imagevecs"] @ text_features.t()

            l1 = _compute_label_loss(scores, d["target"])

            rank_scores = (all_imagevecs @ text_features.t())[:, 0]
            l2 = rank_loss(
                rank_scores, marked_accepted, margin=self.config["rank_margin"]
            )

            if l1 is None and l2 is None:
                return None

            l1 = l1 if l1 is not None else 0.0
            l2 = l2 if l2 is not None else 0.0

            loss = (1.0 - self.config["image_loss_weight"]) * l1 + self.config[
                "image_loss_weight"
            ] * l2
            return loss, l1, l2

        for _ in range(self.config["rounds"] + self.config["num_warmup_steps"]):
            opt.zero_grad()
            loss = opt_closure()
            if loss is None:
                print("no loss yet: no qualified labels")
                break
            (loss, l1, l2) = loss
            loss.backward()
            print(f"loss:{loss:.02f} label_loss: {l1:.02f} rank_loss: {l2:.02f}")
            opt.step()
            lr_scheduler.step()


import os
import sys
#nopep8
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.getcwd())
import argparse
from src.resemblyzer.voice_encoder import VoiceEncoder
from sklearn.metrics import roc_curve
from scipy.interpolate import interp1d
from scipy.optimize import brentq
import librosa
import matplotlib.pyplot as plt
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from math import floor
from torch.utils.tensorboard import SummaryWriter
from yaml import safe_load

import soundfile

# from src.data.data_utils import ImbalancedDatasetSampler
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.getcwd())
from src.datasets.data_utils import H_L, STEP_SIZE_EM, split_audio_ixs, mel_spectrogram, preprocess_aud, HDF5TorchDataset  # noqa
np.random.seed(42)

class SPEAKER_ENCODER(nn.Module):
    """
    Attributes
    ----------
    ...

    Methods
    -------
    ...

    """

    def __init__(self,
                 input_shape=(160, 40),
                 load_model=False,
                 epoch=0,
                 dataset_train='',
                 dataset_val='',
                 device=torch.device('cpu'),
                 loss_=None, mode='train'):
        super(SPEAKER_ENCODER, self).__init__()
        self.input_shape = input_shape
        with open('src/config.yaml', 'r') as f:
            self.config_yml = safe_load(f.read())

        model_save_dir = os.path.join(self.config_yml['model_save_dir'],
                                      'SPEAKER_ENCODER')
        if not os.path.exists(model_save_dir):
            os.makedirs(model_save_dir)

        self.model_save_string = os.path.join(
            model_save_dir, self.__class__.__name__ + '_Epoch_{}.pt')

        log_dir = os.path.join(
            self.config_yml['runs_dir'],
            '{}_{}'.format(dataset_train, self.__class__.__name__))
        self.writer = SummaryWriter(log_dir=os.path.join(
            log_dir, "run_{}".format(
                len(os.listdir(log_dir)) if os.path.exists(log_dir) else 0)))

        self.device = device

        if mode == 'train':
            self.dataset_train = HDF5TorchDataset(dataset_train, device=device)
            if dataset_val == '':
                self.dataset_val = self.dataset_train
            else:
                self.dataset_val = HDF5TorchDataset(
                    args.dataset_val, device=device)

        self.lstm = nn.LSTM(self.config_yml['n_mels'],
                            self.config_yml['hidden_size'],
                            self.config_yml['num_layers'],
                            batch_first=True,
                            bidirectional=False)
        self.linear = nn.Linear(1*self.config_yml['hidden_size'],
                                self.config_yml['embedding_size'])

        self.similarity_weight = torch.nn.Parameter(
            torch.Tensor([10.0]), requires_grad=True)
        self.similarity_bias = torch.nn.Parameter(
            torch.Tensor([-5.0]), requires_grad=True)

        self.dropout = nn.Dropout()
        self.relu = nn.ReLU()
        self.softmax = torch.nn.Softmax()
        self.load_model = load_model
        self.epoch = epoch
        self.loss_ = loss_
        self.opt = None
        self.ce_loss = nn.CrossEntropyLoss(reduction='mean')

    def forward(self, frames):

        _, (x, _) = self.lstm(frames)  # lstm out,hidden,
        # x = x[:, -1]  #last layer -> embeds

        x = self.linear(x[-1])
        x = self.relu(x)
        # x = torch.mean(x,dim=1)
        # x = self.relu(x)

        x = x * torch.reciprocal(torch.norm(x, dim=1, keepdim=True))

        return x

    def loss_fn(self, loss_, embeds, labels):

        lang_count = int(self.config_yml['batch_size'] /
                         self.config_yml['uttr_count'])
        embeds3d = embeds.view(lang_count, self.config_yml['uttr_count'], -1)
        dcl = self.direct_classification_loss(embeds, labels)

        centroids = torch.mean(embeds3d, dim=1)

        centroids_neg = (torch.sum(embeds3d, dim=1, keepdim=True) -
                         embeds3d) / (self.config_yml['uttr_count'] - 1)
        cosim_neg = torch.cosine_similarity(embeds,
                                            centroids_neg.view_as(embeds),
                                            dim=1).view(lang_count, -1)
        centroids = centroids.repeat(
            lang_count * self.config_yml['uttr_count'], 1)

        embeds2de = embeds.unsqueeze(1).repeat_interleave(lang_count, 1).view(
            -1, self.config_yml['embedding_size'])
        cosim = torch.cosine_similarity(embeds2de, centroids)
        cosim_matrix = cosim.view(lang_count, self.config_yml['uttr_count'],
                                  -1)
        neg_ix = list(range(lang_count))

        cosim_matrix[neg_ix, :, neg_ix] = cosim_neg

        sim_matrix = (self.similarity_weight * cosim_matrix) + \
            self.similarity_bias

        sim_matrix = sim_matrix.view(self.config_yml['batch_size'], -1)
        targets = torch.range(
            0, self.config_yml['acc_count'] - 1).repeat_interleave(
                self.config_yml['uttr_count']).long().to(device=self.device)

        ce_loss = loss_(sim_matrix, targets)

        return (ce_loss, dcl), sim_matrix

    def direct_classification_loss(self, embeds, labels):
        labels = labels.reshape(-1, 1).squeeze()
        return self.ce_loss(embeds, labels)

    def train_loop(self,
                   opt,
                   lr_scheduler,
                   loss_,
                   batch_size=1,
                   gaze_pred=None,
                   cpt=0):

        train_iterator = torch.utils.data.DataLoader(self.dataset_train,
                                                     batch_size=batch_size,
                                                     shuffle=True,
                                                     drop_last=True)

        self.val_iterator = torch.utils.data.DataLoader(self.dataset_val,
                                                        batch_size=batch_size,
                                                        shuffle=True,
                                                        drop_last=True)

        if self.load_model:
            self.load_model_cpt(cpt=cpt)

        for epoch in range(self.epoch, 20000):
            for i, (data, labels) in enumerate(train_iterator):
                data = data.view(
                    -1,
                    self.config_yml['MEL_CHANNELS'],
                    self.config_yml['sliding_win_size'],
                ).transpose(1, 2)

                opt.zero_grad()

                embeds = self.forward(data)
                (ce_loss,
                 dcl), sim_matrix = self.loss_fn(loss_, embeds, labels)
                self.loss = ce_loss + dcl

                self.loss.backward()

                torch.nn.utils.clip_grad_norm_(self.parameters(), 3.0)

                opt.step()
                self.writer.add_scalar('ce_loss', ce_loss.data.item(), epoch)
                self.writer.add_scalar('dcl', dcl.data.item(), epoch)
                self.writer.add_scalar('Loss', self.loss.data.item(), epoch)
                # self.writer.add_scalar('ValLoss', self.val_loss(), epoch)

                self.writer.add_scalar('EER', self.eer(sim_matrix), epoch)

            if epoch % 1 == 0:
                # self.writer.add_scalar('Loss', loss.data.item(), epoch)
                # self.writer.add_scalar('Val Loss', self.val_loss(), epoch)
                # self.writer.add_scalar('EER', self.eer(sim_matrix), epoch)
                # self.writer.add_histogram('sim', sim_matrix, epoch)

                torch.save(
                    {
                        'epoch': epoch,
                        'model_state_dict': self.state_dict(),
                        'optimizer_state_dict': opt.state_dict(),
                        'loss': self.loss,
                    }, self.model_save_string.format(epoch))

    def embed(self, aud, group=True):
        aud_splits, mel_splits = split_audio_ixs(len(aud))
        max_aud_length = aud_splits[-1].stop
        if max_aud_length >= len(aud):
            aud = np.pad(aud, (0, max_aud_length - len(aud)), "constant")

        mel = mel_spectrogram(aud).astype(np.float32).T
        mels = np.array([mel[s] for s in mel_splits])
        print(mels.shape)
        mels = torch.from_numpy(mels).to(self.device)
        embeds_all = []
        with torch.no_grad():
            for i in range(0, mels.shape[0], 64):
                embeds = self.forward(
                    mels[i:min(i+64, mels.shape[0]), :, :].to(self.device))
                embeds = embeds / torch.norm(embeds, dim=1, keepdim=True)
                embeds_all.append(embeds)
        embeds = torch.cat(embeds_all)

        if group:
            embeds = torch.mean(embeds, dim=0)
            embeds = embeds / torch.norm(embeds)

        embeds = embeds.cpu().data.numpy()
        # self.train()

        return embeds, (aud_splits, mel_splits)

    def accuracy(self):
        acc = 0
        ix = 0
        for i, data in enumerate(self.val_data):
            uttrs = data[0]
            embeds = self.forward(uttrs)
            # TODO
        return (acc / ix)

    def eer(self, sim_matrix=None):
        with torch.no_grad():

            targets = F.one_hot(
                torch.arange(0, self.config_yml['acc_count']),
                num_classes=self.config_yml['acc_count']).repeat_interleave(
                    self.config_yml['uttr_count'], 1).long().T

            fpr, tpr, thresholds = roc_curve(
                targets.flatten(),
                sim_matrix.detach().flatten().cpu().numpy())

            eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)

        return eer

    def val_loss(self):
        with torch.no_grad():
            val_loss = []

            for ix, (datum, labels) in enumerate(self.val_iterator):
                datum = datum.view(
                    -1,
                    self.config_yml['MEL_CHANNELS'],
                    self.config_yml['sliding_win_size'],
                ).transpose(1, 2)
                embeds = self.forward(datum)

                (ce_loss,
                 dcl), sim_matrix = self.loss_fn(self.loss_, embeds, labels)
                loss = ce_loss + dcl

                val_loss.append(loss)

                if ix == self.config_yml['VAL_LOSS_COUNT']:
                    break

        return torch.mean(torch.stack(val_loss)).data.item()

    def load_model_cpt(self, cpt=0, opt=None, device=torch.device('cuda')):
        self.epoch = cpt

        model_pickle = torch.load(self.model_save_string.format(self.epoch),
                                  map_location=device)
        self.load_state_dict(model_pickle['model_state'])
        if opt:
            self.opt.load_state_dict(model_pickle['optimizer_state'])
        self.global_step = model_pickle['step']
        # self.loss = model_pickle['loss']
        self.loss = None
        print("Loaded Model at epoch {},with loss {}".format(
            self.epoch, self.loss))

    def infer(self, fname, cpt=None):
        aud = preprocess(fname)
        embeds = self.embed(aud, group=True)
        return embeds

    def similarity(self, embed1, embed2):
        sim = torch.cosine_similarity(
            torch.tensor(embed1).unsqueeze(0),
            torch.tensor(embed2).unsqueeze(0)).data.item()
        # sim = embed1 @ embed2
        return sim

    def sim_matrix_infer(self, fnames, cpt):
        sim_matrix = np.zeros((len(fnames), len(fnames)))
        for i, f1 in enumerate(fnames):
            for j, f2 in enumerate(fnames):
                sim_matrix[i, j] = self.similarity(self.infer(f1, cpt),
                                                   self.infer(f2, cpt))

        return sim_matrix

    def diarize(self, aud, ref_embed):
        embed = self.embed(aud, group=True)
        sim = embed@ref_embed
        return sim

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--device",
                        help="cpu or cuda",
                        default='cuda',
                        choices=['cpu', 'cuda'])
    parser.add_argument("--dataset_train",
                        help="path to train_dataset",
                        required=False, default='timit_4_train')
    parser.add_argument("--dataset_val",
                        help="path to val_dataset",
                        required=False, default='timit_4_train')
    parser.add_argument("--mode",
                        help="train or eval",
                        required=True,
                        choices=['train', 'eval'])
    parser.add_argument(
        "--filedir",
        help="dir with fnames to run similiarity eval,atleast 2, separted by a comma",
        type=str)
    parser.add_argument("--load_model",
                        help="to load previously saved model checkpoint",
                        default=False)
    parser.add_argument(
        "--cpt",
        help="# of the save model cpt to load, only valid if valid_cpt is true"
    )

    args = parser.parse_args()

    device = torch.device(args.device)
    loss_ = torch.nn.CrossEntropyLoss(reduction='sum')
    encoder = SPEAKER_ENCODER(dataset_train=args.dataset_train,
                              dataset_val=args.dataset_val,
                              device=device,
                              loss_=loss_,
                              load_model=args.load_model, mode=args.mode).to(device=device)
    optimizer = torch.optim.SGD(encoder.parameters(), lr=1e-2)

    encoder.opt = optimizer
    cpt = args.cpt
    if args.load_model:
        encoder.load_model_cpt(cpt=cpt, device=device)
    lr_scheduler = None
    if args.mode == 'train':
        encoder.train_loop(optimizer,
                           lr_scheduler,
                           loss_,
                           batch_size=encoder.config_yml['acc_count'],
                           cpt=cpt)
    elif args.mode == 'vis':

        fnames = [
            os.path.join(args.filedir, x)
            for x in sorted(os.listdir(args.filedir))
        ]

        sim_matrix = encoder.sim_matrix_infer(fnames, cpt)
    elif args.mode == 'eval':
        aud1, _ = preprocess_aud(os.path.join(
            encoder.config_yml['vis_dir'], 'speaker_male.wav'))
        aud2, _ = preprocess_aud(os.path.join(
            encoder.config_yml['vis_dir'], 'speaker_female.wav'))

        mel1 = mel_spectrogram(aud1, mel_type='simple_40mels')
        mel2 = mel_spectrogram(aud2, mel_type='simple_40mels')
        rate = 16000

        mel1 = librosa.db_to_power(mel1)
        mel2 = librosa.db_to_power(mel2)

        aud1 = librosa.feature.inverse.mel_to_audio(mel1,
                                            sr=rate,
                                            n_fft=400,
                                            hop_length=160,
                                            win_length=400)
        aud2 = librosa.feature.inverse.mel_to_audio(mel2,
                                            sr=rate,
                                            n_fft=400,
                                            hop_length=160,
                                            win_length=400)

        # aud1 = librosa.feature.inverse.mel_to_audio(mel1,
        #                                     sr=rate,
        #                                     n_fft=1024,
        #                                     hop_length=256,
        #                                     win_length=1024)
        # aud2 = librosa.feature.inverse.mel_to_audio(mel2,
        #                                     sr=rate,
        #                                     n_fft=1024,
        #                                     hop_length=256,
        #                                     win_length=1024)

        soundfile.write(os.path.join(
            encoder.config_yml['vis_dir'], 'aud1.wav'),
                        aud1,
                        samplerate=16000)
        soundfile.write(os.path.join(
            encoder.config_yml['vis_dir'], 'aud2.wav'),
                        aud2,
                        samplerate=16000)

        embed1, _ = encoder.embed(aud1, group=True)
        embed2, _ = encoder.embed(aud2, group=True)

        print(embed1@embed2)
        # print(embed1.shape,embed2.shape)

        encoder = VoiceEncoder()
        embed1 = encoder.embed_utterance(aud1)
        embed2 = encoder.embed_utterance(aud2)
        print(embed1@embed2)

        # fig, (ax1, ax2) = plt.subplots(2, 1)
        # fig.suptitle('A tale of 2 subplots')

        # ax1.plot(range(embed1.shape[1]),embed1.squeeze(), 'o-')
        # ax1.set_xlabel('S1')

        # ax2.plot(range(embed2.shape[1]),embed2.squeeze(), '.-')
        # ax2.set_xlabel('S2`')

        # plt.scatter(range(embed1.shape[1]),embed1.squeeze(),embed2.squeeze())
        # plt.savefig(os.path.join(
        #     encoder.config_yml['vis_dir'], 'plt.png'))