# -*- coding: utf-8 -*-
# Part of the psychopy.iohub library.
# Copyright (C) 2012-2016 iSolver Software Solutions
# Distributed under the terms of the GNU General Public License (GPL).
"""iohub wintab util objects / functions for stylus, position traces,
and validation process psychopy graphics.
"""
from __future__ import division, absolute_import

import math
from collections import OrderedDict

import numpy as np

from psychopy import visual, core
from psychopy.visual.basevisual import MinimalStim

class PenPositionStim(MinimalStim):
    """Draws the current pen x,y position with graphics that represent the
    pressure, z axis, and tilt data for the wintab sample used."""
    def __init__(self, win, min_opacity=0.0, hover_color=(255,0,0),
                 touching_color=(0,255,0), tiltline_color=(255,255, 0),
                 tiltline_width=2,
                 min_size=0.033, size_range=0.1666, tiltline_scalar=1.0,
                 name=None, autoLog=None, depth=-10000, colorSpace='rgb255'):
        self.win = win
        self.depth = depth
        super(PenPositionStim, self).__init__(name, autoLog)

        # Pen Hovering Related

        # Opaticy is changed based on pen's z axis if data for z axis
        # is available. Opacity of min_opacity is used when pen is at the
        # furthest hover distance (z value) supported by the device.
        # Opacity of 1.0 is used when z value == 0, meaning pen is touching
        # digitizer surface.
        self.min_opacity = min_opacity
        # If z axis is supported, hover_color specifies the color of the pen
        # position dot when z val > 0.
        self.hover_color = hover_color

        # Pen Pressure Related

        # Smallest radius (in norm units) that the pen position gaussian blob
        # will have, which occurs when pen pressure value is 0
        self.min_size = min_size
        # As pen pressure value increases, so does position gaussian blob
        # radius (in norm units). Max radius is reached when pressure is
        # at max device pressure value, and is equal to min_size+size_range
        self.size_range = size_range
        # Color of pen position blob when pressure > 0.
        self.touching_color = touching_color

        # Pen tilt Related

        # Color of line graphic used to represent the pens tilt relative to
        # the digitizer surface.
        self.tiltline_color = tiltline_color
        self.tiltline_width = tiltline_width
        self.tiltline_scalar = tiltline_scalar
        # Create a Gausian blob stim to use for pen position graphic
        self.pen_guass = visual.PatchStim(win, units='norm', tex='none',
                                          mask='gauss', pos=(0, 0),
                                          colorSpace='rgb255',
                                          size=(self.min_size,self.min_size),
                                          color=self.hover_color,
                                          autoLog=False,
                                          opacity=0.0)

        # Create a line stim to use for pen position graphic
        self.pen_tilt_line = visual.Line(win, units='norm', start=[0, 0],
                                         lineWidth=self.tiltline_width,
                                         end=[0, 0],
                                         colorSpace='rgb255',
                                         lineColor=self.tiltline_color,
                                         opacity=0.0)
        #self.pen_tilt_line.opacity=0.0

    def updateFromEvent(self, evt):
        """Update the pen position and tilt graphics based on the data from
        a wintab sample event.

        :param evt: iohub wintab sample event
        :return:
        """
        # update the pen position stim based on
        # the last tablet event's data
        if evt.pressure > 0:
            # pen is touching tablet surface
            self.pen_guass.color = self.touching_color
        else:
            # pen is hovering just above tablet surface
            self.pen_guass.color = self.hover_color

        if evt.device.axis['pressure']['supported']:
            # change size of pen position blob based on samples pressure
            # value
            pnorm = evt.pressure / evt.device.axis['pressure']['range']
            self.pen_guass.size = self.min_size + pnorm * self.size_range

        # set the position of the gauss blob to be the pen x,y value converted
        # to norm screen coords.
        self.pen_guass.pos = evt.getNormPos()

        # if supported, update all graphics opacity based on the samples z value
        # otherwise opacity is always 1.0
        if evt.device.axis['z']['supported']:
            z = evt.device.axis['z']['range'] - evt.z
            znorm = z / evt.device.axis['z']['range']
            sopacity = self.min_opacity + znorm * (1.0 - self.min_opacity)
            self.pen_guass.opacity = self.pen_tilt_line.opacity = sopacity
        else:
            self.pen_guass.opacity = self.pen_tilt_line.opacity = 1.0

        # Change the tilt line start position to == pen position
        self.pen_tilt_line.start = self.pen_guass.pos

        # Change the tilt line end position based on samples tilt value
        # If tilt is not supported, it will always return 0,0
        # so no line is drawn.
        t1, t2 = evt.tilt
        pen_tilt_xy = 0, 0
        if t1 != t2 != 0:
            pen_tilt_xy = t1 * math.sin(t2), t1 * math.cos(t2)

        pen_pos = self.pen_guass.pos
        tiltend = (pen_pos[0] + pen_tilt_xy[0]*self.tiltline_scalar,
                   pen_pos[1] + pen_tilt_xy[1]*self.tiltline_scalar)
        self.pen_tilt_line.end = tiltend

    def draw(self):
        """Draw the PenPositionStim to the opengl back buffer. This needs
        to be called prior to calling win.flip() for the stim is to be
        displayed.

        :return: None

        """
        self.pen_guass.draw()
        self.pen_tilt_line.draw()

    def clear(self):
        """Hide the graphics on the screen, even if they are drawn, by
        setting opacity to 0.

        :return: None

        """
        self.pen_guass.opacity = 0.0
        self.pen_tilt_line.opacity = 0.0

    def __del__(self):
        self.win = None

class PenTracesStim(MinimalStim):
    """Graphics representing where the pen has been moved on the digitizer
    surface. Positions where sample pressure > 0 are included.

    Implemented as a list of visual.ShapeStim, each representing a
    single pen trace/segment (series on pen samples with pressure >
    0). For improved performance, a single pen trace can have
    max_trace_len points before a new ShapeStim is created and made
    the 'current' pen trace'.
    """
    def __init__( self, win, lineWidth=2, lineColor=(0, 0, 0), opacity=1.0,
                 maxlen=256, name=None, autoLog=None, depth=-1000):
        self.depth = depth
        self.win = win
        super(PenTracesStim, self).__init__(name, autoLog)
        # A single pen trace can have at most max_trace_len points.
        self.max_trace_len = maxlen
        # The list of ShapeStim representing pen traces
        self.pentracestim = []
        # The ShapeStim state new / upcoming position points will be added to.
        self.current_pentrace = None
        # A list representation of the current_pentrace.vertices
        self.current_points = []
        # The last pen position added to a pen trace.
        self.last_pos = [0, 0]

        self.lineWidth=lineWidth
        self.lineColor=lineColor
        self.opacity=opacity

    @property
    def traces(self):
        """List of np arrays, each np array is the set of vertices for one
        pen trace.

        :return: list

        """
        return [pts.vertices for pts in self.pentracestim]

    def updateFromEvents(self, sample_events):
        """
        Update the stim graphics based on 0 - n pen sample events.
        :param sample_events:
        :return: None
        """
        for pevt in sample_events:
            if 'FIRST_ENTER' in pevt.status:
                self.end()
            if pevt.pressure > 0:
                lpx, lpy = self.last_pos
                px, py = pevt.getPixPos(self.win)
                if lpx != px or lpy != py:
                    if len(self.current_points) >= self.max_trace_len:
                        self.end()
                        self.append((lpx, lpy))
                    self.last_pos = (px, py)
                    self.append(self.last_pos)
            else:
                self.end()

    def draw(self):
        """Draws each pen trace ShapeStim to the opengl back buffer. This
        method must be called prior to calling win.flip() if it is to
        appear on the screen.

        :return: None
        """
        for pts in self.pentracestim:
            pts.draw()

    def start(self, first_point):
        """Start a new pen trace, by creating a new ShapeStim, adding it to
        the pentracestim list, and making it the current_pentrace.

        :param first_point: the first point in the ShapStim being craeted.
        :return: None
        """
        self.end()
        self.current_points.append(first_point)
        self.current_pentrace = visual.ShapeStim(self.win,
                                                 units='pix',
                                                 lineWidth=self.lineWidth,
                                                 color=self.lineColor,
                                                 lineColorSpace='rgb255',
                                                 vertices=self.current_points,
                                                 closeShape=False,
                                                 pos=(0, 0),
                                                 size=1,
                                                 ori=0.0,
                                                 opacity=self.opacity,
                                                 interpolate=True)
        self.pentracestim.append(self.current_pentrace)

    def end(self):
        """Stop using the current_pentrace ShapeStim. Next time a pen
        sample position is added to the PenTracesStim instance, a new
        ShapeStim will created and added to the pentracestim list.

        :return: None
        """
        self.current_pentrace = None
        self.current_points = []
        self.last_pos = [0, 0]

    def append(self, pos):
        """Add a pen position (in pix coords) to the current_pentrace
        ShapeStim vertices.

        :param pos: (x,y) tuple
        :return: None
        """
        if self.current_pentrace is None:
            self.start(pos)
        else:
            self.current_points.append(pos)
            self.current_pentrace.vertices = self.current_points

    def clear(self):
        """Remove all ShapStim being used. Next time this stim is drawn, no
        pen traces will exist.

        :return:
        """
        self.end()
        del self.pentracestim[:]

    def __del__(self):
        self.clear()
        self.win = None

#
# Pen position validation process code
#

class ScreenPositionValidation(object):
    NUM_VALID_SAMPLES_PER_TARG = 100
    TARGET_TIMEOUT = 10.0

    def __init__(self, win, io, target_stim=None, pos_grid=None,
                 display_pen_pos=True, force_quit=True, intro_title=None,
                 intro_text1=None, intro_text2=None, intro_target_pos=None):
        """ScreenPositionValidation is used to perform a pen position
        accuracy test for an iohub wintab device.

        :param win: psychopy Window instance to ude for the validation graphics
        :param io: iohub connection instance
        :param target_stim: None to use default, or  psychopy.iohub.util.targetpositionsequence.TargetStim instance
        :param pos_grid: None to use default, or  psychopy.iohub.util.targetpositionsequence.PositionGrid instance
        :param display_pen_pos: True to add calculated pen position graphic
        :param force_quit: Not Used
        :param intro_title: None to use default, str or unicode to set the text used for the introduction screen title, or an instance of psychopy.visual.TextStim
        :param intro_text1: None to use default, str or unicode to set the text used for the introduction text part 1, or an instance of psychopy.visual.TextStim
        :param intro_text2: None to use default, str or unicode to set the text used for the introduction text part 2, or an instance of psychopy.visual.TextStim
        :param intro_target_pos: None to use default, or (x,y) position to place the target graphic on the introduction screen. (x,y) position must be specified in 'norm' coordinate space.
        :return:
        """

        from psychopy.iohub.util.targetpositionsequence import TargetStim, PositionGrid

        self.win = win
        self.io = io
        self._lastPenSample = None
        self._targetStim = target_stim
        self._positionGrid = pos_grid
        self._forceQuit = force_quit
        self._displayPenPosition = display_pen_pos

        # IntroScreen Graphics
        intro_graphics = self._introScreenGraphics = OrderedDict()

        # Title Text
        title_stim = visual.TextStim(self.win, units='norm',
                                     pos=(0, .9),
                                     height=0.1,
                                     text='Pen Position Validation')
        if isinstance(intro_title, str):
            title_stim.setText(intro_title)
        elif isinstance(intro_title, visual.TextStim):
            title_stim = intro_title
        intro_graphics['title'] = title_stim

        # Intro Text part 1
        text1_stim = visual.TextStim(self.win, units='norm',
                                     pos=(0, .65),
                                     height=0.05,
                                     text='On the following screen, '
                                     'press the pen on the target '
                                     'graphic when it appears, '
                                     'as accurately as '
                                     'possible, until the target '
                                     'moves to a different '
                                     'location. Then press at the '
                                     'next target location. '
                                     'Hold the stylus in exactly '
                                     'the same way as you would '
                                     'hold a pen for normal '
                                     'handwriting.',
                                     wrapWidth=1.25
                                     )

        if isinstance(intro_text1, str):
            text1_stim.setText(intro_text1)
        elif isinstance(intro_text1, visual.TextStim):
            text1_stim = intro_text1
        intro_graphics['text1'] = text1_stim

        # Intro Text part 2
        text2_stim = visual.TextStim(self.win, units='norm',
                                     pos=(0, -0.2),
                                     height=0.066,
                                     color='green',
                                     text='Press the pen on the above '
                                     'target to start the '
                                     'validation, or the ESC key '
                                     'to skip the procedure.')
        if isinstance(intro_text2, str):
            text2_stim.setText(intro_text2)
        elif isinstance(intro_text2, visual.TextStim):
            text2_stim = intro_text2
        intro_graphics['text2'] = text2_stim

        self._penStim = None
        if self._displayPenPosition:
            # Validation Screen Graphics
            self._penStim = visual.Circle(self.win,
                                          radius=4,
                                          fillColor=[255, 0, 0],
                                          lineColor=[255, 0, 0],
                                          lineWidth=0,
                                          edges=8,  # int(np.pi*radius),
                                          units='pix',
                                          colorSpace='rgb255',
                                          opacity=0.9,
                                          contrast=1,
                                          interpolate=True,
                                          autoLog=False)

        if self._targetStim is None:
            self._targetStim = TargetStim(win,
                                          radius=16,
                                          fillcolor=[64, 64, 64],
                                          edgecolor=[192, 192, 192],
                                          edgewidth=1,
                                          dotcolor=[255, 255, 255],
                                          dotradius=3,
                                          units='pix',
                                          colorspace='rgb255',
                                          opacity=1.0,
                                          contrast=1.0
                                          )
        if intro_target_pos:
            self._targetStim.setPos(intro_target_pos)

        intro_graphics['target'] = self._targetStim

        if self._positionGrid is None:
            self._positionGrid = PositionGrid(
                winSize=win.monitor.getSizePix(),
                shape=[
                    3,
                    3],
                scale=0.9,
                posList=None,
                noiseStd=None,
                firstposindex=0,
                repeatfirstpos=True)

        # IntroScreen Graphics
        finished_graphics = self._finsihedScreenGraphics = OrderedDict()

        finished_graphics['title'] = visual.TextStim(
            self.win, units='norm', pos=(
                0, .9), height=0.1, text='Validation Complete')
        finished_graphics['result_status'] = visual.TextStim(
            self.win, units='norm', pos=(
                0, .7), height=0.07, color='blue', text='Result: {}')
        finished_graphics['result_stats'] = visual.TextStim(self.win, units='norm', pos=(
            0, .6), height=0.05, text='{}/{} Points Validated. Min, Max, Mean Errors: {}, {}, {}')
        finished_graphics['exit_text'] = visual.TextStim(
            self.win, units='norm', pos=(
                0, .5), height=0.05, text='Press any key to continue...')

    @property
    def targetStim(self):
        return self._targetStim

    @targetStim.setter
    def targetStim(self, ts):
        self._targetStim = ts

    @property
    def positionGrid(self):
        return self._positionGrid

    @positionGrid.setter
    def positionGrid(self, ts):
        self._positionGrid = ts

    def _enterIntroScreen(self):
        kb = self.io.devices.keyboard
        pen = self.io.devices.tablet

        exit_screen = False
        hitcount = 0
        pen.reporting = True
        kb.getPresses()

        while exit_screen is False:
            for ig in self._introScreenGraphics.values():
                ig.draw()
            samples = pen.getSamples()
            if samples:
                self._drawPenStim(samples[-1])
                spos = samples[-1].getPixPos(self.win)
                if samples[-1].pressure > 0 and \
                        self._introScreenGraphics['target'].contains(spos):
                    if hitcount > 10:
                        exit_screen = True
                    hitcount = hitcount + 1
                else:
                    hitcount = 0
            self.win.flip()
            if 'escape' in kb.getPresses():
                exit_screen = True
                pen.reporting = False
                break

        pen.reporting = False
        return True

    def _enterValidationSequence(self):
        val_results = dict(target_data=dict(), avg_err=0, min_err=1000,
                           max_err=-1000, status='PASSED', point_count=0,
                           ok_point_count=0)

        self._lastPenSample = None

        #kb = self.io.devices.keyboard
        pen = self.io.devices.pen

        self._positionGrid.randomize()

        pen.reporting = True
        for tp in self._positionGrid:
            self._targetStim.setPos(tp)
            self._targetStim.draw()
            targ_onset_time = self.win.flip()

            pen.clearEvents()

            val_sample_list = []

            while len(val_sample_list) < self.NUM_VALID_SAMPLES_PER_TARG:
                if core.getTime() - targ_onset_time > self.TARGET_TIMEOUT:
                    break
                self._targetStim.draw()

                samples = pen.getSamples()
                for s in samples:
                    spos = s.getPixPos(self.win)
                    if s.pressure > 0 and self.targetStim.contains(spos):
                        dx = math.fabs(tp[0] - spos[0])
                        dy = math.fabs(tp[1] - spos[1])
                        perr = math.sqrt(dx * dx + dy * dy)
                        val_sample_list.append((spos[0], spos[1], perr))
                    else:
                        val_sample_list = []

                if samples:
                    self._drawPenStim(samples[-1])
                    self._lastPenSample = samples[-1]
                elif self._lastPenSample:
                    self._drawPenStim(self._lastPenSample)
                self.win.flip()

            tp = int(tp[0]), int(tp[1])
            val_results['target_data'][tp] = None
            val_results['point_count'] = val_results['point_count'] + 1

            if val_sample_list:
                pos_acc_array = np.asarray(val_sample_list)
                serr_array = pos_acc_array[:, 2]

                targ_err_stats = val_results['target_data'][tp] = dict()
                targ_err_stats['samples'] = pos_acc_array
                targ_err_stats['count'] = len(val_sample_list)
                targ_err_stats['min'] = serr_array.min()
                targ_err_stats['max'] = serr_array.max()
                targ_err_stats['mean'] = serr_array.mean()
                targ_err_stats['median'] = np.median(serr_array)
                targ_err_stats['stdev'] = serr_array.std()

                val_results['min_err'] = min(
                    val_results['min_err'], targ_err_stats['min'])
                val_results['max_err'] = max(
                    val_results['max_err'], targ_err_stats['max'])

                val_results['avg_err'] = val_results[
                    'avg_err'] + targ_err_stats['mean']
                val_results['ok_point_count'] = val_results[
                    'ok_point_count'] + 1
            else:
                val_results['status'] = 'FAILED'

            self._lastPenSample = None

        if val_results['ok_point_count'] > 0:
            val_results['avg_err'] = val_results[
                'avg_err'] / val_results['ok_point_count']

        pen.reporting = False

        return val_results

    def _enterFinishedScreen(self, results):
        kb = self.io.devices.keyboard

        status = results['status']
        ok_point_count = results['ok_point_count']
        min_err = results['min_err']
        max_err = results['max_err']
        avg_err = results['avg_err']
        point_count = results['point_count']
        self._finsihedScreenGraphics['result_status'].setText(
            'Result: {}'.format(status))

        self._finsihedScreenGraphics['result_stats'].setText(
            '%d/%d '
            'Points Validated.'
            'Min, Max, Mean '
            'Errors: '
            '%.3f, %.3f, %.3f'
            '' %
            (ok_point_count, point_count, min_err, max_err, avg_err))
        for ig in self._finsihedScreenGraphics.values():
            ig.draw()
        self.win.flip()
        kb.clearEvents()

        while not kb.getPresses():
            for ig in self._finsihedScreenGraphics.values():
                ig.draw()
            self.win.flip()

    def _drawPenStim(self, s):
        if self._displayPenPosition:
            spos = s.getPixPos(self.win)
            if spos:
                self._penStim.setPos(spos)
                if s.pressure == 0:
                    self._penStim.setFillColor([255, 0, 0])
                    self._penStim.setLineColor([255, 0, 0])
                else:
                    self._penStim.setFillColor([0, 0, 255])
                    self._penStim.setLineColor([0, 0, 255])

                self._penStim.draw()

    def run(self):
        """Starts the validation process. This function will not return
        until the validation is complete. The validation results are
        returned in dict format.

        :return: dist containing validation results.

        """

        continue_val = self._enterIntroScreen()

        if continue_val is False:
            return None

        # delay about 0.5 sec before staring validation
        ftime = self.win.flip()
        while core.getTime() - ftime < 0.5:
            self.win.flip()
            self.io.clearEvents()

        val_results = self._enterValidationSequence()

        # delay about 0.5 sec before showing validation end screen
        ftime = self.win.flip()
        while core.getTime() - ftime < 0.5:
            self.win.flip()
            self.io.clearEvents()

        self._enterFinishedScreen(val_results)
        self.io.clearEvents()
        self.win.flip()

        return val_results

        # returning None indicates to experiment that the vaidation process
        # was terminated by the user
        # return None

    def free(self):
        self.win = None
        self.io = None
        self._finsihedScreenGraphics.clear()
        self._introScreenGraphics.clear()
        self._targetStim = None
        self._penStim = None

    def __del__(self):
        self.free()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Aug 19 16:02:24 2021

@author: af1tang
"""
import random
import numpy as np
import torch, torch.nn as nn
from convogym.utils._visualization import display_dialog_history
from convogym.prefixes import load_persona_facts

class RankingLoss(nn.Module):
    """
    Implements triplet loss with negative sampling.

    Parameters
    -------
    state_estimator : StateEstimator object or Callable
        A state estimator model that maps from dialog history text -> embedding vector to represent state information.

    swaps : bool, optional
        Whether to swap anchor and positive sample roles during triplet loss.

    margin : float, optional
        Margin parameter used for triplet loss.

    k : int, optional
        Number of total candidates to use during negative sampling. The default is 20.

    reduction : string, optional
        See https://pytorch.org/docs/stable/generated/torch.nn.TripletMarginLoss.html for details.
        > Specifies the reduction to apply to the output:
        >    ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
        >    ``'mean'``: the sum of the output will be divided by the number of
        >    elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
        >    and :attr:`reduce` are in the process of being deprecated, and in the meantime,
        >    specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
    """
    def __init__(self, state_estimator, swap=False, margin=1.0, k = 20, reduction='sum'):
        super(RankingLoss, self).__init__()
        self.ranking_loss = nn.TripletMarginLoss(margin=margin, swap=swap, reduction=reduction)
        self.state_estimator = state_estimator
        self.k = k

    def _calculate_loss(self, inp, pos_cands, neg_cands):
        """
        Calculates the triplet loss for a single sample.

        Parameters
        ----------
        inp : torch.Tensor
            A single state embedding.

        pos_cands : torch.Tensor
            The set of persona embeddings that co-occur with the dialog.

        neg_cands : torch.Tensor
            A set of persona embeddings negatively sampled from a list of candidates.

        Returns
        -------
        total_loss : torch.Tensor
            Loss on single sample over candidate comparisons.

        """
        num_pos, num_neg = pos_cands.size(0), neg_cands.size(0)
        total_loss = 0.
        count = 0
        for i in range(0, num_neg, num_pos):
            if pos_cands.size(0) == neg_cands[i:i+num_pos].size(0):
                total_loss = total_loss + self.ranking_loss(inp, pos_cands, neg_cands[i:i+num_pos])
                count +=1
        return total_loss

    def forward(self, states, positives, candidates):
        """
        Calculates the dialog history embeddings and outputs the triplet loss on batch.

        Parameters
        ----------
        states : torch.Tensor
            Embeddings of input dialog trajectories.

        positives : List of list strings
            List of person 2 persona facts corresponding to each conversation in the batch.

        candidates : List of strings
            List of supporting set of persona facts to generate candidates from.

        Returns
        -------
        loss : float
            Triplet loss on batch.

        """
        # generate candidates using negative sampling
        pos_targets, neg_targets = self._generate_candidates(candidates, positives)
        # calculate ranking loss over candidates
        loss = sum([self._calculate_loss(states[i].unsqueeze(0), pos_targets[i], neg_targets[i])
                                for i in range(len(states))])
        return loss

    def score(self, states, positives, candidates):
        """
        Calculates precision and recall metrics based on embeddings of dialog trajectories.

        Parameters
        ----------
        states : torch.Tensor
            Embeddings of input dialog trajectories.

        positives : List of list strings
            List of person 2 persona facts corresponding to each conversation in the batch.

        candidates : List of strings
            List of supporting set of persona facts to generate candidates from.

        Returns
        -------
        prec1 : float
            # hits with the top guess / # samples in the batch. Measures quality of the top ranked guess.

        prec5 : float
            Average precsion@5 score, defined as # hits in top 5 guesses / 5.

        rec5 : float
            Average recall@5 score, defined as # hits in top 5 guesses / # positives. Measures diversity of guesses.

        rec10 : float
            Average recall@5 score, defined as # hits in top 10 guesses / # positives. Measures diversity of guesses.

        Examples
        -------
        >>> from convogym.representers import StateEstimator
        >>> from convogym._decoder import model
        >>> dialog_history = ["hi how are you doing?",  # corresponds to person 2
                              "good, just got back from playing baseball.",
                              "i like it. i also like to play the piano."]
        >>> persona = ["i play the piano.", "i'm a little league all-star."]
        >>> state_estimator = StateEstimator(model)
        >>> ranking_loss = RankingLoss()
        >>> states = state_estimator(dialog_history) # outputs a (1 x 1024) tensor
        >>> prec1, prec5, rec5, rec10 = ranking_loss.score(states, [persona])
        >>> print("precision@5: %.2f, recall@5: %.2f" %(prec5, rec5))
        precision@5: 0.40, recall@5: 1.00

        """
        pos_targets, neg_targets = self._generate_candidates(candidates, positives)
        # calculate precision@k and recall@k metrics
        prec1 = [self._score_triplet(states[i].unsqueeze(0), pos_targets[i], neg_targets[i], at_k=1)
                 for i in range(len(states))]
        prec5 = [self._score_triplet(states[i].unsqueeze(0), pos_targets[i], neg_targets[i], at_k=5)
                 for i in range(len(states))]
        rec5 = [self._score_triplet(states[i].unsqueeze(0), pos_targets[i],
                neg_targets[i], at_k=5, use_recall=True) for i in range(len(states))]
        rec10 = [self._score_triplet(states[i].unsqueeze(0), pos_targets[i],
                neg_targets[i], at_k=10, use_recall=True) for i in range(len(states))]

        return prec1, prec5, rec5, rec10

    def _generate_candidates(self, candidates, positives):
        """
        Generates negative samples for triplet rankings.
        """
        negatives = [random.sample( set(candidates) - set(pos_set), self.k) for pos_set in positives]
        pos_samples = [self.state_estimator._get_persona_embeddings(p) for p in positives]
        neg_samples = [self.state_estimator._get_persona_embeddings(n) for n in negatives]
        return pos_samples, neg_samples

    def _score_triplet(self, x, pos,neg, at_k, use_recall=False):
        """
        Scores the ranking model based on the ranked list of generated candidates.

        If use_recall then recall@k is used. Otherwise precision@k is used.
        """
        dist = nn.CosineSimilarity()
        scores = []
        for k,v in enumerate(pos):
            scores.append((k,1, dist(x,v.unsqueeze(0)).item()))
        for k,v in enumerate(neg):
            scores.append((k+len(pos), 0, dist(x,v.unsqueeze(0)).item()))
        scores = sorted(scores, key = lambda x: x[-1], reverse=True)
        picks = [s[1] for s in scores[0:at_k]]
        # precision or recall @ k
        score = sum(picks) / len(pos) if use_recall else sum(picks) / len(picks)
        return score

# interface with gym callback objects
class Reward:
    """
    A helper object for callbacks that outputs turn-level rewards based on state information at each turn.

    In the base Reward, a custom criterion function is

    Parameters
    -------
    state_estimator : StateEstimator object or Callable
        A state estimator model that maps from dialog history text -> embedding vector to represent state information.

    criterion : nn.Module or Callable
        Maps from state vector to a scalar reward. The default is None.

        If want (state, action) -> reward, simply concatenate action to state vector as input.

    """
    def __init__(self, state_estimator, criterion):
        self.state_estimator = state_estimator
        self.criterion = criterion

    def calculate_reward(self, scb, mcb):
        """
        Obtains reward from given persona and state information and passes it to the state callback.

        Attributes
        -------
        scb.reward : float
            A scalar reward for current turn.

        scb.rewards: List of floats
            List of reward for each turn.

        """
        loss = self.criterion(scb.state.view(1,-1))
        scb.reward = - loss.item()
        scb.rewards.append(scb.reward)
        return scb, mcb

    def score_trajectory(self, scb, mcb):
        """
        Not implemented in base class.
        """
        raise NotImplementedError

class ManualReward(Reward):
    """
    Implements a wrapper on the Reward base object.

    At each turn, the user is manually prompted to score the dialog up to the current turn.
    (state, action, next_state, reward) samples collected this way can be used to learn reward function that approximates user scoring (i.e., a criterion).
    """
    def __init__(self, state_estimator, *args, **kwargs):
        super().__init__(state_estimator=state_estimator, criterion=None)

    def calculate_reward(self, scb, mcb):
        """
        Obtains reward from given persona and state information and passes it to the state callback.

        Attributes
        -------
        scb.reward : float
            A scalar reward for current turn.

        scb.rewards: List of floats
            List of reward for each turn.

        """
        reward = None
        while reward is None:
            display_dialog_history(mcb.dialog_hx, self.state_estimator.tokenizer)
            try:
                reward = float(input( "Enter a reward (float) for current dialog: " ))
            except:
                reward = None
        scb.reward = reward
        scb.rewards.append(scb.reward)
        return scb, mcb

    def score_trajectory(self, scb, mcb):
        """
        No additional metrics for this wrapper.
        """
        return scb, mcb

class RankingReward(Reward):
    """
    Implements a wrapper on the Reward base object.

    Calculates various precision@k and recall@k metrics based on the RankingLoss criterion.

    Parameters
    -------
    state_estimator : StateEstimator object or Callable
        A state estimator model that maps from dialog history text -> embedding vector to represent state information.

    path_to_persona_facts : os.path or string, optional
        Path to list of persona facts. Used to generate candidates for ranking / scoring.
        If none, a list of persona facts is constructed from the PersonaChat data set and saved locally.
        The default is none.
    """
    def __init__(self, state_estimator, path_to_persona_facts=None):
        super().__init__(state_estimator, None)
        self.criterion = RankingLoss(state_estimator)
        self.candidates = load_persona_facts(path_to_persona_facts)
        self.metrics = {'prec1s': [], 'prec5s': [],
                        'rec5s': [], 'rec10s': [], 'rewards': []}
    def calculate_reward(self, scb, mcb):
        """
        Obtains reward from given persona and state information and passes it to the state callback.

        Attributes
        -------
        scb.reward : float
            A scalar reward for current turn.

        scb.rewards: List of floats
            List of reward for each turn.

        """
        loss = self.criterion(states=scb.state.view(1,-1),
                              positives=[scb.personas],
                              candidates=self.candidates)
        scb.reward = - loss.item()
        scb.rewards.append(scb.reward)
        return scb, mcb

    def score_trajectory(self, scb, mcb):
        """
        Calculates precision@k and recall@k for current dialog trajectory.

        Parameters
        ----------
        scb : StateCb object
            State callback object.

        Metrics
        -------
        prec1 : float
            # hits with the top guess / # samples in the batch. Measures quality of the top ranked guess.

        prec5 : float
            Average precsion@5 score, defined as # hits in top 5 guesses / 5.

        rec5 : float
            Average recall@5 score, defined as # hits in top 5 guesses / # positives. Measures diversity of guesses.

        rec10 : float
            Average recall@5 score, defined as # hits in top 10 guesses / # positives. Measures diversity of guesses.

        """
        prec1, prec5, rec5, rec10 = self.criterion.score(states=scb.state.view(1,-1),
                                                        positives=[scb.personas],
                                                        candidates=self.candidates)
        # log data
        self.metrics['prec1s'].extend(prec1)
        self.metrics['prec5s'].extend(prec5)
        self.metrics['rec5s'].extend(rec5)
        self.metrics['rec10s'].extend(rec10)
        self.metrics['rewards'].append(scb.rewards)
        print("prec@1: %.2f | prec@5: %.2f | rec@5: %.2f | rec@10: %.2f" % (prec1[0], prec5[0], rec5[0], rec10[0]))
        print("avg p@1: %.2f | avg p@5: %.2f | avg r@5: %.2f | avg. r@10: %.2f" % (np.mean(self.metrics['prec1s']),
                                                                                np.mean(self.metrics['prec5s']),
                                                                                np.mean(self.metrics['rec5s']),
                                                                                np.mean(self.metrics['rec10s'])))
        print()


# coding=utf-8
# Copyright 2022 The Deeplab2 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Implements Axial-Attention layers proposed in Axial-DeepLab.

Axial-Attention factorizes 2D self-attention into two 1D self-attentions, so
that it can be applied on large inputs. Axial-Attention is typically used to
replace 3x3 convolutions in a bottleneck residual block.

[1] Axial-Deeplab: Stand-Alone Axial-Attention for Panoptic Segmentation,
    ECCV 2020 Spotlight.
      Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille,
      Liang-Chieh Chen.
"""

import numpy as np
import tensorflow as tf

from deeplab2.model import utils
from deeplab2.model.layers import activations
from deeplab2.model.layers import positional_encodings

class AxialAttention(tf.keras.layers.Layer):
  """An axial-attention layer."""

  def __init__(self,
               query_shape=129,
               memory_flange=32,
               total_key_depth=512,
               total_value_depth=1024,
               num_heads=8,
               name='axial_attention',
               use_query_rpe_similarity=True,
               use_key_rpe_similarity=True,
               use_content_similarity=True,
               retrieve_value_rpe=True,
               retrieve_value_content=True,
               initialization_std_for_query_key_rpe=1.0,
               initialization_std_for_value_rpe=1.0,
               self_attention_activation='softmax',
               bn_layer=tf.keras.layers.BatchNormalization,
               conv_kernel_weight_decay=0.0):
    """Initializes an axial-attention layer.

    This function is designed to support both global and local axial-attention
    in a unified way. If query_shape is larger than the length of input, a
    global attention is applied. If query_shape is smaller than the length of
    input, a local attention is applied. In this case, the input is divided into
    blocks of length query_shape, padded by memory_flange on both sides. Then,
    local attention is applied within each query block. The choice of
    query_shape does not affect the output value but affects computation
    efficiency and memory usage. In general, use global attention (large
    query_shape) if possible. Local axial-attention has not been supported yet.

    Args:
      query_shape: An integer, the block size for local axial attention.
        Defaults to 129 since 129 is usually the largest feature map where we do
        global attention (1025 with stride 8, or 2049 with stride 16).
      memory_flange: An integer, the memory flange padded to each query block in
        local attention. It has no effect in global attention. Defaults to 32,
        which is equivalent to a span of 65 in Aixal-DeepLab paper -- A pixel
        can see 32 pixels on its left and 32 pixels on its right.
      total_key_depth: An integer, the total depth of keys, which is also the
        depth of queries and the depth of key (query) positional encodings.
      total_value_depth: An integer, the total depth of the values, which is
        also the depth of value positional encodings.
      num_heads: An integer, the number of heads in multi-head attention.
      name: A string, the name of this axial attention layer.
      use_query_rpe_similarity: A boolean, whether to use the attention
        similarity between the queries and the relative positional encodings.
      use_key_rpe_similarity: A boolean, whether to use the attention similarity
        between the keys and the relative positional encodings.
      use_content_similarity: A boolean, whether to use the content similarity
        between the queries and the keys.
      retrieve_value_rpe: A boolean, whether to retrieve the relative positional
        encodings of the values.
      retrieve_value_content: A boolean, whether to retrieve the content of the
        values.
      initialization_std_for_query_key_rpe: A float, the initialization std for
        the relative positional encodings of the queries and keys.
      initialization_std_for_value_rpe: A float, the initialization std for the
        relative positional encodings of the values.
      self_attention_activation: A string, type of activation function for
        self-attention. Support 'sigmoid' and 'softmax'.
      bn_layer: A tf.keras.layers.Layer that computes the normalization
        (default: tf.keras.layers.BatchNormalization).
      conv_kernel_weight_decay: A float, the weight decay for convolution
        kernels.

    Returns:
      output: A [batch, length, total_value_depth] tensor.

    Raises:
      ValueError: If none of the three similarities (use_query_rpe_similarity,
        use_key_rpe_similarity, use_content_similarity) is used.
      ValueError: If neither of value content or value rpe is retrieved.
      ValueError: If self_attention_activation is not supported.
      ValueError: If total_key_depth is not divisible by num_heads.
      ValueError: If total_value_depth is not divisible by num_heads.
    """
    # Validate the attention similarity choices.
    if not any([
        use_content_similarity, use_key_rpe_similarity, use_query_rpe_similarity
    ]):
      raise ValueError(
          'Should use at least one similarity to compute attention.')

    # Validate the retrieve value choices.
    if not retrieve_value_content and not retrieve_value_rpe:
      raise ValueError('Should retrieve at least one of content or rpe.')

    if total_key_depth % num_heads:
      raise ValueError('Total_key_depth should be divisible by num_heads.')

    if total_value_depth % num_heads:
      raise ValueError('Total_value_depth should be divisible by num_heads.')

    super(AxialAttention, self).__init__(name=name)
    self._query_shape = query_shape
    self._memory_flange = memory_flange
    self._total_key_depth = total_key_depth
    self._total_value_depth = total_value_depth
    self._num_heads = num_heads
    self._use_query_rpe_similarity = use_query_rpe_similarity
    self._use_key_rpe_similarity = use_key_rpe_similarity
    self._use_content_similarity = use_content_similarity
    self._retrieve_value_rpe = retrieve_value_rpe
    self._retrieve_value_content = retrieve_value_content
    self._initialization_std_for_query_key_rpe = (
        initialization_std_for_query_key_rpe)
    self._initialization_std_for_value_rpe = initialization_std_for_value_rpe
    self._self_attention_activation = self_attention_activation
    self._conv_kernel_weight_decay = conv_kernel_weight_decay

    self._batch_norm_qkv = bn_layer(axis=-1, name='batch_norm_qkv')
    self._batch_norm_similarity = bn_layer(
        axis=[0, 2], name='batch_norm_similarity')
    self._batch_norm_retrieved_output = bn_layer(
        axis=[0, 2, 4], name='batch_norm_retrieved_output')

    self._key_depth_per_head = total_key_depth // num_heads
    self._attention_activate_fn = activations.get_activation(
        self_attention_activation)

  def build(self, input_shape):
    """Builds axial-attention layer weights.

    Args:
      input_shape: An integer list of length 3, the shape of the input tensor.

    Raises:
      NotImplementedError: Local axial-attention has not been implemented. It is
        triggered if query_shape is less than input_shape.
    """

    # Apply global attention if query_shape is larger than input_shape[1].
    if self._query_shape >= input_shape[1]:
      self._query_shape = input_shape[1]
      self._memory_flange = 0
    else:
      raise NotImplementedError('Local axial attention has not been '
                                'implemented yet.')
    self._memory_shape = self._query_shape + 2 * self._memory_flange

    # Compute query key value with one convolution and an optional batch norm.
    # The initialization std is standard transformer initialization (without
    # batch norm), as used in SASA and ViT. In our case, we use batch norm by
    # default, so it does not require careful tuning. If one wants to remove
    # all batch norms in axial attention, this standard initialization should
    # still be good, but a more careful initialization is encouraged.
    self.qkv_kernel = self.add_weight(
        name='qkv_kernel',
        shape=[input_shape[-1],
               self._total_key_depth * 2 + self._total_value_depth],
        initializer=tf.keras.initializers.TruncatedNormal(
            stddev=input_shape[-1]**-0.5),
        regularizer=tf.keras.regularizers.l2(self._conv_kernel_weight_decay))

    if self._use_query_rpe_similarity:
      self._query_rpe = positional_encodings.RelativePositionalEncoding(
          self._query_shape,
          self._memory_shape,
          self._key_depth_per_head,
          'query_rpe',
          initialization_std=self._initialization_std_for_query_key_rpe,
          conv_kernel_weight_decay=self._conv_kernel_weight_decay)

    if self._use_key_rpe_similarity:
      self._key_rpe = positional_encodings.RelativePositionalEncoding(
          self._query_shape,
          self._memory_shape,
          self._key_depth_per_head,
          'key_rpe',
          initialization_std=self._initialization_std_for_query_key_rpe,
          conv_kernel_weight_decay=self._conv_kernel_weight_decay)

    if self._retrieve_value_rpe:
      self._value_rpe = positional_encodings.RelativePositionalEncoding(
          self._query_shape,
          self._memory_shape,
          self._total_value_depth // self._num_heads,
          'value_rpe',
          initialization_std=self._initialization_std_for_value_rpe,
          conv_kernel_weight_decay=self._conv_kernel_weight_decay)

  def call(self, input_tensor, training=False):
    """Performs a forward pass.

    Args:
      input_tensor: An input [batch, length, channel] tensor.
      training: A boolean flag indicating whether training behavior should be
        used (default: False).

    Returns:
      output: An output [batch, length, total_value_depth] tensor.
    """
    # Alternatively, the einsum can be implemented as a 1x1 convolution.
    # However, it is not obvious which implementation is more efficient (without
    # careful benchmarking), so we use einsum for its flexibility and
    # consistency with other parts of the function.
    query_key_value = tf.einsum(
        'nlc,cd->nld', input_tensor, self.qkv_kernel, name='compute_qkv')
    query_key_value = self._batch_norm_qkv(query_key_value, training=training)

    # Split query key value.
    query, key, value = tf.split(
        query_key_value,
        [self._total_key_depth, self._total_key_depth, self._total_value_depth],
        axis=-1)

    # Reshape the query, key, and value.
    query = tf.reshape(query, [-1, self._query_shape, self._num_heads,
                               self._key_depth_per_head])
    query = tf.transpose(a=query, perm=[0, 2, 1, 3])
    key = tf.reshape(key, [-1, np.prod(self._memory_shape), self._num_heads,
                           self._key_depth_per_head])
    key = tf.transpose(a=key, perm=[0, 2, 1, 3])
    value = tf.reshape(value, [-1, np.prod(self._memory_shape), self._num_heads,
                               self._total_value_depth // self._num_heads])

    # Gather all similarity logits into a list.
    similarity_logits = []

    # Compute the content similarity term: q * k.
    if self._use_content_similarity:
      content_similarity = tf.einsum(
          'bhld,bhmd->bhlm', query, key, name='content_similarity')
      similarity_logits.append(content_similarity)

    # Compute the query rpe similarity term: q * rpe.
    if self._use_query_rpe_similarity:
      query_rpe = self._query_rpe(None)
      query_rpe_similarity = tf.einsum(
          'bhld,lmd->bhlm', query, query_rpe, name='query_rpe_similarity')
      similarity_logits.append(query_rpe_similarity)

    # Compute the key rpe similarity term: k * rpe.
    if self._use_key_rpe_similarity:
      key_rpe = self._key_rpe(None)
      key_rpe_similarity = tf.einsum(
          'bhmd,lmd->bhlm', key, key_rpe, name='key_rpe_similarity')
      similarity_logits.append(key_rpe_similarity)

    # Apply an optional batch norm to the similarities and sum them.
    similarity_logits = tf.stack(similarity_logits)
    similarity_logits = self._batch_norm_similarity(similarity_logits,
                                                    training=training)
    similarity_logits = tf.reduce_sum(input_tensor=similarity_logits, axis=0)

    # Apply an attention activation function, e.g. softmax.
    weights = self._attention_activate_fn(similarity_logits)

    # Gather retrieved values or rpes into a list.
    retrieve_list = []

    # Retrieve the content of the attended value.
    if self._retrieve_value_content:
      retrieved_content = tf.einsum(
          'bhlm,bmhd->bhld', weights, value, name='retrieve_value_content')
      retrieve_list.append(retrieved_content)

    # Retrieve the relative position of the attended value.
    if self._retrieve_value_rpe:
      value_rpe = self._value_rpe(None)
      retrieved_rpe = tf.einsum(
          'bhlm,lmd->bhld', weights, value_rpe, name='retrieve_value_rpe')
      retrieve_list.append(retrieved_rpe)

    # Apply batch norms to retrieved contents and rpes respectively.
    retrieved_output = tf.stack(retrieve_list)
    retrieved_output = self._batch_norm_retrieved_output(retrieved_output,
                                                         training=training)
    # Additive contents and rpes.
    retrieved_output = tf.reduce_sum(input_tensor=retrieved_output, axis=0)

    # Combine the heads by transposing and reshaping the tensor.
    retrieved_output = utils.transpose_and_reshape_for_attention_operation(
        retrieved_output)

    return retrieved_output

class AxialAttention2D(tf.keras.layers.Layer):
  """Sequentially applies height-axis and width-axis axial-attention."""

  def __init__(self,
               strides=1,
               filters=512,
               name='attention',
               key_expansion=1,
               value_expansion=2,
               query_shape=(129, 129),
               memory_flange=(32, 32),
               **kwargs):
    """Initializes an AxialAttention2D layer.

    Args:
      strides: An integer, the stride for the output, usually 1 or 2.
      filters: An integer, the base number of channels for the layer.
      name: A string, the name of the attention layer.
      key_expansion: A float, the channel expansion ratio for keys.
      value_expansion: A float, the channel expansion ratio for values.
      query_shape: An integer, the maximum query shape for both the height axis
        and the width axis.
      memory_flange: An integer list of length 2. The memory flange for the
        height axis and the width axis.
      **kwargs: A dictionary of keyword arguments passed to height-axis,
        width-axis, and 2D global AxialAttention.

    Returns:
      output: A [batch, strided height, strided width, output_channels] tensor.
    """
    super(AxialAttention2D, self).__init__(name=name)
    total_key_depth = int(round(filters * key_expansion))
    total_value_depth = int(round(filters * value_expansion))
    self._strides = strides
    self._total_key_depth = total_key_depth
    self._total_value_depth = total_value_depth
    self._height_axis = AxialAttention(
        total_key_depth=total_key_depth,
        total_value_depth=total_value_depth,
        query_shape=query_shape[0],
        memory_flange=memory_flange[0],
        name='height_axis',
        **kwargs)
    self._width_axis = AxialAttention(
        total_key_depth=total_key_depth,
        total_value_depth=total_value_depth,
        query_shape=query_shape[1],
        memory_flange=memory_flange[1],
        name='width_axis',
        **kwargs)

  def call(self, inputs, training=False):
    """Performs a forward pass.

    Args:
      inputs: An input [batch, height, width, channel] tensor.
      training: A boolean flag indicating whether training behavior should be
        used (default: False).

    Returns:
      output: An output [batch, strided_height, strided_width,
        filters * value_expansion] tensor.
    """
    _, height, width, channel = inputs.get_shape().as_list()

    # Transpose and reshape the width axis to the batch dimension.
    x = tf.transpose(a=inputs, perm=[0, 2, 1, 3])
    x = tf.reshape(x, [-1, height, channel])
    x = self._height_axis(x, training=training)
    # Reshape and transpose back to a 4D tensor.
    x = tf.reshape(x, [-1, width, height, self._total_value_depth])
    x = tf.transpose(a=x, perm=[0, 2, 1, 3])
    # Height axis striding.
    if self._strides > 1:
      x = x[:, ::self._strides, :, :]

    # Reshape the height axis to the batch dimension.
    _, strided_height, _, _ = x.get_shape().as_list()
    x = tf.reshape(x, [-1, width, self._total_value_depth])
    x = self._width_axis(x, training=training)
    # Reshape back to a 4D tensor.
    x = tf.reshape(x, [-1, strided_height, width, self._total_value_depth])
    # Width axis striding.
    if self._strides > 1:
      x = x[:, :, ::self._strides, :]

    return x

class GlobalAttention2D(tf.keras.layers.Layer):
  """A 2D global attention layer."""

  def __init__(self,
               strides=1,
               filters=512,
               name='attention',
               key_expansion=1,
               value_expansion=2,
               query_shape=(129, 129),
               memory_flange=(32, 32),
               double_global_attention=False,
               **kwargs):
    """Initializes a GlobalAttention2D layer.

    Args:
      strides: An integer, the stride for the output, usually 1 or 2.
      filters: An integer, the base number of channels for the layer.
      name: A string, the name of the attention layer.
      key_expansion: A float, the channel expansion ratio for keys.
      value_expansion: A float, the channel expansion ratio for values.
      query_shape: An integer, the maximum query shape for both the height axis
        and the width axis.
      memory_flange: An integer list of length 2. The memory flange for the
        height axis and the width axis.
      double_global_attention: A boolean, whether to use two global attention
        layers. Two global attention layers match the parameter count to a
        seqentially applied height and width axial attention layer.
      **kwargs: A dictionary of keyword arguments passed to height-axis,
        width-axis, and 2D global AxialAttention.

    Returns:
      output: A [batch, strided height, strided width, output_channels] tensor.

    Raises:
      ValueError: If relative positional encoding is enforced in kwargs.
    """
    if any([kwargs.get('use_query_rpe_similarity', False),
            kwargs.get('use_key_rpe_similarity', False),
            kwargs.get('retrieve_value_rpe', False)]):
      raise ValueError('GlobalAttention2D does not support relative positional '
                       'encodings.')

    super(GlobalAttention2D, self).__init__(name=name)
    total_key_depth = int(round(filters * key_expansion))
    total_value_depth = int(round(filters * value_expansion))
    self._strides = strides
    self._double_global_attention = double_global_attention
    self._total_key_depth = total_key_depth
    self._total_value_depth = total_value_depth

    # Global attention does not support relative positional encodings.
    kwargs['use_query_rpe_similarity'] = False
    kwargs['use_key_rpe_similarity'] = False
    kwargs['retrieve_value_rpe'] = False
    self._kwargs = kwargs

  def build(self, input_shape):
    """Builds global attention layers according to the 4D input_shape."""
    _, height, width, _ = input_shape
    # Implement 2D global attention as 1D axial-attention by flattening the 2D
    # inputs into 1D. We also disable the relative positional encodings in
    # axial attention, so that only content-based attention is used. The query
    # shape is set to height * width, so that the axial attention is global.
    self._global = AxialAttention(
        total_key_depth=self._total_key_depth,
        total_value_depth=self._total_value_depth,
        query_shape=height*width,
        memory_flange=0,
        name='global',
        **self._kwargs)

    # Use two global attention layers in one residual block. This option
    # ensures that global attention models have similar number of layers and
    # parameters as axial-attention models.
    if self._double_global_attention:
      self._global2 = AxialAttention(
          total_key_depth=self._total_key_depth,
          total_value_depth=self._total_value_depth,
          query_shape=height*width,
          memory_flange=0,
          name='global2',
          **self._kwargs)

  def call(self, inputs, training=False):
    """Performs a forward pass.

    Args:
      inputs: An input [batch, height, width, channel] tensor.
      training: A boolean flag indicating whether training behavior should be
        used (default: False).

    Returns:
      output: An output [batch, strided_height, strided_width,
        filters * value_expansion] tensor.
    """
    _, height, width, channel = inputs.get_shape().as_list()

    # Reshape the inputs so that the attention is global 2D.
    x = tf.reshape(inputs, [-1, height * width, channel])

    # Implement 2D global attention as 1D axial-attention by flattening the 2D
    # inputs into 1D. We also disable the relative positional encodings in
    # axial attention, so that only content-based attention is used.
    x = self._global(x, training=training)

    # Use two global attention layers in one residual block. This option
    # ensures that global attention models have the same number of layers and
    # parameters as axial-attention models.
    if self._double_global_attention:
      x = self._global2(x, training=training)
    x = tf.reshape(x, [-1, height, width, self._total_value_depth])
    if self._strides > 1:
      x = x[:, ::self._strides, ::self._strides, :]

    return x


import logging
from typing import List, Tuple, Dict
import os
import random
from pathlib import Path
import shutil
import sys
from datetime import datetime
from tqdm import tqdm
import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, Subset, Dataset
from torch.utils.data.distributed import DistributedSampler

from apex.parallel import DistributedDataParallel as DDP

from tensorboardX import SummaryWriter

from transformers import AutoTokenizer, BertTokenizer

from vilbert.optimization import AdamW, WarmupLinearSchedule
from airbert import Airbert, BERT_CONFIG_FACTORY

from utils.cli import get_parser
from utils.distributed import set_cuda, wrap_distributed_model, get_local_rank
from utils.misc import set_seed, get_output_dir
from utils.dataset import PanoFeaturesReader
from utils.dataset.common import pad_packed
from utils.dataset.beam_dataset import BeamDataset
from utils.dataset.hard_mining import HardMiningDataset
from utils.dataset.perturbate_dataset import PerturbateDataset

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
    stream=sys.stdout,
)
logger = logging.getLogger(__name__)

@torch.no_grad()
def get_score(data_loader: DataLoader, model: nn.Module, default_gpu: bool):
    if isinstance(data_loader.dataset, PerturbateDataset) and isinstance(data_loader.dataset.dataset, HardMiningDataset):
        dataset: HardMiningDataset = data_loader.dataset.dataset
    elif isinstance(data_loader.dataset, HardMiningDataset):
        dataset = data_loader.dataset
    else:
        raise ValueError(f"Unexpected dataset type ({type(data_loader.dataset)})")

    device = next(model.parameters()).device
    is_training = dataset._training

    batch: Tuple[torch.Tensor, ...]
    for batch in data_loader:
        batch = tuple(t.cuda(device=device, non_blocking=True) for t in batch)
        output = model(*get_model_input(batch))
        dataset.post_step(output, list(batch))

    dataset._training = is_training

    if default_gpu and hasattr(dataset, "save"):
        dataset.save()

def compute_metrics(batch: List[torch.Tensor], outputs: Dict[str, torch.Tensor], args) -> Tuple[torch.Tensor, Dict[str, float]]:
    # B, num_cand
    opt_mask = get_mask_options(batch)

    local_rank = get_local_rank(args)
    batch_size = get_batch_size(batch)
    device = opt_mask.device

    # calculate the masked vision loss
    vision_loss = torch.tensor(0, device=device)
    if args.masked_vision:
        # num_image_regions, im_feature
        vision_predictions = outputs["vision"]
        vision_predictions = vision_predictions.view(-1, vision_predictions.shape[2])

        vision_target, vision_target_mask = get_vision_target(batch)

        vision_loss = F.kl_div(
            F.log_softmax(vision_predictions, dim=-1),
            vision_target,
            reduction="none",
        )

        vision_loss *= vision_target_mask.unsqueeze(-1).float()
        numel = max(1, torch.sum(vision_target_mask).item())
        vision_loss = torch.sum(vision_loss) / numel

    # calculate the masked language loss
    linguistic_loss = torch.tensor(0, device=device)
    if args.masked_language:
        # (num_words * num_cand * B), voc_size
        voc_size = outputs["language"].shape[-1]
        linguistic_predictions = outputs["language"].view(-1, voc_size )

        linguistic_target = get_linguistic_target(batch)

        linguistic_loss = F.cross_entropy(
            linguistic_predictions, linguistic_target, ignore_index=-1
        )

    # calculate the trajectory re-ranking loss
    ranking_loss = torch.tensor(0, device=device)
    correct = torch.tensor(0, device=device)
    if not args.no_ranking:
        target = get_target(batch)
        vil_logit = pad_packed(outputs["action"].squeeze(1), opt_mask)
        ranking_loss = F.cross_entropy(vil_logit, target, ignore_index=-1)
        # calculate accuracy
        correct = torch.sum(torch.argmax(vil_logit, 1) == target).float()

    # calculate the final loss
    loss = ranking_loss + vision_loss + linguistic_loss
    if args.gradient_accumulation_steps > 1:
        loss = loss / args.gradient_accumulation_steps

    # calculate accumulated stats
    reduced_vision_loss = vision_loss.detach().float()
    reduced_linguistic_loss = linguistic_loss.detach().float()
    reduced_ranking_loss = ranking_loss.detach().float()
    reduced_loss = loss.detach().float() * args.gradient_accumulation_steps
    reduced_correct = correct.detach().float()
    reduced_batch_size = torch.tensor(batch_size, device=device).float()

    # TODO: skip this `all_reduce` to speed-up runtime
    if local_rank != -1:
        world_size = float(dist.get_world_size())
        reduced_vision_loss /= world_size
        reduced_linguistic_loss /= world_size
        reduced_ranking_loss /= world_size
        reduced_loss /= world_size
        dist.all_reduce(reduced_vision_loss, op=dist.ReduceOp.SUM) # type: ignore
        dist.all_reduce(reduced_linguistic_loss, op=dist.ReduceOp.SUM) # type: ignore
        dist.all_reduce(reduced_ranking_loss, op=dist.ReduceOp.SUM) # type: ignore
        dist.all_reduce(reduced_loss, op=dist.ReduceOp.SUM) # type: ignore
        dist.all_reduce(reduced_correct, op=dist.ReduceOp.SUM) # type: ignore
        dist.all_reduce(reduced_batch_size, op=dist.ReduceOp.SUM) # type: ignore

    reduced_metrics = {
        "loss/train": reduced_loss,
        "loss/ranking": reduced_ranking_loss,
        "loss/vision": reduced_loss,
        "loss/linguistic": reduced_vision_loss,
        "accuracy/train": reduced_correct / reduced_batch_size,
    }

    # if default_gpu and args.debug:
    #     logger.info(
    #         f"[train] step: {step + 1} "
    #         f"vision loss: {reduced_vision_loss:0.2f} "
    #         f"language loss: {reduced_linguistic_loss:0.2f} "
    #         f"ranking loss: {reduced_ranking_loss:0.2f} "
    #         f"loss: {reduced_loss:0.2f} "
    #         f"accuracy: {reduced_correct / reduced_batch_size:0.2f} "
    #         f"lr: {scheduler.get_lr()[0]:0.1e}"
    #     )

    return loss, reduced_metrics

def main():
    # ----- #
    # setup #
    # ----- #

    # command line parsing
    parser = get_parser(training=True)
    args = parser.parse_args()

    # validate command line arguments
    if not (args.masked_vision or args.masked_language) and args.no_ranking:
        parser.error(
            "No training objective selected, add --masked_vision, "
            "--masked_language, or remove --no_ranking"
        )

    set_seed(args)

    # get device settings
    default_gpu, rank, device = set_cuda(args)

    # create output directory
    save_folder = get_output_dir(args)
    if default_gpu:
        save_folder.parent.mkdir(exist_ok=True, parents=True)

        print(args)

    # ------------ #
    # data loaders #
    # ------------ #

    # tokenizer = AutoTokenizer.from_pretrained(args.bert_tokenizer)
    tokenizer = BertTokenizer.from_pretrained(args.bert_tokenizer)
    if not isinstance(tokenizer, BertTokenizer):
        raise ValueError("fix mypy")
    features_reader = PanoFeaturesReader(args.img_feature)
    train_params = {}
    vln_path = f"data/task/{args.prefix}R2R_train.json"

    if args.training_mode == "provided":
        train_params["beam_path"] = f"data/beamsearch/{args.beam_prefix}beams_train.json"
        if default_gpu:
            logger.info("using provided training trajectories")
            logger.info(f"VLN path: {vln_path}")
            logger.info(f"Beam path: {train_params['beam_path']}")

        if args.hard_mining:
            TrainDataset = HardMiningDataset
            train_params["save_folder"] = str(save_folder)
        else:
            TrainDataset = BeamDataset
    else:
        raise ValueError(f"Unknown training_mode for {args.training_mode}")

    if default_gpu:
        logger.info("Loading train dataset")

    train_dataset: Dataset = TrainDataset(
        **train_params,
        vln_path=vln_path,
        tokenizer=tokenizer,
        features_reader=features_reader,
        max_instruction_length=args.max_instruction_length,
        max_path_length=args.max_path_length,
        max_num_boxes=args.max_num_boxes,
        num_beams=args.num_beams_train,
        num_beams_strict=False,
        training=True,
        masked_vision=args.masked_vision,
        masked_language=args.masked_language,
        default_gpu=default_gpu,
        ground_truth_trajectory=args.ground_truth_trajectory,
        highlighted_language=args.highlighted_language,
        shuffle_visual_features=args.shuffle_visual_features,
        num_negatives=args.num_negatives,
        shuffler=args.shuffler,
    )

    #TODO add perturbation inside BeamDataset?
    if args.perturbation:
        if not isinstance(train_dataset, BeamDataset):
            raise ValueError(f"Unexpected dataset type ({type(train_dataset)})")
        train_dataset = PerturbateDataset(
            dataset=train_dataset,
            num_negatives=args.num_negatives,
            perturbate_path=f"data/task/{args.prefix}R2R_train.json",
            shortest_path=args.shortest_path,
            highlighted_perturbations=args.highlighted_perturbations,
        )

    val_seen_params = {}
    val_unseen_params = {}

    TestDataset = BeamDataset

    if default_gpu:
        logger.info("Loading val datasets")

    val_seen_dataset = TestDataset(
        **val_seen_params,
        vln_path=f"data/task/{args.prefix}R2R_val_seen.json",
        beam_path=f"data/beamsearch/{args.beam_prefix}beams_val_seen.json",
        tokenizer=tokenizer,
        features_reader=features_reader,
        max_instruction_length=args.max_instruction_length,
        max_path_length=args.max_path_length,
        max_num_boxes=args.max_num_boxes,
        num_beams=args.num_beams,
        num_negatives=args.num_negatives,
        num_beams_strict=False,
        training=False,
        masked_vision=False,
        masked_language=False,
        default_gpu=default_gpu,
        ground_truth_trajectory=False,
        highlighted_language=args.highlighted_language,
        shuffle_visual_features=False,
    )

    val_unseen_dataset = TestDataset(
        **val_unseen_params,
        vln_path=f"data/task/{args.prefix}R2R_val_unseen.json",
        beam_path=f"data/beamsearch/{args.beam_prefix}beams_val_unseen.json",
        tokenizer=tokenizer,
        features_reader=features_reader,
        max_instruction_length=args.max_instruction_length,
        max_path_length=args.max_path_length,
        max_num_boxes=args.max_num_boxes,
        num_beams=args.num_beams,
        num_beams_strict=False,
        training=False,
        masked_vision=False,
        masked_language=False,
        default_gpu=default_gpu,
        ground_truth_trajectory=False,
        highlighted_language=args.highlighted_language,
        shuffle_visual_features=False,
        num_negatives=args.num_negatives,
    )

    # in debug mode only run on a subset of the datasets
    if args.debug:
        train_dataset = Subset(
            train_dataset,
            np.random.choice(range(len(train_dataset)), size=128, replace=False), # type: ignore
        )
        val_seen_dataset = Subset(
            val_seen_dataset,
            np.random.choice(range(len(val_seen_dataset)), size=64, replace=False), # type: ignore
        )
        val_unseen_dataset = Subset(
            val_unseen_dataset,
            np.random.choice(range(len(val_unseen_dataset)), size=64, replace=False), # type: ignore
        )

    local_rank = get_local_rank(args)

    if local_rank == -1:
        train_sampler = RandomSampler(train_dataset)
        val_seen_sampler = SequentialSampler(val_seen_dataset)
        val_unseen_sampler = SequentialSampler(val_unseen_dataset)
    else:
        train_sampler = DistributedSampler(train_dataset)
        val_seen_sampler = DistributedSampler(val_seen_dataset)
        val_unseen_sampler = DistributedSampler(val_unseen_dataset)

    # adjust the batch size for distributed training
    batch_size = args.batch_size // args.gradient_accumulation_steps
    if local_rank != -1:
        batch_size = batch_size // dist.get_world_size()
    if default_gpu:
        logger.info(f"batch_size: {batch_size}")

    if default_gpu:
        logger.info(f"Creating dataloader")

    # create data loaders
    train_data_loader = DataLoader(
        train_dataset,
        sampler=train_sampler,
        batch_size=batch_size,
        num_workers=args.num_workers,
        pin_memory=True,
    )
    val_seen_data_loader = DataLoader(
        val_seen_dataset,
        sampler=val_seen_sampler,
        shuffle=False,
        batch_size=batch_size,
        num_workers=args.num_workers,
        pin_memory=True,
    )
    val_unseen_data_loader = DataLoader(
        val_unseen_dataset,
        sampler=val_unseen_sampler,
        shuffle=False,
        batch_size=batch_size,
        num_workers=args.num_workers,
        pin_memory=True,
    )

    # ----- #
    # model #
    # ----- #
    if default_gpu:
        logger.info(f"Loading model")

    config = BERT_CONFIG_FACTORY[args.model_name].from_json_file(args.config_file)
    config.cat_highlight = args.cat_highlight # type: ignore
    config.no_ranking = args.no_ranking # type: ignore
    config.masked_language = args.masked_language # type: ignore
    config.masked_vision = args.masked_vision # type: ignore
    config.model_name = args.model_name
    start_epoch = 0

    if len(args.from_pretrained) == 0:  # hack for catching --from_pretrained ""
        model = Airbert(config)
    else:
        model = Airbert.from_pretrained(
            args.from_pretrained, config, default_gpu=default_gpu
        )

    if default_gpu:
        logger.info(
            f"number of parameters: {sum(p.numel() for p in model.parameters())}"
        )

    # move/distribute model to device
    model.to(device)
    model = wrap_distributed_model(model, local_rank)

    if default_gpu:
        with open(save_folder / "model.txt", "w") as fid:
            fid.write(str(model))

    # ------------ #
    # optimization #
    # ------------ #

    # set parameter specific weight decay
    no_decay = ["bias", "LayerNorm.weight", "LayerNorm.bias"]
    optimizer_grouped_parameters = [
        {"params": [], "weight_decay": 0.0},
        {"params": [], "weight_decay": args.weight_decay},
    ]
    for name, param in model.named_parameters():
        if any(nd in name for nd in no_decay):
            optimizer_grouped_parameters[0]["params"].append(param)
        else:
            optimizer_grouped_parameters[1]["params"].append(param)
    # optimizer
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate,)

    # calculate learning rate schedule
    t_total = (
        len(train_data_loader) // args.gradient_accumulation_steps
    ) * args.num_epochs
    warmup_steps = args.warmup_proportion * t_total
    adjusted_t_total = warmup_steps + args.cooldown_factor * (t_total - warmup_steps)
    scheduler = (
        WarmupLinearSchedule(
            optimizer,
            warmup_steps=warmup_steps,
            t_total=adjusted_t_total,
            last_epoch=-1,
        )
        if not args.no_scheduler
        else MultiplicativeLR(optimizer, lr_lambda=lambda epoch: 1.0) # type: ignore
    )

    # load checkpoint of the optimizer
    weights_path = Path(args.from_pretrained)
    if weights_path.exists():
        state_dict = torch.load(weights_path, map_location="cpu")
        if 'optimizer_state_dict' in state_dict:
            state_dict = state_dict["optimizer_state_dict"]
            optimizer.load_state_dict(state_dict)
        if 'scheduler_state_dict' in state_dict:
            state_dict = state_dict["scheduler_state_dict"]
            scheduler.load_state_dict(state_dict)
        if "epoch" in state_dict:
            start_epoch = state_dict["epoch"]

    # --------------- #
    # before training #
    # --------------- #

    # save the parameters
    if default_gpu:
        with open(save_folder /  "config.txt", "w") as fid:
            print(f"{datetime.now()}", file=fid)
            print("\n", file=fid)
            print(vars(args), file=fid)
            print("\n", file=fid)
            print(config, file=fid)

    # loggers
    if default_gpu:
        writer = SummaryWriter(
            logdir=save_folder / "logging", flush_secs=30
        )
    else:
        writer = None

    # -------- #
    # training #
    # -------- #

    # run training
    if default_gpu:
        logger.info("starting training...")

    best_seen_success_rate, best_unseen_success_rate = 0, 0
    for epoch in range(start_epoch, args.num_epochs):
        if default_gpu and args.debug:
            logger.info(f"epoch {epoch}")

        # Modified:
        if isinstance(train_data_loader.sampler, DistributedSampler):
            train_data_loader.sampler.set_epoch(epoch)

        if args.hard_mining:
            if default_gpu:
                logger.info("setting the beam scores")
            get_score(train_data_loader, model, default_gpu) # type: ignore
            if default_gpu:
                logger.info("the beam scores are set")

        # train for one epoch
        train_epoch(
            epoch,
            model,
            optimizer,
            scheduler,
            train_data_loader,
            writer,
            default_gpu,
            args,
        )

        if default_gpu and args.debug:
            logger.info("saving the model")

        # save the model every epoch
        model_path = save_folder / f"pytorch_model_{epoch + 1}.bin"
        if default_gpu:
            model_state = (
                    model.module.state_dict() # type: ignore
                if hasattr(model, "module")
                else model.state_dict()
            )
            torch.save(model_state, model_path)

        if default_gpu and args.debug:
            logger.info(f"running validation")

        # run validation
        if not args.no_ranking:
            global_step = (epoch + 1) * len(train_data_loader)

            # run validation on the "val seen" split
            with torch.no_grad():
                seen_success_rate = val_epoch(
                    epoch,
                    model,
                    "val_seen",
                    val_seen_data_loader,
                    writer,
                    default_gpu,
                    args,
                    global_step,
                )
                if default_gpu:
                    logger.info(
                        f"[val_seen] epoch: {epoch + 1} success_rate: {seen_success_rate.item():.3f}"
                    )

            # save the model that performs the best on val seen
            if seen_success_rate > best_seen_success_rate:
                best_seen_success_rate = seen_success_rate
                if default_gpu:
                    best_seen_path = save_folder / "pytorch_model_best_seen.bin"
                    shutil.copyfile(model_path, best_seen_path) # type: ignore

            # run validation on the "val unseen" split
            with torch.no_grad():
                unseen_success_rate = val_epoch(
                    epoch,
                    model,
                    "val_unseen",
                    val_unseen_data_loader,
                    writer,
                    default_gpu,
                    args,
                    global_step,
                )
                if default_gpu:
                    logger.info(
                        f"[val_unseen] epoch: {epoch + 1} success_rate: {unseen_success_rate.item():.3f}"
                    )

            # save the model that performs the best on val unseen
            if unseen_success_rate > best_unseen_success_rate:
                best_unseen_success_rate = unseen_success_rate
                if default_gpu:
                    best_unseen_path = save_folder / "pytorch_model_best_unseen.bin"
                    shutil.copyfile(model_path, best_unseen_path)

    # -------------- #
    # after training #
    # -------------- #

    if default_gpu:
        writer.close()

def train_epoch(
    epoch, model, optimizer, scheduler, data_loader, writer, default_gpu, args
) -> None:
    device = next(model.parameters()).device

    model.train()
    model.zero_grad()

    for step, batch in enumerate(tqdm(data_loader, disable= not (default_gpu))):
        # load batch on gpu
        batch = tuple(
            t.cuda(device=device, non_blocking=True) if hasattr(t, "cuda") else t
            for t in batch
        )

        # get the model output
        outputs = model(*get_model_input(batch))

        loss, reduced_metrics = compute_metrics(batch, outputs, args) # type: ignore

        # backward pass
        loss.backward()
        if (step + 1) % args.gradient_accumulation_steps == 0:
            optimizer.step()
            scheduler.step()
            model.zero_grad()

        # write stats to tensorboard
        if default_gpu:
            global_step = step + epoch * len(data_loader)
            reduced_metrics["learning_rate/train"] = float(scheduler.get_lr()[0])
            for name, reduced_loss in reduced_metrics.items():
                writer.add_scalar(name, reduced_loss, global_step=global_step)

def val_epoch(epoch: int, model, tag, data_loader, writer, default_gpu, args, global_step):
    device = next(model.parameters()).device
    local_rank = get_local_rank(args)

    # validation
    model.eval()
    stats = torch.zeros(3, device=device).float()
    for step, batch in enumerate(data_loader):
        # load batch on gpu
        batch = tuple(t.cuda(device=device, non_blocking=True) for t in batch)
        batch_size = get_batch_size(batch)
        target = get_target(batch)

        # get the model output
        inputs = get_model_input(batch)
        outputs = model(*inputs)
        opt_mask = batch[13]

        if not args.no_ranking:
            vil_logit = pad_packed(outputs["action"].squeeze(1), opt_mask)

            # calculate loss
            loss = F.binary_cross_entropy_with_logits(vil_logit, target.float())

            # calculate success rate of the top scoring beam
            correct = torch.sum(
                target.gather(1, torch.argmax(vil_logit, 1).view(-1, 1))
            ).float()

            # accumulate
            stats[0] += loss
            stats[1] += correct
            stats[2] += batch_size

        if default_gpu and args.debug:
            logger.info(
                f"[{tag}] step: {step + 1} "
                f"running loss: {stats[0] / stats[2]:0.2f} "
                f"running success rate: {stats[1] / stats[2]:0.2f}"
            )

    if local_rank != -1:
        dist.all_reduce(stats, op=dist.ReduceOp.SUM) # type: ignore

    # write stats to tensorboard
    if default_gpu:
        writer.add_scalar(
            f"loss/bce_{tag}", stats[0] / stats[2], global_step=global_step
        )
        writer.add_scalar(
            f"accuracy/sr_{tag}", stats[1] / stats[2], global_step=global_step
        )

    return stats[1] / stats[2]

# ------------- #
# batch parsing #
# ------------- #

# batch format:
# 0:target, 1:image_features, 2:image_locations, 3:image_mask, 4:image_targets,
# 5:image_targets_mask, 6:instr_tokens, 7:instr_mask, 8:instr_targets, 9:instr_highlights, 10:segment_ids,
# 11:co_attention_mask, 12:item_id

def get_model_input(batch):
    (
        _,
        image_features,
        image_locations,
        image_mask,
        _,
        _,
        instr_tokens,
        instr_mask,
        _,
        instr_highlights,
        segment_ids,
        co_attention_mask,
        _,
        opt_mask,
    ) = batch

    # remove padding samples
    image_features = image_features[opt_mask]
    image_locations = image_locations[opt_mask]
    image_mask = image_mask[opt_mask]
    instr_tokens = instr_tokens[opt_mask]
    instr_mask = instr_mask[opt_mask]
    instr_highlights = instr_highlights[opt_mask]
    segment_ids = segment_ids[opt_mask]

    # transform batch shape
    co_attention_mask = co_attention_mask.view(
        -1, co_attention_mask.size(2), co_attention_mask.size(3)
    )

    return (
        instr_tokens,
        image_features,
        image_locations,
        segment_ids,
        instr_mask,
        image_mask,
        co_attention_mask,
        instr_highlights,
    )

def get_batch_size(batch):
    return batch[1].size(0)

def get_target(batch):
    return batch[0]

def get_mask_options(batch) -> torch.Tensor:
    return batch[13]

def get_num_options(batch):
    mask = get_mask_options(batch)
    return batch[6][mask].size(0)

def get_linguistic_target(batch):
    opt_mask = get_mask_options(batch)
    return batch[8][opt_mask].flatten()

def get_highlights(batch):
    opt_mask = get_mask_options(batch)
    return batch[9][opt_mask]

def get_vision_target(batch):
    opt_mask = get_mask_options(batch)
    return (
        batch[4][opt_mask].flatten(0, 1),
        batch[5][opt_mask].flatten(),
    )

if __name__ == "__main__":
    main()


import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
import cv2#
cap = cv2.VideoCapture(0)
sys.path.append("..")
import label_map_util
import visualization_utils as vis_util
MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')
NUM_CLASSES = 90
opener = urllib.request.URLopener()
opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
tar_file = tarfile.open(MODEL_FILE)
for file in tar_file.getmembers():
    file_name = os.path.basename(file.name)
    if 'frozen_inference_graph.pb' in file_name:
        tar_file.extract(file, os.getcwd())
        detection_graph = tf.Graph()
        with detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
                label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
                categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
                category_index = label_map_util.create_category_index(categories)
                def load_image_into_numpy_array(image):
                    (im_width, im_height) = image.size
                    return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
                PATH_TO_TEST_IMAGES_DIR = 'test_images'
                TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 3) ]
                IMAGE_SIZE = (12, 8)
                with detection_graph.as_default():
                    with tf.Session(graph=detection_graph) as sess:
                        while True:
                            ret, image_np = cap.read()
                            image_np_expanded = np.expand_dims(image_np, axis=0)
                            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
                            boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
                            scores = detection_graph.get_tensor_by_name('detection_scores:0')
                            classes = detection_graph.get_tensor_by_name('detection_classes:0')
                            num_detections = detection_graph.get_tensor_by_name('num_detections:0')
                            (boxes, scores, classes, num_detections) = sess.run([boxes, scores, classes, num_detections],feed_dict={image_tensor: image_np_expanded})
                            vis_util.visualize_boxes_and_labels_on_image_array(image_np, np.squeeze(boxes),np.squeeze(classes).astype(np.int32), np.squeeze(scores),category_index,use_normalized_coordinates=True,line_thickness=8)
                            cv2.imshow('object detection', cv2.resize(image_np, (800,600)))
                            if cv2.waitKey(25) & 0xFF == ord('q'):
                                cv2.destroyAllWindows()
                                break


import numpy as np
import argparse
import glob
import os
from functools import partial
import vispy
import scipy.misc as misc
from tqdm import tqdm
import yaml
import time
import sys
from mesh import write_ply, read_ply, output_3d_photo
from utils import get_MiDaS_samples, read_MiDaS_depth
import torch
import cv2
from skimage.transform import resize
import imageio
import copy
from networks import Inpaint_Color_Net, Inpaint_Depth_Net, Inpaint_Edge_Net
from MiDaS.run import run_depth
from MiDaS.monodepth_net import MonoDepthNet
import MiDaS.MiDaS_utils as MiDaS_utils
from bilateral_filtering import sparse_bilateral_filtering

parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='argument.yml',help='Configure of post processing')
args = parser.parse_args()
config = yaml.load(open(args.config, 'r'))
if config['offscreen_rendering'] is True:
    vispy.use(app='egl')
os.makedirs(config['mesh_folder'], exist_ok=True)
os.makedirs(config['video_folder'], exist_ok=True)
os.makedirs(config['depth_folder'], exist_ok=True)
sample_list = get_MiDaS_samples(config['src_folder'], config['depth_folder'], config, config['specific'])
normal_canvas, all_canvas = None, None

if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0):
    device = config["gpu_ids"]
else:
    device = "cpu"

print(f"running on device {device}")

for idx in tqdm(range(len(sample_list))):
    depth = None
    sample = sample_list[idx]
    print("Current Source ==> ", sample['src_pair_name'])
    mesh_fi = os.path.join(config['mesh_folder'], sample['src_pair_name'] +'.ply')
    image = imageio.imread(sample['ref_img_fi'])

    print(f"Running depth extraction at {time.time()}")
    if config['require_midas'] is True:
        image_width = image.shape[1]
        run_depth([sample['ref_img_fi']], config['src_folder'], config['depth_folder'],
                  config['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=image_width)
    if 'npy' in config['depth_format']:
        config['output_h'], config['output_w'] = np.load(sample['depth_fi']).shape[:2]
    else:
        config['output_h'], config['output_w'] = imageio.imread(sample['depth_fi']).shape[:2]
    frac = 1
    config['output_h'], config['output_w'] = int(config['output_h'] * frac), int(config['output_w'] * frac)
    config['original_h'], config['original_w'] = config['output_h'], config['output_w']
    if image.ndim == 2:
        image = image[..., None].repeat(3, -1)
    if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum(np.abs(image[..., 1] - image[..., 2])) == 0:
        config['gray_image'] = True
    else:
        config['gray_image'] = False
    image = cv2.resize(image, (config['output_w'], config['output_h']), interpolation=cv2.INTER_AREA)
    depth = read_MiDaS_depth(sample['depth_fi'], 3.0, config['output_h'], config['output_w'])
    mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2]
    if not(config['load_ply'] is True and os.path.exists(mesh_fi)):
        vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), config, num_iter=config['sparse_iter'], spdb=False)
        depth = vis_depths[-1]
        model = None
        torch.cuda.empty_cache()
        print("Start Running 3D_Photo ...")
        print(f"Loading edge model at {time.time()}")
        depth_edge_model = Inpaint_Edge_Net(init_weights=True)
        depth_edge_weight = torch.load(config['depth_edge_model_ckpt'],
                                       map_location=torch.device(device))
        depth_edge_model.load_state_dict(depth_edge_weight)
        depth_edge_model = depth_edge_model.to(device)
        depth_edge_model.eval()

        print(f"Loading depth model at {time.time()}")
        depth_feat_model = Inpaint_Depth_Net()
        depth_feat_weight = torch.load(config['depth_feat_model_ckpt'],
                                       map_location=torch.device(device))
        depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
        depth_feat_model = depth_feat_model.to(device)
        depth_feat_model.eval()
        depth_feat_model = depth_feat_model.to(device)
        print(f"Loading rgb model at {time.time()}")
        rgb_model = Inpaint_Color_Net()
        rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'],
                                     map_location=torch.device(device))
        rgb_model.load_state_dict(rgb_feat_weight)
        rgb_model.eval()
        rgb_model = rgb_model.to(device)
        graph = None

        print(f"Writing depth ply (and basically doing everything) at {time.time()}")
        rt_info = write_ply(image,
                              depth,
                              sample['int_mtx'],
                              mesh_fi,
                              config,
                              rgb_model,
                              depth_edge_model,
                              depth_edge_model,
                              depth_feat_model)

        if rt_info is False:
            continue
        rgb_model = None
        color_feat_model = None
        depth_edge_model = None
        depth_feat_model = None
        torch.cuda.empty_cache()
    if config['save_ply'] is True or config['load_ply'] is True:
        verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi)
    else:
        verts, colors, faces, Height, Width, hFov, vFov = rt_info

    print(f"Making video at {time.time()}")
    videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name']
    top = (config.get('original_h') // 2 - sample['int_mtx'][1, 2] * config['output_h'])
    left = (config.get('original_w') // 2 - sample['int_mtx'][0, 2] * config['output_w'])
    down, right = top + config['output_h'], left + config['output_w']
    border = [int(xx) for xx in [top, down, left, right]]
    normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov),
                        copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(config['video_folder']),
                        image.copy(), copy.deepcopy(sample['int_mtx']), config, image,
                        videos_poses, video_basename, config.get('original_h'), config.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas,
                        mean_loc_depth=mean_loc_depth)