# -*- coding: utf-8 -*-

"""
*------------------------- PageViewVectorField.py ----------------------------*
显示矢量场数据的页面。

部件内容：
    - WidgetPlotVector 用于使用 matplotlib 显示矢量场及背景
    - WidgetPlotBackground 用于使用 matplotlib 显示全域背景

提升部件
    - 提升类名 PageViewVectorField
    - 头文件 bin.Widgets.PageViewVectorField

作者:           胡一鸣
创建日期:       2022年5月18日

The GUI page to view Vector Field dataset.

Contents:
    - WidgetPlotVector, to view the vector fields and backgroud by matplotlib
    - WidgetPlotBackground, to view the full-scale background by matplotlib

Promoted Widget:
    - name of widget class: PageViewVectorField
    - header file: bin.Widgets.PageViewVectorField

author:         Hu Yiming
date:           Mar 30, 2022
*------------------------- PageViewVectorField.py ----------------------------*
"""

from logging import Logger
import os
from PySide6.QtWidgets import QWidget
from PySide6.QtWidgets import QMessageBox
from PySide6.QtWidgets import QInputDialog
from PySide6.QtWidgets import QDialog
# from PySide6.QtWidgets import QMenu
# from PySide6.QtCore import QPoint
from matplotlib.backend_bases import MouseEvent
from matplotlib.backends.backend_qtagg import (
    FigureCanvasQTAgg as FigureCanvas)
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from matplotlib.image import AxesImage
from matplotlib.quiver import Quiver
from matplotlib.patches import Rectangle
from matplotlib.colorbar import Colorbar, make_axes
import h5py
import numpy as np

from bin.BlitManager import BlitManager
from bin.HDFManager import HDFDataNode, HDFHandler
from bin.TaskManager import TaskManager
from bin.Widgets.DialogChooseItem import DialogHDFChoose
from bin.Widgets.PageVirtualImage import DialogSaveImage
from lib.TaskVectorFieldProcess import TaskCurl, TaskFlipVectorField
from lib.TaskVectorFieldProcess import TaskDivergence
from lib.TaskVectorFieldProcess import TaskPotential
from lib.TaskVectorFieldProcess import TaskRotateVectorAngle
from lib.TaskVectorFieldProcess import TaskSliceI
from lib.TaskVectorFieldProcess import TaskSliceJ
from lib.TaskVectorFieldProcess import TaskSubtractVectorOffset
from ui import uiDialogCreateImage
from ui import uiPageViewVectorField
from ui import uiDialogAdjustQuiverEffect
from ui import uiDialogVectorProcessing

class PageViewVectorField(QWidget):
    """
    显示二维矢量场的部件类。使用 Quiver。

    Ui 文件地址：ROOT_PATH/ui/uiPageVieweVectorField.ui

    Widget to view vector fields by quiver plot.

    The path of the ui file: ROOT_PATH/ui/uiPageViewVectorField.ui
    """
    def __init__(self, parent: QWidget = None):
        super().__init__(parent)
        self.ui = uiPageViewVectorField.Ui_Form()
        self.ui.setupUi(self)

        self._data_path = ''
        self._background_path = ''
        self._image_ax = None
        self._colorbar_ax = None
        self._background_ax = None
        self._quiver_object = None
        self._colorbar_object = None
        self._background_object = None
        self._image_object = None
        self._image_max = 0
        self._image_min = 0

        self.ui.lineEdit_vector_path.setReadOnly(True)
        self.ui.lineEdit_background_path.setReadOnly(True)

        self._initUi()
        self._createAxes()

    @property
    def hdf_handler(self) -> HDFHandler:
        global qApp
        return qApp.hdf_handler

    @property
    def data_object(self) -> h5py.Dataset:
        return self.hdf_handler.file[self._data_path]

    @property
    def data_path(self) -> str:
        return self._data_path

    @property
    def logger(self) -> Logger:
        global qApp
        return qApp.logger

    @property
    def image_canvas(self) -> FigureCanvas:
        return self.ui.widget_quiver.canvas

    @property
    def background_canvas(self) -> FigureCanvas:
        return self.ui.widget_background.canvas

    @property
    def image_figure(self) -> Figure:
        return self.ui.widget_quiver.figure

    @property
    def background_figure(self) -> Figure:
        return self.ui.widget_background.figure

    @property
    def image_ax(self) -> Axes:
        return self._image_ax

    @property
    def background_ax(self) -> Axes:
        return self._background_ax

    @property
    def colorbar_ax(self) -> Axes:
        return self._colorbar_ax

    @property
    def image_object(self) -> AxesImage:
        return self._image_object

    @property
    def quiver_object(self) -> Quiver:
        return self._quiver_object

    @property
    def colorbar_object(self) -> Colorbar:
        return self._colorbar_object

    @property
    def background_object(self) -> AxesImage:
        return self._background_object

    @property
    def background_path(self) -> str:
        return self._background_path

    @property
    def image_blit_manager(self) -> BlitManager:
        return self.ui.widget_quiver.blit_manager

    @property
    def background_blit_manager(self) -> BlitManager:
        return self.ui.widget_background.blit_manager

    @property
    def background_visible(self) -> bool:
        return self.ui.checkBox_background_visible.isChecked()

    @property
    def task_manager(self) -> TaskManager:
        global qApp
        return qApp.task_manager

    def setVectorField(self, data_path: str):
        """
        Set the data path in HDF5 file, to show the vector field.

        Will set the data_path attribute. The vector field must be a 3D matrix.

        arguments:
            data_path: (str) the path of the vector field data.

        raises:
            TypeError, KeyError, ValueError
        """
        if not isinstance(data_path, str):
            raise TypeError('data_path must be a str, not '
                '{0}'.format(type(data_path).__name__))

        img_node = self.hdf_handler.getNode(data_path)
        # May raise KeyError if the path does not exist
        if not isinstance(img_node, HDFDataNode):
            raise ValueError('Item {0} must be a Dataset'.format(data_path))

        data_obj = self.hdf_handler.file[data_path]
        if not len(data_obj.shape) == 3:
            raise ValueError('Data must be a 3D matrix (2, i, j)')

        self._data_path = data_path
        self.ui.lineEdit_vector_path.setText(self.data_path)
        self.setWindowTitle('{0} - Vector'.format(img_node.name))

        # We must first render the background image, in order not to
        # cover the vector field.
        if 'background_path' in self.data_object.attrs:
            background_path = self.data_object.attrs['background_path']
            try:
                self.setBackground(background_path)
            except (KeyError, ValueError, TypeError):
                new_background_path = self._createNewBackground()
                self.setBackground(new_background_path)
        else:
            new_background_path = self._createNewBackground()
            self.setBackground(new_background_path)

        self._createQuiver()

        self.image_canvas.draw()
        self.image_canvas.flush_events()

        self.ui.widget_quiver.setProcessingActionItemPath(self.data_path)

    def _createAxes(self):
        """
        Create the axes that contains the quiver, colorbar and the background
        respectively.
        """
        if self._image_ax is None:
            self._image_ax = self.image_figure.add_subplot()
            self.image_blit_manager.addArtist('image_axes', self._image_ax)

        if self._background_ax is None:
            self._background_ax = self.background_figure.add_subplot()
            self.background_blit_manager.addArtist(
                'background_axes',
                self._background_ax,
            )

        if self._colorbar_ax is None:
            self._colorbar_ax, _kw = make_axes(
                self.background_ax,
                location = 'right',
                orientation = 'vertical',
            )
            self._colorbar_ax.xaxis.set_visible(False)
            self._colorbar_ax.yaxis.tick_right()

            self.background_blit_manager['colorbar_axes'] = self._colorbar_ax

    def _createQuiver(self):
        """
        Read the vector field and its attributes, and show it.

        TODO: read and save attributes, like length unit.
        """
        if self._quiver_object in self.image_ax.collections:
            _index = self.image_ax.collections.index(self._quiver_object)
            self.image_ax.collections.pop(_index)

        _, height, width = self.data_object.shape
        array_i = np.linspace(0, height - 1, height)
        array_j = np.linspace(0, width - 1, width)
        coord_i, coord_j = np.meshgrid(array_i, array_j, indexing = 'ij')
        vec_i, vec_j = self.data_object

        if 'quiver_scale' in self.data_object.attrs:
            quiver_scale = self.data_object.attrs['quiver_scale']
        else:
            self.data_object.attrs['quiver_scale'] = 1
            quiver_scale = 1

        if 'quiver_width' in self.data_object.attrs:
            quiver_width = self.data_object.attrs['quiver_width']
        else:
            self.data_object.attrs['quiver_width'] = 0.15
            quiver_width = 0.15

        if 'quiver_color' in self.data_object.attrs:
            quiver_color = self.data_object.attrs['quiver_color']
        else:
            self.data_object.attrs['quiver_color'] = 'black'
            quiver_color = 'black'

        X, Y = coord_j, coord_i
        U, V = vec_j, vec_i

        self._quiver_object = self.image_ax.quiver(
            X, Y, U, V,
            units = 'xy',
            scale = quiver_scale,
            width = quiver_width,
            angles = 'xy',
            pivot = 'mid',
            color = quiver_color,
        )

        self.image_blit_manager['quiver'] = self._quiver_object

    def setBackground(self, background_path: str):
        """
        Set the background path in HDF5 file, to show the background image.

        Will set the background_path attribute. The background image must be
        a 2D matrix, RGB images (3 channels) are not supported here.

        The background image's shape must be equal to one of the channel of
        the vector field. For example, if the shape of the vector field is
            2 x 256 x 256
        the shape of the background image must be 256 x 256 then.

        arguments:
            background_path: (str) the path of the image or data.

        raises:
            TypeError, KeyError, ValueError
        """
        if not isinstance(background_path, str):
            raise TypeError('background_path must be a str, not '
                '{0}'.format(type(background_path).__name__))

        background_node = self.hdf_handler.getNode(background_path)
        # May raise KeyError if the path does not exist
        if not isinstance(background_node, HDFDataNode):
            raise ValueError('Item {0} must be a '
                'Dataset'.format(background_path))

        background_data_obj = self.hdf_handler.file[background_path]
        if not len(background_data_obj.shape) == 2:
            raise ValueError('Data must be a 2D matrix (single channel image)')

        if self.data_path == '':
            raise ValueError('Must set vector field before set background')
        _, height, width = self.data_object.shape
        if (background_data_obj.shape[0] != height
                or background_data_obj.shape[1] != width):
            raise ValueError('The background image\'s shape must be equal to'
                ' one of the channel of the vector field.')

        self._background_path = background_path
        self.ui.lineEdit_background_path.setText(self.background_path)
        self.data_object.attrs['background_path'] = background_path

        self._createBackgroundImage()
        self._createColorbar()

        self.background_canvas.draw()
        self.background_canvas.flush_events()
        self.image_canvas.draw()
        self.image_canvas.flush_events()

    def _createBackgroundImage(self):
        """
        Create the axes and images of the background images.

        TODO: read and save attributes, like norm, cmap, alpha, etc.
        """
        if self._background_object in self.background_ax.images:
            # clear current objects in the axes
            _index = self.background_ax.images.index(self._background_object)
            self.background_ax.images.pop(_index)

        self._background_object = self.background_ax.imshow(
            self.hdf_handler.file[self.background_path]
        )
        self.background_blit_manager['background_image'] = self._background_object

        if self._image_object in self.image_ax.images:
            # clear current objects in the axes
            _index = self.image_ax.images.index(self._image_object)
            self.image_ax.images.pop(_index)

        self._image_object = self.image_ax.imshow(
            self.hdf_handler.file[self.background_path],
            interpolation = 'hermite',
        )
        self._image_object.set_visible(self.background_visible)
        self.image_blit_manager['image'] = self._image_object

    def _createColorbar(self):
        """
        Create the colorbar according to the background image.
        """
        if self._colorbar_object is None:
            self._colorbar_object = Colorbar(
                ax = self.colorbar_ax,
                mappable = self.background_object
            )
        else:
            self.colorbar_object.update_normal(self.background_object)

    def _initUi(self):
        """
        Initialize UI.
        """
        self.ui.checkBox_background_visible.setChecked(False)
        self.ui.checkBox_background_visible.stateChanged.connect(
            self._setImageVisible
        )

        self.ui.pushButton_browse.clicked.connect(
            self._browse
        )
        self.ui.pushButton_browse_background.clicked.connect(
            self._browseBackground
        )
        # self.ui.pushButton_vector_processing.clicked.connect(
        #     self._vectorProcessing
        # )
        self.ui.pushButton_adjust_effects.clicked.connect(
            self._adjustEffects
        )

        self.ui.pushButton_show_color_wheel.setVisible(False)   # TODO
        self.ui.pushButton_show_color_wheel.setVisible(False)   # TODO

    def _adjustEffects(self):
        """
        Open a dialog to adjust quiver display effects.
        """
        dialog = DialogAdjustQuiverEffect(self)
        dialog.setVectorField(self.data_path)
        dialog_code = dialog.exec()
        if dialog_code == DialogAdjustQuiverEffect.Rejected:
            return

        self.data_object.attrs['quiver_scale'] = dialog.getScale()
        self.data_object.attrs['quiver_width'] = dialog.getWidth()
        self.data_object.attrs['quiver_color'] = dialog.getColor()

        self.setVectorField(self.data_path)

    def _browse(self):
        """
        Open a dialog to browse which vector field to be opened.
        """
        dialog = DialogHDFChoose(self)
        dialog_code = dialog.exec()
        if dialog_code == dialog.Accepted:
            current_path = dialog.getCurrentPath()
        else:
            return

        try:
            self.setVectorField(current_path)
        except (KeyError, ValueError, TypeError,) as e:
            self.logger.error('{0}'.format(e), exc_info = True)
            msg = QMessageBox(parent = self)
            msg.setWindowTitle('Warning')
            msg.setIcon(QMessageBox.Warning)
            msg.setStandardButtons(QMessageBox.Ok)
            msg.setText('Cannot open this data: {0}'.format(e))
            msg.exec()

    def _browseBackground(self):
        """
        Open a dialog to browse which image to be opened as the background.
        """
        dialog = DialogHDFChoose(self)
        dialog_code = dialog.exec()
        if dialog_code == dialog.Accepted:
            current_path = dialog.getCurrentPath()
        else:
            return

        try:
            self.setBackground(current_path)
        except (KeyError, ValueError, TypeError,) as e:
            self.logger.error('{0}'.format(e), exc_info = True)
            msg = QMessageBox(parent = self)
            msg.setWindowTitle('Warning')
            msg.setIcon(QMessageBox.Warning)
            msg.setStandardButtons(QMessageBox.Ok)
            msg.setText('Cannot open this data: {0}'.format(e))
            msg.exec()

    def _setImageVisible(self):
        """
        Set whether the background image to be visible or not.
        """
        self.image_object.set_visible(self.background_visible)
        self.image_blit_manager.update()

    def _createNewBackground(self):
        """
        Create a new background image.

        In default, it will be /.../[vector name]_bkgrd.img . If there has been
        a background at this path, try to use it. Otherwise, add an index and
        try again: /.../[vector name]_bkgrd_1.img . The new path will be under
        the same group as the vector field.

        Users can calculate a new background and reset it themselves.

        returns:
            (str) the path of the new background image.
        """

        # Get a valid new name of the new background image
        data_node = self.hdf_handler.getNode(self.data_path)
        if '.' in data_node.name:
            name_array = data_node.name.split('.')
            name_array.pop()
            original_name = '.'.join(name_array)
        else:
            original_name = data_node.name

        bkgrd_name = original_name + '_bkgrd.img'
        _count = 0
        while bkgrd_name in data_node.parent:
            _count += 1
            bkgrd_name = original_name + '_bkgrd_{0}.img'.format(_count)

        # Get the shape of the new background image
        _, height, width = self.data_object.shape

        self.hdf_handler.addNewData(
            parent_path = data_node.parent.path,
            name = bkgrd_name,
            shape = (height, width),
            dtype = 'float64',
        )

        # Get the path of the new background image
        if data_node.parent.path == '/':
            bkgrd_path = '/' + bkgrd_name
        else:
            bkgrd_path = data_node.parent.path + '/' + bkgrd_name

        self.data_object.attrs['background_path'] = bkgrd_path

        return bkgrd_path

    def _vectorProcessing(self):
        """
        Open a dialog to show vector processing methods.

        Including change angles, subtracting mean vector, and calculate curl,
        divergence, potential, ...
        """
        dialog = DialogVectorProcessing(self)
        dialog_code = dialog.exec()
        if not dialog_code == dialog.Accepted:
            return
        result = dialog.getResult()
        if result == 'rotate':
            self._vectorRotateAngle()
        if result == 'subtract':
            self._vectorSubtract()
        if result == 'flip':
            self._vectorFlip()
        if result == 'potential':
            self._vectorPotential()
        if result == 'divergence':
            self._vectorDivergence()
        if result == 'curl':
            self._vectorCurl()
        if result == 'vec_i':
            self._vectorSliceI()
        if result == 'vec_j':
            self._vectorSliceJ()

    def _vectorRotateAngle(self):
        """
        Rotate every vector an angle.

        In practical experiment, there may exist some angular shift between
        the scanning array and the pixelized camera, but we usually cannot
        be aware of this by pure 4D-STEM dataset. Fortunately, we know that
        some vector fields (like Electric field) keeps non-curl, so we use
        this property to get correct vector fields.
        """
        angle, is_accepted = QInputDialog.getDouble(
            self,
            'Input rotation angle',
            'Here input a rotation angle of every vector. Unit: degree',
            0,
            minValue = -360,
            maxValue = 360,
            decimals = 1,
            step = 1,
        )

        if not is_accepted:
            return

        dialog_save = DialogSaveVectorField(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskRotateVectorAngle(
            self.data_path,
            image_parent_path,
            image_name,
            angle = angle,
            parent = self,
            **meta,
        )

        self.task_manager.addTask(self.task)

    def _vectorSubtract(self):
        """
        Subtract every vector by their mean vector.

        In practical experiment, there may exist some shift between the center
        of the diffraction pattern and the origin of the diffraction plane. In
        this case, the calculated center of mass will have an global offset.
        Here this function will be able to recover the fine architecture of the
        vector field, by set the vector offset to be zero.
        """
        dialog_save = DialogSaveVectorField(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskSubtractVectorOffset(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

    def _vectorFlip(self):
        """
        Exchange every vector's i, j components.

        In practical experiment, the coordinate of the diffraction plane may
        differ from the convention of 4D-Explorer. For example, both typical
        i-j indexing and x-y coordinates are right-handed system:

            ┌------------> j        ^y
            |                       |
            |                       |
            |                       |
            |                       |
            v                       |
            i                       └----------------> x

        However, in some conventions, the coordinate is left-handed:

            ┌------------> x
            |
            |
            |
            |
            v
            y

        In this case, the calculated vector field's components should be
        exchanged in order to reveal correct electromagnetic field.
        """
        dialog_save = DialogSaveVectorField(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskFlipVectorField(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

    def _vectorPotential(self):
        """
        Calculate potential for the vector field.

        The vector field should be a non-curl field, otherwise the result is
        invalid in physics.
        """
        dialog_save = DialogSaveImage(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskPotential(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

    def _vectorDivergence(self):
        """
        Calculate divergence for the vector field.
        """
        dialog_save = DialogSaveImage(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskDivergence(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

    def _vectorCurl(self):
        """
        Calculate curl for the vector field.
        """
        dialog_save = DialogSaveImage(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskCurl(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

    def _vectorSliceI(self):
        """
        Slice i-component for the vector field.
        """
        dialog_save = DialogSaveImage(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskSliceI(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

    def _vectorSliceJ(self):
        """
        Slice j-component for the vector field.
        """
        dialog_save = DialogSaveImage(self)
        dialog_save.setParentPath(self.data_path)
        dialog_code = dialog_save.exec()
        if not dialog_code == dialog_save.Accepted:
            return
        image_name = dialog_save.getNewName()
        image_parent_path = dialog_save.getParentPath()
        meta = self.data_object.attrs

        self.task = TaskSliceJ(
            self.data_path,
            image_parent_path,
            image_name,
            parent = self,
            **meta,
        )
        self.task_manager.addTask(self.task)

class DialogSaveVectorField(DialogSaveImage):
    """
    选择在 HDF 文件中保存矢量场的路径。

    Dialog to choose where to save the reconstructed image in the HDF file.
    """
    def __init__(self, parent: QWidget = None):
        super().__init__(parent)

    def getNewName(self) -> str:
        """
        returns the new name of the imported dataset.

        Will add '.img' automatically as the extension.
        """
        name = self.ui.lineEdit_name.text()
        if '.' in name:
            if name.split('.')[-1] == 'vec':
                return name
        return name + '.vec'

class DialogAdjustQuiverEffect(QDialog):
    """
    用于调整矢量场箭头图的效果的对话框。

    Dialog to adjust quiver display effect.
    """
    def __init__(self, parent: QWidget = None):
        super().__init__(parent)
        self.ui = uiDialogAdjustQuiverEffect.Ui_Dialog()
        self.ui.setupUi(self)
        self._item_path = None
        self._colors = {
            'blue': 0,
            'cyan': 1,
            'green': 2,
            'black': 3,
            'magenta': 4,
            'red': 5,
            'white': 6,
            'yellow': 7,
        }
        self.ui.doubleSpinBox_scale.setMaximum(1e9)
        self.ui.doubleSpinBox_scale.setMinimum(0.01)
        self.ui.doubleSpinBox_scale.setValue(1)
        self.ui.doubleSpinBox_width.setMinimum(0.01)
        self.ui.doubleSpinBox_width.setValue(0.15)
        self.ui.comboBox_color.setCurrentIndex(3)

        self.ui.pushButton_cancel.clicked.connect(self.reject)
        self.ui.pushButton_ok.clicked.connect(self.accept)

    @property
    def hdf_handler(self) -> HDFHandler:
        global qApp
        return qApp.hdf_handler

    @property
    def data_object(self) -> h5py.Dataset:
        return self.hdf_handler.file[self._item_path]

    @property
    def data_scale(self) -> float:
        return self.data_object.attrs['quiver_scale']

    @property
    def data_width(self) -> float:
        return self.data_object.attrs['quiver_width']

    @property
    def data_color(self) -> str:
        return self.data_object.attrs['quiver_color']

    def setVectorField(self, item_path: str):
        """
        arguments:
            item_path: (str) the vector field's path in hdf5 file.
        """
        self._item_path = item_path

        self.ui.doubleSpinBox_scale.setValue(self.data_scale)
        self.ui.doubleSpinBox_width.setValue(self.data_width)
        self.ui.comboBox_color.setCurrentIndex(
            self._colors[self.data_color]
        )

    def getScale(self) -> float:
        return self.ui.doubleSpinBox_scale.value()

    def getWidth(self) -> float:
        return self.ui.doubleSpinBox_width.value()

    def getColor(self) -> str:
        return self.ui.comboBox_color.currentText()

class DialogVectorProcessing(QDialog):
    """
    用于计算 Vector Field 相关操作的对话框。

    Dialog to calculate some processings for vector fields.
    """
    def __init__(self, parent: QWidget):
        super().__init__(parent)
        self.ui = uiDialogVectorProcessing.Ui_Dialog()
        self.ui.setupUi(self)
        self.ui.pushButton_rotate_vector_angle.clicked.connect(
            self._rotate_vector_angle
        )
        self.ui.pushButton_subtract_mean_vector.clicked.connect(
            self._subtract_mean_vector
        )
        self.ui.pushButton_flip.clicked.connect(
            self._flip
        )
        self.ui.pushButton_calculate_curl.clicked.connect(
            self._calculate_curl
        )
        self.ui.pushButton_calculate_divergence.clicked.connect(
            self._calculate_divergence
        )
        self.ui.pushButton_calculate_potential.clicked.connect(
            self._calculate_potential
        )
        self.ui.pushButton_vec_i.clicked.connect(
            self._slice_vec_i
        )
        self.ui.pushButton_vec_j.clicked.connect(
            self._slice_vec_j
        )

    def getResult(self):
        return self._result

    def _rotate_vector_angle(self):
        self._result = 'rotate'
        self.accept()

    def _subtract_mean_vector(self):
        self._result = 'subtract'
        self.accept()

    def _calculate_curl(self):
        self._result = 'curl'
        self.accept()

    def _calculate_divergence(self):
        self._result = 'divergence'
        self.accept()

    def _calculate_potential(self):
        self._result = 'potential'
        self.accept()

    def _slice_vec_i(self):
        self._result = 'vec_i'
        self.accept()

    def _slice_vec_j(self):
        self._result = 'vec_j'
        self.accept()

    def _flip(self):
        self._result = 'flip'
        self.accept()

# if 'quiver_scale' in self.data_object.attrs:
#             quiver_scale = self.data_object.attrs['quiver_scale']
#         else:
#             self.data_object.attrs['quiver_scale'] = 1
#             quiver_scale = 1

#         if 'quiver_width' in self.data_object.attrs:
#             quiver_width = self.data_object.attrs['quiver_width']
#         else:
#             self.data_object.attrs['quiver_width'] = 0.15
#             quiver_width = 0.15

#         if 'quiver_color' in self.data_object.attrs:


from typing import Dict, List, Optional, Sequence, Any, Union
import numpy as np
import logging
log = logging.getLogger(__name__)

import zhinst.utils
import qcodes as qc
from qcodes.instrument.base import Instrument
import qcodes.utils.validators as vals

class HF2LI(Instrument):
    """Qcodes driver for Zurich Instruments HF2LI lockin amplifier.

    This driver is meant to emulate a single-channel lockin amplifier,
    so one instance has a single demodulator, a single sigout channel,
    and multiple auxout channels (for X, Y, R, Theta, or an arbitrary manual value).
    Multiple instances can be run simultaneously as independent lockin amplifiers.

    This instrument has a great deal of additional functionality that is
    not currently supported by this driver.

    Args:
        name: Name of instrument.
        device: Device name, e.g. "dev204", used to create zhinst API session.
        demod: Index of the demodulator to use.
        sigout: Index of the sigout channel to use as excitation source.
        auxouts: Dict of the form {output: index},
            where output is a key of HF2LI.OUTPUT_MAPPING, for example {"X": 0, "Y": 3}
            to use the instrument as a lockin amplifier in X-Y mode with auxout channels 0 and 3.
        num_sigout_mixer_channels: Number of mixer channels to enable on the sigouts. Default: 1.
    """
    OUTPUT_MAPPING = {-1: 'manual', 0: 'X', 1: 'Y', 2: 'R', 3: 'Theta'}
    def __init__(self, name: str, device: str, demod: int, sigout: int,
        auxouts: Dict[str, int], num_sigout_mixer_channels: int=1, **kwargs) -> None:
        super().__init__(name, **kwargs)
        instr = zhinst.utils.create_api_session(device, 1, required_devtype='HF2LI')
        self.daq, self.dev_id, self.props = instr
        self.demod = demod
        self.sigout = sigout
        self.auxouts = auxouts
        log.info(f'Successfully connected to {name}.')

        for ch in self.auxouts:
            self.add_parameter(
                name=ch,
                label=f'Scaled {ch} output value',
                unit='V',
                get_cmd=lambda channel=ch: self._get_output_value(channel),
                get_parser=float,
                docstring=f'Scaled and demodulated {ch} value.'
            )
            self.add_parameter(
                name=f'gain_{ch}',
                label=f'{ch} output gain',
                unit='V/Vrms',
                get_cmd=lambda channel=ch: self._get_gain(channel),
                get_parser=float,
                set_cmd=lambda gain, channel=ch: self._set_gain(gain, channel),
                vals=vals.Numbers(),
                docstring=f'Gain factor for {ch}.'
            )
            self.add_parameter(
                name=f'offset_{ch}',
                label=f'{ch} output offset',
                unit='V',
                get_cmd=lambda channel=ch: self._get_offset(channel),
                get_parser=float,
                set_cmd=lambda offset, channel=ch: self._set_offset(offset, channel),
                vals=vals.Numbers(-2560, 2560),
                docstring=f'Manual offset for {ch}, applied after scaling.'
            )
            self.add_parameter(
                name=f'output_{ch}',
                label=f'{ch} outptut select',
                get_cmd=lambda channel=ch: self._get_output_select(channel),
                get_parser=str
            )
            # Making output select only gettable, since we are
            # explicitly mapping auxouts to X, Y, R, Theta, etc.
            self._set_output_select(ch)

        self.add_parameter(
            name='phase',
            label='Phase',
            unit='deg',
            get_cmd=self._get_phase,
            get_parser=float,
            set_cmd=self._set_phase,
            vals=vals.Numbers(-180,180)
        )
        self.add_parameter(
            name='time_constant',
            label='Time constant',
            unit='s',
            get_cmd=self._get_time_constant,
            get_parser=float,
            set_cmd=self._set_time_constant,
            vals=vals.Numbers()
        )
        self.add_parameter(
            name='frequency',
            label='Frequency',
            unit='Hz',
            get_cmd=self._get_frequency,
            get_parser=float
        )
        self.add_parameter(
            name='sigout_range',
            label='Signal output range',
            unit='V',
            get_cmd=self._get_sigout_range,
            get_parser=float,
            set_cmd=self._set_sigout_range,
            vals=vals.Enum(0.01, 0.1, 1, 10)
        )
        self.add_parameter(
            name='sigout_offset',
            label='Signal output offset',
            unit='V',
            get_cmd=self._get_sigout_offset,
            get_parser=float,
            set_cmd=self._set_sigout_offset,
            vals=vals.Numbers(-1, 1),
            docstring='Multiply by sigout_range to get actual offset voltage.'
        )
        for i in range(num_sigout_mixer_channels):
            self.add_parameter(
                name=f'sigout_enable{i}',
                label=f'Signal output mixer {i} enable',
                get_cmd=lambda mixer_channel=i: self._get_sigout_enable(mixer_channel),
                get_parser=float,
                set_cmd=lambda amp, mixer_channel=i: self._set_sigout_enable(mixer_channel, amp),
                vals=vals.Enum(0,1,2,3),
                docstring="""\
                0: Channel off (unconditionally)
                1: Channel on (unconditionally)
                2: Channel off (will be turned off on next change of sign from negative to positive)
                3: Channel on (will be turned on on next change of sign from negative to positive)
                """
            )
            self.add_parameter(
                name=f'sigout_amplitude{i}',
                label=f'Signal output mixer {i} amplitude',
                unit='Gain',
                get_cmd=lambda mixer_channel=i: self._get_sigout_amplitude(mixer_channel),
                get_parser=float,
                set_cmd=lambda amp, mixer_channel=i: self._set_sigout_amplitude(mixer_channel, amp),
                vals=vals.Numbers(-1, 1),
                docstring='Multiply by sigout_range to get actual output voltage.'
            )

    def _get_phase(self) -> float:
        path = f'/{self.dev_id}/demods/{self.demod}/phaseshift/'
        return self.daq.getDouble(path)

    def _set_phase(self, phase: float) -> None:
        path = f'/{self.dev_id}/demods/{self.demod}/phaseshift/'
        self.daq.setDouble(path, phase)

    def _get_gain(self, channel: str) -> float:
        path = f'/{self.devid}/auxouts/{self.auxouts[channel]}/scale/'
        return self.daq.getDouble(path)

    def _set_gain(self, gain: float, channel: str) -> None:
        path = f'/{self.dev_id}/auxouts/{self.auxouts[channel]}/scale/'
        self.daq.setDouble(path, gain)

    def _get_offset(self, channel: str) -> float:
        path = f'/{self.dev_id}/auxouts/{self.auxouts[channel]}/offset/'
        return self.daq.getDouble(path)

    def _set_offset(self, offset: float, channel: str) -> None:
        path = f'/{self.dev_id}/auxouts/{self.auxouts[channel]}/offset/'
        self.daq.setDouble(path, offset)

    def _get_output_value(self, channel: str) -> float:
        path = f'/{self.dev_id}/auxouts/{self.auxouts[channel]}/value/'
        return self.daq.getDouble(path)

    def _get_output_select(self, channel: str) -> str:
        path = f'/{self.dev_id}/auxouts/{self.auxouts[channel]}/outputselect/'
        idx = self.daq.getInt(path)
        return self.OUTPUT_MAPPING[idx]

    def _set_output_select(self, channel: str) -> None:
        path = f'/{self.dev_id}/auxouts/{self.auxouts[channel]}/outputselect/'
        keys = list(self.OUTPUT_MAPPING.keys())
        idx = keys[list(self.OUTPUT_MAPPING.values()).index(channel)]
        self.daq.setInt(path, idx)

    def _get_time_constant(self) -> float:
        path = f'/{self.dev_id}/demods/{self.demod}/timeconstant/'
        return self.daq.getDouble(path)

    def _set_time_constant(self, tc: float) -> None:
        path = f'/{self.dev_id}/demods/{self.demod}/timeconstant/'
        self.daq.setDouble(path, tc)

    def _get_sigout_range(self) -> float:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/range/'
        return self.daq.getDouble(path)

    def _set_sigout_range(self, rng: float) -> None:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/range/'
        self.daq.setDouble(path, rng)

    def _get_sigout_offset(self) -> float:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/offset/'
        return self.daq.getDouble(path)

    def _set_sigout_offset(self, offset: float) -> None:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/offset/'
        self.daq.setDouble(path, offset)

    def _get_sigout_amplitude(self, mixer_channel: int) -> float:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/amplitudes/{mixer_channel}/'
        return self.daq.getDouble(path)

    def _set_sigout_amplitude(self, mixer_channel: int, amp: float) -> None:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/amplitudes/{mixer_channel}/'
        self.daq.setDouble(path, amp)

    def _get_sigout_enable(self, mixer_channel: int) -> int:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/enables/{mixer_channel}/'
        return self.daq.getInt(path)

    def _set_sigout_enable(self, mixer_channel: int, val: int) -> None:
        path = f'/{self.dev_id}/sigouts/{self.sigout}/enables/{mixer_channel}/'
        self.daq.setInt(path, val)

    def _get_frequency(self) -> float:
        path = f'/{self.dev_id}/demods/{self.demod}/freq/'
        return self.daq.getDouble(path)

    def sample(self) -> dict:
        path = f'/{self.dev_id}/demods/{self.demod}/sample/'
        return self.daq.getSample(path)


import argparse
import os
import re
import sys
import time
import json
import logging
from collections import defaultdict
#from langchain.llms import OpenAI
#import tiktoken
import openai
import numpy as np
from bs4 import BeautifulSoup
from lib.utils import deserialize_dictionary  #TODO: fix the module import from parsing.py
from lib.utils import load_json_from_file
import ast

import nltk
nltk.data.path.append('./nltk_data/')
from nltk.tokenize import word_tokenize

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# json_string = '{"First Name": "firstName", "Middle Name": "middleName", "Middle Initial": "middleName-initial", "Last Name": "lastName", "Individual Taxpayer Identification Number (itin)": "itin",\
#                "Mailing Addressline1": "mailingAddressLine1", "Mailing Addressline2 (optional)": "mailingAddressLine2", "Zipcode/Postcode": "zip/postalCode",\
#                "State/Province/Region": "state/province/region", "City": "city", "Country": "country", "Primary Phone Number": "primaryPhoneNumber", \
#                "Phone Type": "primaryPhoneType", "Employment Status":"employmentStatus", "Income(annual)": "income-annual-usd", "Housing Status": "housingStatus",\
#                "Monthly Rent/Mortgage": "monthlyRent/Mortgage", "Education Degree": "educationDegree", "Working Industry": "workingIndustry", \
#                "Professional Title": "professionalTitle", "Professional Tenure":"professionalTenure", "Social Security Number (ssn)": "ssn", \
#                "Date of Birth": "dob", "Email Address": "email"}'

json_string = '''{
  "First Name": "firstName",
  "Middle Name": "middleName",
  "Middle Initial": "middleName-initial",
  "Last Name": "lastName",
  "Suffix": "suffix",
  "Tax ID Type": "taxIDType",
  "Individual Taxpayer Identification Number (itin)": "itin",
  "Social Security Number (ssn)": "ssn",
  "Date of Birth (mm/dd/yyyy)": "dob-mm/dd/yyyy",
  "Mother Maiden Name": "motherMaidenName",
  "Address Type": "addressType",
  "Mailing Addressline1": "mailingAddressLine1",
  "Mailing Addressline2 (optional)": "mailingAddressLine2",
  "Zipcode/Postcode": "zip/postalCode",
  "State/Province/Region": "state/province/region",
  "City": "city",
  "Country": "country",
  "Primary Phone Type": "primaryPhoneType",
  "Primary Phone Number": "primaryPhoneNumber",
  "Email Address": "email",
  "Employment Status": "employmentStatus",
  "Income(annual)": "income-annual-usd",
  "Type of residence/Housing Status": "housingStatus",
  "Monthly Rent/Mortgage": "monthlyRent/Mortgage",
  "Education Degree": "educationDegree",
  "Working Industry": "workingIndustry",
  "Professional Title": "professionalTitle",
  "Professional Tenure": "professionalTenure",
  "Linkedin Profile": "linkedinProfile",
  "Personal Website": "personalWebsite",
  "Visa Sponsorship": "visaSponsorship"
}'''

# Convert the JSON string to a Python dictionary
name_maps = json.loads(json_string)

input_string = """First Name: Min
Middle Initial: NA
Last Name: Lu
Mailing Addressline1: 15481 Bristol Ridge Ter
Mailing Addressline2 (optional): #54
State/Province/Region: CA
City: San Diego
Zipcode/Postcode: 92127
Email Address: minlu19@gmail.com
Primary Phone Number: 2176399259
Employment Status: employed
Education Degree: Doctoral Degree (Ph.D, Ed.D, M.D.)
Income(annual): 250000
Monthly Rent/Mortgage: 2500
Date of Birth (mm/dd/yyyy): 12/15/1989
Social Security Number (ssn): 316-14-4952
Middle Name: NA
Suffix: None
Tax ID Type: SOCIAL_SECURITY_NUMBER
Individual Taxpayer Identification Number (itin): 316-88-8888
Country: USA
Primary Phone Type: CELL
Type of residence/Housing Status: OWN
Working Industry: Tech
Professional Title: Engineer
Professional Tenure: 10 Years
Address Type: DOMESTIC
Mother Maiden Name: A
Linkedin Profile: https://www.linkedin.com/in/abc/
Personal Website: https://www.abc.com
Visa Sponsorship: NO"""

label_values_map = {
                         "veteranStatus": {"Filler":["please select"],"Veteran":["Veteran"], "Not Veteran":["Not Veteran"], "Other":["Other", "I don't want to declare"]},
                         "gender": {"Filler":["please select"],"male":["male","m"],"female":["female","f"],"other":["other", "x", "don't want to specify", "decline to self identify"]},
                         "visaSponsorship": {"Filler":["please select"],"Yes":["Yes"], "No":["No"]},
                         "disabilityStatus": {"Filler":["please select"],"Yes":["Yes"], "No":["No"]},
                         "isHispanicLatino": {"Filler":["please select"],"Yes":["Yes"], "No":["No"]},
                         # equalent to "residency status":["US citizen", "Green card (Permanent Resident)", "Foreign (Non-resident)"],
                         "race": {"Filler":["please select"],"American Indian or Alaskan Native":["American Indian or Alaskan Native"], "Asian":["Asian"], "Black or African American":["Black or African American"], "Hispanic or Latino":["Hispanic or Latino"], "White":["White"], "Native Hawaiian or Other Pacific Islander":["Native Hawaiian or Other Pacific Islander"], "Two or More Races":["Two or More Races"], "Decline To Self Identify":["Decline To Self Identify"]},

}

# input_string = """First Name: Xin
# Middle Name: X
# Last Name: Xu
# Suffix: None
# Tax ID Type: SSN
# Individual Taxpayer Identification Number (itin): 316-14-8882
# Social Security Number (ssn): 316-14-8882
# Date of Birth (mm/dd/yyyy): 12/15/1989
# Mother's Maiden Name: A
# Address Type: DOMESTIC
# Mailing Addressline1: 15469 Bristol Ridge Ter
# Mailing Addressline2: #58
# Zipcode/Postcode: 92122
# State/Province/Region: CA
# City: San Diego
# Country: USA
# Primary Phone Type: CELL
# Primary Phone Number: 2176399999
# Email Address: minlu8@gmail.com
# Employment Status: Employed
# Income(annual): 250000
# Type of residence/Housing Status: OWN
# Monthly Rent/Mortgage: 2500
# Education Degree: Doctoral Degree (Ph.D, Ed.D, M.D.)
# Working Industry: Tech
# Professional Title: Engineer
# Professional Tenure: 10 Years
# """

# Split the string into lines
lines = input_string.splitlines()

# Initialize an empty dictionary
pii_dict = {}

# Iterate through the lines and add them to the dictionary
for line in lines:
    key, value = line.split(": ", 1)  # split at the first occurrence of ': '
    pii_dict[key] = value
piikeys = pii_dict.keys()
piistr = str(pii_dict)

print(piistr)

def cosine_similarity_matrix(embeddings1, embeddings2):
    # Normalize the embeddings
    embeddings1_normalized = embeddings1 / np.linalg.norm(embeddings1, axis=1, keepdims=True)
    embeddings2_normalized = embeddings2 / np.linalg.norm(embeddings2, axis=1, keepdims=True)

    # Compute the similarity matrix
    similarity = np.dot(embeddings1_normalized, embeddings2_normalized.T)

    return similarity

def load_html_files_from_directory(directory_path):
    # List all files in the directory
    all_files = [f for f in os.listdir(directory_path) if os.path.isfile(os.path.join(directory_path, f))]

    # Filter for .html files
    html_files = [f for f in all_files if f.endswith('.html')]

    # Load each HTML file into memory
    html_contents = {}
    for html_file in html_files:
        with open(os.path.join(directory_path, html_file), 'r', encoding='utf-8') as file:
            html_contents[html_file] = file.read()

    return html_contents

def nltk_tokenize(input_string):
    tokens_nltk = word_tokenize(input_string.lower())
    tokens_nltk = [token for token in tokens_nltk if not no_alphanumeric_string(token)]
    return ' '.join(tokens_nltk)

def no_alphanumeric_string(s):
    return bool(re.match(r'^[^\w]', s))

def remove_trailing_non_alnum_regex(s):
    return re.sub(r'[^a-zA-Z0-9]*$', '', s)

def map_select_values(name, label_values, form_texts):

    value_tokens = []
    form_texts_list = []
    for form_text in form_texts:
        name_form_text = ' '.join([name.lower(), form_text.lower()])
        value_tokens.append(name_form_text)
        form_texts_list.append(form_text)
    print(value_tokens)
    openai_api_key = os.environ["OPENAI_API_KEY"]
    resp = openai.Embedding.create(
                input=value_tokens,
                engine="text-embedding-ada-002")
    value_token_embedding = dict()
    for i in range(len(value_tokens)):
        #"first name", embeddding vector
        value_token_embedding[value_tokens[i]] = resp['data'][i]['embedding']

    value_token_embedding_array_2d = []
    value_token_labels = []
    for k, v in value_token_embedding.items():
        value_token_embedding_array_2d.append(v)
        value_token_labels.append(k)
    value_token_embedding_array_2d = np.vstack(value_token_embedding_array_2d)

    synonym_label = []
    synonym_label_tokens = []
    label_synonym_embedding_array_2d = []
    label_value_synoyms = label_values_map.get(name)
    if label_value_synoyms:
        for label in label_values:
            synonyms = label_value_synoyms.get(label)
            for synonym in synonyms:
                name_value = ' '.join([name.lower(), synonym.lower()])
                print(name_value)
                synonym_label_tokens.append(name_value)
                synonym_label.append(label)
    print(synonym_label_tokens)
    openai_api_key = os.environ["OPENAI_API_KEY"]
    resp = openai.Embedding.create(
                input=synonym_label_tokens,
                engine="text-embedding-ada-002")

    for i in range(len(synonym_label_tokens)):
        #"first name", embeddding vector
        label_synonym_embedding_array_2d.append(resp['data'][i]['embedding'])
    label_synonym_embedding_array_2d = np.vstack(label_synonym_embedding_array_2d)

    simmatrix = cosine_similarity_matrix(value_token_embedding_array_2d, label_synonym_embedding_array_2d)
    label_dict = dict()
    for ind, value in enumerate(value_tokens):
        name_value = ' '.join([name.lower(), value.lower()])
        dot_products = simmatrix[ind,:]
        max_index = np.argmax(dot_products)
        max_similarity = dot_products[max_index]
        print(f'{value} is mapped to {synonym_label[max_index]} with score: {max_similarity}')
        #if max_similarity > .9:
        label_dict[synonym_label[max_index]] = (form_texts_list[ind], max_similarity)
    return label_dict

def generate_xpath(element):
    path_parts = []
    current = element
    while current is not None and current.name is not None and current.parent is not None:
        siblings = list(current.parent.children)
        tag_siblings = [sibling for sibling in siblings if sibling.name == current.name]
        if len(tag_siblings) > 1:
            index = tag_siblings.index(current) + 1  # XPath is 1-indexed
            path_parts.append(f"{current.name}[{index}]")
        else:
            path_parts.append(current.name)
        current = current.parent
    return '/' + '/'.join(reversed(path_parts))

def get_inverted_index(html_data):
    inverted = dict()
    for filename, content in html_data.items():
        print(f"Contents of {filename}:")
        #print(content)
        page = BeautifulSoup(content, 'html.parser')
        vocab = set()
        # Extracting text strings for all elements in the soup
        for s in page.stripped_strings:
            vocab.add(s.lower())

        # Extracting the 'name' attributes for all elements in the soup
        element_name_attributes = [tag['name'] for tag in page.find_all() if tag.has_attr('name')]
        for s in element_name_attributes:
            vocab.add(s.lower())

        #print(vocab)
        for v in vocab:
            if v not in inverted:
                inverted[v] = [filename]
            else:
                inverted[v].append(filename)
    return inverted

def scan_for_inputs(html_data, label_dict, output_folder):
    for filename, content in html_data.items():
        print(f"Contents of {filename}:")
        #print(content)
        soup, inputs = get_all_inputs(content)
        data = get_label_for_inputs(inputs, soup, label_dict)
        json_output = output_folder + '/' + filename[:-5] + ".json"
        with open(json_output, 'w') as json_file:
            json.dump(data, json_file)

def get_all_inputs(html_content):
    soup = BeautifulSoup(html_content, 'html.parser')
    inputs  = soup.find_all('input')
    selects = soup.find_all('select')
    textareas = soup.find_all('textarea')
    radios = soup.find_all(attrs={'role': 'radio'})
    checkboxes = soup.find_all(attrs={'role': 'checkbox'})
    submit_buttons = soup.find_all(attrs={'role': 'submit'})
    radiogroups = soup.find_all(attrs={'role': 'radiogroup'})
    input_elements = inputs + selects + textareas + radiogroups + radios + checkboxes + submit_buttons
    return soup, input_elements

def check_ground_truth(groundTruth, prediction):
    target_dict = {}
    for ele in groundTruth:
        if len(ele[0]) > 1:
            target_dict[ele[0]] = ele[2]
    pred_dict = {}
    for ele in prediction:
        if len(ele[0]) > 1:
            pred_dict[ele[0]] = ele[2]
    #print(target_dict)
    #print(pred_dict)
    miss = 0
    true_det = 0
    false_det = 0
    match_cat = {}
    for key, value in target_dict.items():
        if value != "no match":
            if key not in pred_dict:
                miss = miss + 1
                match_cat[key] = "miss|" + value
            else:
                pred_label = pred_dict[key]
                if pred_label == value:
                    true_det = true_det + 1
                    match_cat[key] = pred_label + "|" + value
                else:
                    false_det = false_det + 1
        else:
            if key in pred_dict:
                pred_label = pred_dict[key]
                if pred_label != "no match":
                    false_det = false_det + 1
                    match_cat[key] = pred_label + "|" + value
    return true_det, false_det, miss, match_cat

def get_label_for_fields(fields, label_dict):
    data = []
    for field in fields:
        elem_id = ""
        elem_type = ""
        label_text = ""
        label_conf = 0.0
        #if we can find a good match using name, skip
        foundLabel = False
        if field.get('htmlName'):
            s =  field.get('htmlName').lower()
            s = nltk_tokenize(s)
            if s in label_dict:
                label_text = label_dict[s][0] #label category
                label_conf = label_dict[s][1] #confidence score
                foundLabel = True
        if not foundLabel:
            if field.get('label-tag'):
                s =  field.get('label-tag').lower()
                s = nltk_tokenize(s)
                if s in label_dict:
                    label_text = label_dict[s][0] #label category
                    label_conf = label_dict[s][1] #confidence score
                    foundLabel = True
        if field.get('htmlID'):
            elem_id = "#"+field.get('htmlID')

        if field.get('tagName'):
            elem_type = field.get('tagName')

        mappings = dict()
        print(label_text)
        if len(label_text) > 1 and field.get('tagName') == 'select':
            options = field.get('options')
            text_value_map = dict()
            for op in options:
                text_value_map[op.get("text")] = op.get("value")
            form_option_texts = text_value_map.keys()
            label_values = label_values_map[label_text]
            result = map_select_values(label_text, label_values, form_option_texts)
            mappings = result
        xpath = field.get('htmlSelector')
        data_item = (elem_id, xpath, label_text, label_conf, elem_type, mappings)
        data.append(data_item)
    return data

def get_label_for_inputs(non_hidden_inputs, soup, label_dict):

    data = []
    for input_elem in non_hidden_inputs:

        elem_id = ""
        elem_type = ""
        label_text = ""
        label_conf = 0.0

        #if we can find a good match using name, skip
        foundLabel = False
        if input_elem.get('name'):
            s =  input_elem.get('name').lower()
            s = nltk_tokenize(s)
            if s in label_dict:
                label_text = label_dict[s][0] #label category
                label_conf = label_dict[s][1] #confidence score
                foundLabel = True
            #print(f"name is {s}, found label {foundLabel}")

        if not foundLabel:
            # Find associated label using 'for' attribute
            associated_label = soup.find('label', {'for': input_elem.get('id')})

            # If input is a child of label, the parent is the associated label
            if not associated_label and input_elem.find_parent('label'):
                associated_label = input_elem.find_parent('label')

            if associated_label:
                #print(associated_label)

                for s in associated_label.stripped_strings:
                    token = nltk_tokenize(s)
                    #print(f'normalized to {token}')
                    if token in label_dict:
                        print(f'{token} is matched to {label_dict[token]}')
                        label_text = label_dict[token][0] #label category
                        label_conf = round(label_dict[token][1], 2)#confidence score
                        foundLabel = True
                        break
                    #print(f"label is {s}, found label {foundLabel}")

        #print(f'label is {label_text}')
        #if 'id' in input_elem:
        if input_elem.get('id'):
            elem_id = "#"+input_elem.get('id')

        if input_elem.get('type'):
            elem_type = input_elem.get('type')

        xpath = generate_xpath(input_elem)
        data_item = (elem_id, xpath, label_text, label_conf, elem_type)
        data.append(data_item)

    return data

def chunk_list(input_list, chunk_size):
    return [input_list[i:i + chunk_size] for i in range(0, len(input_list), chunk_size)]

def get_embedding(input_tokens):
    print(input_tokens)
    token_embedding = dict()
    sub_tokens_list = chunk_list(input_tokens, 20)
    for sub_tokens in sub_tokens_list:
        print(sub_tokens)
        resp = openai.Embedding.create(
                input=sub_tokens,
                engine="text-embedding-ada-002")
        for i in range(len(sub_tokens)):
            token_embedding[sub_tokens[i]] = resp['data'][i]['embedding']
        print(sub_tokens)
    return token_embedding

def classifier_function(fieldName, sources, targets, model = "gpt-4"):
    messages = [{"role": "user", "content": f"You are domain expert in data schema matching for {fieldName}, you are generating map between two value lists, do not generate python code! Now you are acting like mapper function from a list specified by ```{sources}``` on to one of predefined list specified in: ```{targets}```, meaning the response should be one of the element in ```{targets}```.\n\n Only respond with your `return` python dictionary object. Do not include any other explanatory text or python code in your response."}]

    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0
    )

    return response.choices[0].message["content"]

def ai_function(function, args, description, model = "gpt-4"):
    # parse args to comma separated string
    args = ", ".join(args)
    messages = [{"role": "system", "content": f"You are now the following python function: ```# {description}\n{function}```\n\nOnly respond with your `return` value. Do not include any other explanatory text in your response."},{"role": "user", "content": args}]

    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0
    )

    return response.choices[0].message["content"]

def cosine_similarity_matrix(embeddings1, embeddings2):
    # Normalize the embeddings
    embeddings1_normalized = embeddings1 / np.linalg.norm(embeddings1, axis=1, keepdims=True)
    embeddings2_normalized = embeddings2 / np.linalg.norm(embeddings2, axis=1, keepdims=True)

    # Compute the similarity matrix
    similarity = np.dot(embeddings1_normalized, embeddings2_normalized.T)

    return similarity

def matching_simple(formfields, openai_api_key, model = "gpt-3.5-turbo"):
    openai.api_key = openai_api_key
    pii_names = []
    for readable_name, machine_name in name_maps.items():
        pii_names.append((readable_name, machine_name))

    form_names = []
    for element in formfields:
        form_name = ""
        if 'name' in element:
            form_name = form_name + ' ' + element['name']
        if 'label' in element:
            form_name = form_name + ' ' + element['label']
        if 'context' in element:
            form_name = form_name + ' ' + element['context']
        if len(form_name) < 1:
            continue
        form_names.append(form_name)

    function_string = "def map_names(form_names: list, pii_names: list) -> dict:"
    description_string = """Based on the semantic of list elements, build a map from the first list to the second list. please make sure the value of the map are one of tuple in second list or 'no match' if no good match can be found in second list"""
    args = [str(pii_names), str(form_names)]
    result_string = ai_function(function_string, args, description_string, model)
    logging.info(result_string)
    result = ast.literal_eval(result_string)
    print(result)

def matching_simple_bak(formfields, openai_api_key, model = "gpt-3.5-turbo"):
    openai.api_key = openai_api_key

    pii_names = []
    for readable_name, machine_name in name_maps.items():
        pii_names.append(readable_name)

    resp = openai.Embedding.create(
                input=pii_names,
                engine="text-embedding-ada-002")

    pii_name_embedding_array_2d = []
    for i in range(len(pii_names)):
        pii_name_embedding_array_2d.append(resp['data'][i]['embedding'])
    pii_name_embedding_array_2d = np.vstack(pii_name_embedding_array_2d)

    form_names = []
    for element in formfields:
        form_name = ""
        if 'name' in element:
            form_name = form_name + ' ' + element['name']
        elif 'label' in element:
            form_name = form_name + ' ' + element['label']
        elif 'context' in element:
            form_name = form_name + ' ' + element['context']
        if len(form_name) < 1:
            continue
        form_names.append(form_name)

    resp = openai.Embedding.create(
            input=form_names,
            engine="text-embedding-ada-002")

    form_name_embedding_array_2d = []
    for i in range(len(form_names)):
        form_name_embedding_array_2d.append(resp['data'][i]['embedding'])
    form_name_embedding_array_2d = np.vstack(form_name_embedding_array_2d)

    simmatrix = cosine_similarity_matrix(pii_name_embedding_array_2d, form_name_embedding_array_2d)

    form_name_candidates = {}
    for ind, form_name in enumerate(form_names):
        dot_products = simmatrix[:,ind]
        max_index = np.argmax(dot_products)
        max_similarity = dot_products[max_index]
        top_k_indices = np.argsort(dot_products)[-5:]
        candidates = []
        for index in top_k_indices:
            pii_name = pii_names[index]
            similarity = dot_products[index]
            if similarity > .6 and similarity > max_similarity - .05:
                candidates.append(pii_name)
        #verify the match with classification
        form_name_candidates[form_name] = candidates

    form_name_to_pii_name = []
    for element in formfields:
        form_name = ""
        machine_name = "no match"
        sub_dict = ""
        if 'name' in element:
            form_name = form_name + ' ' + element['name']
            sub_dict = element['name']
        elif 'label' in element:
            form_name = form_name + ' ' + element['label']
        elif 'context' in element:
            form_name = form_name + ' ' + element['context']
        if len(form_name) >= 1:
            #print('\n\n')
            #print(form_name)
            sub_dict = sub_dict + ' ' +element['label'] + ' ' + element['context']
            candidates = form_name_candidates[form_name]
            result = classifier_function(sub_dict, candidates,  "gpt-3.5-turbo")
            element['matched'] = result
            if result in name_maps:
                machine_name = name_maps[result]
            else:
                logging.warning(f"Key not found in name_maps: {result}")
            print(sub_dict)
            print(candidates)
            print(result)
        xid = ''
        if 'id' in element:
            xid = '#'+element['id']
        xpath = form_name
        input_type = 'NA'
        if 'type' in element:
            input_type = element['type']
        pii_name = machine_name
        mapped_value = 'N/A'
        form_name_to_pii_name.append((xid, xpath, pii_name, mapped_value, input_type))
    return form_name_to_pii_name

def matching(formfields, openai_api_key, model = "gpt-3.5-turbo"):
    openai.api_key = openai_api_key
    form_key_values = []
    form_key_texts = defaultdict(dict)
    for key, values in formfields.items():
        if (key[1] is None) or (len(key[1]) < 1):
            continue
        form_key_values.append(key[1])
        text_to_value = dict()
        for v in values:
            print(v)
            # key[1] is the actual name / context of the input group
            # v[0], v[1] is the id and xpath of input element,
            # v[2] is the value of the input element
            text_to_value[(key[1], str(v[2]))] = (v[0],v[1],v[3])

        form_key_texts[key[1]] = text_to_value

    unique_form_keys = set(form_key_values)
    if len(unique_form_keys) < len(form_key_values):
        logging.warning('This means some form field names have duplication')
        exit(1)

    #print(formfields['input group 5  Social Security number'])

    pii_key_values = []
    for key, values in pii_dict.items():
        pii_key_values.append(key)

    #print(form_key_values)
    #print(pii_key_values)

    function_string = "def map_names(form_key_values: list, pii_key_values: list) -> dict:"
    description_string = """Based on the semantic of list elements, build a map from the first list to the second list."""
    args = [str(pii_key_values), str(form_key_values)]
    result_string = ai_function(function_string, args, description_string, model)
    logging.info(result_string)
    result = ast.literal_eval(result_string)
    print(result)
    filtered_result = {}
    for pii_name_raw, form_name in result.items():
        if (form_name is None) or (len(form_name) == 0):
            continue
        print(pii_name_raw)
        print(form_name)
        # Compute the cosine similarity
        resp = openai.Embedding.create(
            input=[pii_name_raw, form_name],
            engine="text-similarity-davinci-001")

        embedding_a = resp['data'][0]['embedding']
        embedding_b = resp['data'][1]['embedding']

        similarity = np.dot(embedding_a, embedding_b)
        #similarity = nlp(pii_name_raw).similarity(nlp(form_name))
        print(similarity)
        if similarity >=.6:
            filtered_result[pii_name_raw] = form_name
    result = filtered_result

    form_name_to_pii_name = []

    matched_form_fields = set(result.values())
    for form_name in form_key_values:
        if len(form_name) < 1:
            continue
        if form_name not in matched_form_fields:
            form_value = []
            for (k,v) in form_key_texts[form_name]:
                form_value.append(v)
            if len(form_value) > 1:
                for k, v in form_key_texts[form_name]:
                    xid = v[0]
                    xpath = v[1]
                    input_type = v[2] if len(v) > 2 else 'unknown'
                    print(input_type)
                    form_name_to_pii_name.append((xid, xpath, 'unknown', 'na', input_type))
            if len(form_value) == 1:
                default_value = form_value[0]
                xid, xpath, input_type = form_key_texts[form_name][(form_name, default_value)]
                print(input_type)
                form_name_to_pii_name.append((xid, xpath, 'unknown', 'na', input_type))

    for pii_name_raw, form_name in result.items():

        if pii_name_raw in name_maps:
            pii_name = name_maps[pii_name_raw]
        else:
            logging.warning(f"Key not found in name_maps: {pii_name_raw}")
            continue

        form_value = []
        for (k,v) in form_key_texts[form_name]:
            form_value.append(v)
        print(pii_name_raw)
        print(pii_name)
        print(form_name)
        print(form_value)
        pii_value = pii_dict[pii_name_raw]

        if len(form_value) > 1:
            function_string = "def value_select(values: list, query: str) -> str:"
            args = [str(form_value), pii_value]
            description_string = """Based on the query string to select the best value in the list. Please output the whole result, no abbreviation"""
            mapped_value = ai_function(function_string, args, description_string, model)
            if (form_name, mapped_value) in form_key_texts[form_name]:
                xid, xpath, input_type = form_key_texts[form_name][(form_name, mapped_value)]
                # logging form_name_values
                logging.info((xid, xpath, mapped_value))
                form_name_to_pii_name.append((xid, xpath, pii_name, mapped_value))
            else:
                logging.warning(f"Key not found in form key texts map: {(form_name, mapped_value)}")
            for k, v in form_key_texts[form_name]:
                if (k[0] == form_name) and (k[1] == mapped_value):
                    continue
                xid = v[0]
                xpath = v[1]
                input_type = v[2]
                print(input_type)
                form_name_to_pii_name.append((xid, xpath, 'unselected', 'na', input_type))

        if len(form_value) == 1:
            #check if there is format string availale
            default_value = form_value[0]
            xid, xpath, input_type = form_key_texts[form_name][(form_name, default_value)]
            mapped_value = ''
            if len(default_value) > 1:
                function_string = "def value_norm(format: str, value: str) -> str:"
                args = [default_value, pii_value]
                print(default_value)
                description_string = """Based on the format string, normalized the value string."""
                mapped_value = ai_function(function_string, args, description_string, model)
            else:
                mapped_value = pii_value
            # logging form_name_values
            logging.info((xid, xpath, mapped_value))
            form_name_to_pii_name.append((xid, xpath, pii_name, mapped_value, input_type))

        if len(form_value) == 0:
            mapped_value = pii_value
            # logging form_name_values
            logging.info((xid, xpath, mapped_value))
            form_name_to_pii_name.append((xid, xpath, pii_name, mapped_value, input_type))
    return form_name_to_pii_name

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Matching pii data to extracted form fields.")
    parser.add_argument("--method", type=str, default="simple", help="The pii data json file.")
    parser.add_argument("--formfields", type=str, help="The extraced form fields pickle file.")
    parser.add_argument("--output", type=str, help="The filled information for #id/#xpath of input elements, in json format.")

    args = parser.parse_args()

    if args.method == 'simple':
        formfields = load_json_from_file(args.formfields)
        openai_api_key = os.environ["OPENAI_API_KEY"]
        json_output = args.output
        model = "gpt-3.5-turbo"
        form_name_to_pii_name = matching_simple(formfields, openai_api_key, model)
        with open(json_output, 'w') as json_file:
            json.dump(form_name_to_pii_name, json_file)

    else:

        formfields = deserialize_dictionary(args.formfields)
        openai_api_key = os.environ["OPENAI_API_KEY"]
        json_output = args.output

        model = "gpt-3.5-turbo"
        form_name_to_pii_name = matching(formfields, openai_api_key, model)

        with open(json_output, 'w') as json_file:
            json.dump(form_name_to_pii_name, json_file)


import torch
import torch.nn as nn
import torch.distributions as dist
import torch.nn.functional as F

class LinearVariational(nn.Module):
    """
    Mean field approximation of nn.Linear
    """
    def __init__(self, in_features, out_features, parent, n_batches, bias=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.include_bias = bias
        self.parent = parent
        self.n_batches = n_batches

        if getattr(parent, 'accumulated_kl_div', None) is None:
            parent.accumulated_kl_div = 0

        # Initialize the variational parameters.
        # Q(w)=N(mu_theta,sigma2_theta)
        # Do some random initialization with sigma=0.001
        self.w_mu = nn.Parameter(
            torch.FloatTensor(in_features, out_features).normal_(mean=0, std=0.1)
        )
        # proxy for variance
        # log(1 + exp(ρ))* eps
        self.w_p = nn.Parameter(
            torch.FloatTensor(in_features, out_features).normal_(mean=-3., std=0.1)
        )
        # # Prior mean
        # self.w_mu_pr = nn.Parameter(
        #     torch.FloatTensor(in_features, out_features).normal_(mean=0, std=0.1)
        # )

        if self.include_bias:
            self.b_mu = nn.Parameter(
                #torch.zeros(out_features)
                torch.FloatTensor(out_features).normal_(mean=0., std=0.1)
            )
            # proxy for variance
            self.b_p = nn.Parameter(
                #torch.zeros(out_features)
                torch.FloatTensor(out_features).normal_(mean=-3., std=0.1)
            )
            # # bias prior mean
            # self.b_mu_pr = nn.Parameter(
            #     #torch.zeros(out_features)
            #     torch.FloatTensor(out_features).normal_(mean=0., std=0.1)
            # )

    def reparameterize(self, mu, p):
        sigma = torch.log(1 + torch.exp(p))
        eps = torch.randn_like(sigma)
        return mu + (eps * sigma)

    def kl_divergence(self, mu_theta, p_theta, mu_prior, prior_sd=1.):
        # log_prior = dist.Normal(mu_prior, prior_sd).log_prob(z)
        # log_p_q = dist.Normal(mu_theta, torch.log(1 + torch.exp(p_theta))).log_prob(z)

        # return (log_p_q - log_prior).sum() / self.n_batches
        postr = dist.Normal(mu_theta, torch.log(1 + torch.exp(p_theta)))
        prior = dist.Normal(mu_prior, prior_sd)
        return dist.kl.kl_divergence(postr,prior).sum() / self.n_batches

    # def forward(self, x):
    #     w = self.reparameterize(self.w_mu, self.w_p)

    #     if self.include_bias:
    #         b = self.reparameterize(self.b_mu, self.b_p)
    #     else:
    #         b = 0

    #     z = x @ w + b

    #     z_mu, z_var = self.calcPredDist(x)

    #     self.parent.accumulated_kl_div += self.kl_divergence(w, self.w_mu, self.w_p, 0., 1.)
    #     if self.include_bias:
    #         self.parent.accumulated_kl_div += self.kl_divergence(b, self.b_mu, self.b_p, 0., 1.)
    #     return z, z_mu, z_var

    def forward(self, x):
        # sampling delta_W
        sigma_weight = torch.log1p(torch.exp(self.w_p))
        delta_weight = (sigma_weight * torch.randn_like(sigma_weight))

        if self.include_bias:
            sigma_bias = torch.log1p(torch.exp(self.b_p))
            bias = (sigma_bias * torch.randn_like(sigma_bias))
        else:
            bias = torch.zeros(self.out_features)

        # get kl divergence
        self.parent.accumulated_kl_div += self.kl_divergence(self.w_mu, self.w_p, 0., 1.)
        if self.include_bias:
            self.parent.accumulated_kl_div += self.kl_divergence(self.b_mu, self.b_p, 0., 1.)

        #print(self.w_mu,self.b_mu)
        #print(self.w_mu.size(),self.b_mu.size())
        # linear outputs
        outputs = F.linear(x, self.w_mu.transpose(0,1), self.b_mu)

        sign_input = x.clone().uniform_(-1, 1).sign()
        sign_output = outputs.clone().uniform_(-1, 1).sign()

        perturbed_outputs = F.linear(x * sign_input, delta_weight.transpose(0,1), bias) * sign_output

        z_mu, z_var = self.calcPredDist(x)

        # returning outputs + perturbations
        return outputs + perturbed_outputs, z_mu, z_var

    def calcPredDist(self, x):
        if self.include_bias:
            z_mu = x @ self.w_mu + self.b_mu
            #z_var = (x @ self.w_p)**2 + self.b_p**2
            z_var = x**2 @ torch.log(1 + torch.exp(self.w_p))**2 + torch.log(1 + torch.exp(self.b_p))**2
        else:
            z_mu = x @ self.w_mu
            z_var = x**2 @ torch.log(1 + torch.exp(self.w_p))**2
        return z_mu, z_var


#!/usr/bin/env python
# coding=utf-8
import sys
sys.path.append('../')
import torch
import torch.nn as nn
from transformer.attention import MultiHeadAttention
from transformer.positionwise_feedforward import PositionwiseFeedForward

class EncoderLayer(nn.Module):
    """
    EncoderLayer = self_attention with sublayerconnection + feedforward with sublayerconnection
    """
    def __init__(self, d_model, d_hidden, n_heads, dropout=0.3):
        """
        Args:
            d_model: hidden size of transformer
            d_hidden: feed forward hidden size, usually 4 * d_model
            n_heads: number of heads in multi-head attention
            dropout: dropout rate
        """
        super(EncoderLayer, self).__init__()
        self.self_attention = MultiHeadAttention(n_heads=n_heads, d_model=d_model, dropout=dropout)
        self.pos_feedforward = PositionwiseFeedForward(d_model=d_model, d_hidden=d_hidden, dropout=dropout)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, encoder_input, mask=None):
        encoder_input, encoder_attn = self.self_attention(encoder_input, encoder_input, encoder_input, mask=mask)
        encoder_output = self.pos_feedforward(encoder_input)
        return encoder_output, encoder_attn

class DecoderLayer(nn.Module):
    """
    DecoderLayer = self_attention + enc_attention + feedforward
    """
    def __init__(self, d_model, d_hidden, n_heads, dropout=0.3):
        """
        Args:
            d_model: hidden size of transformer
            d_hidden: feed forward hidden size, usually 4 * d_model
            n_heads: number of heads in multi-head attention
            dropout: dropout rate
        """
        super(DecoderLayer, self).__init__()
        self.self_attention = MultiHeadAttention(n_heads=n_heads, d_model=d_model, dropout=dropout)
        self.enc_attention = MultiHeadAttention(n_heads=n_heads, d_model=d_model, dropout=dropout)
        self.pos_feedforward = PositionwiseFeedForward(d_model=d_model, d_hidden=d_hidden, dropout=dropout)

    def forward(self, encoder_output, decoder_input, self_attn_mask=None, enc_attn_mask=None):
        decoder_output, dec_self_attn = self.self_attention(decoder_input, decoder_input, decoder_input, mask=self_attn_mask)
        decoder_output, dec_enc_attn = self.enc_attention(decoder_output, encoder_output, encoder_output, mask=enc_attn_mask)
        decoder_output = self.pos_feedforward(decoder_output)
        return decoder_output, dec_self_attn, dec_enc_attn

if __name__ == '__main__':
    d_model = 512
    d_hidden = 2048
    n_heads = 8
    batch_size = 16
    len_frames = 20

    encoder_input = torch.ones(batch_size, len_frames, d_model)
    decoder_output = torch.ones(batch_size, len_frames // 2, d_model)
    encoder_mask = torch.ones(batch_size, 1, len_frames)
    self_attn_mask = (1 - torch.triu(torch.ones(1, len_frames // 2, len_frames // 2), diagonal=1))
    dec_attn_mask = encoder_mask.clone()

    temp_encoder = EncoderLayer(d_model=d_model, d_hidden=d_hidden, n_heads=n_heads)
    temp_decoder = DecoderLayer(d_model=d_model, d_hidden=d_hidden, n_heads=n_heads)

    encoder_output = temp_encoder(encoder_input=encoder_input, mask=encoder_mask)
    print(encoder_output.shape)
    decoder_output, *_ = temp_decoder(encoder_output=encoder_output, decoder_input=decoder_output,
                                      self_attn_mask=self_attn_mask, dec_attn_mask=dec_attn_mask)
    print(decoder_output.shape)


# ----------------------------------------------------------------------------
# "34 Amazing Games in Python from the 1K ZX81"
# Python Code redevelopment Kevin Phillips, 2015
#
# Recreation of classic ZX81 games using python and pygame
# ----------------------------------------------------------------------------
# Requires:
# ----------------------------------------------------------------------------
# Python 2.7	-	www.python.org
# pygame 1.91	-	www.pygame.org
# ZXBasic.py	-	ZX Basic class (import module)
# ZX81.ttf		-	ZX81 Font (http://www.dafont.com/zx81.font) - installed
# g1.png, etc	-	A collection of 20 .png files representing key graphics
# ----------------------------------------------------------------------------
#
# GAME:			UFO, page 8
# DESCRIPTION:	Use 0 and 8 to move your laser beneath the descending UFO
#				You score when you hit it by pressing 1. Game over when it lands
#
# From the book "34 Amazing games for the 1K ZX81"
# Original book and game design (C) 1982, Alastair Gourlay, Mark Ramshaw and
# Interface publishing.
#
# ----------------------------------------------------------------------------
import ZXBasic
from random import randint
from pygame.locals import *
import pygame.time

# Note that all pygame functionality is managed in class.  We do however need
# to import the pygame.time module here to control the speed of the game
fpsClock = pygame.time.Clock()

zx = ZXBasic.ZXBasic()
zx.createScreen()
zx.CLS()

# Define game as a function
# We will use functions for all games, as some of the games have
# options to restart.  This makes more sense if called from a games loop
# in the main program.
def ufo():
	# define variables.  P is player position, L is length of river
	P = 1
	Q = 7
	S = 0
	X = 1

	zx.PRINT_AT(0,0,"|i               ")
	for x in range(1,15):
		zx.PRINT_AT(x,0,"|i |t             |i ")

	# Refresh the screen
	zx.screenRefresh()

	while True:

		# wipe the player and UFO
		zx.PRINT_AT(int(P),int(Q),'|t ')
		zx.PRINT_AT(16,X,'|t ')

		# Work out the movement of our barrel by checking the keys
		X = X + (zx.INKEY('0') - zx.INKEY('8'))

		# Move the alien
		P = P + zx.RND() / 3
		Q = Q + 2 * zx.RND() - 1

		# if the UFO has landed, we exit
		if P > 15:
			return True

		# Check for bounds
		if Q < 1 or Q > 13:
			Q = 7
		if X < 1 or X > 13:
			X = 1

		# Print our game graphics
		zx.PRINT_AT(int(P),int(Q),'|tX')
		zx.PRINT_AT(16,X,'|g7')

		# Check if we fired and hit
		if zx.INKEY('1') and X == int(Q + 0.5):
			S = S + 10
			zx.PRINT_AT(int(P),int(Q),'|iX')
			zx.PRINT_AT(0,5,'|t%s' % S)

		# Refresh the screen and loop
		zx.screenRefresh()
		fpsClock.tick(4)

while True:
	ufo()

	# Refresh the display
	zx.screenRefresh()

	# When the game has completed, wait here until the space bar is pressed.
	# Because of the key press repeat, we may find that allowing any key (like the ZX81)
	# will instantly restart the game.  We need to also keep an eye on events here for
	# QUIT as well.

	# Lets just set a flag and use a while loop to keep the game paused
	replay = True

	while replay:
		for event in pygame.event.get():
			if event.type == KEYDOWN:
				if event.key == K_SPACE:
					replay = False
			if event.type == QUIT:
				pygame.quit()
				sys.exit()

	zx.CLS()
	# The game will now return to the top of the while loop (ie start a new game)
	# Easily done


import numpy as np
import scipy.ndimage

# fh = open('ex.txt', 'r')
# shape = (5,10)
fh = open('input.txt', 'r')
shape = (100,100)

lines = fh.readlines()

seafloor = np.ndarray(shape=shape, dtype=int)
for i,line in enumerate(lines):
    seafloor[i] =[int(x) for x in line.strip()]

neighbors = np.array(
      [[0, 1, 0],
       [1, 0, 1],
       [0, 1, 0]])
lowpoints = seafloor < scipy.ndimage.minimum_filter(seafloor, footprint=neighbors, mode='constant', cval=999)
# >>> ans * np.ones(shape=ans.shape) + seafloor*ans
low_values = lowpoints * np.ones(shape=lowpoints.shape) + seafloor*lowpoints
print(low_values.sum())


from farm.builders.interfaces.abstract_habitat_builder import AbstractHabitatBuilder
from farm.domain.collections.pig_pen import PigPen
from farm.domain.single_entities.pig import Pig
from farm.visitors.collections.add_visitor import AddVisitor
from farm.visitors.collections.get_animals_visitor import GetAnimalsVisitor

class ConcretePigPenBuilder(AbstractHabitatBuilder):
    """
    ConcretePigPenBuilder Subclass

    Concrete Builder Class to set attributes, add animal, and get details
    to a pig pen instance

    Arguments:
        AbstractHabitatBuilder -- AbstractHabitatBuilder
    """

    def __init__(self) -> None:
        """
        __init__ method

        Creates an instance of Pig pen
        """
        self.pig_pen: PigPen = PigPen()

    def build_capacity(self, capacity: int) -> None:
        """
        build_capacity method

        Setter/builder method to set capacity attribute
        of pig pen instance

        Arguments:
            capacity -- int
        """
        self.pig_pen.capacity = capacity

    def build_material(self, material: str) -> None:
        """
        build_material method

        Setter/builder method to set material attribute
        of pig pen instance

        Arguments:
            material -- str
        """
        self.pig_pen.material = material

    def build_add_animal(self, pig: Pig) -> None:
        """
        build_add_animal method

        Adds Pig instance to pig pen instance

        Arguments:
            pig -- Pig
        """
        add_visitor = AddVisitor()
        self.pig_pen.accept(add_visitor, pig)

    def get_habitat(self) -> str:
        get_animals_builder = GetAnimalsVisitor()
        pigs_list = []
        if self.pig_pen.pigs[0] and self.pig_pen.pigs[0].name:
            pigs_list = [cow.name for cow in self.pig_pen.accept(get_animals_builder)]
        return f"""
        Pig Pen:
        Capacity: {self.pig_pen.capacity}
        Material: {self.pig_pen.material}
        Number: {len(self.pig_pen.accept(get_animals_builder))}
        Animals: {pigs_list}
        """


import numpy as np
from scipy.special import expit
from scipy.misc import logsumexp

def sigmoid(x):
    return expit(x)

def sigmoidForData(X, coef, intercept=None, returnFull=False):
    if coef.ndim == 1:  # coef is 1d array
        if intercept is None:
            linearTerms = X.dot(coef)
        else:
            if isinstance(intercept, float):
                linearTerms = X.dot(coef) + intercept  # intercept should be a float
            else:
                raise("'coef' is 1-dimensional, which requires a scalar intercept."
                      "Non-scalar intercept is not yet implemented.")

        sigmoids = sigmoid(linearTerms)
        if returnFull:
            return np.vstack((sigmoids, 1-sigmoids)).T
        else:
            return sigmoids
    else:  # coef is 2d array
        if intercept is None:
            linearTerms = X.dot(coef)
        else:
            if len(intercept) == coef.shape[1]:
                linearTerms = X.dot(coef) + intercept
            else:
                ValueError("The shapes of 'coef' and 'intercept' are not conformable. "
                           "'coef' has shape %s and 'intercept' has shape %s." % (coef.shape, intercept.shape))
        return sigmoid(linearTerms)

def softmax(x, addDefaultBase=True, axis=1, returnLog=False):
    # x is a 1d or 2d array
    if x.ndim == 1:
        if addDefaultBase:
            x = np.append(x, 0)
        logSumExp = logsumexp(x)
        logSoftMax = x - logSumExp
        if returnLog:
            return logSoftMax
        else:
            return np.exp(logSoftMax)
    else:
        if addDefaultBase:
            x = np.append(x, np.zeros((x.shape[0], 1)), axis=1)
        logSumExp = logsumexp(x, axis, keepdims=True)
        logSoftMax = x - logSumExp
        if returnLog:
            return logSoftMax
        else:
            return np.exp(logSoftMax)

def softmaxForData(X, Coef, intercepts=None, addDefaultBase=True, returnLog=False):
    if intercepts is None:
        linearTerms = X.dot(Coef)
    else:
        linearTerms = X.dot(Coef) + intercepts

    return softmax(linearTerms, addDefaultBase, returnLog=returnLog)


import numpy as np
import matplotlib.pyplot as plt
import keras
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout
from keras.optimizers import RMSprop

#自作パッケージ
from data_manager import DataManager as dm

#各種パラメータ

batch_size = 128  # 訓練データを128ずつのデータに分けて学習させる
num_classes = 24 # 分類させる数。今回のコードの総数は24種類
epochs = 20 # 訓練データを繰り返し学習させる数
x_output_path = "./output/x/"
y_output_path = "./output/y/"

codeFiles = [
  "c",
  "c#",
  "d",
  "d#",
  "e",
  "f",
  "f#",
  "g",
  "g#",
  "a",
  "a#",
  "b",
  "cm",
  "c#m",
  "dm",
  "d#m",
  "em",
  "fm",
  "f#m",
  "gm",
  "g#m",
  "am",
  "a#m",
  "bm"
]

# クロマベクトルのデータとその教師データを取得する
dm = dm()
train_range = [1,2,3,4,5,6,7,8,9,10]
test_range = [1,2,3,4,5,6,7,8,9,10]
x_train, y_train, x_test, y_test = dm.get_chromas(train_range=train_range, test_range=test_range)

# y(教師データ)にはコードの名前が文字列として入っているが、Kerasで扱いやすい形(0 or 1)に変換する
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# モデルの作成
model = Sequential()
model.add(Dense(48, activation='sigmoid', input_shape=(None,12)))
model.add(Dense(48, activation='sigmoid'))
model.add(Dense(num_classes, activation='softmax'))

model.compile(loss='categorical_crossentropy',
 optimizer=RMSprop(),
 metrics=['accuracy'])

model.summary()

# 学習
history = model.fit(x_train, y_train,
 batch_size=batch_size,
 epochs=epochs,
 verbose=1,
 validation_data=(x_test, y_test))

# 評価はevaluateで行う
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

# 結果をプロットする
fig, (loss_plot, accuracy_plot) = plt.subplots(ncols=2)

accuracy_plot.plot(history.history['accuracy'])
accuracy_plot.plot(history.history['val_accuracy'])
accuracy_plot.set_title('Model accuracy')
accuracy_plot.set_ylabel('Accuracy')
accuracy_plot.set_xlabel('Epoch')
accuracy_plot.set_ylim(0, 1.00)
accuracy_plot.legend(['Train', 'Test'], loc='upper left')

loss_plot.plot(history.history['loss'])
loss_plot.plot(history.history['val_loss'])
loss_plot.set_title('Model loss')
loss_plot.set_ylabel('Loss')
loss_plot.set_xlabel('Epoch')
loss_plot.legend(['Train', 'Test'], loc='upper left')

fig.tight_layout()
plt.show()


import streamlit as st
import yfinance as yf

def stock_price():
    st.write("""
            Simple stock price App.
            Shown are the stock price and volume of Google
            """)
    # Символ тикера
    ticker_symbol = 'GOOGL'

    ticker_data = yf.Ticker(ticker_symbol)

    ticker_df = ticker_data.history(period='1d', start='2019-01-01')

    return st.line_chart(ticker_df.Close), st.line_chart(ticker_df.Volume)

if __name__ == '__main__':
    stock_price()


from mock import MagicMock, patch
from django.test import TestCase

from auth_backend.resources import base
from auth_backend.resources.base import ObjectResource, Action

from auth_backend.tests.mock_path import *  # noqa

class ObjectResourceTestCase(TestCase):
    def setUp(self):
        self.rtype = 'type_token'
        self.name = 'name_token'
        self.scope_type = 'scope_type_token'
        self.scope_name = 'scope_name_token'
        self.actions = [Action(id='view', name='view', is_instance_related=True),
                        Action(id='edit', name='edit', is_instance_related=True)]
        self.inspect = MagicMock()
        self.scope_id = 'scope_id_token'
        self.parent = MagicMock()
        self.parent.rtype = 'parent_type_token'
        self.operations = [{
            'operate_id': 'view',
            'actions_id': ['view'],
        }, {
            'operate_id': 'edit',
            'actions_id': ['view', 'edit']
        }]
        self.backend = MagicMock()

    def tearDown(self):
        base.resource_type_lib = {}

    @patch(RESOURCE_CLEAN_INSTANCE, MagicMock())
    def test_clean_instances__is_resource_cls(self):
        resource = ObjectResource(rtype=self.rtype,
                                  name=self.name,
                                  scope_type=self.scope_type,
                                  scope_name=self.scope_name,
                                  actions=self.actions,
                                  inspect=self.inspect,
                                  scope_id=self.scope_id,
                                  parent=self.parent,
                                  operations=self.operations,
                                  backend=self.backend,
                                  resource_cls=str)

        instance = 'instance'
        self.assertEqual(resource.clean_instances(instance), instance)
        super(ObjectResource, resource).clean_instances.assert_not_called()

    @patch(RESOURCE_CLEAN_INSTANCE, MagicMock())
    def test_clean_instances__fall_back_to_super(self):
        resource = ObjectResource(rtype=self.rtype,
                                  name=self.name,
                                  scope_type=self.scope_type,
                                  scope_name=self.scope_name,
                                  actions=self.actions,
                                  inspect=self.inspect,
                                  scope_id=self.scope_id,
                                  parent=self.parent,
                                  operations=self.operations,
                                  backend=self.backend,
                                  resource_cls=list)

        instance = 'instance'
        self.assertIsNotNone(resource.clean_instances(instance))
        super(ObjectResource, resource).clean_instances.assert_called_once_with(instance)